Normunpol batch raytrace

CHANGELOG.md

@@ -13,6 +13,7 @@ ZOS-API can also be added in patch releases.
 
 - Support for Python 3.14 (#221)
 - New `zospy.tools` submodule with `open_tool` (a context manager to open a tool and close it automatically after use) and tool wrappers (#226)
+- New Batch Ray Trace tool: `zospy.tools.BatchRayTraceNormUnpol` (#225)
 
 ### Changed
 

tests/tools/test_batch_raytrace.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from io import StringIO
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+import zospy as zp
+from zospy.tools import BatchRayTraceNormUnpol
+
+if TYPE_CHECKING:
+    from zospy.zpcore import OpticStudioSystem
+
+
+RAY_COORDINATES = [(0, 0), (0.5, 0.5), (-0.5, 0.5), (1, 1), (-1, -1)]
+PUPIL_COORDINATES = [(0, 0), (0, 0), (-1, -1), (0.5, 0.5), (0, 0)]
+EXPECTED_DATA = """
+rayNumber,ErrorCode,vignetteCode,X,Y,Z,L,M,N,l2,m2,n2,opd,intensity
+1,0,0,0.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,-1.0,21.292,1.0
+2,0,0,0.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,-1.0,21.292,1.0
+3,0,0,0.0061213962871815575,0.0061213962871815575,0.0,0.04982748996703953,0.04982748996703953,0.997514131472416,0.0,0.0,-1.0,21.29250706701466,1.0
+4,0,0,-1.180861838340208e-05,-1.180861838340208e-05,0.0,-0.024821950654654583,-0.024821950654654583,0.9993836808410451,0.0,0.0,-1.0,21.292012741140418,1.0
+5,0,0,0.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,-1.0,21.292,1.0
+"""
+
+
+def test_batch_ray_trace_norm_unpol(simple_system: OpticStudioSystem):
+    tool = BatchRayTraceNormUnpol(
+        hx=[c[0] for c in RAY_COORDINATES],
+        hy=[c[1] for c in RAY_COORDINATES],
+        px=[c[0] for c in PUPIL_COORDINATES],
+        py=[c[1] for c in PUPIL_COORDINATES],
+    )
+
+    result = tool.run(simple_system)
+
+    # Compare the actual result with the expected data
+    assert_frame_equal(result.data, pd.read_csv(StringIO(EXPECTED_DATA)))
+
+
+@pytest.mark.parametrize(
+    "surface, expected_coordinate",
+    [
+        (1, (0.0, 0.5, 0.0)),  # Stop surface
+        (2, (0.0, 0.50124164265831, 0.006282066217268977)),  # Lens front surface
+        (3, (0.0, 0.6209696175613508, -0.009642406048217225)),  # Lens back surface
+        (4, (0.0, 3.944854065030327, 0.0)),  # Image surface
+        ("Image", (0.0, 3.944854065030327, 0.0)),  # Image surface
+    ],
+)
+def test_batch_ray_trace_norm_unpol_surface(
+    decentered_system: OpticStudioSystem,
+    surface: int | Literal["Image"],
+    expected_coordinate: tuple[float, float, float],
+):
+    tool = BatchRayTraceNormUnpol(
+        hx=[0.0],
+        hy=[0.5],
+        px=[0.0],
+        py=[0.5],
+        surface=surface,
+    )
+
+    result = tool.run(decentered_system)
+
+    # Check that the coordinates of the first ray match the expected values
+    assert np.allclose(result.data.loc[0, ["X", "Y", "Z"]].values, expected_coordinate)
+
+
+@pytest.mark.parametrize(
+    "wavelength, expected_coordinate",
+    [
+        (1, (0.0, 3.9220149310491603, 0)),  # Wavelength 1
+        (2, (0.0, 3.926593029532985, 0)),  # Wavelength 2
+    ],
+)
+def test_batch_ray_trace_norm_unpol_wavelength(
+    decentered_system: OpticStudioSystem, wavelength: int, expected_coordinate: tuple[float, float, float]
+):
+    decentered_system.SystemData.Wavelengths.AddWavelength(0.6328, 1.0)  # Add a second wavelength to the system
+    zp.solvers.fixed(
+        decentered_system.LDE.GetSurfaceAt(2).MaterialCell
+    )  # Use a material that has a different refractive index at the second wavelength
+    decentered_system.LDE.GetSurfaceAt(2).Material = "BK7"
+
+    tool = BatchRayTraceNormUnpol(
+        hx=[0.0],
+        hy=[0.5],
+        px=[0.0],
+        py=[0.5],
+        wavelength=wavelength,
+    )
+
+    result = tool.run(decentered_system)
+
+    # Check that the coordinates of the first ray match the expected values for the specified wavelength
+    assert np.allclose(result.data.loc[0, ["X", "Y", "Z"]].values, expected_coordinate)

zospy/analyses/base.py

@@ -138,10 +138,16 @@ def _serialize_analysis_data_type(data: AnalysisData) -> _TypeInfo:
 
 
 def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo, module: str) -> object:
-    if typeinfo["module"].startswith(module):
+    module_name = typeinfo.get("module")
+    member_name = typeinfo.get("name")
+
+    if not module_name or not member_name:
+        raise ValueError("Missing module information for zospy class deserialization.")
+
+    if module_name.startswith(module):
         try:
-            m = import_module(typeinfo["module"])
-            t = getattr(m, typeinfo["name"])
+            m = import_module(module_name)
+            t = getattr(m, member_name)
 
             return TypeAdapter(t).validate_python(data)
         except (ModuleNotFoundError, AttributeError):
@@ -236,7 +242,9 @@ def _deserialize_types(cls, data: Any, handler):
             if "__analysis_data__" in data:
                 data["data"] = _deserialize_analysis_data(data["data"], data.pop("__analysis_data__"))
             if "__analysis_settings__" in data:
-                data["settings"] = _deserialize_zospy_class(data["settings"], data.pop("__analysis_settings__"))
+                data["settings"] = _deserialize_zospy_class(
+                    data["settings"], data.pop("__analysis_settings__"), module="zospy.analyses"
+                )
 
         return handler(data)
 

zospy/analyses/decorators.py

@@ -19,7 +19,7 @@
 
 def _default_config_dataclass(default_config: ConfigDict, cls=None, config: ConfigDict | None = None, **kwargs):
     """Pydantic dataclass with default configuration."""
-    config = default_config if config is None else default_config.update(config)
+    config = default_config if config is None else default_config | config
 
     if cls is None:
         # Called with parentheses

zospy/analyses/raysandspots/single_ray_trace.py

@@ -32,8 +32,10 @@ def ray_trace_data_table(self, args):
         # Fill empty columns with NaN
         for row in rows:
             if (row_length := len(row)) < header_length:
-                # Warning is only raised once for the full loop
-                warn("Header and row length mismatch. Empty columns will be filled with NaN.")
+                # Avoid warning just because surface comment is missing.
+                if not isinstance(row[-1], str) and header[-1] != "Comment":
+                    # Warning is only raised once for the full loop
+                    warn("Header and row length mismatch. Empty columns will be filled with NaN.")
 
                 # Check if the last row value is a comment
                 if isinstance(row[-1], str):

zospy/tools/__init__.py

@@ -9,6 +9,16 @@
 from __future__ import annotations
 
 from zospy.tools.base import open_tool
+from zospy.tools.batch_raytrace import (
+    BatchRayTraceNormUnpol,
+    BatchRayTraceNormUnpolSettings,
+)
 from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings
 
-__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool")
+__all__ = (
+    "BatchRayTraceNormUnpol",
+    "BatchRayTraceNormUnpolSettings",
+    "QuickFocus",
+    "QuickFocusSettings",
+    "open_tool",
+)

zospy/tools/batch_raytrace.py

@@ -0,0 +1,193 @@
+"""Batch Ray Trace tool."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Annotated, Literal
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+from pydantic import BeforeValidator, ConfigDict, Field, PositiveInt, model_validator
+
+from zospy.analyses.decorators import analysis_settings
+from zospy.analyses.parsers.types import ZOSAPIConstant  # ruff: ignore[typing-only-first-party-import]
+from zospy.api import constants
+from zospy.tools.base import BaseToolWrapper
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zospy.api import _ZOSAPI
+
+__all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings")
+
+
+def ndarray_to_list(v: Sequence | NDArray) -> Sequence:
+    """Convert ndarray to list."""
+    if isinstance(v, np.ndarray):
+        return v.tolist()
+
+    return v
+
+
+# Constrained types
+NormalizedCoordinate = Annotated[float, Field(le=1.0, ge=-1.0)]
+CoordinateVector = Annotated[
+    Sequence[NormalizedCoordinate] | NDArray[np.floating],
+    BeforeValidator(ndarray_to_list),
+]
+PositiveIntVector = Annotated[
+    Sequence[PositiveInt] | NDArray[np.integer],
+    BeforeValidator(ndarray_to_list),
+]
+
+
+@analysis_settings(config=ConfigDict(arbitrary_types_allowed=True))
+class BatchRayTraceNormUnpolSettings:
+    """Settings for the Batch Ray Trace of unpolarized light.
+
+    Attributes
+    ----------
+    hx : Sequence[float] | np.ndarray
+        Sequence of normalized X field coordinates. Defaults to (0,).
+    hy : Sequence[float] | np.ndarray
+        Sequence of normalized Y field coordinates. Defaults to (0,).
+    px : Sequence[float] | np.ndarray
+        Sequence of normalized X pupil coordinates. Defaults to (0,).
+    py : Sequence[float] | np.ndarray
+        Sequence of normalized Y pupil coordinates. Defaults to (0,).
+    wavelength : int | Sequence[int] | np.ndarray
+        The wavelength number that is to be used. Must be an integer or a sequence of integers specifying the wavelength number.
+        Defaults to 1.
+    surface : str | int
+        Surface up to which the rays will be traced. Either 'Image' or an integer specifying the surface number.
+        Defaults to 'Image'.
+    rays_type : zospy.api.constants.Tools.RayTrace.RaysType
+        Type of rays to trace ('Real' or 'Paraxial'). Defaults to 'Real'.
+    opd_mode : zospy.api.constants.Tools.RayTrace.OPDMode
+        Mode of optical path difference for rays (e.g. 'None'). Defaults to 'None'.
+    """
+
+    hx: CoordinateVector = Field(default=(0,), description="Normalized X field coordinate")
+    hy: CoordinateVector = Field(default=(0,), description="Normalized Y field coordinate")
+    px: CoordinateVector = Field(default=(0,), description="Normalized X pupil coordinate")
+    py: CoordinateVector = Field(default=(0,), description="Normalized Y pupil coordinate")
+    wavelength: PositiveInt | PositiveIntVector = Field(default=1, description="Wavelength number")
+    surface: Literal["Image"] | Annotated[int, Field(ge=0)] = Field(default="Image", description="Surface number")
+    rays_type: ZOSAPIConstant("Tools.RayTrace.RaysType") = Field(default="Real", description="Type of rays to trace")
+    opd_mode: ZOSAPIConstant("Tools.RayTrace.OPDMode") = Field(
+        default="None", description="Mode of optical path difference"
+    )
+
+    @model_validator(mode="after")
+    def validate_lengths(self):
+        """Validate that hx, hy, px, py (and wavelength) have the same lengths."""
+        if not len(self.hx) == len(self.hy) == len(self.px) == len(self.py):
+            raise ValueError("Hx, Hy, Px, Py must all have the same length.")
+        expected_len = len(self.hx)
+
+        if isinstance(self.wavelength, Sequence) and len(self.wavelength) != expected_len:
+            raise ValueError(
+                f"`wavelength` sequence length ({len(self.wavelength)}) must match "
+                f"the length of Hx, Hy, Px, Py ({expected_len})."
+            )
+
+        return self
+
+
+class BatchRayTraceNormUnpol(BaseToolWrapper[pd.DataFrame, BatchRayTraceNormUnpolSettings]):
+    """Batch Ray Trace of unpolarized light."""
+
+    def __init__(
+        self,
+        *,
+        hx: Sequence[float] | NDArray[np.floating] = (0,),
+        hy: Sequence[float] | NDArray[np.floating] = (0,),
+        px: Sequence[float] | NDArray[np.floating] = (0,),
+        py: Sequence[float] | NDArray[np.floating] = (0,),
+        wavelength: int | Sequence[int] | NDArray[np.integer] = 1,
+        surface: Literal["Image"] | int = "Image",
+        rays_type: str | constants.Tools.RayTrace.RaysType = "Real",
+        opd_mode: str | constants.Tools.RayTrace.OPDMode = "None",
+    ):
+        """Create a new Batch Ray Trace of unpolarized light.
+
+        This tool returns a DataFrame with the results of the ray trace in long format. The DataFrame contains the following columns:
+
+        - `rayNumber`: The ray number.
+        - `ErrorCode`: The error code for the ray trace.
+        - `vignetteCode`: The vignette code for the ray trace.
+        - `X`, `Y`, `Z`: The X, Y, Z coordinates of the ray.
+        - `L`, `M`, `N`: The direction cosines of the ray.
+        - `l2`, `m2`, `n2`: The direction cosines of the ray in the local coordinate system.
+        - `opd`: The optical path difference of the ray.
+        - `intensity`: The intensity of the ray.
+
+        See Also
+        --------
+        BatchRayTraceNormUnpolSettings : Settings for the Batch Ray Trace of unpolarized light.
+        """
+        super().__init__(settings_kws=locals())
+
+    def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.RayTrace.IBatchRayTrace]:
+        """Get a callable that opens the batch raytrace tool in OpticStudio and returns the tool object."""
+        return oss.Tools.OpenBatchRayTrace
+
+    def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame:
+        """Run the Batch Ray Trace of unpolarized light.
+
+        Returns
+        -------
+        DataFrame
+            The data in long format.
+        """
+        number_of_rays = len(self.settings.hx)
+        wavelengths = (
+            [self.settings.wavelength] * number_of_rays
+            if isinstance(self.settings.wavelength, int)
+            else self.settings.wavelength
+        )
+        rays_type = constants.process_constant(constants.Tools.RayTrace.RaysType, self.settings.rays_type)
+        opd_mode = constants.process_constant(constants.Tools.RayTrace.OPDMode, self.settings.opd_mode)
+        surface = -1 if self.settings.surface == "Image" else self.settings.surface
+
+        # Initiate batch ray trace
+        norm_unpol_data = tool.CreateNormUnpol(
+            number_of_rays,
+            rays_type,
+            surface,
+        )
+
+        # Add rays
+        for wavelength, hx, hy, px, py in zip(
+            wavelengths, self.settings.hx, self.settings.hy, self.settings.px, self.settings.py, strict=True
+        ):
+            norm_unpol_data.AddRay(wavelength, hx, hy, px, py, opd_mode)
+
+        # Run ray trace and read results
+        tool.RunAndWaitForCompletion()
+        norm_unpol_data.StartReadingResults()
+
+        # Read all results and append to outputs
+        outputs = []
+        columns = [
+            "rayNumber",
+            "ErrorCode",
+            "vignetteCode",
+            "X",
+            "Y",
+            "Z",
+            "L",
+            "M",
+            "N",
+            "l2",
+            "m2",
+            "n2",
+            "opd",
+            "intensity",
+        ]
+        outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(number_of_rays)]
+
+        # Convert to DataFrame and return
+        return pd.DataFrame(outputs, columns=columns)