Added optika.sensors.AbstractImagingSensor.readout() method (#60)

optika/sensors/_sensors.py

@@ -2,12 +2,13 @@
 Models of light sensors that can be used in optical systems.
 """
 
-from typing import TypeVar
+from typing import TypeVar, Sequence
 import abc
 import dataclasses
 import astropy.units as u
 import named_arrays as na
 import optika
+from . import AbstractImagingSensorMaterial
 
 __all__ = [
     "AbstractImagingSensor",
@@ -16,7 +17,7 @@
 ]
 
 
-MaterialT = TypeVar("MaterialT", bound=optika.materials.AbstractMaterial)
+MaterialT = TypeVar("MaterialT", bound=AbstractImagingSensorMaterial)
 
 
 @dataclasses.dataclass(eq=False, repr=False)
@@ -49,6 +50,14 @@ def width_pixel(self) -> u.Quantity | na.AbstractCartesian2dVectorArray:
         The physical size of each pixel on the sensor.
         """
 
+    @property
+    @abc.abstractmethod
+    def axis_pixel(self) -> na.Cartesian2dVectorArray[str, str]:
+        """
+        The names of the logical axes corresponding to the rows and
+        columns of the pixel grid.
+        """
+
     @property
     @abc.abstractmethod
     def num_pixel(self) -> na.Cartesian2dVectorArray[int, int]:
@@ -72,6 +81,68 @@ def aperture(self):
             half_width=self.width_pixel * self.num_pixel / 2,
         )
 
+    def readout(
+        self,
+        rays: optika.rays.RayVectorArray,
+        timedelta: None | u.Quantity | na.AbstractScalar = None,
+        axis: None | str | Sequence[str] = None,
+        where: bool | na.AbstractScalar = True,
+    ) -> na.FunctionArray[
+        na.Cartesian2dVectorArray,
+        na.AbstractScalar,
+    ]:
+        """
+        Given a set of rays incident on the sensor surface,
+        where each ray represents an expected number of photons per unit time,
+        simulate the number of electrons that would be measured by the sensor.
+
+        Parameters
+        ----------
+        rays
+            A set of incident rays in global coordinates to measure.
+        timedelta
+            The exposure time of the measurement.
+            If :obj:`None` (the default), the value in :attr:`timedelta_exposure`
+            will be used.
+        axis
+            The logical axes along which to collect photons.
+        where
+            A boolean mask used to indicate which photons should be considered
+            when calculating the signal measured by the sensor.
+        """
+        if timedelta is None:
+            timedelta = self.timedelta_exposure
+
+        if self.transformation is not None:
+            rays = self.transformation.inverse(rays)
+
+        where = where & rays.unvignetted
+
+        rays = dataclasses.replace(
+            rays,
+            intensity=rays.intensity * timedelta,
+        )
+
+        electrons = self.material.electrons_measured(
+            rays=rays,
+            normal=self.sag.normal(rays.position),
+        )
+
+        hist = na.histogram2d(
+            x=na.as_named_array(rays.position.x),
+            y=rays.position.y,
+            bins={
+                self.axis_pixel.x: self.num_pixel.x,
+                self.axis_pixel.y: self.num_pixel.y,
+            },
+            axis=axis,
+            min=self.aperture.bound_lower.xy,
+            max=self.aperture.bound_upper.xy,
+            weights=electrons * where,
+        )
+
+        return hist
+
 
 @dataclasses.dataclass(eq=False, repr=False)
 class IdealImagingSensor(
@@ -87,6 +158,12 @@ class IdealImagingSensor(
     width_pixel: u.Quantity | na.AbstractCartesian2dVectorArray = 0 * u.um
     """The physical size of each pixel on the sensor."""
 
+    axis_pixel: na.Cartesian2dVectorArray[str, str] = None
+    """
+    The names of the logical axes corresponding to the rows and 
+    columns of the pixel grid.
+    """
+
     num_pixel: na.Cartesian2dVectorArray[int, int] = None
     """The number of pixels along each axis of the sensor."""
 

optika/sensors/_tests/test_sensors.py

@@ -19,15 +19,48 @@ def test_timedelta_exposure(self, a: optika.sensors.AbstractImagingSensor):
         result = a.timedelta_exposure
         assert result >= 0 * u.s
 
+    @pytest.mark.parametrize(
+        argnames="rays",
+        argvalues=[
+            optika.rays.RayVectorArray(
+                intensity=100 * u.photon / u.s,
+                position=na.Cartesian3dVectorArray() * u.mm,
+            ),
+            optika.rays.RayVectorArray(
+                intensity=na.random.poisson(100, shape_random=dict(t=11)) * u.erg / u.s,
+                wavelength=500 * u.nm,
+                position=na.Cartesian3dVectorArray(
+                    x=na.random.uniform(-1, 1, shape_random=dict(t=11)) * u.mm,
+                    y=na.random.uniform(-1, 1, shape_random=dict(t=11)) * u.mm,
+                    z=0 * u.mm,
+                ),
+            ),
+        ],
+    )
+    def test_readout(
+        self,
+        a: optika.sensors.AbstractImagingSensor,
+        rays: optika.rays.RayVectorArray,
+    ):
+        result = a.readout(rays)
+        assert isinstance(result, na.FunctionArray)
+        assert isinstance(result.inputs, na.Cartesian2dVectorArray)
+        assert isinstance(result.outputs, na.AbstractScalar)
+        assert result.outputs.unit.is_equivalent(u.electron)
+        assert a.axis_pixel.x in result.outputs.shape
+        assert a.axis_pixel.y in result.outputs.shape
+
 
 @pytest.mark.parametrize(
     argnames="a",
     argvalues=[
         optika.sensors.IdealImagingSensor(
             name="test sensor",
             width_pixel=15 * u.um,
+            axis_pixel=na.Cartesian2dVectorArray("detector_x", "detector_y"),
             num_pixel=na.Cartesian2dVectorArray(2048, 1024),
-        )
+            transformation=na.transformations.Cartesian3dTranslation(x=1 * u.mm),
+        ),
     ],
 )
 class TestIdealImagingSensor(