Add support for using toymodel in telescope frame

ctapipe/image/tests/test_toy.py

@@ -6,11 +6,13 @@
 from pytest import approx
 from scipy.stats import norm, poisson, skewnorm
 
+from ctapipe.coordinates.telescope_frame import TelescopeFrame
 from ctapipe.image.toymodel import WaveformModel, obtain_time_image
 
 
+@pytest.mark.parametrize("frame", ["telescope", "camera"])
 @pytest.mark.parametrize("seed", [None, 0])
-def test_intensity(seed, monkeypatch, prod5_lst):
+def test_intensity(seed, frame, monkeypatch, prod5_lst):
     """
     Test generation of the toymodel roughly follows the given intensity.
 
@@ -19,13 +21,19 @@ def test_intensity(seed, monkeypatch, prod5_lst):
     """
     from ctapipe.image import toymodel
 
-    geom = prod5_lst.camera.geometry
+    if frame == "camera":
+        geom = prod5_lst.camera.geometry
+        unit = u.m
+    else:
+        geom = prod5_lst.camera.geometry.transform_to(TelescopeFrame())
+        unit = u.deg
+
+    x, y = u.Quantity([0.2, 0.3], unit)
+    width = 0.05 * unit
+    length = 0.15 * unit
 
-    x, y = u.Quantity([0.2, 0.3], u.m)
-    width = 0.05 * u.m
-    length = 0.15 * u.m
     intensity = 200
-    psi = "30d"
+    psi = 30 * u.deg
 
     # make sure we set a fixed seed for this test even when testing the
     # API without giving the rng
@@ -43,33 +51,43 @@ def test_intensity(seed, monkeypatch, prod5_lst):
         )
 
     # test if signal reproduces given cog values
-    assert np.average(geom.pix_x.to_value(u.m), weights=signal) == approx(0.2, rel=0.15)
-    assert np.average(geom.pix_y.to_value(u.m), weights=signal) == approx(0.3, rel=0.15)
+    assert np.average(geom.pix_x.to_value(unit), weights=signal) == approx(
+        0.2, rel=0.15
+    )
+    assert np.average(geom.pix_y.to_value(unit), weights=signal) == approx(
+        0.3, rel=0.15
+    )
 
     # test if signal reproduces given width/length values
     cov = np.cov(geom.pix_x.value, geom.pix_y.value, aweights=signal)
     eigvals, _ = np.linalg.eigh(cov)
 
-    assert np.sqrt(eigvals[0]) == approx(width.to_value(u.m), rel=0.15)
-    assert np.sqrt(eigvals[1]) == approx(length.to_value(u.m), rel=0.15)
+    assert np.sqrt(eigvals[0]) == approx(width.to_value(unit), rel=0.15)
+    assert np.sqrt(eigvals[1]) == approx(length.to_value(unit), rel=0.15)
 
     # test if total intensity is inside in 99 percent confidence interval
     assert poisson(intensity).ppf(0.05) <= signal.sum() <= poisson(intensity).ppf(0.95)
 
 
-def test_skewed(prod5_lst):
+@pytest.mark.parametrize("frame", ["telescope", "camera"])
+def test_skewed(frame, prod5_lst):
     from ctapipe.image.toymodel import SkewedGaussian
 
     # test if the parameters we calculated for the skew normal
     # distribution produce the correct moments
     rng = np.random.default_rng(0)
-    geom = prod5_lst.camera.geometry
+    if frame == "camera":
+        geom = prod5_lst.camera.geometry
+        unit = u.m
+    else:
+        geom = prod5_lst.camera.geometry.transform_to(TelescopeFrame())
+        unit = u.deg
 
-    x, y = u.Quantity([0.2, 0.3], u.m)
-    width = 0.05 * u.m
-    length = 0.15 * u.m
+    x, y = u.Quantity([0.2, 0.3], unit)
+    width = 0.05 * unit
+    length = 0.15 * unit
     intensity = 50
-    psi = "30d"
+    psi = 30 * u.deg
     skewness = 0.3
 
     model = SkewedGaussian(
@@ -81,20 +99,26 @@ def test_skewed(prod5_lst):
     mean, var, skew = skewnorm(a=a, loc=loc, scale=scale).stats(moments="mvs")
 
     assert np.isclose(mean, 0)
-    assert np.isclose(var, length.to_value(u.m) ** 2)
+    assert np.isclose(var, length.to_value(unit) ** 2)
     assert np.isclose(skew, skewness)
 
 
-def test_compare(prod5_lst):
+@pytest.mark.parametrize("frame", ["telescope", "camera"])
+def test_compare(frame, prod5_lst):
     from ctapipe.image.toymodel import Gaussian, SkewedGaussian
 
-    geom = prod5_lst.camera.geometry
+    if frame == "camera":
+        geom = prod5_lst.camera.geometry
+        unit = u.m
+    else:
+        geom = prod5_lst.camera.geometry.transform_to(TelescopeFrame())
+        unit = u.deg
 
-    x, y = u.Quantity([0.2, 0.3], u.m)
-    width = 0.05 * u.m
-    length = 0.15 * u.m
+    x, y = u.Quantity([0.2, 0.3], unit)
+    width = 0.05 * unit
+    length = 0.15 * unit
     intensity = 50
-    psi = "30d"
+    psi = 30 * u.deg
 
     skewed = SkewedGaussian(x=x, y=y, width=width, length=length, psi=psi, skewness=0)
     normal = Gaussian(x=x, y=y, width=width, length=length, psi=psi)
@@ -105,14 +129,29 @@ def test_compare(prod5_lst):
     assert np.isclose(signal_skewed, signal_normal).all()
 
 
-def test_obtain_time_image(prod5_sst):
+@pytest.mark.parametrize("frame", ["telescope", "camera"])
+def test_obtain_time_image(frame, prod5_sst):
     geom = prod5_sst.camera.geometry
 
-    centroid_x = u.Quantity(0.05, u.m)
-    centroid_y = u.Quantity(0.05, u.m)
+    if frame == "camera":
+        geom = prod5_sst.camera.geometry
+        unit = u.m
+        scale = 1.0
+    else:
+        geom_cam_frame = prod5_sst.camera.geometry
+        geom = geom_cam_frame.transform_to(TelescopeFrame())
+        unit = u.deg
+        # further down we test the std deviation, but that scales
+        # with the size of our shower in the camera
+        r_tel_frame = geom.guess_radius().to_value(u.deg)
+        r_cam_frame = geom_cam_frame.guess_radius().to_value(u.m)
+        scale = r_tel_frame / r_cam_frame
+
+    centroid_x = u.Quantity(0.05, unit)
+    centroid_y = u.Quantity(0.05, unit)
     psi = u.Quantity(70, u.deg)
 
-    time_gradient = u.Quantity(0, u.ns / u.m)
+    time_gradient = u.Quantity(0, u.ns / unit)
     time_intercept = u.Quantity(0, u.ns)
     time = obtain_time_image(
         geom.pix_x,
@@ -125,7 +164,7 @@ def test_obtain_time_image(prod5_sst):
     )
     np.testing.assert_allclose(time, 0)
 
-    time_gradient = u.Quantity(0, u.ns / u.m)
+    time_gradient = u.Quantity(0, u.ns / unit)
     time_intercept = u.Quantity(40, u.ns)
     time = obtain_time_image(
         geom.pix_x,
@@ -138,7 +177,7 @@ def test_obtain_time_image(prod5_sst):
     )
     np.testing.assert_allclose(time, 40)
 
-    time_gradient = u.Quantity(20, u.ns / u.m)
+    time_gradient = u.Quantity(20 / scale, u.ns / unit)
     time_intercept = u.Quantity(40, u.ns)
     time = obtain_time_image(
         geom.pix_x,
@@ -151,7 +190,7 @@ def test_obtain_time_image(prod5_sst):
     )
     np.testing.assert_allclose(time.std(), 1.710435, rtol=0.1)
 
-    time_gradient = u.Quantity(20, u.ns / u.m)
+    time_gradient = u.Quantity(20, u.ns / unit)
     time_intercept = u.Quantity(40, u.ns)
     time = obtain_time_image(
         centroid_x,
@@ -165,13 +204,20 @@ def test_obtain_time_image(prod5_sst):
     np.testing.assert_allclose(time, 40)
 
 
-def test_waveform_model(prod5_sst):
+@pytest.mark.parametrize("frame", ["telescope", "camera"])
+def test_waveform_model(frame, prod5_sst):
     from ctapipe.image.toymodel import Gaussian
 
     prod5_sst = deepcopy(prod5_sst)
-    geom = prod5_sst.camera.geometry
     readout = prod5_sst.camera.readout
 
+    if frame == "camera":
+        geom = prod5_sst.camera.geometry
+        unit = u.m
+    else:
+        geom = prod5_sst.camera.geometry.transform_to(TelescopeFrame())
+        unit = u.deg
+
     ref_duration = 67
     n_ref_samples = 100
     pulse_sigma = 3
@@ -184,12 +230,12 @@ def test_waveform_model(prod5_sst):
     )
     readout.sampling_rate = u.Quantity(2, u.GHz)
 
-    centroid_x = u.Quantity(0.05, u.m)
-    centroid_y = u.Quantity(0.05, u.m)
-    length = u.Quantity(0.03, u.m)
-    width = u.Quantity(0.008, u.m)
+    centroid_x = u.Quantity(0.05, unit)
+    centroid_y = u.Quantity(0.05, unit)
+    length = u.Quantity(0.03, unit)
+    width = u.Quantity(0.008, unit)
     psi = u.Quantity(70, u.deg)
-    time_gradient = u.Quantity(50, u.ns / u.m)
+    time_gradient = u.Quantity(50, u.ns / unit)
     time_intercept = u.Quantity(20, u.ns)
 
     _, charge, _ = Gaussian(