Repository: cta-observatory/ctapipe
Pull request #2349 (2023-06-13 16:27:16)

Add support for using toymodel in telescope frame
missing tests

maxnoe: Just because I had the script already from #2347 


https://github.com/cta-observatory/ctapipe/assets/5488440/ad475013-53d6-4e23-a765-84db08397f69



Patch:
diff --git a/ctapipe/image/toymodel.py b/ctapipe/image/toymodel.py
index bb9a9f5a..7c1575e7 100644
--- a/ctapipe/image/toymodel.py
+++ b/ctapipe/image/toymodel.py
@@ -18,15 +18,16 @@ Examples:
     >>> print(image.shape)
     (400,)
 """
-import numpy as np
-from ctapipe.utils import linalg
-from ctapipe.image.hillas import camera_to_shower_coordinates
-import astropy.units as u
-from astropy.coordinates import Angle
-from scipy.stats import multivariate_normal, skewnorm, norm
-from scipy.ndimage import convolve1d
 from abc import ABCMeta, abstractmethod
+
+import astropy.units as u
+import numpy as np
 from numpy.random import default_rng
+from scipy.ndimage import convolve1d
+from scipy.stats import multivariate_normal, norm, skewnorm
+
+from ctapipe.image.hillas import camera_to_shower_coordinates
+from ctapipe.utils import linalg
 
 __all__ = [
     "WaveformModel",
@@ -40,15 +41,6 @@ __all__ = [
 TOYMODEL_RNG = default_rng(0)
 
 
-@u.quantity_input(
-    x=u.m,
-    y=u.m,
-    centroid_x=u.m,
-    centroid_y=u.m,
-    psi=u.deg,
-    time_gradient=u.ns / u.m,
-    time_intercept=u.ns,
-)
 def obtain_time_image(x, y, centroid_x, centroid_y, psi, time_gradient, time_intercept):
     """Create a pulse time image for a toymodel shower. Assumes the time development
     occurs only along the longitudinal (major) axis of the shower, and scales
@@ -68,7 +60,7 @@ def obtain_time_image(x, y, centroid_x, centroid_y, psi, time_gradient, time_int
         Y camera coordinate for the centroid of the shower
     psi : convertible to `astropy.coordinates.Angle`
         rotation angle about the centroid (0=x-axis)
-    time_gradient : u.Quantity[time/length]
+    time_gradient : u.Quantity[time/angle]
         Rate at which the time changes with distance along the shower axis
     time_intercept : u.Quantity[time]
         Pulse time at the shower centroid
@@ -79,11 +71,17 @@ def obtain_time_image(x, y, centroid_x, centroid_y, psi, time_gradient, time_int
         Pulse time in nanoseconds at (x, y)
 
     """
+    unit = x.unit
+    x = x.to_value(unit)
+    y = y.to_value(unit)
+    centroid_x = centroid_x.to_value(unit)
+    centroid_y = centroid_y.to_value(unit)
+    psi = psi.to_value(u.rad)
+    time_gradient = time_gradient.to_value(u.ns / unit)
+    time_intercept = time_intercept.to_value(u.ns)
+
     longitudinal, _ = camera_to_shower_coordinates(x, y, centroid_x, centroid_y, psi)
-    longitudinal_m = longitudinal.to_value(u.m)
-    time_gradient_ns_m = time_gradient.to_value(u.ns / u.m)
-    time_intercept_ns = time_intercept.to_value(u.ns)
-    return longitudinal_m * time_gradient_ns_m + time_intercept_ns
+    return longitudinal * time_gradient + time_intercept
 
 
 class WaveformModel:
@@ -183,11 +181,9 @@ class WaveformModel:
 
 
 class ImageModel(metaclass=ABCMeta):
-    @u.quantity_input(x=u.m, y=u.m)
     @abstractmethod
     def pdf(self, x, y):
-        """Probability density function.
-        """
+        """Probability density function."""
 
     def generate_image(self, camera, intensity=50, nsb_level_pe=20, rng=None):
         """Generate a randomized DL1 shower image.
@@ -244,7 +240,6 @@ class ImageModel(metaclass=ABCMeta):
 
 
 class Gaussian(ImageModel):
-    @u.quantity_input(x=u.m, y=u.m, length=u.m, width=u.m)
     def __init__(self, x, y, length, width, psi):
         """Create 2D Gaussian model for a shower image in a camera.
 
@@ -256,7 +251,7 @@ class Gaussian(ImageModel):
             width of shower (minor axis)
         length: u.Quantity[length]
             length of shower (major axis)
-        psi : convertable to `astropy.coordinates.Angle`
+        psi : u.Quantity[angle]
             rotation angle about the centroid (0=x-axis)
 
         Returns
@@ -264,32 +259,36 @@ class Gaussian(ImageModel):
         a `scipy.stats` object
 
         """
+        self.unit = x.unit
         self.x = x
         self.y = y
         self.width = width
         self.length = length
         self.psi = psi
 
-    @u.quantity_input(x=u.m, y=u.m)
-    def pdf(self, x, y):
-        """2d probability for photon electrons in the camera plane"""
         aligned_covariance = np.array(
-            [[self.length.to_value(u.m) ** 2, 0], [0, self.width.to_value(u.m) ** 2]]
+            [
+                [self.length.to_value(self.unit) ** 2, 0],
+                [0, self.width.to_value(self.unit) ** 2],
+            ]
         )
         # rotate by psi angle: C' = R C R+
         rotation = linalg.rotation_matrix_2d(self.psi)
         rotated_covariance = rotation @ aligned_covariance @ rotation.T
+        self.dist = multivariate_normal(
+            mean=u.Quantity([self.x, self.y]).to_value(self.unit),
+            cov=rotated_covariance,
+        )
 
-        return multivariate_normal(
-            mean=[self.x.to_value(u.m), self.y.to_value(u.m)], cov=rotated_covariance
-        ).pdf(np.column_stack([x.to_value(u.m), y.to_value(u.m)]))
+    def pdf(self, x, y):
+        """2d probability for photon electrons in the camera plane"""
+        X = np.column_stack([x.to_value(self.unit), y.to_value(self.unit)])
+        return self.dist.pdf(X)
 
 
 class SkewedGaussian(ImageModel):
-    """A shower image that has a skewness along the major axis.
-    """
+    """A shower image that has a skewness along the major axis."""
 
-    @u.quantity_input(x=u.m, y=u.m, length=u.m, width=u.m)
     def __init__(self, x, y, length, width, psi, skewness):
         """Create 2D skewed Gaussian model for a shower image in a camera.
         Skewness is only applied along the main shower axis.
@@ -305,7 +304,7 @@ class SkewedGaussian(ImageModel):
             width of shower (minor axis)
         length: u.Quantity[length]
             length of shower (major axis)
-        psi : convertable to `astropy.coordinates.Angle`
+        psi : u.Quantity[angle]
             rotation angle about the centroid (0=x-axis)
 
         Returns
@@ -313,6 +312,7 @@ class SkewedGaussian(ImageModel):
         a `scipy.stats` object
 
         """
+        self.unit = x.unit
         self.x = x
         self.y = y
         self.width = width
@@ -320,6 +320,12 @@ class SkewedGaussian(ImageModel):
         self.psi = psi
         self.skewness = skewness
 
+        a, loc, scale = self._moments_to_parameters()
+        self.long_dist = skewnorm(a=a, loc=loc, scale=scale)
+        self.trans_dist = norm(loc=0, scale=self.width.to_value(self.unit))
+        self.rotation = linalg.rotation_matrix_2d(-self.psi)
+        self.mu = u.Quantity([self.x, self.y]).to_value(self.unit)
+
     def _moments_to_parameters(self):
         """Returns loc and scale from mean, std and skewnewss."""
         # see https://en.wikipedia.org/wiki/Skew_normal_distribution#Estimation
@@ -327,45 +333,36 @@ class SkewedGaussian(ImageModel):
         delta = np.sign(self.skewness) * np.sqrt(
             (np.pi / 2 * skew23) / (skew23 + (0.5 * (4 - np.pi)) ** (2 / 3))
         )
-        a = delta / np.sqrt(1 - delta ** 2)
-        scale = self.length.to_value(u.m) / np.sqrt(1 - 2 * delta ** 2 / np.pi)
+        a = delta / np.sqrt(1 - delta**2)
+        scale = self.length.to_value(self.unit) / np.sqrt(1 - 2 * delta**2 / np.pi)
         loc = -scale * delta * np.sqrt(2 / np.pi)
 
         return a, loc, scale
 
-    @u.quantity_input(x=u.m, y=u.m)
     def pdf(self, x, y):
         """2d probability for photon electrons in the camera plane."""
-        mu = u.Quantity([self.x, self.y]).to_value(u.m)
-
-        rotation = linalg.rotation_matrix_2d(-Angle(self.psi))
-        pos = np.column_stack([x.to_value(u.m), y.to_value(u.m)])
-        long, trans = rotation @ (pos - mu).T
-
-        trans_pdf = norm(loc=0, scale=self.width.to_value(u.m)).pdf(trans)
-
-        a, loc, scale = self._moments_to_parameters()
-
-        return trans_pdf * skewnorm(a=a, loc=loc, scale=scale).pdf(long)
+        pos = np.column_stack([x.to_value(self.unit), y.to_value(self.unit)])
+        long, trans = self.rotation @ (pos - self.mu).T
+        return self.trans_dist.pdf(trans) * self.long_dist.pdf(long)
 
 
 class RingGaussian(ImageModel):
     """A shower image consisting of a ring with gaussian radial profile.
-    Simplified model for a muon ring.
 
+    Simplified model for a muon ring.
     """
 
-    @u.quantity_input(x=u.m, y=u.m, radius=u.m, sigma=u.m)
     def __init__(self, x, y, radius, sigma):
+        self.unit = x.unit
         self.x = x
         self.y = y
         self.sigma = sigma
         self.radius = radius
+        self.dist = norm(
+            self.radius.to_value(self.unit), self.sigma.to_value(self.unit)
+        )
 
-    @u.quantity_input(x=u.m, y=u.m)
     def pdf(self, x, y):
         """2d probability for photon electrons in the camera plane."""
-
         r = np.sqrt((x - self.x) ** 2 + (y - self.y) ** 2)
-
-        return norm(self.radius.to_value(u.m), self.sigma.to_value(u.m)).pdf(r)
+        return self.dist.pdf(r)
diff --git a/ctapipe/io/toymodel.py b/ctapipe/io/toymodel.py
index 6e24d04d..05477f4e 100644
--- a/ctapipe/io/toymodel.py
+++ b/ctapipe/io/toymodel.py
@@ -145,7 +145,7 @@ class ToyEventSource(TelescopeComponent, EventSource):
                 y=y,
                 length=length * u.m,
                 width=width * u.m,
-                psi=f"{psi}d",
+                psi=psi * u.deg,
                 skewness=skewness,
             )
             image, _, _ = model.generate_image(cam, intensity)
diff --git a/docs/changes/2349.feature.rst b/docs/changes/2349.feature.rst
new file mode 100644
index 00000000..8eb063a7
--- /dev/null
+++ b/docs/changes/2349.feature.rst
@@ -0,0 +1,1 @@
+Add support for using ``ctapipe.image.toymodel`` features in ``TelescopeFrame``.


Test patch:
diff --git a/ctapipe/image/tests/test_toy.py b/ctapipe/image/tests/test_toy.py
index 85f76242..8229163f 100644
--- a/ctapipe/image/tests/test_toy.py
+++ b/ctapipe/image/tests/test_toy.py
@@ -6,11 +6,13 @@ import pytest
 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(