diff --git a/CHANGELOG.md b/CHANGELOG.md index b5f12997..cb7cfa18 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/tools/test_batch_raytrace.py b/tests/tools/test_batch_raytrace.py new file mode 100644 index 00000000..c4f2886e --- /dev/null +++ b/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) diff --git a/zospy/analyses/base.py b/zospy/analyses/base.py index fea1b107..cf22873e 100644 --- a/zospy/analyses/base.py +++ b/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) diff --git a/zospy/analyses/decorators.py b/zospy/analyses/decorators.py index 718c4d5a..75ef552b 100644 --- a/zospy/analyses/decorators.py +++ b/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 diff --git a/zospy/analyses/raysandspots/single_ray_trace.py b/zospy/analyses/raysandspots/single_ray_trace.py index ae33f2ab..acaa1be2 100644 --- a/zospy/analyses/raysandspots/single_ray_trace.py +++ b/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): diff --git a/zospy/tools/__init__.py b/zospy/tools/__init__.py index 10e37d35..26b8bda6 100644 --- a/zospy/tools/__init__.py +++ b/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", +) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py new file mode 100644 index 00000000..2507d696 --- /dev/null +++ b/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)