Adds optika.systems.SequentialSystem class

optika/__init__.py

@@ -5,3 +5,5 @@
 from . import apertures
 from . import rulings
 from . import surfaces
+from . import propagators
+from . import systems

optika/_tests/test_sags.py

@@ -181,7 +181,7 @@ class TestConicSag(
     argnames="sag",
     argvalues=[
         optika.sags.ParabolicSag(
-            radius=radius,
+            focal_length=radius / 2,
             transformation=transformation,
         )
         for radius in radius_parameterization()

optika/_tests/test_systems.py

@@ -0,0 +1,106 @@
+import pytest
+import astropy.units as u
+import named_arrays as na
+import optika
+from . import test_mixins
+from . import test_plotting
+
+
+class AbstractTestAbstractSystem(
+    test_plotting.AbstractTestPlottable,
+    test_mixins.AbstractTestTransformable,
+):
+    pass
+
+
+class AbstractTestAbstractSequentialSystem(
+    AbstractTestAbstractSystem,
+):
+    def test_object(self, a: optika.systems.AbstractSequentialSystem):
+        if a.object is not None:
+            assert isinstance(a.object, optika.surfaces.AbstractSurface)
+
+    def test_surfaces(self, a: optika.systems.AbstractSequentialSystem):
+        for surface in a.surfaces:
+            assert isinstance(surface, optika.surfaces.AbstractSurface)
+
+    def test_axis_surface(self, a: optika.systems.AbstractSequentialSystem):
+        assert isinstance(a.axis_surface, str)
+
+    def test_surfaces_all(self, a: optika.systems.AbstractSequentialSystem):
+        for surface in a.surfaces_all:
+            assert isinstance(surface, optika.surfaces.AbstractSurface)
+
+    def grid_input_normalized(self, a: optika.systems.AbstractSequentialSystem):
+        assert isinstance(a.grid_input_normalized, optika.vectors.ObjectVectorArray)
+
+    def test_index_field_stop(self, a: optika.systems.AbstractSequentialSystem):
+        assert isinstance(a.index_field_stop, int)
+        assert a.surfaces_all[a.index_field_stop].is_field_stop
+
+    def test_index_pupil_stop(self, a: optika.systems.AbstractSequentialSystem):
+        assert isinstance(a.index_pupil_stop, int)
+        assert a.surfaces_all[a.index_pupil_stop].is_pupil_stop
+
+    def test_field_stop(self, a: optika.systems.AbstractSequentialSystem):
+        assert a.field_stop.is_field_stop
+
+    def test_pupil_stop(self, a: optika.systems.AbstractSequentialSystem):
+        assert a.pupil_stop.is_pupil_stop
+
+    def test_raytrace(self, a: optika.systems.AbstractSequentialSystem):
+        raytrace = a.raytrace
+        assert isinstance(raytrace, optika.rays.RayFunctionArray)
+        assert isinstance(raytrace.inputs, optika.vectors.ObjectVectorArray)
+        assert isinstance(raytrace.outputs, optika.rays.RayVectorArray)
+        assert a.axis_surface in raytrace.shape
+
+    def test_rayfunction(self, a: optika.systems.AbstractSequentialSystem):
+        rayfunction = a.rayfunction
+        assert isinstance(rayfunction, optika.rays.RayFunctionArray)
+        assert isinstance(rayfunction.inputs, optika.vectors.ObjectVectorArray)
+        assert isinstance(rayfunction.outputs, optika.rays.RayVectorArray)
+        assert a.axis_surface not in rayfunction.shape
+
+
+@pytest.mark.parametrize(
+    argnames="a",
+    argvalues=[
+        optika.systems.SequentialSystem(
+            surfaces=[
+                optika.surfaces.Surface(
+                    name="mirror",
+                    sag=optika.sags.SphericalSag(-200 * u.mm),
+                    material=optika.materials.Mirror(),
+                    aperture=optika.apertures.CircularAperture(20 * u.mm),
+                    is_pupil_stop=True,
+                    transformation=na.transformations.Cartesian3dTranslation(
+                        z=100 * u.mm
+                    ),
+                ),
+                optika.surfaces.Surface(
+                    name="detector",
+                    aperture=optika.apertures.RectangularAperture(10 * u.mm),
+                    is_field_stop=True,
+                ),
+            ],
+            grid_input_normalized=optika.vectors.ObjectVectorArray(
+                wavelength=500 * u.nm,
+                field=na.Cartesian2dVectorLinearSpace(
+                    start=0,
+                    stop=1,
+                    axis=na.Cartesian2dVectorArray("field_x", "field_y"),
+                    num=5,
+                ),
+                pupil=na.Cartesian2dVectorLinearSpace(
+                    start=0,
+                    stop=1,
+                    axis=na.Cartesian2dVectorArray("pupil_x", "pupil_y"),
+                    num=5,
+                ),
+            ),
+        )
+    ],
+)
+class TestSequentialSystem(AbstractTestAbstractSequentialSystem):
+    pass

optika/propagators.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+from typing import Sequence
 import abc
 import dataclasses
+import named_arrays as na
 import optika
 
 __all__ = [
@@ -8,6 +11,63 @@
 ]
 
 
+def propagate_rays(
+    propagators: AbstractRayPropagator | Sequence[AbstractRayPropagator],
+    rays: optika.rays.RayVectorArray,
+) -> optika.rays.RayVectorArray:
+    """
+    Iterate through a sequence of ray propagators, calling
+    :meth:`optika.propagators.AbstractRayPropagator.propagate_rays` on the given
+    set of input rays.
+
+    Parameters
+    ----------
+    propagators
+        a sequence of ray propagators to interact with ``rays``
+    rays
+        the input rays to propagate through the sequence
+    """
+    if isinstance(propagators, AbstractRayPropagator):
+        propagators = [propagators]
+
+    for propagator in propagators:
+        rays = propagator.propagate_rays(rays)
+
+    return rays
+
+
+def accumulate_rays(
+    propagators: AbstractRayPropagator | Sequence[AbstractRayPropagator],
+    rays: optika.rays.RayVectorArray,
+    axis: str,
+) -> optika.rays.RayVectorArray:
+    """
+    Iterate through a sequence of ray propagators, calling
+    :meth:`optika.propagators.AbstractRayPropagator.propagate_rays` on the given
+    set of input rays, and store the resulting the rays at every propagator.
+
+    Parameters
+    ----------
+    propagators
+        a sequence of ray propagators to interact with ``rays``
+    rays
+        the input rays to propagate through the sequence
+    axis
+        the new axis representing the sequence of propagators
+    """
+    if isinstance(propagators, AbstractRayPropagator):
+        propagators = [propagators]
+
+    result = [rays]
+    for propagator in propagators:
+        rays = propagator.propagate_rays(rays)
+        result.append(rays)
+
+    result = na.stack(result, axis=axis)
+
+    return result
+
+
 @dataclasses.dataclass(eq=False, repr=False)
 class AbstractPropagator(
     abc.ABC,

optika/rays/__init__.py

@@ -1 +1,2 @@
 from ._ray_vectors import *
+from ._ray_functions import *

optika/rays/_ray_functions.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+import dataclasses
+import named_arrays as na
+import optika
+from . import RayVectorArray
+
+__all__ = [
+    "AbstractRayFunctionArray",
+    "RayFunctionArray",
+]
+
+
+@dataclasses.dataclass(eq=False, repr=False)
+class AbstractRayFunctionArray(
+    na.AbstractFunctionArray,
+):
+    @property
+    def type_explicit(self) -> type[RayFunctionArray]:
+        return RayFunctionArray
+
+
+@dataclasses.dataclass(eq=False, repr=False)
+class RayFunctionArray(
+    AbstractRayFunctionArray,
+    na.FunctionArray[
+        optika.vectors.ObjectVectorArray,
+        RayVectorArray,
+    ],
+):
+    pass

optika/sags.py

@@ -28,6 +28,10 @@
     "RadiusOfRotationT",
     bound=float | u.Quantity | na.AbstractScalar,
 )
+FocalLengthT = TypeVar(
+    "FocalLengthT",
+    bound=float | u.Quantity | na.AbstractScalar,
+)
 
 
 @dataclasses.dataclass
@@ -375,9 +379,13 @@ class ConicSag(
 class ParabolicSag(
     AbstractConicSag[RadiusT, int],
 ):
-    radius: RadiusT = np.inf * u.mm
+    focal_length: FocalLengthT = np.inf * u.mm
     transformation: None | na.transformations.AbstractTransformation = None
 
+    @property
+    def radius(self) -> RadiusT:
+        return 2 * self.focal_length
+
     @property
     def conic(self) -> int:
         return -1