From d6d4b7197dc32c74b97adc0fbd2948545d30e626 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Tue, 2 Jun 2026 18:20:48 +0200 Subject: [PATCH 01/35] Add Batch Ray Trace NormUnpol as analysis. --- .../analyses/raysandspots/batch_ray_trace.py | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 zospy/analyses/raysandspots/batch_ray_trace.py diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py new file mode 100644 index 00000000..e99de709 --- /dev/null +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -0,0 +1,154 @@ +"""Batch Ray Trace analysis.""" + +from __future__ import annotations + +from typing import Literal, Annotated, Union + +import pandas as pd +from pydantic import confloat, conint, Field, model_validator +from System import Int32, Double + +from zospy.analyses.base import BaseAnalysisWrapper +from zospy.analyses.decorators import analysis_settings +from zospy.api import constants + +__all__ = ("NormUnpol", "NormUnpolSettings") + + +# Constrained types +FloatInRange = Annotated[float, confloat(le=1.0, ge=-1.0)] +PositiveInt = Annotated[int, conint(gt=0)] + +@analysis_settings +class NormUnpolSettings: + """Settings for the Batch Ray Trace of unpolarized light. + + Attributes + ---------- + hx : list[float] + List of normalized X field coordinates. Defaults to [0]. + hy : list[float] + List of normalized Y field coordinates. Defaults to [0]. + px : list[float] + List of normalized X pupil coordinates. Defaults to [0]. + py : list[float] + List of normalized Y pupil coordinates. Defaults to [0]. + wavelength : int | list[int] + The wavelength number that is to be used. Must be an integer or a list 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: list[FloatInRange] = Field(default=[0], description="Normalized X field coordinate") + hy: list[FloatInRange] = Field(default=[0], description="Normalized Y field coordinate") + px: list[FloatInRange] = Field(default=[0], description="Normalized X pupil coordinate") + py: list[FloatInRange] = Field(default=[0], description="Normalized Y pupil coordinate") + wavelength: Union[PositiveInt, list[PositiveInt]] = Field(default=1, description="Wavelength number") + surface: Literal["Image"] | Annotated[int, Field(ge=0)] = Field(default="Image", description="Surface number") + rays_type: Literal["Real", "Paraxial"] = Field(default="Real", description="Type of rays to trace") + opd_mode: Literal["None", "Current", "CurrentAndChief"] = 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.""" + lengths = [len(self.hx), len(self.hy), len(self.px), len(self.py)] + if len(set(lengths)) != 1: + raise ValueError( + f"Hx, Hy, Px, Py must all have the same length; got lengths {lengths}." + ) + expected_len = lengths[0] + + if isinstance(self.wavelength, list): + if len(self.wavelength) != expected_len: + raise ValueError( + f"`wavelength` list length ({len(self.wavelength)}) must match " + f"the length of Hx, Hy, Px, Py ({expected_len})." + ) + + return self + + +class NormUnpol(BaseAnalysisWrapper[pd.DataFrame | None, NormUnpolSettings], + analysis_type="RayTrace", +): + """Batch Ray Trace of unpolarized light.""" + + def __init__( + self, + *, + hx: list[float] = [0], + hy: list[float] = [0], + px: list[float] = [0], + py: list[float] = [0], + wavelength: int | list[int] = 1, + surface: Literal["Image"] | int = "Image", + rays_type: str = "Real", + opd_mode: str = "None", + ): + """Create a new Batch Ray Trace of unpolarized light. + + See Also + -------- + NormUnpolSettings : Settings for the Batch Ray Trace of unpolarized light. + """ + super().__init__(settings_kws=locals()) + + def run_analysis(self) -> pd.DataFrame | None: + """Run the Batch Ray Trace of unpolarized light.""" + number_of_rays = len(self.settings.hx) + if isinstance(self.settings.wavelength, int): + wvls = [self.settings.wavelength]*number_of_rays + else: + wvls = 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 + raytrace = self.oss.Tools.OpenBatchRayTrace() + normUnPolData = raytrace.CreateNormUnpol( + number_of_rays, + rays_type, + surface, + ) + + # Add rays + for ii in range(number_of_rays): + normUnPolData.AddRay( + wvls[ii], + self.settings.hx[ii], + self.settings.hy[ii], + self.settings.px[ii], + self.settings.py[ii], + opd_mode, + ) + + # Run ray trace and read results + self.analysis.ApplyAndWaitForCompletion() + raytrace.RunAndWaitForCompletion() + normUnPolData.StartReadingResults() + + # Create placeholders for all output arguments + placeholders = [Int32(1)]*3 + [Double(1.0)]*11 + + # Read all results and append to outputs + outputs = [] + columns = ["rayNumber","ErrorCode","vignetteCode", + "X","Y","Z","L","M","N", + "l2","m2","n2","opd","Intensity"] + for ii in range(number_of_rays): + outputs.append(normUnPolData.ReadNextResult(*placeholders)[1:]) + + # Convert to DataFrame and return + return pd.DataFrame(outputs, columns=columns) + From 44d0618a65787a4ada8b75ade6b675fb9aac149e Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Tue, 2 Jun 2026 18:25:33 +0200 Subject: [PATCH 02/35] Add batch ray trace to __init__.py. --- zospy/analyses/raysandspots/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/zospy/analyses/raysandspots/__init__.py b/zospy/analyses/raysandspots/__init__.py index a6ed31e0..f5614847 100644 --- a/zospy/analyses/raysandspots/__init__.py +++ b/zospy/analyses/raysandspots/__init__.py @@ -7,5 +7,8 @@ SingleRayTrace, SingleRayTraceSettings, ) +from zospy.analyses.raysandspots.batch_ray_trace import NormUnpol as BatchRayTraceNormUnpol +from zospy.analyses.raysandspots.batch_ray_trace import NormUnpolSettings as BatchRayTraceNormUnpolSettings -__all__ = ("RayFan", "RayFanSettings", "SingleRayTrace", "SingleRayTraceSettings") +__all__ = ("RayFan", "RayFanSettings", "SingleRayTrace", "SingleRayTraceSettings", + "BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") From 35329cda2d1599b745263173171d0c8b79e819d4 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:01:41 +0200 Subject: [PATCH 03/35] Close batch raytrace. --- zospy/analyses/raysandspots/batch_ray_trace.py | 1 + 1 file changed, 1 insertion(+) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index e99de709..03a514c7 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -148,6 +148,7 @@ def run_analysis(self) -> pd.DataFrame | None: "l2","m2","n2","opd","Intensity"] for ii in range(number_of_rays): outputs.append(normUnPolData.ReadNextResult(*placeholders)[1:]) + raytrace.Close() # Convert to DataFrame and return return pd.DataFrame(outputs, columns=columns) From ce8feca356c87ad5238b5e4f00a3aee1b837b91a Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:03:38 +0200 Subject: [PATCH 04/35] Avoid warning just because surface comment is missing. --- zospy/analyses/raysandspots/single_ray_trace.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zospy/analyses/raysandspots/single_ray_trace.py b/zospy/analyses/raysandspots/single_ray_trace.py index b6450f9b..983e7315 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): From 3e3b88831248817dd41408ab98cbf2c660c3ad2a Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:17:28 +0200 Subject: [PATCH 05/35] Add context manager for managing tools --- tests/functions/test_tools.py | 79 +++++++++++++++++++++++++++++++++++ zospy/functions/__init__.py | 5 ++- zospy/functions/tools.py | 63 ++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 2 deletions(-) create mode 100644 tests/functions/test_tools.py create mode 100644 zospy/functions/tools.py diff --git a/tests/functions/test_tools.py b/tests/functions/test_tools.py new file mode 100644 index 00000000..93a40047 --- /dev/null +++ b/tests/functions/test_tools.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from zospy.functions.tools import open_tool + +if TYPE_CHECKING: + from pytest_mock import MockerFixture + + from zospy.zpcore import OpticStudioSystem + + +class TestOpenTool: + def test_open_tool(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + assert oss.Tools.CurrentTool == tool + + assert oss.Tools.CurrentTool is None + + def test_close_current(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + oss.Tools.OpenLocalOptimization() + + with open_tool(oss, oss.Tools.OpenBatchRayTrace, close_current=True) as tool: + # Testing that the previous tool was closed messes up the OpticStudio connection, + # but since only one tool can be open at a time, this should be sufficient to verify that the previous tool was closed. + assert oss.Tools.CurrentTool == tool + + assert oss.Tools.CurrentTool is None + + def test_no_close_current(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + local_optimization = oss.Tools.OpenLocalOptimization() + + with ( + pytest.raises(RuntimeError, match="Cannot open tool because another tool is already open"), + open_tool(oss, oss.Tools.OpenBatchRayTrace, close_current=False), + ): + pass + + assert oss.Tools.CurrentTool == local_optimization + + def test_closes_after_exception(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with ( + pytest.raises(ValueError, match="An error occurred while using the tool"), + open_tool(oss, oss.Tools.OpenBatchRayTrace), + ): + raise ValueError("An error occurred while using the tool") + + assert oss.Tools.CurrentTool is None + + def test_does_not_close_if_already_closed(self, oss: OpticStudioSystem, mocker: MockerFixture): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + tool.Close() + + spy = mocker.spy(tool, "Close") + + assert oss.Tools.CurrentTool is None + spy.assert_not_called() + + def test_does_not_close_other_tool(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + tool.Close() + + tool2 = oss.Tools.OpenLocalOptimization() + + assert oss.Tools.CurrentTool == tool2 diff --git a/zospy/functions/__init__.py b/zospy/functions/__init__.py index 29ed5907..3046a76c 100644 --- a/zospy/functions/__init__.py +++ b/zospy/functions/__init__.py @@ -4,10 +4,11 @@ - **`zospy.functions.lde`** provides helper functions for the Lens Data Editor (LDE); - **`zospy.functions.nce`** provides helper functions for the Non-sequential Component Editor (NCE). +- **`zospy.functions.tools`** provides helper functions for Tools in OpticStudio. """ from __future__ import annotations -from zospy.functions import lde, nce +from zospy.functions import lde, nce, tools -__all__ = ("lde", "nce") +__all__ = ("lde", "nce", "tools") diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py new file mode 100644 index 00000000..17f537ea --- /dev/null +++ b/zospy/functions/tools.py @@ -0,0 +1,63 @@ +"""Utility functions for Tools in OpticStudio.""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from zospy.api import _ZOSAPI + from zospy.zpcore import OpticStudioSystem + +logger = logging.getLogger(__name__) + + +@contextmanager +def open_tool( + oss: OpticStudioSystem, tool: Callable[..., _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False +) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: + """Context manager for opening a tool in OpticStudio. + + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it can + will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system to use for opening the tool. + tool : Callable[..., _ZOSAPI.Tools.ISystemTool] + A callable that returns an instance of the tool to be opened. + close_current : bool + Whether to close the currently open tool before opening the new one. Defaults to False. + + Yields + ------ + _ZOSAPI.Tools.ISystemTool + An instance of the opened tool. + + Raises + ------ + RuntimeError + If there is already an open tool and `close_current` is False. + """ + if oss.Tools.CurrentTool is not None: + if close_current: + logger.warning( + "A tool is already open. Closing the currently open tool (%s) before opening the new one.", + oss.Tools.CurrentTool.__class__.__name__, + ) + oss.Tools.CurrentTool.Close() + else: + raise RuntimeError("Cannot open tool because another tool is already open.") + + new_tool = tool() + + try: + yield new_tool + finally: + if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: + logger.info("Closing tool %s.", new_tool.__class__.__name__) + new_tool.Close() From e91317c273b51d1b64c9cc2f0bc1a0a366906419 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:22:36 +0200 Subject: [PATCH 06/35] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 080d916e..82c1ab27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ ZOS-API can also be added in patch releases. ### Added - Support for Python 3.14 (#221) +- `zospy.functions.tools.open_tool`, a context manager to open a tool and close it automatically after use (#226) ### Changed From a00fc924e3a7b43fbf4b5559be8df0c9dccc202c Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:26:40 +0200 Subject: [PATCH 07/35] Fix Copilot review comments --- zospy/functions/tools.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py index 17f537ea..72e2e49a 100644 --- a/zospy/functions/tools.py +++ b/zospy/functions/tools.py @@ -12,16 +12,18 @@ from zospy.api import _ZOSAPI from zospy.zpcore import OpticStudioSystem +__all__ = ("open_tool",) + logger = logging.getLogger(__name__) @contextmanager def open_tool( - oss: OpticStudioSystem, tool: Callable[..., _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False + oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False ) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: """Context manager for opening a tool in OpticStudio. - Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it can + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. Parameters From 5ec9de1dc7dc50ce5c1a6793045a0db093ae3bb4 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:54:24 +0200 Subject: [PATCH 08/35] Iterate over rays with zip Co-authored-by: crnh <30109443+crnh@users.noreply.github.com> --- zospy/analyses/raysandspots/batch_ray_trace.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 03a514c7..2b400890 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -123,15 +123,14 @@ def run_analysis(self) -> pd.DataFrame | None: ) # Add rays - for ii in range(number_of_rays): - normUnPolData.AddRay( - wvls[ii], - self.settings.hx[ii], - self.settings.hy[ii], - self.settings.px[ii], - self.settings.py[ii], - opd_mode, - ) + for wavelength, hx, hy, px, py in zip( + wavelengths, + self.settings.hx, + self.settings.hy, + self.settings.px, + self.settings.py, + ): + normUnPolData.AddRay(wavelength, hx, hy, px, py, opd_mode) # Run ray trace and read results self.analysis.ApplyAndWaitForCompletion() From d4dabe0315e6b52cc1d53c1aca876a3f12b46c04 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 11:17:25 +0200 Subject: [PATCH 09/35] Rename classes at import --- zospy/analyses/raysandspots/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zospy/analyses/raysandspots/__init__.py b/zospy/analyses/raysandspots/__init__.py index f5614847..f0b5f372 100644 --- a/zospy/analyses/raysandspots/__init__.py +++ b/zospy/analyses/raysandspots/__init__.py @@ -7,8 +7,10 @@ SingleRayTrace, SingleRayTraceSettings, ) -from zospy.analyses.raysandspots.batch_ray_trace import NormUnpol as BatchRayTraceNormUnpol -from zospy.analyses.raysandspots.batch_ray_trace import NormUnpolSettings as BatchRayTraceNormUnpolSettings +from zospy.analyses.raysandspots.batch_ray_trace import ( + BatchRayTraceNormUnpol, + BatchRayTraceNormUnpolSettings, +) __all__ = ("RayFan", "RayFanSettings", "SingleRayTrace", "SingleRayTraceSettings", "BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") From a9c74f40d4f470d890be130b18723fb15c9f15d4 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 11:19:54 +0200 Subject: [PATCH 10/35] Rename batch raytrace classes --- zospy/analyses/raysandspots/batch_ray_trace.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 2b400890..33b5b5d9 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -12,7 +12,7 @@ from zospy.analyses.decorators import analysis_settings from zospy.api import constants -__all__ = ("NormUnpol", "NormUnpolSettings") +__all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") # Constrained types @@ -20,7 +20,7 @@ PositiveInt = Annotated[int, conint(gt=0)] @analysis_settings -class NormUnpolSettings: +class BatchRayTraceNormUnpolSettings: """Settings for the Batch Ray Trace of unpolarized light. Attributes @@ -74,7 +74,7 @@ def validate_lengths(self): return self -class NormUnpol(BaseAnalysisWrapper[pd.DataFrame | None, NormUnpolSettings], +class BatchRayTraceNormUnpol(BaseAnalysisWrapper[pd.DataFrame | None, BatchRayTraceNormUnpolSettings], analysis_type="RayTrace", ): """Batch Ray Trace of unpolarized light.""" @@ -95,7 +95,7 @@ def __init__( See Also -------- - NormUnpolSettings : Settings for the Batch Ray Trace of unpolarized light. + BatchRayTraceNormUnpolSettings : Settings for the Batch Ray Trace of unpolarized light. """ super().__init__(settings_kws=locals()) From d5e15859977ffaa71147a3bcdcb2f92cce50c742 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 11:23:29 +0200 Subject: [PATCH 11/35] Rename wavelengths variable --- zospy/analyses/raysandspots/batch_ray_trace.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 33b5b5d9..039bcda3 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -103,9 +103,9 @@ def run_analysis(self) -> pd.DataFrame | None: """Run the Batch Ray Trace of unpolarized light.""" number_of_rays = len(self.settings.hx) if isinstance(self.settings.wavelength, int): - wvls = [self.settings.wavelength]*number_of_rays + wavelengths = [self.settings.wavelength]*number_of_rays else: - wvls = self.settings.wavelength + wavelengths = self.settings.wavelength rays_type = constants.process_constant( constants.Tools.RayTrace.RaysType, self.settings.rays_type ) From cf6949ac7b99c8550253dda5f5a5ef5cdabe37e0 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 11:42:23 +0200 Subject: [PATCH 12/35] Use Sequence instead of list --- .../analyses/raysandspots/batch_ray_trace.py | 62 +++++++++---------- 1 file changed, 29 insertions(+), 33 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 039bcda3..7229ac7e 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Literal, Annotated, Union +from typing import Literal, Annotated, Union, Sequence import pandas as pd from pydantic import confloat, conint, Field, model_validator @@ -16,7 +16,7 @@ # Constrained types -FloatInRange = Annotated[float, confloat(le=1.0, ge=-1.0)] +NormalizedCoordinate = Annotated[float, confloat(le=1.0, ge=-1.0)] PositiveInt = Annotated[int, conint(gt=0)] @analysis_settings @@ -25,16 +25,16 @@ class BatchRayTraceNormUnpolSettings: Attributes ---------- - hx : list[float] - List of normalized X field coordinates. Defaults to [0]. - hy : list[float] - List of normalized Y field coordinates. Defaults to [0]. - px : list[float] - List of normalized X pupil coordinates. Defaults to [0]. - py : list[float] - List of normalized Y pupil coordinates. Defaults to [0]. - wavelength : int | list[int] - The wavelength number that is to be used. Must be an integer or a list of integers specifying the wavelength number. + hx : Sequence[float] + Sequence of normalized X field coordinates. Defaults to [0]. + hy : Sequence[float] + Sequence of normalized Y field coordinates. Defaults to [0]. + px : Sequence[float] + Sequence of normalized X pupil coordinates. Defaults to [0]. + py : Sequence[float] + Sequence of normalized Y pupil coordinates. Defaults to [0]. + wavelength : int | Sequence[int] + 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. @@ -45,11 +45,11 @@ class BatchRayTraceNormUnpolSettings: Mode of optical path difference for rays (e.g. 'None'). Defaults to 'None'. """ - hx: list[FloatInRange] = Field(default=[0], description="Normalized X field coordinate") - hy: list[FloatInRange] = Field(default=[0], description="Normalized Y field coordinate") - px: list[FloatInRange] = Field(default=[0], description="Normalized X pupil coordinate") - py: list[FloatInRange] = Field(default=[0], description="Normalized Y pupil coordinate") - wavelength: Union[PositiveInt, list[PositiveInt]] = Field(default=1, description="Wavelength number") + hx: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized X field coordinate") + hy: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized Y field coordinate") + px: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized X pupil coordinate") + py: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized Y pupil coordinate") + wavelength: PositiveInt | Sequence[PositiveInt] = Field(default=1, description="Wavelength number") surface: Literal["Image"] | Annotated[int, Field(ge=0)] = Field(default="Image", description="Surface number") rays_type: Literal["Real", "Paraxial"] = Field(default="Real", description="Type of rays to trace") opd_mode: Literal["None", "Current", "CurrentAndChief"] = Field(default="None", description="Mode of optical path difference") @@ -57,24 +57,23 @@ class BatchRayTraceNormUnpolSettings: @model_validator(mode="after") def validate_lengths(self): """Validate that hx, hy, px, py (and wavelength) have the same lengths.""" - lengths = [len(self.hx), len(self.hy), len(self.px), len(self.py)] - if len(set(lengths)) != 1: + if not len(self.hx) == len(self.hy) == len(self.px) == len(self.py): raise ValueError( - f"Hx, Hy, Px, Py must all have the same length; got lengths {lengths}." + f"Hx, Hy, Px, Py must all have the same length." ) - expected_len = lengths[0] + expected_len = len(self.hx) - if isinstance(self.wavelength, list): + if isinstance(self.wavelength, Sequence): if len(self.wavelength) != expected_len: raise ValueError( - f"`wavelength` list length ({len(self.wavelength)}) must match " + f"`wavelength` sequence length ({len(self.wavelength)}) must match " f"the length of Hx, Hy, Px, Py ({expected_len})." ) return self -class BatchRayTraceNormUnpol(BaseAnalysisWrapper[pd.DataFrame | None, BatchRayTraceNormUnpolSettings], +class BatchRayTraceNormUnpol(BaseAnalysisWrapper[pd.DataFrame, BatchRayTraceNormUnpolSettings], analysis_type="RayTrace", ): """Batch Ray Trace of unpolarized light.""" @@ -82,11 +81,11 @@ class BatchRayTraceNormUnpol(BaseAnalysisWrapper[pd.DataFrame | None, BatchRayTr def __init__( self, *, - hx: list[float] = [0], - hy: list[float] = [0], - px: list[float] = [0], - py: list[float] = [0], - wavelength: int | list[int] = 1, + hx: Sequence[float] = (0,), + hy: Sequence[float] = (0,), + px: Sequence[float] = (0,), + py: Sequence[float] = (0,), + wavelength: int | Sequence[int] = 1, surface: Literal["Image"] | int = "Image", rays_type: str = "Real", opd_mode: str = "None", @@ -102,10 +101,7 @@ def __init__( def run_analysis(self) -> pd.DataFrame | None: """Run the Batch Ray Trace of unpolarized light.""" number_of_rays = len(self.settings.hx) - if isinstance(self.settings.wavelength, int): - wavelengths = [self.settings.wavelength]*number_of_rays - else: - wavelengths = self.settings.wavelength + 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 ) From c016b02095c29b3ee3d33ddd1f71b739a9637c4a Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 12:13:48 +0200 Subject: [PATCH 13/35] Remove out parameters for ray trace data and use list comprehension. --- zospy/analyses/raysandspots/batch_ray_trace.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 7229ac7e..6b690394 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -6,7 +6,6 @@ import pandas as pd from pydantic import confloat, conint, Field, model_validator -from System import Int32, Double from zospy.analyses.base import BaseAnalysisWrapper from zospy.analyses.decorators import analysis_settings @@ -133,16 +132,12 @@ def run_analysis(self) -> pd.DataFrame | None: raytrace.RunAndWaitForCompletion() normUnPolData.StartReadingResults() - # Create placeholders for all output arguments - placeholders = [Int32(1)]*3 + [Double(1.0)]*11 - # Read all results and append to outputs outputs = [] columns = ["rayNumber","ErrorCode","vignetteCode", "X","Y","Z","L","M","N", "l2","m2","n2","opd","Intensity"] - for ii in range(number_of_rays): - outputs.append(normUnPolData.ReadNextResult(*placeholders)[1:]) + outputs = [normUnPolData.ReadNextResult()[1:] for _ in range(number_of_rays)] raytrace.Close() # Convert to DataFrame and return From c6b8f00511271e64f30629a3eef5678d217374f6 Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 12:22:16 +0200 Subject: [PATCH 14/35] Change pydantic types and remove None as output --- zospy/analyses/raysandspots/batch_ray_trace.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 6b690394..0433de0d 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -2,10 +2,11 @@ from __future__ import annotations -from typing import Literal, Annotated, Union, Sequence +from typing import Literal, Annotated +from collections.abc import Sequence import pandas as pd -from pydantic import confloat, conint, Field, model_validator +from pydantic import Field, PositiveInt, model_validator from zospy.analyses.base import BaseAnalysisWrapper from zospy.analyses.decorators import analysis_settings @@ -15,8 +16,7 @@ # Constrained types -NormalizedCoordinate = Annotated[float, confloat(le=1.0, ge=-1.0)] -PositiveInt = Annotated[int, conint(gt=0)] +NormalizedCoordinate = Annotated[float, Field(le=1.0, ge=-1.0)] @analysis_settings class BatchRayTraceNormUnpolSettings: @@ -97,7 +97,7 @@ def __init__( """ super().__init__(settings_kws=locals()) - def run_analysis(self) -> pd.DataFrame | None: + def run_analysis(self) -> pd.DataFrame: """Run the Batch Ray Trace of unpolarized light.""" number_of_rays = len(self.settings.hx) wavelengths = [self.settings.wavelength] * number_of_rays if isinstance(self.settings.wavelength, int) else self.settings.wavelength From 8cd0a4033d37c9c8f6897e33bc48700d2b43d00f Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 12:51:04 +0200 Subject: [PATCH 15/35] Use ZOSAPIConstant and fix case. --- zospy/analyses/raysandspots/batch_ray_trace.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index 0433de0d..b9be3f15 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -10,6 +10,7 @@ from zospy.analyses.base import BaseAnalysisWrapper from zospy.analyses.decorators import analysis_settings +from zospy.analyses.parsers.types import ZOSAPIConstant from zospy.api import constants __all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") @@ -50,8 +51,8 @@ class BatchRayTraceNormUnpolSettings: py: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized Y pupil coordinate") wavelength: PositiveInt | Sequence[PositiveInt] = Field(default=1, description="Wavelength number") surface: Literal["Image"] | Annotated[int, Field(ge=0)] = Field(default="Image", description="Surface number") - rays_type: Literal["Real", "Paraxial"] = Field(default="Real", description="Type of rays to trace") - opd_mode: Literal["None", "Current", "CurrentAndChief"] = Field(default="None", description="Mode of optical path difference") + 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): @@ -86,8 +87,8 @@ def __init__( py: Sequence[float] = (0,), wavelength: int | Sequence[int] = 1, surface: Literal["Image"] | int = "Image", - rays_type: str = "Real", - opd_mode: str = "None", + 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. @@ -111,7 +112,7 @@ def run_analysis(self) -> pd.DataFrame: # Initiate batch ray trace raytrace = self.oss.Tools.OpenBatchRayTrace() - normUnPolData = raytrace.CreateNormUnpol( + norm_unpol_data = raytrace.CreateNormUnpol( number_of_rays, rays_type, surface, @@ -125,19 +126,19 @@ def run_analysis(self) -> pd.DataFrame: self.settings.px, self.settings.py, ): - normUnPolData.AddRay(wavelength, hx, hy, px, py, opd_mode) + norm_unpol_data.AddRay(wavelength, hx, hy, px, py, opd_mode) # Run ray trace and read results self.analysis.ApplyAndWaitForCompletion() raytrace.RunAndWaitForCompletion() - normUnPolData.StartReadingResults() + 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 = [normUnPolData.ReadNextResult()[1:] for _ in range(number_of_rays)] + outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(number_of_rays)] raytrace.Close() # Convert to DataFrame and return From 142bceaca6e803b250238385ba487afb836c1b72 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:29:12 +0200 Subject: [PATCH 16/35] Add base classes for tool wrappers and implement quick focus tool --- tests/tools/__init__.py | 0 .../test_tools.py => tools/test_base.py} | 2 +- zospy/__init__.py | 3 +- zospy/analyses/__init__.py | 2 +- zospy/functions/tools.py | 65 ---- zospy/tools/__init__.py | 14 + zospy/tools/base.py | 353 ++++++++++++++++++ zospy/tools/quick_focus.py | 65 ++++ 8 files changed, 436 insertions(+), 68 deletions(-) create mode 100644 tests/tools/__init__.py rename tests/{functions/test_tools.py => tools/test_base.py} (98%) delete mode 100644 zospy/functions/tools.py create mode 100644 zospy/tools/__init__.py create mode 100644 zospy/tools/base.py create mode 100644 zospy/tools/quick_focus.py diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/functions/test_tools.py b/tests/tools/test_base.py similarity index 98% rename from tests/functions/test_tools.py rename to tests/tools/test_base.py index 93a40047..acb12892 100644 --- a/tests/functions/test_tools.py +++ b/tests/tools/test_base.py @@ -4,7 +4,7 @@ import pytest -from zospy.functions.tools import open_tool +from zospy.tools import open_tool if TYPE_CHECKING: from pytest_mock import MockerFixture diff --git a/zospy/__init__.py b/zospy/__init__.py index 6d5be33a..aba7cc1d 100644 --- a/zospy/__init__.py +++ b/zospy/__init__.py @@ -5,7 +5,7 @@ import logging from importlib.metadata import version -from zospy import analyses, functions, solvers +from zospy import analyses, functions, solvers, tools from zospy.api import config, constants from zospy.zpcore import ZOS @@ -17,6 +17,7 @@ "constants", "functions", "solvers", + "tools", ) config.set_decimal_point_and_thousands_separator() diff --git a/zospy/analyses/__init__.py b/zospy/analyses/__init__.py index 2853daa3..c04e995a 100644 --- a/zospy/analyses/__init__.py +++ b/zospy/analyses/__init__.py @@ -2,7 +2,7 @@ Analyses in OpticStudio are available through `zospy.analyses`. This module provides basic classes and functions for interacting with -analyses (in `zospy.analyses.base`), as well as a procedural interface +analyses (in `zospy.analyses.base`), as well as an object-oriented interface to several analyses. Examples diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py deleted file mode 100644 index 72e2e49a..00000000 --- a/zospy/functions/tools.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Utility functions for Tools in OpticStudio.""" - -from __future__ import annotations - -import logging -from contextlib import contextmanager -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable, Generator - - from zospy.api import _ZOSAPI - from zospy.zpcore import OpticStudioSystem - -__all__ = ("open_tool",) - -logger = logging.getLogger(__name__) - - -@contextmanager -def open_tool( - oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False -) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: - """Context manager for opening a tool in OpticStudio. - - Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it - will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. - - Parameters - ---------- - oss : OpticStudioSystem - The OpticStudio system to use for opening the tool. - tool : Callable[..., _ZOSAPI.Tools.ISystemTool] - A callable that returns an instance of the tool to be opened. - close_current : bool - Whether to close the currently open tool before opening the new one. Defaults to False. - - Yields - ------ - _ZOSAPI.Tools.ISystemTool - An instance of the opened tool. - - Raises - ------ - RuntimeError - If there is already an open tool and `close_current` is False. - """ - if oss.Tools.CurrentTool is not None: - if close_current: - logger.warning( - "A tool is already open. Closing the currently open tool (%s) before opening the new one.", - oss.Tools.CurrentTool.__class__.__name__, - ) - oss.Tools.CurrentTool.Close() - else: - raise RuntimeError("Cannot open tool because another tool is already open.") - - new_tool = tool() - - try: - yield new_tool - finally: - if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: - logger.info("Closing tool %s.", new_tool.__class__.__name__) - new_tool.Close() diff --git a/zospy/tools/__init__.py b/zospy/tools/__init__.py new file mode 100644 index 00000000..10e37d35 --- /dev/null +++ b/zospy/tools/__init__.py @@ -0,0 +1,14 @@ +"""OpticStudio Tools. + +Tools in OpticStudio are available through `zospy.tools`. +This module provides basic classes and functions for interacting with +tools (in `zospy.tools.base`), as well as an object-oriented interface +to several tools. +""" + +from __future__ import annotations + +from zospy.tools.base import open_tool +from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings + +__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool") diff --git a/zospy/tools/base.py b/zospy/tools/base.py new file mode 100644 index 00000000..c9f88861 --- /dev/null +++ b/zospy/tools/base.py @@ -0,0 +1,353 @@ +"""Base functionality for tools.""" + +from __future__ import annotations + +import dataclasses +import logging +import weakref +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import is_dataclass +from types import NoneType +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, get_args + +import numpy as np +import pandas as pd +import pydantic + +from zospy.analyses.base import _deserialize_analysis_data, _deserialize_zospy_class, _serialize_analysis_data_type +from zospy.analyses.parsers.types import ValidatedDataFrame + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from zospy.api import _ZOSAPI + from zospy.zpcore import OpticStudioSystem + + +__all__ = ("BaseToolWrapper", "ToolResult", "open_tool") + +logger = logging.getLogger(__name__) + + +@contextmanager +def open_tool( + oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False +) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: + """Context manager for opening a tool in OpticStudio. + + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it + will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system to use for opening the tool. + tool : Callable[..., _ZOSAPI.Tools.ISystemTool] + A callable that returns an instance of the tool to be opened. + close_current : bool + Whether to close the currently open tool before opening the new one. Defaults to False. + + Yields + ------ + _ZOSAPI.Tools.ISystemTool + An instance of the opened tool. + + Raises + ------ + RuntimeError + If there is already an open tool and `close_current` is False. + """ + if oss.Tools.CurrentTool is not None: + if close_current: + logger.warning( + "A tool is already open. Closing the currently open tool (%s) before opening the new one.", + oss.Tools.CurrentTool.__class__.__name__, + ) + oss.Tools.CurrentTool.Close() + else: + raise RuntimeError("Cannot open tool because another tool is already open.") + + new_tool = tool() + + try: + yield new_tool + finally: + if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: + logger.info("Closing tool %s.", new_tool.__class__.__name__) + new_tool.Close() + + +ToolOutputData = TypeVar("ToolOutputData") +ToolSettings = TypeVar("ToolSettings") + + +@pydantic.dataclasses.dataclass(frozen=True, config=pydantic.ConfigDict(ser_json_inf_nan="constants")) +class ToolResult(Generic[ToolOutputData, ToolSettings]): + """Zemax OpticStudio tool result. + + Attributes + ---------- + data : ToolOutputData + The data of the tool. Can be a `pandas.DataFrame`, `numpy.ndarray`, or a tool-specific dataclass. + settings : ToolSettings | None + The settings of the tool. + error_message : str | None + Error message from the analysis. If the tool ran successfully, this will be None. + """ + + data: ToolOutputData + settings: ToolSettings | None + error_message: str | None + + def to_json(self): + """Convert the result to a JSON string.""" + return pydantic.RootModel(self).model_dump_json(indent=4) + + @classmethod + def from_json(cls, data: str): + """Create a result from a JSON string.""" + return pydantic.TypeAdapter(cls).validate_json(data) + + @pydantic.field_serializer("data", mode="wrap", when_used="json") + @staticmethod + def _serialize_data( + value: ToolOutputData, + nxt: pydantic.SerializerFunctionWrapHandler, + info, # noqa: ARG004 + ): + if isinstance(value, pd.DataFrame): + return pydantic.TypeAdapter( + ValidatedDataFrame, config=pydantic.ConfigDict(ser_json_inf_nan="constants") + ).dump_python(value, mode="json") + + if isinstance(value, np.ndarray): + return value.tolist() + + return nxt(value) + + @pydantic.model_serializer(mode="wrap", when_used="json") + def _serialize_types(self, nxt: pydantic.SerializerFunctionWrapHandler): + data = nxt(self) + data["__tool_data__"] = _serialize_analysis_data_type(self.data) + data["__tool_settings__"] = { + "data_type": "zospy_class", + "name": type(self.settings).__name__, + "module": type(self.settings).__module__, + } + + return data + + @pydantic.model_validator(mode="wrap") + @classmethod + def _deserialize_types(cls, data: Any, handler): + if isinstance(data, dict): + if "__tool_data__" in data: + data["data"] = _deserialize_analysis_data(data["data"], data.pop("__tool_data__")) + if "__tool_settings__" in data: + data["settings"] = _deserialize_zospy_class(data["settings"], data.pop("__tool_settings__")) + + return handler(data) + + +class BaseToolWrapper(ABC, Generic[ToolOutputData, ToolSettings]): + """Base class for tool wrappers. + + This class provides a common interface for all tool wrappers. It defines the methods and properties that all + tool wrappers should implement. + + Attributes + ---------- + CONNECTION_MODE : {"standalone", "extension"} | None + The connection mode required for the tool. If None, the tool can be used in any mode. + """ + + CONNECTION_MODE: Literal["standalone", "extension"] | None = None + + def __init__(self, *, settings_kws: dict[str, Any] | None = None): + """Create a new tool wrapper. + + Settings can be changed by passing the settings as keyword arguments. Use the `with_settings` method to specify + the settings using a settings object. + + Parameters + ---------- + settings_kws : dict[str, Any] + Arguments to set the settings of the tool. + + Raises + ------ + ValueError + If `settings` is not a dataclass. + """ + self._settings = self._default_settings() + self.update_settings(settings_kws=settings_kws) + + self._oss = None + + def __init_subclass__( + cls, + *, + connection_mode: Literal["standalone", "extension"] | None = None, + **kwargs, + ): + """Determine the settings type and class-level configuration of the tool.""" + cls.CONNECTION_MODE = connection_mode + + if not hasattr(cls, "_settings_type"): + if hasattr(cls, "__orig_bases__"): + base = cls.__orig_bases__[0] + cls._settings_type: type[ToolSettings] = get_args(base)[1] + else: + cls._settings_type = NoneType + + super().__init_subclass__(**kwargs) + + @abstractmethod + def _get_tool_opener(self, oss: OpticStudioSystem) -> Callable[[], _ZOSAPI.Tools.ISystemTool]: + """Callable that opens the tool in OpticStudio and returns the tool object.""" + raise NotImplementedError("Subclasses of BaseToolWrapper must implement the `tool_opener` property.") + + def update_settings( + self, + *, + settings: ToolSettings | None = None, + settings_kws: dict[str, Any] | None = None, + ) -> None: + """Update the settings of the tool using a settings object or keyword arguments. + + Settings can be specified as an object and as keyword arguments. If both are specified, the keyword arguments + take precedence. If no settings are specified, the default settings are used. Furthermore, instead of using + a reference to the settings object, a new settings object is created with the specified parameters. This is done + to avoid modifying the original settings object. + + Parameters + ---------- + settings : ToolSettings + Tool settings object. + settings_kws + Dictionary with the settings parameters. + + Raises + ------ + ValueError + If `settings` is not a dataclass. + """ + # Use the existing settings if no settings are specified + settings = settings or self.settings + + if settings is None: + # Analysis does not have settings + return + + if not is_dataclass(settings): + raise TypeError("settings should be a dataclass.") + + # Create a new settings object with the specified parameters. If no parameters are specified, this creates a + # copy of the settings object. This is done to avoid modifying the original settings object. + self._settings = dataclasses.replace(settings, **(settings_kws or {})) + + @classmethod + def _default_settings(cls) -> ToolSettings: + """Get the default settings of the tool. + + Returns + ------- + ToolSettings + The default settings. + """ + return cls._settings_type() + + @classmethod + def with_settings(cls, settings: ToolSettings): + """Create a new tool with the specified settings. + + Parameters + ---------- + settings : ToolSettings + Settings of the tool. + + Returns + ------- + BaseAnalysisWrapper + The tool wrapper. + """ + instance = cls() + instance.update_settings(settings=settings) + + return instance + + @property + def settings(self) -> ToolSettings: + """Settings of the tool.""" + return self._settings + + @property + def oss(self) -> OpticStudioSystem: + """The OpticStudio system. This property is set when the tool is run.""" + if self._oss is None: + raise ValueError("OpticStudioSystem has not been set.") + + return self._oss + + @abstractmethod + def _run_tool(self, tool: _ZOSAPI.Tools.ISystemTool, *args, **kwargs) -> ToolOutputData: + """Run the tool and return the results.""" + + def _check_mode(self): + connection_modes: dict[str, str] = { + "Server": "standalone", + "Plugin": "extension", + } + + if self.CONNECTION_MODE is None: + return + + if str(self.oss.ZOS.Application.Mode) not in connection_modes: + raise ValueError(f"Unknown connection mode: {self.oss.ZOS.Application.Mode}") + + current_mode = connection_modes[str(self.oss.ZOS.Application.Mode)] + + if current_mode != self.CONNECTION_MODE: + raise ValueError(f"The tool requires {self.CONNECTION_MODE} mode, got {current_mode}.") + + def run( + self, + oss: OpticStudioSystem, + *, + close_current: bool = False, + ) -> ToolResult[ToolOutputData, ToolSettings]: + """Run the tool and return the results. + + This method opens and runs the tool in OpticStudio with the specified settings. + If another tool is already open, it will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system. + close_current : bool + Whether to close the current tool if one is already open. + + Returns + ------- + ToolResult + The tool result. + """ + self._oss = weakref.proxy(oss) + self._check_mode() + + with open_tool(oss, self._get_tool_opener(oss), close_current=close_current) as tool: + data = self._run_tool(tool) + + error_message = tool.ErrorMessage + + return ToolResult( + data=data, + settings=self.settings, + error_message=error_message, + ) + + def __call__(self, oss: OpticStudioSystem, *args, **kwargs): + """Run the tool and return the results.""" + return self.run(oss, *args, **kwargs) diff --git a/zospy/tools/quick_focus.py b/zospy/tools/quick_focus.py new file mode 100644 index 00000000..580965a5 --- /dev/null +++ b/zospy/tools/quick_focus.py @@ -0,0 +1,65 @@ +"""Quick Focus Tool.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pydantic import Field + +from zospy.analyses.decorators import analysis_settings +from zospy.analyses.parsers.types import ZOSAPIConstant # noqa: TC001 +from zospy.api import constants +from zospy.api.constants import process_constant +from zospy.tools.base import BaseToolWrapper + +if TYPE_CHECKING: + from collections.abc import Callable + + from zospy.api import _ZOSAPI + +__all__ = ("QuickFocus", "QuickFocusSettings") + + +@analysis_settings +class QuickFocusSettings: + """Settings for the Quick Focus tool. + + Attributes + ---------- + criterion : constants.Tools.General.QuickFocusCriterion | str + The criterion to use for quick focusing. Defaults to 'SpotSizeRadial'. + use_centroid : bool + Reference all calculations to the image centroid rather than the chief ray. Defaults to True. + """ + + criterion: ZOSAPIConstant("Tools.General.QuickFocusCriterion") = Field(default="SpotSizeRadial") + use_centroid: bool = Field(default=True) + + +class QuickFocus(BaseToolWrapper[None, QuickFocusSettings]): + """Wrapper for the Quick Focus tool.""" + + def __init__( + self, + *, + criterion: str | constants.Tools.General.QuickFocusCriterion = "SpotSizeRadial", + use_centroid: bool = True, + ) -> None: + """Initialize the Quick Focus tool. + + See Also + -------- + QuickFocusSettings : Settings for the Quick Focus tool. + """ + super().__init__(settings_kws=locals()) + + def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.General.IQuickFocus]: + """Get a callable that opens the Quick Focus tool in OpticStudio and returns the tool object.""" + return oss.Tools.OpenQuickFocus + + def _run_tool(self, tool: _ZOSAPI.Tools.General.IQuickFocus) -> None: + """Run the Quick Focus tool.""" + tool.Criterion = process_constant(constants.Tools.General.QuickFocusCriterion, self.settings.criterion) + tool.UseCentroid = self.settings.use_centroid + + tool.RunAndWaitForCompletion() From 162fbec53f89c8771442a09a3d2f65873f8a7021 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:34:01 +0200 Subject: [PATCH 17/35] Add test for quick focus tool --- tests/tools/test_quick_focus.py | 19 +++++++++++++++++++ zospy/functions/__init__.py | 4 ++-- 2 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 tests/tools/test_quick_focus.py diff --git a/tests/tools/test_quick_focus.py b/tests/tools/test_quick_focus.py new file mode 100644 index 00000000..69b57f67 --- /dev/null +++ b/tests/tools/test_quick_focus.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from zospy.tools import QuickFocus + +if TYPE_CHECKING: + from zospy.zpcore import OpticStudioSystem + + +def test_quick_focus(simple_system: OpticStudioSystem): + """Test the Quick Focus tool.""" + simple_system.LDE.GetSurfaceAt(3).Thickness = 30 + + QuickFocus().run(simple_system) + + assert simple_system.LDE.GetSurfaceAt(3).Thickness == pytest.approx(19.777975) diff --git a/zospy/functions/__init__.py b/zospy/functions/__init__.py index 3046a76c..e974a3bb 100644 --- a/zospy/functions/__init__.py +++ b/zospy/functions/__init__.py @@ -9,6 +9,6 @@ from __future__ import annotations -from zospy.functions import lde, nce, tools +from zospy.functions import lde, nce -__all__ = ("lde", "nce", "tools") +__all__ = ("lde", "nce") From 1d28bbb42641c0de14f90da0115cec2965038307 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 15:18:30 +0200 Subject: [PATCH 18/35] Add tests for tools base functionality --- tests/analyses/conftest.py | 5 +- tests/helpers.py | 4 + tests/tools/test_base.py | 219 ++++++++++++++++++++++++++++++++++++- zospy/analyses/base.py | 8 +- zospy/tools/base.py | 8 +- 5 files changed, 232 insertions(+), 12 deletions(-) diff --git a/tests/analyses/conftest.py b/tests/analyses/conftest.py index 1fe2ea16..cbd5989e 100644 --- a/tests/analyses/conftest.py +++ b/tests/analyses/conftest.py @@ -7,6 +7,7 @@ import pytest from tests.config import CONFIG_DATA_FOLDER, REFERENCE_DATA_FOLDER, REFERENCE_VERSION +from tests.helpers import _all_subclasses from zospy.analyses.base import AnalysisResult, BaseAnalysisWrapper from zospy.analyses.psf.huygens_psf import BaseHuygensPSF from zospy.analyses.systemviewers.base import SystemViewerWrapper @@ -56,10 +57,6 @@ def reference_data(request) -> AnalysisResult: return AnalysisResult.from_json(data_file.read_text(encoding="utf-8")) -def _all_subclasses(cls): - return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in _all_subclasses(c)]) - - _analysis_wrapper_classes = [ c for c in _all_subclasses(BaseAnalysisWrapper) if c not in {BaseHuygensPSF, SystemViewerWrapper} ] diff --git a/tests/helpers.py b/tests/helpers.py index 01281c5c..abcb02f1 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -66,3 +66,7 @@ def assert_dataclass_equal( return True return None + + +def _all_subclasses(cls): + return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in _all_subclasses(c)]) diff --git a/tests/tools/test_base.py b/tests/tools/test_base.py index acb12892..135a8342 100644 --- a/tests/tools/test_base.py +++ b/tests/tools/test_base.py @@ -1,12 +1,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import inspect +import json +from dataclasses import dataclass, fields +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +import numpy as np import pytest +from pandas import DataFrame +from pydantic import Field +from pydantic.fields import FieldInfo +from tests.helpers import _all_subclasses +from zospy.analyses.decorators import analysis_settings +from zospy.analyses.parsers.types import ValidatedDataFrame from zospy.tools import open_tool +from zospy.tools.base import BaseToolWrapper, ToolResult, ToolSettings +from zospy.tools.quick_focus import QuickFocusSettings if TYPE_CHECKING: + from collections.abc import Callable + from pytest_mock import MockerFixture from zospy.zpcore import OpticStudioSystem @@ -77,3 +92,205 @@ def test_does_not_close_other_tool(self, oss: OpticStudioSystem): tool2 = oss.Tools.OpenLocalOptimization() assert oss.Tools.CurrentTool == tool2 + + +@dataclass +class MockToolOutputData: + int_data: int = 1 + string_data: str = "a" + + +@analysis_settings +class MockToolSettings: + int_setting: int = Field(default=1, description="An integer setting") + string_setting: str = Field(default="a", description="A string setting") + + +class MockTool(BaseToolWrapper[MockToolOutputData, MockToolSettings]): + def __init__( + self, + *, + int_setting: int = 1, + string_setting: str = "a", + ): + super().__init__(settings_kws=locals()) + + def _get_tool_opener(self, oss: OpticStudioSystem) -> Callable[[], Any]: # noqa: ARG002 + return lambda: SimpleNamespace( + Close=lambda: None, ErrorMessage=None + ) # This tool does not actually open anything in OpticStudio + + def _run_tool(self) -> MockToolOutputData: + return MockToolOutputData() + + +_tool_wrapper_classes = [c for c in _all_subclasses(BaseToolWrapper) if c not in {MockTool}] + + +@pytest.fixture(scope="module", params=_tool_wrapper_classes) +def tool_wrapper_class(request): + return request.param + + +class TestToolWrapper: + @staticmethod + def get_settings_defaults(settings_class): + result = {} + + for field in fields(settings_class): + if isinstance(field.default, FieldInfo): + result[field.name] = field.default.default + else: + result[field.name] = field.default + + return result + + def test_get_settings_type(self): + assert MockTool._settings_type == MockToolSettings # noqa: SLF001 + + def test_settings_type_is_specified(self): + assert MockTool._settings_type is not ToolSettings # noqa: SLF001 + + def test_tools_correct_tool_opener(self, oss: OpticStudioSystem, tool_wrapper_class): + instance = tool_wrapper_class() + + try: + instance._get_tool_opener(oss) # noqa: SLF001 + except Exception: # noqa: BLE001 + pytest.fail(f"{tool_wrapper_class.__name__} does not implement _get_tool_opener correctly.") + + def test_init_all_keyword_only_parameters(self, tool_wrapper_class): + assert all(p.kind.name == "KEYWORD_ONLY" for p in inspect.signature(tool_wrapper_class).parameters.values()) + + def test_init_contains_all_settings(self, tool_wrapper_class): + if tool_wrapper_class().settings is None: + return + + init_signature = inspect.signature(tool_wrapper_class.__init__) + settings_fields = fields(tool_wrapper_class().settings) + + assert all(field.name in init_signature.parameters for field in settings_fields) + + def test_tools_default_values(self, tool_wrapper_class): + if tool_wrapper_class().settings is None: + return + + settings_defaults = self.get_settings_defaults(type(tool_wrapper_class().settings)) + init_signature = inspect.signature(tool_wrapper_class.__init__) + + for field_name, default_value in settings_defaults.items(): + assert field_name in init_signature.parameters + assert init_signature.parameters[field_name].default == default_value + + def test_change_settings_from_parameters(self): + tool = MockTool(int_setting=2, string_setting="b") + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_change_settings_from_object(self): + settings = MockToolSettings(int_setting=2, string_setting="b") + tool = MockTool().with_settings(settings) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_settings_object_is_copied(self): + settings = MockToolSettings(int_setting=2, string_setting="b") + tool = MockTool().with_settings(settings) + + assert tool.settings is not settings + assert tool.settings == settings + + def test_update_settings_object(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings(settings=MockToolSettings(int_setting=2, string_setting="b")) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_dictionary(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings(settings_kws={"int_setting": 2, "string_setting": "b"}) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_object_and_dictionary(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings( + settings=MockToolSettings(int_setting=2, string_setting="a"), settings_kws={"string_setting": "b"} + ) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_no_dataclass_raises_type_error(self): + with pytest.raises(TypeError, match="settings should be a dataclass"): + MockTool().update_settings(settings=123) + + +class TestAnalysisResultJSONConversion: + # Only test for non-dataclass results, because dataclass results are tested separately in the corresponding + # analysis' tests. + @pytest.mark.parametrize( + "result_type,result_value,type_info", + [ + (DataFrame, DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), {"data_type": "dataframe"}), + (np.ndarray, np.array([1, 2, 3]), {"data_type": "ndarray"}), + ], + ) + def test_data_to_json(self, result_type, result_value, type_info): + result = ToolResult[result_type, MockToolSettings]( + data=result_value, + settings=MockToolSettings(), + error_message=None, + ) + + result_json = result.to_json() + result_dict = json.loads(result_json) + + assert "__tool_data__" in result_dict + assert result_dict["__tool_data__"] == type_info + + def test_settings_to_json(self): + result = ToolResult[ValidatedDataFrame, MockToolSettings]( + data=DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), + settings=MockToolSettings(), + error_message=None, + ) + + result_json = result.to_json() + result_dict = json.loads(result_json) + + assert "__tool_settings__" in result_dict + assert result_dict["__tool_settings__"] == { + "data_type": "zospy_class", + "name": "MockToolSettings", + "module": "tests.tools.test_base", + } + + @pytest.mark.parametrize( + "result_type,result_value", + [ + (ValidatedDataFrame, DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})), + (np.ndarray, np.array([1, 2, 3])), + ], + ) + def test_roundtrip(self, result_type, result_value): + result = ToolResult[result_type, QuickFocusSettings]( + data=result_value, + settings=QuickFocusSettings(), + error_message=None, + ) + + result_json = result.to_json() + + result_roundtrip = ToolResult.from_json(result_json) + assert all(result_roundtrip.data == result.data) + assert str(result_roundtrip.settings.criterion) == str(result.settings.criterion) + assert result_roundtrip.settings.use_centroid == result.settings.use_centroid + assert result_roundtrip.error_message == result.error_message diff --git a/zospy/analyses/base.py b/zospy/analyses/base.py index 4a6507e3..63e42d6b 100644 --- a/zospy/analyses/base.py +++ b/zospy/analyses/base.py @@ -137,8 +137,8 @@ def _serialize_analysis_data_type(data: AnalysisData) -> _TypeInfo: raise ValueError(f"Cannot serialize data type: {type(data)}") -def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo) -> AnalysisData: - if typeinfo["module"].startswith("zospy.analyses"): +def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo, module: str) -> object: + if typeinfo["module"].startswith(module): try: m = import_module(typeinfo["module"]) t = getattr(m, typeinfo["name"]) @@ -150,7 +150,7 @@ def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo) -> AnalysisData: return data -def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> AnalysisData: +def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> object: if typeinfo["data_type"] == "none": return None @@ -161,7 +161,7 @@ def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> Analys return np.array(data) if typeinfo["data_type"] == "zospy_class": - return _deserialize_zospy_class(data, typeinfo) + return _deserialize_zospy_class(data, typeinfo, module="zospy.analyses") raise ValueError(f"Cannot deserialize data type: {typeinfo['data_type']}") diff --git a/zospy/tools/base.py b/zospy/tools/base.py index c9f88861..c99dfba7 100644 --- a/zospy/tools/base.py +++ b/zospy/tools/base.py @@ -90,14 +90,14 @@ class ToolResult(Generic[ToolOutputData, ToolSettings]): ---------- data : ToolOutputData The data of the tool. Can be a `pandas.DataFrame`, `numpy.ndarray`, or a tool-specific dataclass. - settings : ToolSettings | None + settings : ToolSettings The settings of the tool. error_message : str | None Error message from the analysis. If the tool ran successfully, this will be None. """ data: ToolOutputData - settings: ToolSettings | None + settings: ToolSettings error_message: str | None def to_json(self): @@ -145,7 +145,9 @@ def _deserialize_types(cls, data: Any, handler): if "__tool_data__" in data: data["data"] = _deserialize_analysis_data(data["data"], data.pop("__tool_data__")) if "__tool_settings__" in data: - data["settings"] = _deserialize_zospy_class(data["settings"], data.pop("__tool_settings__")) + data["settings"] = _deserialize_zospy_class( + data["settings"], data.pop("__tool_settings__"), module="zospy.tools" + ) return handler(data) From fb98499e38ccb07b9b200a3f35c7f77b85d0313b Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 17:10:42 +0200 Subject: [PATCH 19/35] Use BeforeValidator to accept ndarray inputs. --- .../analyses/raysandspots/batch_ray_trace.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/analyses/raysandspots/batch_ray_trace.py index b9be3f15..94779b62 100644 --- a/zospy/analyses/raysandspots/batch_ray_trace.py +++ b/zospy/analyses/raysandspots/batch_ray_trace.py @@ -5,8 +5,9 @@ from typing import Literal, Annotated from collections.abc import Sequence +import numpy as np import pandas as pd -from pydantic import Field, PositiveInt, model_validator +from pydantic import Field, PositiveInt, BeforeValidator, model_validator from zospy.analyses.base import BaseAnalysisWrapper from zospy.analyses.decorators import analysis_settings @@ -16,8 +17,22 @@ __all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") +def ndarray_to_list(v): + """Function to 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], + BeforeValidator(ndarray_to_list), +] +PositiveIntVector = Annotated[ + Sequence[PositiveInt], + BeforeValidator(ndarray_to_list), +] @analysis_settings class BatchRayTraceNormUnpolSettings: @@ -45,11 +60,11 @@ class BatchRayTraceNormUnpolSettings: Mode of optical path difference for rays (e.g. 'None'). Defaults to 'None'. """ - hx: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized X field coordinate") - hy: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized Y field coordinate") - px: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized X pupil coordinate") - py: Sequence[NormalizedCoordinate] = Field(default=[0], description="Normalized Y pupil coordinate") - wavelength: PositiveInt | Sequence[PositiveInt] = Field(default=1, description="Wavelength number") + 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") From f9cbec6a2d3df3061133f268c79d0499c4154652 Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Fri, 5 Jun 2026 18:16:32 +0200 Subject: [PATCH 20/35] Move batch raytrace to tools. --- zospy/analyses/raysandspots/__init__.py | 7 +------ zospy/tools/__init__.py | 7 ++++++- .../batch_ray_trace.py => tools/batch_raytrace.py} | 0 3 files changed, 7 insertions(+), 7 deletions(-) rename zospy/{analyses/raysandspots/batch_ray_trace.py => tools/batch_raytrace.py} (100%) diff --git a/zospy/analyses/raysandspots/__init__.py b/zospy/analyses/raysandspots/__init__.py index f0b5f372..a6ed31e0 100644 --- a/zospy/analyses/raysandspots/__init__.py +++ b/zospy/analyses/raysandspots/__init__.py @@ -7,10 +7,5 @@ SingleRayTrace, SingleRayTraceSettings, ) -from zospy.analyses.raysandspots.batch_ray_trace import ( - BatchRayTraceNormUnpol, - BatchRayTraceNormUnpolSettings, -) -__all__ = ("RayFan", "RayFanSettings", "SingleRayTrace", "SingleRayTraceSettings", - "BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") +__all__ = ("RayFan", "RayFanSettings", "SingleRayTrace", "SingleRayTraceSettings") diff --git a/zospy/tools/__init__.py b/zospy/tools/__init__.py index 10e37d35..0ad84d8c 100644 --- a/zospy/tools/__init__.py +++ b/zospy/tools/__init__.py @@ -10,5 +10,10 @@ from zospy.tools.base import open_tool from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings +from zospy.tools.batch_raytrace import ( + BatchRayTraceNormUnpol, + BatchRayTraceNormUnpolSettings, +) -__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool") +__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool", + "BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") diff --git a/zospy/analyses/raysandspots/batch_ray_trace.py b/zospy/tools/batch_raytrace.py similarity index 100% rename from zospy/analyses/raysandspots/batch_ray_trace.py rename to zospy/tools/batch_raytrace.py From a036f2dcb47afd7fad9375a92ea68f7903bb18cc Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Fri, 5 Jun 2026 21:10:08 +0200 Subject: [PATCH 21/35] Switch to BaseToolWrapper. --- zospy/tools/batch_raytrace.py | 63 ++++++++++++++++++++--------------- 1 file changed, 36 insertions(+), 27 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 94779b62..61162dab 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -1,19 +1,24 @@ -"""Batch Ray Trace analysis.""" +"""Batch Ray Trace tool.""" from __future__ import annotations -from typing import Literal, Annotated -from collections.abc import Sequence +from typing import TYPE_CHECKING import numpy as np import pandas as pd from pydantic import Field, PositiveInt, BeforeValidator, model_validator -from zospy.analyses.base import BaseAnalysisWrapper +from zospy.tools.base import BaseToolWrapper from zospy.analyses.decorators import analysis_settings from zospy.analyses.parsers.types import ZOSAPIConstant from zospy.api import constants +if TYPE_CHECKING: + from typing import Literal, Annotated + from collections.abc import Sequence, Callable + + from zospy.api import _ZOSAPI + __all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") @@ -88,9 +93,7 @@ def validate_lengths(self): return self -class BatchRayTraceNormUnpol(BaseAnalysisWrapper[pd.DataFrame, BatchRayTraceNormUnpolSettings], - analysis_type="RayTrace", -): +class BatchRayTraceNormUnpol(BaseToolWrapper[pd.DataFrame, BatchRayTraceNormUnpolSettings]): """Batch Ray Trace of unpolarized light.""" def __init__( @@ -113,39 +116,46 @@ def __init__( """ super().__init__(settings_kws=locals()) - def run_analysis(self) -> pd.DataFrame: + 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_analysis(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame: """Run the Batch Ray Trace of unpolarized light.""" - 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( + tool.Hx = self.settings.hx + tool.Hy = self.settings.hy + tool.Px = self.settings.px + tool.Py = self.settings.py + tool.number_of_rays = len(tool.Hx) + tool.wavelengths = [self.settings.wavelength] * number_of_rays if isinstance(self.settings.wavelength, int) else self.settings.wavelength + tool.rays_type = constants.process_constant( constants.Tools.RayTrace.RaysType, self.settings.rays_type ) - opd_mode = constants.process_constant( + tool.opd_mode = constants.process_constant( constants.Tools.RayTrace.OPDMode, self.settings.opd_mode ) - surface = -1 if self.settings.surface == "Image" else self.settings.surface + tool.surface = -1 if self.settings.surface == "Image" else self.settings.surface # Initiate batch ray trace - raytrace = self.oss.Tools.OpenBatchRayTrace() - norm_unpol_data = raytrace.CreateNormUnpol( - number_of_rays, - rays_type, - surface, + norm_unpol_data = tool.CreateNormUnpol( + tool.number_of_rays, + tool.rays_type, + tool.surface, ) # Add rays for wavelength, hx, hy, px, py in zip( - wavelengths, - self.settings.hx, - self.settings.hy, - self.settings.px, - self.settings.py, + tool.wavelengths, + tool.Hx, + tool.Hy, + tool.Px, + tool.Py, ): - norm_unpol_data.AddRay(wavelength, hx, hy, px, py, opd_mode) + norm_unpol_data.AddRay(wavelength, hx, hy, px, py, tool.opd_mode) # Run ray trace and read results - self.analysis.ApplyAndWaitForCompletion() - raytrace.RunAndWaitForCompletion() + tool.RunAndWaitForCompletion() norm_unpol_data.StartReadingResults() # Read all results and append to outputs @@ -154,7 +164,6 @@ def run_analysis(self) -> pd.DataFrame: "X","Y","Z","L","M","N", "l2","m2","n2","opd","Intensity"] outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(number_of_rays)] - raytrace.Close() # Convert to DataFrame and return return pd.DataFrame(outputs, columns=columns) From d46a3f09dd0999a2b3bdb5a470b64cc6f4b0b363 Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Fri, 5 Jun 2026 21:21:49 +0200 Subject: [PATCH 22/35] Fix typos and small issues. --- zospy/tools/batch_raytrace.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 61162dab..399eba5d 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -121,7 +121,7 @@ def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.RayTrace.IBatchRay return oss.Tools.OpenBatchRayTrace - def run_analysis(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame: + def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame: """Run the Batch Ray Trace of unpolarized light.""" tool.Hx = self.settings.hx tool.Hy = self.settings.hy @@ -163,7 +163,7 @@ def run_analysis(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFr 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)] + outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(tool.number_of_rays)] # Convert to DataFrame and return return pd.DataFrame(outputs, columns=columns) From 885724be45b1d6da39a0a1908949494c11986415 Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Fri, 5 Jun 2026 21:34:00 +0200 Subject: [PATCH 23/35] Format with hatch. --- zospy/tools/__init__.py | 11 ++++++++--- zospy/tools/batch_raytrace.py | 27 +++++++++++++-------------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/zospy/tools/__init__.py b/zospy/tools/__init__.py index 0ad84d8c..26b8bda6 100644 --- a/zospy/tools/__init__.py +++ b/zospy/tools/__init__.py @@ -9,11 +9,16 @@ from __future__ import annotations from zospy.tools.base import open_tool -from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings from zospy.tools.batch_raytrace import ( BatchRayTraceNormUnpol, BatchRayTraceNormUnpolSettings, ) +from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings -__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool", - "BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") +__all__ = ( + "BatchRayTraceNormUnpol", + "BatchRayTraceNormUnpolSettings", + "QuickFocus", + "QuickFocusSettings", + "open_tool", +) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 399eba5d..39331ef5 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -6,17 +6,17 @@ import numpy as np import pandas as pd -from pydantic import Field, PositiveInt, BeforeValidator, model_validator +from pydantic import BeforeValidator, Field, PositiveInt, model_validator -from zospy.tools.base import BaseToolWrapper from zospy.analyses.decorators import analysis_settings from zospy.analyses.parsers.types import ZOSAPIConstant from zospy.api import constants +from zospy.tools.base import BaseToolWrapper if TYPE_CHECKING: - from typing import Literal, Annotated - from collections.abc import Sequence, Callable - + from collections.abc import Callable, Sequence + from typing import Annotated, Literal + from zospy.api import _ZOSAPI __all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") @@ -79,19 +79,18 @@ 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( - f"Hx, Hy, Px, Py must all have the same length." + "Hx, Hy, Px, Py must all have the same length." ) expected_len = len(self.hx) - if isinstance(self.wavelength, Sequence): - if 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})." - ) + 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.""" @@ -119,7 +118,7 @@ def __init__( 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.""" From 316aeddee45b2c741ee8ef7cf5cb938a4ea671fd Mon Sep 17 00:00:00 2001 From: Andi Barg <47257059+andibarg@users.noreply.github.com> Date: Fri, 5 Jun 2026 21:39:32 +0200 Subject: [PATCH 24/35] Move back some typing imports. --- zospy/tools/batch_raytrace.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 39331ef5..cb612972 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -3,6 +3,8 @@ from __future__ import annotations from typing import TYPE_CHECKING +from typing import Annotated, Literal +from collections.abc import Sequence import numpy as np import pandas as pd @@ -14,8 +16,7 @@ from zospy.tools.base import BaseToolWrapper if TYPE_CHECKING: - from collections.abc import Callable, Sequence - from typing import Annotated, Literal + from collections.abc import Callable from zospy.api import _ZOSAPI From 4a5d28df1a4d6e735f0a8a2b6f11d891141f2ba5 Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Tue, 16 Jun 2026 22:37:08 +0200 Subject: [PATCH 25/35] Merge changelog correctly. --- CHANGELOG.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c1be08d..b5f12997 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,11 +12,7 @@ ZOS-API can also be added in patch releases. ### Added - Support for Python 3.14 (#221) -<<<<<<< HEAD -- `zospy.functions.tools.open_tool`, a context manager to open a tool and close it automatically after use (#226) -======= - New `zospy.tools` submodule with `open_tool` (a context manager to open a tool and close it automatically after use) and tool wrappers (#226) ->>>>>>> upstream/main ### Changed From bb024e1fcfbbed5498f793407c38f21962bd677c Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Tue, 16 Jun 2026 23:36:14 +0200 Subject: [PATCH 26/35] Address review comments. --- zospy/tools/batch_raytrace.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index cb612972..1eb17e6d 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -46,15 +46,15 @@ class BatchRayTraceNormUnpolSettings: Attributes ---------- - hx : Sequence[float] + hx : Sequence[float] | np.ndarray Sequence of normalized X field coordinates. Defaults to [0]. - hy : Sequence[float] + hy : Sequence[float] | np.ndarray Sequence of normalized Y field coordinates. Defaults to [0]. - px : Sequence[float] + px : Sequence[float] | np.ndarray Sequence of normalized X pupil coordinates. Defaults to [0]. - py : Sequence[float] + py : Sequence[float] | np.ndarray Sequence of normalized Y pupil coordinates. Defaults to [0]. - wavelength : int | Sequence[int] + 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 @@ -66,10 +66,10 @@ class BatchRayTraceNormUnpolSettings: 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") + 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") @@ -122,7 +122,13 @@ def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.RayTrace.IBatchRay def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame: - """Run the Batch Ray Trace of unpolarized light.""" + """Run the Batch Ray Trace of unpolarized light. + + Returns + ------- + DataFrame + The data in long format. + """ tool.Hx = self.settings.hx tool.Hy = self.settings.hy tool.Px = self.settings.px From 2c93899c7cd9ebf9174f6c86ebe941743e88bc2c Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Wed, 17 Jun 2026 00:20:45 +0200 Subject: [PATCH 27/35] Avoid setting tool attributes. --- zospy/tools/batch_raytrace.py | 34 +++++++++++++++------------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 1eb17e6d..f3e38742 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -129,36 +129,32 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame DataFrame The data in long format. """ - tool.Hx = self.settings.hx - tool.Hy = self.settings.hy - tool.Px = self.settings.px - tool.Py = self.settings.py - tool.number_of_rays = len(tool.Hx) - tool.wavelengths = [self.settings.wavelength] * number_of_rays if isinstance(self.settings.wavelength, int) else self.settings.wavelength - tool.rays_type = constants.process_constant( + 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 ) - tool.opd_mode = constants.process_constant( + opd_mode = constants.process_constant( constants.Tools.RayTrace.OPDMode, self.settings.opd_mode ) - tool.surface = -1 if self.settings.surface == "Image" else self.settings.surface + surface = -1 if self.settings.surface == "Image" else self.settings.surface # Initiate batch ray trace norm_unpol_data = tool.CreateNormUnpol( - tool.number_of_rays, - tool.rays_type, - tool.surface, + number_of_rays, + rays_type, + surface, ) # Add rays for wavelength, hx, hy, px, py in zip( - tool.wavelengths, - tool.Hx, - tool.Hy, - tool.Px, - tool.Py, + wavelengths, + self.settings.hx, + self.settings.hy, + self.settings.px, + self.settings.py, ): - norm_unpol_data.AddRay(wavelength, hx, hy, px, py, tool.opd_mode) + norm_unpol_data.AddRay(wavelength, hx, hy, px, py, opd_mode) # Run ray trace and read results tool.RunAndWaitForCompletion() @@ -169,7 +165,7 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame columns = ["rayNumber","ErrorCode","vignetteCode", "X","Y","Z","L","M","N", "l2","m2","n2","opd","Intensity"] - outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(tool.number_of_rays)] + outputs = [norm_unpol_data.ReadNextResult()[1:] for _ in range(number_of_rays)] # Convert to DataFrame and return return pd.DataFrame(outputs, columns=columns) From f99aaf8c50e5dc719117ae787f3531912e942b87 Mon Sep 17 00:00:00 2001 From: Andi Barg Date: Wed, 17 Jun 2026 00:25:26 +0200 Subject: [PATCH 28/35] Small fix in docstring. --- zospy/tools/batch_raytrace.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index f3e38742..5efe1be7 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -47,13 +47,13 @@ class BatchRayTraceNormUnpolSettings: Attributes ---------- hx : Sequence[float] | np.ndarray - Sequence of normalized X field coordinates. Defaults to [0]. + Sequence of normalized X field coordinates. Defaults to (0,). hy : Sequence[float] | np.ndarray - Sequence of normalized Y field coordinates. Defaults to [0]. + Sequence of normalized Y field coordinates. Defaults to (0,). px : Sequence[float] | np.ndarray - Sequence of normalized X pupil coordinates. Defaults to [0]. + Sequence of normalized X pupil coordinates. Defaults to (0,). py : Sequence[float] | np.ndarray - Sequence of normalized Y pupil coordinates. Defaults to [0]. + 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. From 54aa05321970f0c370e69cec8edb61b4d0d632b8 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Wed, 17 Jun 2026 12:02:24 +0200 Subject: [PATCH 29/35] Fix dictionary combination error in analysis_settings decorator --- zospy/analyses/decorators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From cf154eff1ee77d36b9f2688ce6aac45ae6a73afc Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Wed, 17 Jun 2026 12:03:23 +0200 Subject: [PATCH 30/35] Add NDArrays to type annotations and document DataFrame columns --- zospy/tools/batch_raytrace.py | 81 +++++++++++++++++++++++------------ 1 file changed, 53 insertions(+), 28 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 5efe1be7..a9225fd7 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -2,13 +2,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING -from typing import Annotated, Literal from collections.abc import Sequence +from typing import TYPE_CHECKING, Annotated, Literal import numpy as np import pandas as pd -from pydantic import BeforeValidator, Field, PositiveInt, model_validator +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 @@ -23,24 +23,26 @@ __all__ = ("BatchRayTraceNormUnpol", "BatchRayTraceNormUnpolSettings") -def ndarray_to_list(v): - """Function to convert ndarray to list.""" +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], + Sequence[NormalizedCoordinate] | NDArray[np.floating], BeforeValidator(ndarray_to_list), ] PositiveIntVector = Annotated[ - Sequence[PositiveInt], + Sequence[PositiveInt] | NDArray[np.integer], BeforeValidator(ndarray_to_list), ] -@analysis_settings + +@analysis_settings(config=ConfigDict(arbitrary_types_allowed=True)) class BatchRayTraceNormUnpolSettings: """Settings for the Batch Ray Trace of unpolarized light. @@ -73,15 +75,15 @@ class BatchRayTraceNormUnpolSettings: 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") + 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." - ) + 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: @@ -99,17 +101,28 @@ class BatchRayTraceNormUnpol(BaseToolWrapper[pd.DataFrame, BatchRayTraceNormUnpo def __init__( self, *, - hx: Sequence[float] = (0,), - hy: Sequence[float] = (0,), - px: Sequence[float] = (0,), - py: Sequence[float] = (0,), - wavelength: int | Sequence[int] = 1, + 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. @@ -120,7 +133,6 @@ def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.RayTrace.IBatchRay """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. @@ -130,13 +142,13 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.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 + 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 @@ -153,6 +165,7 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame self.settings.hy, self.settings.px, self.settings.py, + strict=True ): norm_unpol_data.AddRay(wavelength, hx, hy, px, py, opd_mode) @@ -162,11 +175,23 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame # Read all results and append to outputs outputs = [] - columns = ["rayNumber","ErrorCode","vignetteCode", - "X","Y","Z","L","M","N", - "l2","m2","n2","opd","Intensity"] + 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) - From 5738a66e94ba270f2547930052c9ce20a43f305b Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Wed, 17 Jun 2026 12:03:59 +0200 Subject: [PATCH 31/35] Format code --- zospy/tools/batch_raytrace.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index a9225fd7..4b5ad678 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -11,7 +11,7 @@ from pydantic import BeforeValidator, ConfigDict, Field, PositiveInt, model_validator from zospy.analyses.decorators import analysis_settings -from zospy.analyses.parsers.types import ZOSAPIConstant +from zospy.analyses.parsers.types import ZOSAPIConstant # noqa: TC001 from zospy.api import constants from zospy.tools.base import BaseToolWrapper @@ -30,6 +30,7 @@ def ndarray_to_list(v: Sequence | NDArray) -> Sequence: return v + # Constrained types NormalizedCoordinate = Annotated[float, Field(le=1.0, ge=-1.0)] CoordinateVector = Annotated[ @@ -160,12 +161,7 @@ def _run_tool(self, tool: _ZOSAPI.Tools.RayTrace.IBatchRayTrace) -> pd.DataFrame # 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 + 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) From d03ef5725d115930f231c0e5ed2b69435d0ab099 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 26 Jun 2026 13:16:17 +0200 Subject: [PATCH 32/35] Add unit tests for batch raytrace tool --- tests/tools/test_batch_raytrace.py | 100 +++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 tests/tools/test_batch_raytrace.py 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) From e48ae0f3e5a16f7384f01a7a6b53ee808e905bbd Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 26 Jun 2026 13:17:07 +0200 Subject: [PATCH 33/35] Fix regression in analysis result deserialization Not all usages of _deserialize_zospy_class correctly specified the module --- zospy/analyses/base.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/zospy/analyses/base.py b/zospy/analyses/base.py index 63e42d6b..7d8be383 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) From cabd447ab080607319df9d28e4137b8c094075e5 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 26 Jun 2026 13:18:26 +0200 Subject: [PATCH 34/35] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) 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 From b5bea86e7c77371033e62b173d470cb3839c2c0c Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:19:24 +0200 Subject: [PATCH 35/35] Lint and format --- zospy/tools/batch_raytrace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zospy/tools/batch_raytrace.py b/zospy/tools/batch_raytrace.py index 4b5ad678..2507d696 100644 --- a/zospy/tools/batch_raytrace.py +++ b/zospy/tools/batch_raytrace.py @@ -11,7 +11,7 @@ from pydantic import BeforeValidator, ConfigDict, Field, PositiveInt, model_validator from zospy.analyses.decorators import analysis_settings -from zospy.analyses.parsers.types import ZOSAPIConstant # noqa: TC001 +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