ml.wrappers

Machine Learning Wrappers

This module contains wrappers for integrating Optiland with PyTorch. The OpticalSystemModule class is a key component that allows users to define and optimize optical systems within the PyTorch ecosystem. The forward method can be customized or overridden to implement specific optical system behaviors.

Kramer Harrison, 2025

Classes

OpticalSystemModule(*args, **kwargs)

A PyTorch nn.Module that wraps an Optiland OptimizationProblem.

class OpticalSystemModule(*args: Any, **kwargs: Any)

A PyTorch nn.Module that wraps an Optiland OptimizationProblem.

This class exposes the optical system’s variables as trainable nn.Parameters, allowing the entire system to be integrated as a differentiable layer into larger machine learning models.

Parameters:

• optic (Optic) – The optical system definition.

• problem (OptimizationProblem) – The optimization problem defining variables and objectives.

• objective_fn (Callable[[], torch.Tensor] | None) – An optional callable tha takes no arguments and returns a scalar PyTorch tensor representing the loss or metric to be optimized. If None, problem.sum_squared() is used as default.

apply_bounds()

Applies the defined bounds to the parameters in-place. This should be called after each optimizer step to enforce constraints.

forward() → torch.Tensor

Defines the forward pass of the optical system.

This implementation synchronizes the PyTorch parameters with the Optiland problem variables, updates the optics, and then computes the loss using either the user-provided objective function or the default sum of squared errors. The output is a differentiable scalar tensor.

Users are encouraged to customize or override this method to suit their specific optimization objectives.