Optics Core

Shared building blocks used by every lens model: base classes, optical surfaces, light/wave representations, and image-simulation utilities.

Base Classes

Base class for all optical objects. Provides device transfer, dtype conversion, and cloning by introspecting instance tensors.

deeplens.DeepObj

DeepObj(dtype=None)

Base class for all differentiable optical objects in DeepLens.

Provides device management, dtype conversion, and deep-copy support via automatic introspection over instance tensors and nested DeepObj sub-objects. All lens, surface, material, ray, and wave objects inherit from this class.

Attributes:

Name	Type	Description
dtype	dtype	Floating-point dtype of all owned tensors.
device	str or device	Compute device, set by to.

Initialize the base object and record its floating-point dtype.

Parameters:

Name	Type	Description	Default
dtype	dtype	Floating-point dtype for owned tensors. Defaults to torch.get_default_dtype() when None.	None

Source code in deeplens-src/deeplens/base.py

def __init__(self, dtype=None):
    """Initialize the base object and record its floating-point dtype.

    Args:
        dtype (torch.dtype, optional): Floating-point dtype for owned
            tensors. Defaults to `torch.get_default_dtype()` when None.
    """
    self.dtype = torch.get_default_dtype() if dtype is None else dtype

__str__

__str__()

Return a multi-line string listing the object's attributes.

Scalars and tensors are printed as key: value; lists and tuples are expanded element-wise; dicts and sets are skipped.

Returns:

Name	Type	Description
text	str	Human-readable summary of the object's attributes.

Source code in deeplens-src/deeplens/base.py

def __str__(self):
    """Return a multi-line string listing the object's attributes.

    Scalars and tensors are printed as `key: value`; lists and tuples are
    expanded element-wise; dicts and sets are skipped.

    Returns:
        text (str): Human-readable summary of the object's attributes.
    """
    lines = [self.__class__.__name__ + ":"]
    for key, val in vars(self).items():
        if val.__class__.__name__ in ["list", "tuple"]:
            for i, v in enumerate(val):
                lines += "{}[{}]: {}".format(key, i, v).split("\n")
        elif val.__class__.__name__ in ["dict", "OrderedDict", "set"]:
            pass
        else:
            lines += "{}: {}".format(key, val).split("\n")

    return "\n    ".join(lines)

__call__

__call__(inp)

Forward the input to the subclass forward method.

Parameters:

Name	Type	Description	Default
inp	Any	Input passed through to self.forward.	required

Returns:

Name	Type	Description
output	Any	Result of self.forward(inp).

Source code in deeplens-src/deeplens/base.py

def __call__(self, inp):
    """Forward the input to the subclass `forward` method.

    Args:
        inp (Any): Input passed through to `self.forward`.

    Returns:
        output (Any): Result of `self.forward(inp)`.
    """
    return self.forward(inp)

clone

clone()

Return a deep copy of this object.

Returns:

Name	Type	Description
obj	DeepObj	A new, independent deep copy of self.

Source code in deeplens-src/deeplens/base.py

def clone(self):
    """Return a deep copy of this object.

    Returns:
        obj (DeepObj): A new, independent deep copy of `self`.
    """
    return copy.deepcopy(self)

to

to(device)

Move all tensors and nested objects to a device.

Recursively walks over every instance attribute and moves tensors, nn.Parameter data, nn.Module sub-objects, nested DeepObj objects, and tensors/DeepObj items inside lists and tuples to the target device.

Parameters:

Name	Type	Description	Default
device	str or device	Target device, e.g. "cuda", "cpu", or a torch.device instance.	required

Returns:

Name	Type	Description
self	DeepObj	The updated object (for chaining).

Examples:

lens = GeoLens(filename="lens.json")
lens.to("cuda")  # move all tensors to GPU

Source code in deeplens-src/deeplens/base.py

def to(self, device):
    """Move all tensors and nested objects to a device.

    Recursively walks over every instance attribute and moves tensors,
    `nn.Parameter` data, `nn.Module` sub-objects, nested `DeepObj` objects,
    and tensors/`DeepObj` items inside lists and tuples to the target device.

    Args:
        device (str or torch.device): Target device, e.g. `"cuda"`, `"cpu"`,
            or a `torch.device` instance.

    Returns:
        self (DeepObj): The updated object (for chaining).

    Example:
        ```python
        lens = GeoLens(filename="lens.json")
        lens.to("cuda")  # move all tensors to GPU
        ```
    """
    self.device = torch.device(device)

    for key, val in list(vars(self).items()):
        if key == "device":
            continue
        setattr(self, key, self._map_state(val, device=self.device))
    return self

astype

astype(dtype)

Convert all floating-point tensors to a target dtype.

Recursively converts owned floating-point and complex tensors, nn.Parameter data, modules, nested DeepObj objects, and values inside lists, tuples, and dictionaries. Conversion is local to this object; it never changes PyTorch's process-wide default dtype.

Parameters:

Name	Type	Description	Default
dtype	dtype or None	Target floating-point dtype, one of torch.float16, torch.float32, or torch.float64. When None, this is a no-op and self is returned unchanged.	required

Returns:

Name	Type	Description
self	DeepObj	The updated object (for chaining).

Raises:

Type	Description
AssertionError	If dtype is not one of the three supported floating-point dtypes.

Examples:

lens = GeoLens(filename="lens.json")
lens.astype(torch.float64)  # switch to double precision

Source code in deeplens-src/deeplens/base.py

def astype(self, dtype):
    """Convert all floating-point tensors to a target dtype.

    Recursively converts owned floating-point and complex tensors,
    `nn.Parameter` data, modules, nested `DeepObj` objects, and values inside
    lists, tuples, and dictionaries. Conversion is local to this object; it
    never changes PyTorch's process-wide default dtype.

    Args:
        dtype (torch.dtype or None): Target floating-point dtype, one of
            `torch.float16`, `torch.float32`, or `torch.float64`. When None,
            this is a no-op and `self` is returned unchanged.

    Returns:
        self (DeepObj): The updated object (for chaining).

    Raises:
        AssertionError: If dtype is not one of the three supported
            floating-point dtypes.

    Example:
        ```python
        lens = GeoLens(filename="lens.json")
        lens.astype(torch.float64)  # switch to double precision
        ```
    """
    if dtype is None:
        return self

    dtype_ls = [torch.float16, torch.float32, torch.float64]
    assert dtype in dtype_ls, f"Data type {dtype} is not supported."

    self.dtype = dtype
    for key, val in list(vars(self).items()):
        if key == "dtype":
            continue
        setattr(self, key, self._map_state(val, dtype=dtype))
    return self

Optical material model (refractive index and dispersion) used by refractive surfaces.

deeplens.Material

Material(name=None, device='cpu')

Bases: DeepObj

Optical material defined by its wavelength-dependent refractive index.

Materials are looked up by name in the bundled CDGM, SCHOTT, or MISC AGF catalogs, in a custom JSON catalog, in the bundled refractiveindex.info catalog (optical glasses from SCHOTT/OHARA/HOYA/HIKARI/SUMITA/CDGM/LZOS/ Crystran plus common substrate crystals), or specified inline as "n/V" (Cauchy approximation from Abbe number V). Names defined by more than one source resolve in that order, so refractiveindex.info only fills gaps and never overrides an existing name.

Supported dispersion models: "sellmeier", "cauchy", "schott", "interp" (lookup table), "rii" (refractiveindex.info dispersion formulas), and "optimizable" (Cauchy with learnable n, V).

Attributes:

Name	Type	Description
name	str	Lowercase material name.
device	str	Compute device for dispersion tensors.
dispersion	str	Dispersion model in use ("sellmeier", "cauchy", "schott", "interp", "rii", or "optimizable").
n	float or Tensor	Refractive index at the d-line (587.6 nm). Becomes a learnable tensor after get_optimizer_params.
V	float or Tensor	Abbe number. Also learnable in "optimizable" mode.

Initialize an optical material.

Parameters:

Name	Type	Description	Default
name	str or None	Material name (case-insensitive). Accepted forms: Glass catalog name, e.g. "N-BK7", "H-K9L" "air" (n = 1, non-dispersive). Legacy names "vacuum" and "occluder" are accepted and normalised to "air". Inline Cauchy "n/V", e.g. "1.5168/64.17" Custom name registered in materials_data.json refractiveindex.info name, e.g. "s-bsl7", "sapphire", "znse" (optical glasses and substrate crystals) Defaults to None (treated as "air").	None
device	str	Compute device. Defaults to "cpu".	'cpu'

Raises:

Type	Description
NotImplementedError	If name is not found in any catalog.

Examples:

mat = Material("N-BK7")
n_green = mat.get_ri(0.587)  # refractive index at 587 nm

Source code in deeplens-src/deeplens/material/materials.py

def __init__(self, name=None, device="cpu"):
    """Initialize an optical material.

    Args:
        name (str or None, optional): Material name (case-insensitive).
            Accepted forms:

            - Glass catalog name, e.g. `"N-BK7"`, `"H-K9L"`
            - `"air"` (n = 1, non-dispersive). Legacy names `"vacuum"` and
              `"occluder"` are accepted and normalised to `"air"`.
            - Inline Cauchy `"n/V"`, e.g. `"1.5168/64.17"`
            - Custom name registered in `materials_data.json`
            - refractiveindex.info name, e.g. `"s-bsl7"`, `"sapphire"`,
              `"znse"` (optical glasses and substrate crystals)

            Defaults to None (treated as `"air"`).
        device (str, optional): Compute device. Defaults to `"cpu"`.

    Raises:
        NotImplementedError: If *name* is not found in any catalog.

    Example:
        ```python
        mat = Material("N-BK7")
        n_green = mat.get_ri(0.587)  # refractive index at 587 nm
        ```
    """
    raw = "air" if name is None else name.lower()
    # Normalise legacy aliases to "air"
    self.name = "air" if raw in ("vacuum", "occluder") else raw
    self.wvln_range = None
    self.load_dispersion()
    self._validated_wavelengths = set()
    self.device = device

get_name

get_name()

Return the material name, or an inline "n/V" string if optimizable.

Returns:

Name	Type	Description
name	str	The catalog name, or a "{n}/{V}" string formatted from the current (n, V) when the dispersion mode is "optimizable".

Source code in deeplens-src/deeplens/material/materials.py

def get_name(self):
    """Return the material name, or an inline `"n/V"` string if optimizable.

    Returns:
        name (str): The catalog name, or a `"{n}/{V}"` string formatted from
            the current (n, V) when the dispersion mode is `"optimizable"`.
    """
    if self.dispersion == "optimizable":
        return f"{self.n.item():.4f}/{self.V.item():.2f}"
    else:
        return self.name

load_dispersion

load_dispersion()

Resolve the material name into a dispersion model and its parameters.

Sets self.dispersion and the corresponding coefficients (Sellmeier k*/l*, Schott a*, Cauchy A/B, or interpolation tables), along with self.n (d-line index) and self.V (Abbe number). Looks up the name in the AGF catalogs, the inline "n/V" form, the custom JSON tables, then the bundled refractiveindex.info catalog (RII_data) as a final fallback, so names from earlier sources keep precedence.

Raises:

Type	Description
NotImplementedError	If the material name is not found in any catalog.
ValueError	If a custom "interp" entry has mismatched wavelength and index table lengths.

Source code in deeplens-src/deeplens/material/materials.py

def load_dispersion(self):
    """Resolve the material name into a dispersion model and its parameters.

    Sets `self.dispersion` and the corresponding coefficients (Sellmeier
    `k*`/`l*`, Schott `a*`, Cauchy `A`/`B`, or interpolation tables), along
    with `self.n` (d-line index) and `self.V` (Abbe number). Looks up the
    name in the AGF catalogs, the inline `"n/V"` form, the custom JSON
    tables, then the bundled refractiveindex.info catalog (`RII_data`) as a
    final fallback, so names from earlier sources keep precedence.

    Raises:
        NotImplementedError: If the material name is not found in any
            catalog.
        ValueError: If a custom `"interp"` entry has mismatched wavelength
            and index table lengths.
    """
    # Air (n=1, non-dispersive)
    if self.name == "air":
        self.dispersion = "sellmeier"
        self.k1, self.l1, self.k2, self.l2, self.k3, self.l3 = 0, 0, 0, 0, 0, 0
        self.n, self.V = 1.0, 1e38

    # Material found in AGF file
    elif self.name.lower() in MATERIAL_data:
        self.set_material_param_agf(MATERIAL_data, self.name.lower())

    # Material is given by a (n, V) string, e.g. "1.5168/64.17"
    elif "/" in self.name:
        self.dispersion = "cauchy"
        self.n = float(self.name.split("/")[0])
        self.V = float(self.name.split("/")[1])
        self.A, self.B = self.nV_to_AB(self.n, self.V)

    # Material found in custom JSON file
    elif self.name in CUSTOM_data["INTERP_TABLE"]:
        self.load_interp_table(CUSTOM_data["INTERP_TABLE"][self.name])

    elif self.name in CUSTOM_data["SELLMEIER_TABLE"]:
        self.dispersion = "sellmeier"
        self.k1, self.l1, self.k2, self.l2, self.k3, self.l3 = CUSTOM_data[
            "SELLMEIER_TABLE"
        ][self.name]
        try:
            self.n = CUSTOM_data["MATERIAL_TABLE"][self.name][0]
            self.V = CUSTOM_data["MATERIAL_TABLE"][self.name][1]
        except KeyError:
            print(
                f"Warning: {self.name} found in SELLMEIER_TABLE but not in MATERIAL_TABLE."
            )

    elif self.name in CUSTOM_data["SCHOTT_TABLE"]:
        self.dispersion = "schott"
        self.a0, self.a1, self.a2, self.a3, self.a4, self.a5 = CUSTOM_data[
            "SCHOTT_TABLE"
        ][self.name]
        try:
            self.n = CUSTOM_data["MATERIAL_TABLE"][self.name][0]
            self.V = CUSTOM_data["MATERIAL_TABLE"][self.name][1]
        except KeyError:
            print(
                f"Warning: {self.name} found in SCHOTT_TABLE but not in MATERIAL_TABLE."
            )

    elif self.name in CUSTOM_data["MATERIAL_TABLE"]:
        self.dispersion = "cauchy"
        self.n, self.V = CUSTOM_data["MATERIAL_TABLE"][self.name]
        self.A, self.B = self.nV_to_AB(self.n, self.V)

    # refractiveindex.info catalog (fallback: only reached for names not
    # defined by any of the sources above, so existing names keep priority).
    elif self.name in RII_data["FORMULA"]:
        entry = RII_data["FORMULA"][self.name]
        self.dispersion = "rii"
        self.rii_formula = entry["formula"]
        self.rii_coeffs = entry["coeffs"]
        self.rii_wvln_range = entry.get("wvln_range")
        self.wvln_range = self.rii_wvln_range
        self.n = entry["nd"]
        self.V = entry["vd"]

    elif self.name in RII_data["INTERP"]:
        self.load_interp_table(RII_data["INTERP"][self.name])

    else:
        raise NotImplementedError(f"Material {self.name} not implemented.")

load_interp_table

load_interp_table(mat_data)

Set up a tabulated-index ("interp") material from a data table.

Stores the reference wavelength/index arrays (and cached tensors), then samples self.n (d-line) and self.V (Abbe number) from the table.

Parameters:

Name	Type	Description	Default
mat_data	dict	Mapping with keys "wvlns" and "n", two equal length lists of wavelengths [µm] and refractive indices.	required

Raises:

Type	Description
ValueError	If the wavelength and index tables differ in length.

Source code in deeplens-src/deeplens/material/materials.py

def load_interp_table(self, mat_data):
    """Set up a tabulated-index (`"interp"`) material from a data table.

    Stores the reference wavelength/index arrays (and cached tensors), then
    samples `self.n` (d-line) and `self.V` (Abbe number) from the table.

    Args:
        mat_data (dict): Mapping with keys `"wvlns"` and `"n"`, two equal
            length lists of wavelengths [µm] and refractive indices.

    Raises:
        ValueError: If the wavelength and index tables differ in length.
    """
    self.dispersion = "interp"
    self.ref_wvlns = mat_data["wvlns"]
    self.ref_n = mat_data["n"]
    if len(self.ref_wvlns) != len(self.ref_n):
        raise ValueError(
            f"Interpolation wavelength and index tables for {self.name} "
            f"have different lengths."
        )
    self._ref_wvlns_t = torch.tensor(self.ref_wvlns)
    self._ref_n_t = torch.tensor(self.ref_n)
    # nd/Vd are defined at the visible He/H lines. Only report them when the
    # table actually spans the F and C lines; otherwise np.interp would clamp
    # to an endpoint and fabricate a meaningless d-line index (e.g. for
    # IR/UV-only crystals). In that case expose the in-band reference index
    # and mark Vd non-applicable (1e38), matching the formula-based path.
    wmin, wmax = min(self.ref_wvlns), max(self.ref_wvlns)
    self.wvln_range = tuple(mat_data.get("wvln_range", (wmin, wmax)))
    if wmin <= 0.4861 and 0.6563 <= wmax:
        nd = float(np.interp(0.58756, self.ref_wvlns, self.ref_n))
        nF = float(np.interp(0.4861, self.ref_wvlns, self.ref_n))
        nC = float(np.interp(0.6563, self.ref_wvlns, self.ref_n))
        self.n = nd
        self.V = (nd - 1) / (nF - nC) if nF != nC else 1e38
    else:
        self.n = float(np.interp(0.5 * (wmin + wmax), self.ref_wvlns, self.ref_n))
        self.V = 1e38

set_material_param_agf

set_material_param_agf(material_data, material_name)

Set dispersion model and coefficients from an AGF catalog entry.

Reads the calculate_mode flag to pick the Schott (mode 1) or Sellmeier (mode 2) model, fills the corresponding coefficients, and sets self.n and self.V from the catalog's nd/vd fields.

Parameters:

Name	Type	Description	Default
material_data	dict	Parsed AGF catalog (name to parameter dict).	required
material_name	str	Lowercase material name to look up.	required

Raises:

Type	Description
NotImplementedError	If the entry's calculate_mode is neither 1 nor 2.

Source code in deeplens-src/deeplens/material/materials.py

def set_material_param_agf(self, material_data, material_name):
    """Set dispersion model and coefficients from an AGF catalog entry.

    Reads the `calculate_mode` flag to pick the Schott (mode 1) or Sellmeier
    (mode 2) model, fills the corresponding coefficients, and sets `self.n`
    and `self.V` from the catalog's `nd`/`vd` fields.

    Args:
        material_data (dict): Parsed AGF catalog (name to parameter dict).
        material_name (str): Lowercase material name to look up.

    Raises:
        NotImplementedError: If the entry's `calculate_mode` is neither 1
            nor 2.
    """
    if material_name in material_data:
        material = material_data[material_name]

        if material["calculate_mode"] == 1:
            self.dispersion = "schott"
            self.a0 = material["a_coeff"]
            self.a1 = material["b_coeff"]
            self.a2 = material["c_coeff"]
            self.a3 = material["d_coeff"]
            self.a4 = material["e_coeff"]
            self.a5 = material["f_coeff"]
        elif material["calculate_mode"] == 2:
            self.dispersion = "sellmeier"
            self.k1 = material["a_coeff"]
            self.l1 = material["b_coeff"]
            self.k2 = material["c_coeff"]
            self.l2 = material["d_coeff"]
            self.k3 = material["e_coeff"]
            self.l3 = material["f_coeff"]
        else:
            raise NotImplementedError(
                f"Error: {material_name} calculate_mode {material['calculate_mode']}"
            )

        self.n = material["nd"]
        self.V = material["vd"]
        self.wvln_range = material.get("wvln_range")
    else:
        print(f"error: not {material_name}")

set_sellmeier_param

set_sellmeier_param(params=None)

Manually set the six Sellmeier coefficients for a custom material.

Switches the dispersion model to "sellmeier" so subsequent ior calls use the newly set parameters.

Parameters:

Name	Type	Description	Default
params	tuple or list or None	The six coefficients (k1, l1, k2, l2, k3, l3). Defaults to None (all zeros).	None

Source code in deeplens-src/deeplens/material/materials.py

def set_sellmeier_param(self, params=None):
    """Manually set the six Sellmeier coefficients for a custom material.

    Switches the dispersion model to `"sellmeier"` so subsequent `ior` calls
    use the newly set parameters.

    Args:
        params (tuple or list or None, optional): The six coefficients
            `(k1, l1, k2, l2, k3, l3)`. Defaults to None (all zeros).
    """
    # Switch the dispersion model so ior() uses the newly set parameters.
    self.dispersion = "sellmeier"
    if params is None:
        self.k1, self.l1, self.k2, self.l2, self.k3, self.l3 = (
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        )
    else:
        self.k1, self.l1, self.k2, self.l2, self.k3, self.l3 = params

refractive_index

refractive_index(wvln)

Compute the refractive index at a given wavelength.

Thin wrapper over ior that accepts a Python float and returns a float, otherwise passes a tensor through unchanged.

Parameters:

Name	Type	Description	Default
wvln	float or Tensor	Wavelength in micrometres [µm].	required

Returns:

Name	Type	Description
n	float or Tensor	Refractive index. A float when wvln is a float, otherwise a tensor matching the input shape.

Source code in deeplens-src/deeplens/material/materials.py

def refractive_index(self, wvln):
    """Compute the refractive index at a given wavelength.

    Thin wrapper over `ior` that accepts a Python float and returns a float,
    otherwise passes a tensor through unchanged.

    Args:
        wvln (float or torch.Tensor): Wavelength in micrometres [µm].

    Returns:
        n (float or torch.Tensor): Refractive index. A float when `wvln` is a
            float, otherwise a tensor matching the input shape.
    """
    if isinstance(wvln, float):
        wvln = torch.tensor(wvln, device=self.device)
        return self.ior(wvln).item()

    return self.ior(wvln)

ior

ior(wvln)

Compute the refractive index from the active dispersion model.

Dispatches on self.dispersion: Sellmeier, Schott, Cauchy, linear interpolation of a lookup table, the refractiveindex.info dispersion formulas ("rii"), or an optimizable Cauchy form with learnable (n, V). The Cauchy branch evaluates \(n = A + B/\lambda^2\) with \(\lambda\) in nanometres; all other branches take \(\lambda\) in micrometres directly.

Parameters:

Name	Type	Description	Default
wvln	Tensor	Wavelength in micrometres [µm]. Must lie in (0.1, 10).	required

Returns:

Name	Type	Description
n	Tensor	Refractive index, same shape as wvln.

Raises:

Type	Description
NotImplementedError	If self.dispersion is unknown.

Source code in deeplens-src/deeplens/material/materials.py

def ior(self, wvln):
    """Compute the refractive index from the active dispersion model.

    Dispatches on `self.dispersion`: Sellmeier, Schott, Cauchy, linear
    interpolation of a lookup table, the refractiveindex.info dispersion
    formulas (`"rii"`), or an optimizable Cauchy form with learnable (n, V).
    The Cauchy branch evaluates $n = A + B/\\lambda^2$ with $\\lambda$ in
    nanometres; all other branches take $\\lambda$ in micrometres directly.

    Args:
        wvln (torch.Tensor): Wavelength in micrometres [µm]. Must lie in
            (0.1, 10).

    Returns:
        n (torch.Tensor): Refractive index, same shape as `wvln`.

    Raises:
        NotImplementedError: If `self.dispersion` is unknown.
    """
    if not torch.is_tensor(wvln):
        wvln = torch.as_tensor(wvln, device=self.device)
    self._validate_wavelength(wvln)

    if self.dispersion == "sellmeier":
        # Sellmeier equation: https://en.wikipedia.org/wiki/Sellmeier_equation
        n2 = (
            1
            + self.k1 * wvln**2 / (wvln**2 - self.l1)
            + self.k2 * wvln**2 / (wvln**2 - self.l2)
            + self.k3 * wvln**2 / (wvln**2 - self.l3)
        )
        n = torch.sqrt(n2)

    elif self.dispersion == "schott":
        # Schott equation: https://johnloomis.org/eop501/notes/matlab/sect1/schott.html
        ws = wvln**2
        n2 = (
            self.a0
            + self.a1 * ws
            + (self.a2 + (self.a3 + (self.a4 + self.a5 / ws) / ws) / ws) / ws
        )
        n = torch.sqrt(n2)

    elif self.dispersion == "cauchy":
        # Cauchy equation: https://en.wikipedia.org/wiki/Cauchy%27s_equation
        n = self.A + self.B / (wvln * 1e3) ** 2

    elif self.dispersion == "interp":
        # Use cached tensors, move to correct device if needed
        if (
            self._ref_wvlns_t.device != wvln.device
            or self._ref_wvlns_t.dtype != wvln.dtype
        ):
            self._ref_wvlns_t = self._ref_wvlns_t.to(
                device=wvln.device, dtype=wvln.dtype
            )
            self._ref_n_t = self._ref_n_t.to(device=wvln.device, dtype=wvln.dtype)
        ref_wvlns = self._ref_wvlns_t
        ref_n = self._ref_n_t

        # Find the lower and upper bracketing wavelengths
        i = torch.searchsorted(ref_wvlns, wvln, side="right")
        num_ref_wvlns = len(ref_wvlns)
        idx_low = torch.clamp(i - 1, 0, num_ref_wvlns - 1)
        idx_high = torch.clamp(i, 0, num_ref_wvlns - 1)

        wvln_ref_low = ref_wvlns[idx_low]
        wvln_ref_high = ref_wvlns[idx_high]
        n_ref_low = ref_n[idx_low]
        n_ref_high = ref_n[idx_high]

        # Interpolate n
        denom = wvln_ref_high - wvln_ref_low
        has_interval = denom != 0
        safe_denom = torch.where(has_interval, denom, torch.ones_like(denom))
        weight_high = torch.where(
            has_interval,
            (wvln - wvln_ref_low) / safe_denom,
            torch.zeros_like(wvln),
        )
        weight_low = 1.0 - weight_high
        n = n_ref_low * weight_low + n_ref_high * weight_high

    elif self.dispersion == "rii":
        # refractiveindex.info dispersion formulas, wavelength in [µm].
        # https://refractiveindex.info -> "Dispersion formulas". The
        # coefficient layout is C1 followed by (numerator, denominator/
        # exponent) pairs, matching the bundled refractiveindex_data.json.
        c = self.rii_coeffs
        ws = wvln**2
        if self.rii_formula == 1:
            # Sellmeier (preferred): squared denominators.
            n2 = 1.0 + c[0]
            for i in range(1, len(c) - 1, 2):
                n2 = n2 + c[i] * ws / (ws - c[i + 1] ** 2)
            n = torch.sqrt(n2)
        elif self.rii_formula == 2:
            # Sellmeier-2: non-squared denominators.
            n2 = 1.0 + c[0]
            for i in range(1, len(c) - 1, 2):
                n2 = n2 + c[i] * ws / (ws - c[i + 1])
            n = torch.sqrt(n2)
        elif self.rii_formula == 3:
            # Polynomial.
            n2 = c[0] + torch.zeros_like(wvln)
            for i in range(1, len(c) - 1, 2):
                n2 = n2 + c[i] * wvln ** c[i + 1]
            n = torch.sqrt(n2)
        else:
            raise NotImplementedError(
                f"refractiveindex.info formula {self.rii_formula} not implemented."
            )

    elif self.dispersion == "optimizable":
        # Cauchy's equation, calculate (A, B) on the fly. Clamp the Abbe
        # number away from zero before dividing: an unconstrained optimizable
        # V can be driven toward 0, which blows up B and the gradients
        # (physical Abbe numbers are well above 1).
        V_safe = torch.clamp(self.V, min=1.0)
        B = (self.n - 1) / V_safe / (1 / 0.486**2 - 1 / 0.656**2)
        A = self.n - B * 1 / 0.587**2
        n = A + B / wvln**2

    else:
        raise NotImplementedError(f"Error: {self.dispersion} not implemented.")

    return n

nV_to_AB staticmethod

nV_to_AB(n, V)

Convert (n, V) to Cauchy coefficients (A, B).

Solves the two-term Cauchy model \(n(\lambda) = A + B/\lambda^2\) for A and B given the d-line index and Abbe number, using the F/d/C lines (486.1 / 587.6 / 656.3 nm). B is in nm², matching the Cauchy branch of ior.

Parameters:

Name	Type	Description	Default
n	float	Refractive index at the d-line.	required
V	float	Abbe number.	required

Returns:

Name	Type	Description
A	float	Cauchy constant term.
B	float	Cauchy dispersion term, in nm².

Source code in deeplens-src/deeplens/material/materials.py

@staticmethod
def nV_to_AB(n, V):
    """Convert (n, V) to Cauchy coefficients (A, B).

    Solves the two-term Cauchy model $n(\\lambda) = A + B/\\lambda^2$ for A
    and B given the d-line index and Abbe number, using the F/d/C lines
    (486.1 / 587.6 / 656.3 nm). B is in nm², matching the Cauchy branch of
    `ior`.

    Args:
        n (float): Refractive index at the d-line.
        V (float): Abbe number.

    Returns:
        A (float): Cauchy constant term.
        B (float): Cauchy dispersion term, in nm².
    """

    def ivs(a):
        return 1.0 / a**2

    lambdas = [656.3, 587.6, 486.1]
    B = (n - 1) / V / (ivs(lambdas[2]) - ivs(lambdas[0]))
    A = n - B * ivs(lambdas[1])
    return A, B

match_material

match_material(mat_table=None)

Snap this material to the closest real glass in a catalog.

Finds the catalog entry minimising the normalised (n, V) distance (n scaled by 0.4, V by 40), renames this material to it, and reloads its dispersion. No-op for air.

Parameters:

Name	Type	Description	Default
mat_table	str or dict or None	Catalog to match against. None or "CDGM" uses the CDGM common glasses, "PLASTIC" uses the plastic table, or pass a name-to-(n, V) dict directly. Defaults to None.	None

Raises:

Type	Description
NotImplementedError	If mat_table is an unrecognised string.

Source code in deeplens-src/deeplens/material/materials.py

def match_material(self, mat_table=None):
    """Snap this material to the closest real glass in a catalog.

    Finds the catalog entry minimising the normalised (n, V) distance
    (n scaled by 0.4, V by 40), renames this material to it, and reloads its
    dispersion. No-op for air.

    Args:
        mat_table (str or dict or None, optional): Catalog to match against.
            `None` or `"CDGM"` uses the CDGM common glasses, `"PLASTIC"` uses
            the plastic table, or pass a name-to-(n, V) dict directly.
            Defaults to None.

    Raises:
        NotImplementedError: If `mat_table` is an unrecognised string.
    """
    if not self.name == "air":
        # Material match table
        if mat_table is None:
            print(
                "No material table provided. Using CDGM common glasses as default."
            )
            mat_table = CUSTOM_data["CDGM_GLASS"]
        elif mat_table == "CDGM":
            # CDGM common glasses
            mat_table = CUSTOM_data["CDGM_GLASS"]
        elif mat_table == "PLASTIC":
            mat_table = CUSTOM_data["PLASTIC_TABLE"]
        else:
            raise NotImplementedError(
                f"Material table {mat_table} not implemented."
            )

        # Find the closest material
        n_range = 0.4  # refractive index range usually [1.5, 1.9]
        V_range = 40.0  # Abbe number range usually [30, 70]
        n_self = float(self.n) if torch.is_tensor(self.n) else self.n
        V_self = float(self.V) if torch.is_tensor(self.V) else self.V
        self.name = min(
            mat_table,
            key=lambda name: (
                abs(mat_table[name][0] - n_self) / n_range
                + abs(mat_table[name][1] - V_self) / V_range
            ),
        )

        # Load the new material parameters
        self.load_dispersion()

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.01])

Make (n, V) learnable and return optimizer parameter groups.

Converts self.n and self.V to gradient-tracking tensors and switches the dispersion model to "optimizable". Optimizing the refractive index matters more than the Abbe number.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates [lr_n, lr_V] for n and V. Defaults to [1e-4, 1e-2].	[0.0001, 0.01]

Returns:

Name	Type	Description
params	list	Two optimizer parameter-group dicts, one for n and one for V, each with its own learning rate.

Source code in deeplens-src/deeplens/material/materials.py

def get_optimizer_params(self, lrs=[1e-4, 1e-2]):
    """Make (n, V) learnable and return optimizer parameter groups.

    Converts `self.n` and `self.V` to gradient-tracking tensors and switches
    the dispersion model to `"optimizable"`. Optimizing the refractive index
    matters more than the Abbe number.

    Args:
        lrs (list, optional): Learning rates `[lr_n, lr_V]` for n and V.
            Defaults to `[1e-4, 1e-2]`.

    Returns:
        params (list): Two optimizer parameter-group dicts, one for n and
            one for V, each with its own learning rate.
    """
    if isinstance(self.n, float):
        self.n = torch.tensor(self.n, device=self.device)
        self.V = torch.tensor(self.V, device=self.device)

    self.n.requires_grad = True
    self.V.requires_grad = True
    self.dispersion = "optimizable"

    params = [
        {"params": [self.n], "lr": lrs[0]},
        {"params": [self.V], "lr": lrs[1]},
    ]
    return params

Surfaces

Axial spacing

Current surface constructors use d_next, the axial distance to the next surface. A few upstream-generated argument descriptions still say d; follow the rendered function signature. Use lens.surf_d(i) for global z.

Geometric Surfaces

Base class for all geometric optical surfaces. Implements surface intersection (Newton's method with one differentiable step) and differentiable vector Snell's law refraction.

deeplens.geometric_surface.Surface

Surface(
    r,
    d_next,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
)

Bases: DeepObj

Base class for all geometric optical surfaces.

Surfaces use sequential geometry: each surface owns only d_next, the axial thickness [mm] from its vertex to the next vertex (or from the last surface to the sensor). Absolute vertex positions are derived by GeoLens.surf_d; they are not stored on individual surfaces. A surface also has an aperture radius r [mm] and separates two optical media. Subclasses override _sag and _dfdxy to define their shape.

Ray-surface interaction is handled in three stages by ray_reaction:

1. Coordinate transform: the ray is brought into the local surface frame.
2. Intersection: solved via Newton's method (newtons_method), using a non-differentiable iteration loop followed by a single differentiable Newton step to enable gradient flow.
3. Refraction / reflection: vector Snell's law (refract) or specular reflection (reflect).

Attributes:

Name	Type	Description
d_next	Tensor	Thickness to the next vertex [mm], scalar tensor.
r	float	Aperture radius [mm]. For a square aperture this is the circumscribed-circle radius (half-diagonal).
mat2	Material	Optical material on the transmission side.
pos_x	Tensor	Lateral x offset of the vertex [mm], scalar tensor.
pos_y	Tensor	Lateral y offset of the vertex [mm], scalar tensor.
vec_local	Tensor	Local surface normal direction, shape [3].
is_square	bool	If True the aperture is square; otherwise circular.
w	float	Square aperture side length [mm] (only set when is_square).
h	float	Square aperture side length [mm] (only set when is_square).

Initialize a generic optical surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius [mm]. For a square aperture this is the circumscribed-circle radius (half-diagonal), so the side length is r * sqrt(2).	required
d_next	float	Axial thickness to the next vertex [mm].	required
mat2	str or Material	Material on the transmission side (e.g. "N-BK7", "air").	required
pos_xy	list[float]	Lateral offset [x, y][mm]. Defaults to [0.0, 0.0].	[0.0, 0.0]
vec_local	list[float]	Local surface normal direction; normalized internally. Defaults to [0.0, 0.0, 1.0] (on-axis).	[0.0, 0.0, 1.0]
is_square	bool	Use a square aperture instead of a circular one. Defaults to False.	False
device	str	Compute device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def __init__(
    self,
    r,
    d_next,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize a generic optical surface.

    Args:
        r (float): Aperture radius [mm]. For a square aperture this is the
            circumscribed-circle radius (half-diagonal), so the side length
            is `r * sqrt(2)`.
        d_next (float): Axial thickness to the next vertex [mm].
        mat2 (str or Material): Material on the transmission side
            (e.g. "N-BK7", "air").
        pos_xy (list[float], optional): Lateral offset [x, y] [mm].
            Defaults to [0.0, 0.0].
        vec_local (list[float], optional): Local surface normal direction;
            normalized internally. Defaults to [0.0, 0.0, 1.0] (on-axis).
        is_square (bool, optional): Use a square aperture instead of a
            circular one. Defaults to False.
        device (str, optional): Compute device. Defaults to "cpu".
    """
    super().__init__()

    self.d_next = (
        d_next.detach().clone()
        if torch.is_tensor(d_next)
        else torch.tensor(d_next, dtype=torch.get_default_dtype())
    )
    if not self.d_next.is_floating_point():
        self.d_next = self.d_next.to(torch.get_default_dtype())

    state_dtype = self.d_next.dtype
    state_device = self.d_next.device
    self.vec_global = torch.tensor(
        [0.0, 0.0, 1.0], dtype=state_dtype, device=state_device
    )
    self.pos_x = torch.as_tensor(pos_xy[0], dtype=state_dtype, device=state_device)
    self.pos_y = torch.as_tensor(pos_xy[1], dtype=state_dtype, device=state_device)
    self.vec_local = F.normalize(
        torch.as_tensor(vec_local, dtype=state_dtype, device=state_device),
        p=2,
        dim=-1,
    )

    self.mat2 = Material(mat2)
    self.r = float(r)
    self.is_square = is_square
    if is_square:
        self.w = self.r * float(np.sqrt(2))
        self.h = self.r * float(np.sqrt(2))

    self.device = device if device is not None else torch.device("cpu")
    self.to(self.device)
    self._cache_rotation_matrices()

    # Newton method parameters
    self.newton_maxiter = 8  # [int], maximum number of Newton iterations
    self.newton_convergence = (
        50.0 * 1e-6
    )  # [mm], Newton method convergence threshold
    self.newton_step_bound = 5.0  # [mm], maximum step size in each iteration

init_from_dict classmethod

init_from_dict(surf_dict)

Initialize a surface from a serialized dict.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters, typically produced by surf_dict.	required

Returns:

Name	Type	Description
surface	Surface	The reconstructed surface instance.

Raises:

Type	Description
NotImplementedError	Always, on the base class; subclasses override.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize a surface from a serialized dict.

    Args:
        surf_dict (dict): Surface parameters, typically produced by
            `surf_dict`.

    Returns:
        surface (Surface): The reconstructed surface instance.

    Raises:
        NotImplementedError: Always, on the base class; subclasses override.
    """
    raise NotImplementedError(
        f"init_from_dict() is not implemented for {cls.__name__}."
    )

paraxial_power

paraxial_power(n1, n2)

Return the paraxial (first-order) optical power of this surface [1/mm].

Only the vertex geometry contributes: conic constants, aspheric terms and freeform departures vanish in the paraxial limit, matching the first-order convention used by Zemax and CODE V. Surfaces with a base curvature or an explicit focal length override this; flat and freeform surfaces contribute pure transfer.

Parameters:

Name	Type	Description	Default
n1	Tensor	Refractive index of the incident medium.	required
n2	Tensor	Refractive index of the transmission medium.	required

Returns:

Name	Type	Description
power	Tensor	Surface power [1/mm], scalar.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def paraxial_power(self, n1, n2):
    """Return the paraxial (first-order) optical power of this surface [1/mm].

    Only the vertex geometry contributes: conic constants, aspheric terms
    and freeform departures vanish in the paraxial limit, matching the
    first-order convention used by Zemax and CODE V. Surfaces with a base
    curvature or an explicit focal length override this; flat and freeform
    surfaces contribute pure transfer.

    Args:
        n1 (torch.Tensor): Refractive index of the incident medium.
        n2 (torch.Tensor): Refractive index of the transmission medium.

    Returns:
        power (torch.Tensor): Surface power [1/mm], scalar.
    """
    return torch.zeros_like(n2)

ray_reaction

ray_reaction(ray, n1, n2, refraction=True)

Compute the output ray after intersection and refraction/reflection.

The input ray is expressed in this surface's reference frame (its vertex is at z=0). This method applies only the surface-local lateral offset/orientation, intersection and refraction/reflection. GeoLens applies the axial d_next frame step between surface interactions.

Parameters:

Name	Type	Description	Default
ray	Ray	Incident ray bundle.	required
n1	float	Refractive index of the incident medium.	required
n2	float	Refractive index of the transmission medium.	required
refraction	bool	If True refract the ray; if False reflect it. Defaults to True.	True

Returns:

Name	Type	Description
ray	Ray	Updated ray bundle after the surface interaction.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def ray_reaction(self, ray, n1, n2, refraction=True):
    """Compute the output ray after intersection and refraction/reflection.

    The input ray is expressed in this surface's reference frame (its
    vertex is at z=0). This method applies only the surface-local lateral
    offset/orientation, intersection and refraction/reflection. `GeoLens`
    applies the axial `d_next` frame step between surface interactions.

    Args:
        ray (Ray): Incident ray bundle.
        n1 (float): Refractive index of the incident medium.
        n2 (float): Refractive index of the transmission medium.
        refraction (bool, optional): If True refract the ray; if False
            reflect it. Defaults to True.

    Returns:
        ray (Ray): Updated ray bundle after the surface interaction.
    """
    # Transform ray to local coordinate system
    ray = self.to_local_coord(ray)

    # Intersection
    ray = self.intersect(ray, n1)

    if refraction:
        old_d = ray.d.clone()
        ray = self.refract(ray, n1 / n2)
        ray = self.bend_penalty(ray, old_d)
    else:
        # Reflection
        ray = self.reflect(ray)

    # Transform ray back to the surface reference frame
    ray = self.to_global_coord(ray)

    return ray

intersect

intersect(ray, n=1.0)

Solve ray-surface intersection in the local coordinate system.

Moves each valid ray origin to the surface and, for coherent rays, accumulates optical path length n * t into ray.opl.

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray bundle in local coordinates.	required
n	float	Refractive index of the medium the ray travels through to reach the surface. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
ray	Ray	Ray with updated origins, validity mask, and (for coherent rays) optical path length.

Raises:

Type	Description
Exception	If a coherent ray has float32 dtype and an intersection distance above 100 mm, which can cause OPL precision problems.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in the local coordinate system.

    Moves each valid ray origin to the surface and, for coherent rays,
    accumulates optical path length `n * t` into `ray.opl`.

    Args:
        ray (Ray): Input ray bundle in local coordinates.
        n (float, optional): Refractive index of the medium the ray
            travels through to reach the surface. Defaults to 1.0.

    Returns:
        ray (Ray): Ray with updated origins, validity mask, and (for
            coherent rays) optical path length.

    Raises:
        Exception: If a coherent ray has float32 dtype and an intersection
            distance above 100 mm, which can cause OPL precision problems.
    """
    # Solve ray-surface intersection time by Newton's method
    t, valid = self.newtons_method(ray)

    # Update ray
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    ray.o = torch.where(valid.unsqueeze(-1), new_o, ray.o)
    ray.is_valid = ray.is_valid * valid

    if ray.is_coherent:
        # Check the actual tensor dtype (mirrors ray.py) rather than the
        # global default, which may not reflect this ray's precision.
        if t.abs().max() > 100 and t.dtype != torch.float64:
            raise Exception(
                "Using float32 may cause precision problem for OPL calculation."
            )
        new_opl = ray.opl + n * t.unsqueeze(-1)
        ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    return ray

newton_initial_t

newton_initial_t(ray)

Return the initial ray parameter for Newton intersection solving.

The generic approximation intersects the ray with the local vertex plane at z = 0. Surfaces with a closer analytic base shape can override this hook without duplicating the Newton iteration itself.

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray bundle in local surface coordinates.	required

Returns:

Name	Type	Description
t	Tensor	Initial intersection parameter [mm], shape [...] matching the ray batch.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def newton_initial_t(self, ray):
    """Return the initial ray parameter for Newton intersection solving.

    The generic approximation intersects the ray with the local vertex
    plane at ``z = 0``. Surfaces with a closer analytic base shape can
    override this hook without duplicating the Newton iteration itself.

    Args:
        ray (Ray): Input ray bundle in local surface coordinates.

    Returns:
        t (torch.Tensor): Initial intersection parameter [mm], shape [...]
            matching the ray batch.
    """
    return -ray.o[..., 2] / ray.d[..., 2]

newtons_method

newtons_method(ray)

Solve the ray-surface intersection by Newton's method (local frame).

Runs newton_maxiter - 1 non-differentiable iterations followed by one differentiable Newton step, so gradients flow only through the final step. The solved \(t\) satisfies sag(x, y) - z = 0 along the ray.

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray bundle in local coordinates.	required

Returns:

Name	Type	Description
t	Tensor	Intersection parameter (distance along the ray) [mm], shape [...] matching the ray batch.
valid	Tensor	Boolean mask of converged, in-range intersections, shape [...].

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def newtons_method(self, ray):
    """Solve the ray-surface intersection by Newton's method (local frame).

    Runs `newton_maxiter - 1` non-differentiable iterations followed by one
    differentiable Newton step, so gradients flow only through the final
    step. The solved $t$ satisfies `sag(x, y) - z = 0` along the ray.

    Args:
        ray (Ray): Input ray bundle in local coordinates.

    Returns:
        t (torch.Tensor): Intersection parameter (distance along the ray)
            [mm], shape [...] matching the ray batch.
        valid (torch.Tensor): Boolean mask of converged, in-range
            intersections, shape [...].
    """
    newton_maxiter = self.newton_maxiter
    newton_convergence = self.newton_convergence
    newton_step_bound = self.newton_step_bound

    # Ray direction components (reused across iterations)
    dxdt, dydt, dzdt = ray.d[..., 0], ray.d[..., 1], ray.d[..., 2]

    # Surface-specific initial guess (the generic default is the z=0 plane)
    t = self.newton_initial_t(ray)

    # 1. Non-differentiable Newton's iterations to find the intersection
    #    Run (maxiter - 1) iterations; the differentiable step below acts as
    #    the final iteration while also enabling gradient flow.
    with torch.no_grad():
        for _ in range(newton_maxiter - 1):
            new_o = ray.o + ray.d * t.unsqueeze(-1)
            new_x, new_y = new_o[..., 0], new_o[..., 1]
            valid = self.is_within_data_range(new_x, new_y) & (ray.is_valid > 0)

            x, y = new_x * valid, new_y * valid
            ft = self._sag(x, y) - new_o[..., 2]
            dfdx, dfdy = self._dfdxy(x, y)
            dfdt = dfdx * dxdt + dfdy * dydt - dzdt
            t = t - torch.clamp(
                ft / (dfdt + EPSILON), -newton_step_bound, newton_step_bound
            )

    # 2. One differentiable Newton step (final iteration + gradient flow)
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    new_x, new_y = new_o[..., 0], new_o[..., 1]
    valid = self.is_valid(new_x, new_y) & (ray.is_valid > 0)

    x, y = new_x * valid, new_y * valid
    ft = self._sag(x, y) - new_o[..., 2]
    dfdx, dfdy = self._dfdxy(x, y)
    dfdt = dfdx * dxdt + dfdy * dydt - dzdt
    t = t - torch.clamp(
        ft / (dfdt + EPSILON), -newton_step_bound, newton_step_bound
    )

    # 3. Re-evaluate the actual final update. The previous implementation
    # checked the residual and aperture before applying the last Newton step,
    # so a divergent final point could still be marked valid.
    with torch.no_grad():
        final_o = ray.o + ray.d * t.unsqueeze(-1)
        final_x, final_y = final_o[..., 0], final_o[..., 1]
        valid = self.is_valid(final_x, final_y) & (ray.is_valid > 0)
        x, y = final_x * valid, final_y * valid
        final_residual = self._sag(x, y) - final_o[..., 2]
        valid = (
            valid
            & torch.isfinite(t)
            & torch.isfinite(final_residual)
            & (final_residual.abs() < newton_convergence)
        )

    return t, valid

bend_penalty

bend_penalty(ray, old_d)

Accumulate a soft per-surface bend penalty onto the ray.

The penalty rises smoothly once the bend angle between old_d and the refracted ray.d exceeds bend_angle_max (degrees, default 30) and stays at zero for milder refractions. It is added into ray.bend_penalty.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray after refraction (ray.d is the new direction).	required
old_d	Tensor	Pre-refraction ray directions, shape [..., 3], same shape as ray.d.	required

Returns:

Name	Type	Description
ray	Ray	Ray with bend_penalty (shape [..., 1]) updated.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def bend_penalty(self, ray, old_d):
    """Accumulate a soft per-surface bend penalty onto the ray.

    The penalty rises smoothly once the bend angle between `old_d` and the
    refracted `ray.d` exceeds `bend_angle_max` (degrees, default 30) and
    stays at zero for milder refractions. It is added into `ray.bend_penalty`.

    Args:
        ray (Ray): Ray after refraction (`ray.d` is the new direction).
        old_d (torch.Tensor): Pre-refraction ray directions, shape [..., 3],
            same shape as `ray.d`.

    Returns:
        ray (Ray): Ray with `bend_penalty` (shape [..., 1]) updated.
    """
    bend_angle_max = getattr(self, "bend_angle_max", 30.0)
    cos_bend_min = math.cos(math.radians(bend_angle_max))
    cos_bend = torch.sum(ray.d * old_d, dim=-1).unsqueeze(-1)
    per_surf_penalty = F.relu(cos_bend_min - cos_bend)
    valid = ray.is_valid > 0
    ray.bend_penalty = (
        ray.bend_penalty + per_surf_penalty * valid.unsqueeze(-1).float()
    )
    return ray

refract

refract(ray, eta)

Refract a ray with vector Snell's law in local coordinates.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def refract(self, ray, eta):
    """Refract a ray with vector Snell's law in local coordinates."""
    normal_vec = self.normal_vec(ray)
    dot_product = (-normal_vec * ray.d).sum(-1).unsqueeze(-1)
    k = 1 - eta**2 * (1 - dot_product**2)

    valid = (k >= 0).squeeze(-1) & (ray.is_valid > 0)
    k = k * valid.unsqueeze(-1)

    new_d = eta * ray.d + (eta * dot_product - torch.sqrt(k + EPSILON)) * normal_vec
    ray.d = torch.where(valid.unsqueeze(-1), new_d, ray.d)
    ray.is_valid = ray.is_valid * valid
    return ray

to_local_coord

to_local_coord(ray)

Transform a ray from the surface reference frame to local coordinates.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def to_local_coord(self, ray):
    """Transform a ray from the surface reference frame to local coordinates."""
    offset = torch.stack(
        [self.pos_x, self.pos_y, torch.zeros_like(self.pos_x)]
    ).expand_as(ray.o)
    ray.o = ray.o - offset

    if self._R_to_local is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_local)
        ray.d = self._apply_rotation(ray.d, self._R_to_local)
        ray.d = F.normalize(ray.d, p=2, dim=-1)
    return ray

to_global_coord

to_global_coord(ray)

Transform a ray from local coordinates to the surface reference frame.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def to_global_coord(self, ray):
    """Transform a ray from local coordinates to the surface reference frame."""
    if self._R_to_global is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_global)
        ray.d = self._apply_rotation(ray.d, self._R_to_global)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    offset = torch.stack(
        [self.pos_x, self.pos_y, torch.zeros_like(self.pos_x)]
    ).expand_as(ray.o)
    ray.o = ray.o + offset
    return ray

reflect

reflect(ray)

Reflect the ray specularly off the surface (local coordinate system).

The surface normal points from the surface toward the side the light comes from. The reflected direction is renormalized.

Parameters:

Name	Type	Description	Default
ray	Ray	Incident ray bundle.	required

Returns:

Name	Type	Description
ray	Ray	Reflected ray with updated direction.

Reference

[1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/reflect.xhtml [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def reflect(self, ray):
    """Reflect the ray specularly off the surface (local coordinate system).

    The surface normal points from the surface toward the side the light
    comes from. The reflected direction is renormalized.

    Args:
        ray (Ray): Incident ray bundle.

    Returns:
        ray (Ray): Reflected ray with updated direction.

    Reference:
        [1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/reflect.xhtml
        [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.
    """
    # Compute surface normal vectors
    normal_vec = self.normal_vec(ray)

    # Reflect
    dot_product = (normal_vec * ray.d).sum(-1).unsqueeze(-1)
    new_d = ray.d - 2 * dot_product * normal_vec
    new_d = F.normalize(new_d, p=2, dim=-1)

    # Update valid rays
    valid_mask = ray.is_valid > 0
    ray.d = torch.where(valid_mask.unsqueeze(-1), new_d, ray.d)

    return ray

normal_vec

normal_vec(ray)

Compute the unit surface normal at the ray intersection point (local frame).

The normal points from the surface toward the side the light comes from (it is flipped to oppose forward-propagating rays).

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray bundle whose origins ray.o lie on the surface.	required

Returns:

Name	Type	Description
n_vec	Tensor	Unit surface normal vectors, shape [..., 3].

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def normal_vec(self, ray):
    """Compute the unit surface normal at the ray intersection point (local frame).

    The normal points from the surface toward the side the light comes from
    (it is flipped to oppose forward-propagating rays).

    Args:
        ray (Ray): Input ray bundle whose origins `ray.o` lie on the surface.

    Returns:
        n_vec (torch.Tensor): Unit surface normal vectors, shape [..., 3].
    """
    x, y = ray.o[..., 0], ray.o[..., 1]
    nx, ny, nz = self.dfdxyz(x, y)
    n_vec = torch.stack((nx, ny, nz), axis=-1)
    n_vec = F.normalize(n_vec, p=2, dim=-1)

    is_forward = ray.d[..., 2].unsqueeze(-1) > 0
    n_vec = torch.where(is_forward, n_vec, -n_vec)
    return n_vec

sag

sag(x, y, valid=None)

Calculate the surface sag \(z = f(x, y)\) [mm] with validity masking.

The valid mask zeroes out-of-range coordinates before calling _sag, avoiding NaN for spherical/aspheric surfaces where \(r = \sqrt{x^2 + y^2}\) is undefined in back-propagation at \(x = y = 0\) (since \(dr/dx = x/r\)).

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required
valid	Tensor or None	Boolean mask of valid points, same shape as x. Defaults to None, in which case it is computed via is_valid.	None

Returns:

Name	Type	Description
z	Tensor	Surface sag [mm], same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def sag(self, x, y, valid=None):
    """Calculate the surface sag $z = f(x, y)$ [mm] with validity masking.

    The `valid` mask zeroes out-of-range coordinates before calling `_sag`,
    avoiding NaN for spherical/aspheric surfaces where $r = \\sqrt{x^2 + y^2}$
    is undefined in back-propagation at $x = y = 0$ (since $dr/dx = x/r$).

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.
        valid (torch.Tensor or None, optional): Boolean mask of valid points,
            same shape as `x`. Defaults to None, in which case it is computed
            via `is_valid`.

    Returns:
        z (torch.Tensor): Surface sag [mm], same shape as `x`.
    """
    if valid is None:
        valid = self.is_valid(x, y)

    x, y = x * valid, y * valid
    return self._sag(x, y)

dfdxyz

dfdxyz(x, y, valid=None)

Compute the gradient of the implicit surface function.

The surface is defined implicitly as \(f(x, y, z) = \mathrm{sag}(x, y) - z = 0\). This gradient is used in Newton's method and normal-vector computation. The analytical implementation here only works for explicit surfaces \(z = \mathrm{sag}(x, y)\); for implicit surfaces one could instead use numerical finite differences or autograd.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required
valid	Tensor or None	Boolean mask of valid points, same shape as x. Defaults to None, computed via is_valid.	None

Returns:

Name	Type	Description
dfdx	Tensor	Partial derivative \(\partial f/\partial x\) [1], same shape as x.
dfdy	Tensor	Partial derivative \(\partial f/\partial y\) [1], same shape as x.
dfdz	Tensor	Partial derivative \(\partial f/\partial z = -1\), same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def dfdxyz(self, x, y, valid=None):
    """Compute the gradient of the implicit surface function.

    The surface is defined implicitly as $f(x, y, z) = \\mathrm{sag}(x, y) - z = 0$.
    This gradient is used in Newton's method and normal-vector computation.
    The analytical implementation here only works for explicit surfaces
    $z = \\mathrm{sag}(x, y)$; for implicit surfaces one could instead use
    numerical finite differences or autograd.

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.
        valid (torch.Tensor or None, optional): Boolean mask of valid points,
            same shape as `x`. Defaults to None, computed via `is_valid`.

    Returns:
        dfdx (torch.Tensor): Partial derivative $\\partial f/\\partial x$ [1], same shape as `x`.
        dfdy (torch.Tensor): Partial derivative $\\partial f/\\partial y$ [1], same shape as `x`.
        dfdz (torch.Tensor): Partial derivative $\\partial f/\\partial z = -1$, same shape as `x`.
    """
    if valid is None:
        valid = self.is_valid(x, y)

    x, y = x * valid, y * valid
    dx, dy = self._dfdxy(x, y)
    return dx, dy, -torch.ones_like(x)

d2fdxyz2

d2fdxyz2(x, y, valid=None)

Compute second-order partial derivatives of the implicit surface function.

The surface function is \(f(x, y, z) = \mathrm{sag}(x, y) - z = 0\), so all second derivatives involving \(z\) vanish. Currently used only for surface constraints.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required
valid	Tensor or None	Boolean mask of valid points, same shape as x. Defaults to None, computed via is_within_data_range.	None

Returns:

Name	Type	Description
d2f_dx2	Tensor	\(\partial^2 f/\partial x^2\), same shape as x.
d2f_dxdy	Tensor	\(\partial^2 f/\partial x\partial y\), same shape as x.
d2f_dy2	Tensor	\(\partial^2 f/\partial y^2\), same shape as x.
d2f_dxdz	Tensor	\(\partial^2 f/\partial x\partial z = 0\), same shape as x.
d2f_dydz	Tensor	\(\partial^2 f/\partial y\partial z = 0\), same shape as x.
d2f_dz2	Tensor	\(\partial^2 f/\partial z^2 = 0\), same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def d2fdxyz2(self, x, y, valid=None):
    """Compute second-order partial derivatives of the implicit surface function.

    The surface function is $f(x, y, z) = \\mathrm{sag}(x, y) - z = 0$, so all
    second derivatives involving $z$ vanish. Currently used only for surface
    constraints.

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.
        valid (torch.Tensor or None, optional): Boolean mask of valid points,
            same shape as `x`. Defaults to None, computed via
            `is_within_data_range`.

    Returns:
        d2f_dx2 (torch.Tensor): $\\partial^2 f/\\partial x^2$, same shape as `x`.
        d2f_dxdy (torch.Tensor): $\\partial^2 f/\\partial x\\partial y$, same shape as `x`.
        d2f_dy2 (torch.Tensor): $\\partial^2 f/\\partial y^2$, same shape as `x`.
        d2f_dxdz (torch.Tensor): $\\partial^2 f/\\partial x\\partial z = 0$, same shape as `x`.
        d2f_dydz (torch.Tensor): $\\partial^2 f/\\partial y\\partial z = 0$, same shape as `x`.
        d2f_dz2 (torch.Tensor): $\\partial^2 f/\\partial z^2 = 0$, same shape as `x`.
    """
    if valid is None:
        valid = self.is_within_data_range(x, y)

    x, y = x * valid, y * valid

    # Compute second-order derivatives of sag(x, y)
    d2f_dx2, d2f_dxdy, d2f_dy2 = self._d2fdxy(x, y)

    # Mixed partial derivatives involving z are zero
    zeros = torch.zeros_like(x)
    d2f_dxdz = zeros  # ∂²f/∂x∂z = 0
    d2f_dydz = zeros  # ∂²f/∂y∂z = 0
    d2f_dz2 = zeros  # ∂²f/∂z² = 0

    return d2f_dx2, d2f_dxdy, d2f_dy2, d2f_dxdz, d2f_dydz, d2f_dz2

is_valid

is_valid(x, y)

Return a mask of points within both the data range and aperture boundary.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required

Returns:

Name	Type	Description
valid	Tensor	Boolean mask, same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def is_valid(self, x, y):
    """Return a mask of points within both the data range and aperture boundary.

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.

    Returns:
        valid (torch.Tensor): Boolean mask, same shape as `x`.
    """
    return self.is_within_data_range(x, y) & self.is_within_boundary(x, y)

is_within_boundary

is_within_boundary(x, y)

Return a mask of points inside the aperture boundary.

For a square aperture the limits are the half-side lengths w/2, h/2; otherwise the circular radius r is used.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required

Returns:

Name	Type	Description
valid	Tensor	Boolean mask, same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def is_within_boundary(self, x, y):
    """Return a mask of points inside the aperture boundary.

    For a square aperture the limits are the half-side lengths `w/2`, `h/2`;
    otherwise the circular radius `r` is used.

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.

    Returns:
        valid (torch.Tensor): Boolean mask, same shape as `x`.
    """
    if self.is_square:
        valid = (torch.abs(x) <= (self.w / 2 + EPSILON)) & (
            torch.abs(y) <= (self.h / 2 + EPSILON)
        )
    else:
        r = self.r
        valid = (x**2 + y**2) <= (r**2 + EPSILON)

    return valid

is_within_data_range

is_within_data_range(x, y)

Return a mask of points inside the sag function's data region.

The base surface has an unbounded data region, so all points are valid; subclasses (e.g. spheric) override this to exclude regions where the sag is undefined.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm], any shape.	required
y	Tensor	Local y coordinate [mm], same shape as x.	required

Returns:

Name	Type	Description
valid	Tensor	Boolean mask, same shape as x (all True here).

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def is_within_data_range(self, x, y):
    """Return a mask of points inside the sag function's data region.

    The base surface has an unbounded data region, so all points are valid;
    subclasses (e.g. spheric) override this to exclude regions where the sag
    is undefined.

    Args:
        x (torch.Tensor): Local x coordinate [mm], any shape.
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.

    Returns:
        valid (torch.Tensor): Boolean mask, same shape as `x` (all True here).
    """
    return torch.ones_like(x, dtype=torch.bool)

max_height

max_height()

Return the maximum valid radial height of the surface [mm].

Returns:

Name	Type	Description
max_height	float	Maximum valid height [mm] (10e3 for the base surface).

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def max_height(self):
    """Return the maximum valid radial height of the surface [mm].

    Returns:
        max_height (float): Maximum valid height [mm] (10e3 for the base surface).
    """
    return 10e3

surface_with_offset

surface_with_offset(x, y, valid_check=True, d=0.0)

Compute the z coordinate sag(x, y) + d.

The surface does not own an absolute axial position, so callers that need global geometry pass the vertex position from GeoLens.surf_d.

Parameters:

Name	Type	Description	Default
x	Tensor or float	Local x coordinate [mm].	required
y	Tensor or float	Local y coordinate [mm], same shape as x.	required
valid_check	bool	If True apply is_valid masking via sag; if False use the raw _sag. Defaults to True.	True
d	float or Tensor	Vertex position to add [mm]. Defaults to 0 (the surface reference frame).	0.0

Returns:

Name	Type	Description
z	Tensor	Global z coordinate [mm], same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def surface_with_offset(self, x, y, valid_check=True, d=0.0):
    """Compute the z coordinate `sag(x, y) + d`.

    The surface does not own an absolute axial position, so callers that
    need global geometry pass the vertex position from `GeoLens.surf_d`.

    Args:
        x (torch.Tensor or float): Local x coordinate [mm].
        y (torch.Tensor or float): Local y coordinate [mm], same shape as `x`.
        valid_check (bool, optional): If True apply `is_valid` masking via
            `sag`; if False use the raw `_sag`. Defaults to True.
        d (float or torch.Tensor, optional): Vertex position to add [mm].
            Defaults to 0 (the surface reference frame).

    Returns:
        z (torch.Tensor): Global z coordinate [mm], same shape as `x`.
    """
    x = (
        x
        if torch.is_tensor(x)
        else torch.tensor(x, device=self.device, dtype=self.dtype)
    )
    y = (
        y
        if torch.is_tensor(y)
        else torch.tensor(y, device=self.device, dtype=self.dtype)
    )
    if valid_check:
        return self.sag(x, y) + d
    else:
        return self._sag(x, y) + d

surface_sag

surface_sag(x, y)

Compute the local surface sag at (x, y) as a Python float.

This function is currently not used.

Parameters:

Name	Type	Description	Default
x	Tensor or float	Local x coordinate [mm].	required
y	Tensor or float	Local y coordinate [mm].	required

Returns:

Name	Type	Description
sag	float	Surface sag [mm] at (x, y).

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def surface_sag(self, x, y):
    """Compute the local surface sag at (x, y) as a Python float.

    This function is currently not used.

    Args:
        x (torch.Tensor or float): Local x coordinate [mm].
        y (torch.Tensor or float): Local y coordinate [mm].

    Returns:
        sag (float): Surface sag [mm] at (x, y).
    """
    x = (
        x
        if torch.is_tensor(x)
        else torch.tensor(x, device=self.device, dtype=self.dtype)
    )
    y = (
        y
        if torch.is_tensor(y)
        else torch.tensor(y, device=self.device, dtype=self.dtype)
    )
    return self.sag(x, y).item()

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Build the per-parameter optimizer parameter groups (subclass-specific).

Parameters:

Name	Type	Description	Default
lrs	list[float]	Learning rates for the surface's differentiable parameters. Defaults to [1e-4].	[0.0001]
optim_mat	bool	Whether to also optimize the material refractive index/dispersion. Defaults to False.	False

Returns:

Name	Type	Description
params	list[dict]	Adam parameter groups (param tensors and lr).

Raises:

Type	Description
NotImplementedError	Always, on the base class; subclasses override.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Build the per-parameter optimizer parameter groups (subclass-specific).

    Args:
        lrs (list[float], optional): Learning rates for the surface's
            differentiable parameters. Defaults to [1e-4].
        optim_mat (bool, optional): Whether to also optimize the material
            refractive index/dispersion. Defaults to False.

    Returns:
        params (list[dict]): Adam parameter groups (param tensors and `lr`).

    Raises:
        NotImplementedError: Always, on the base class; subclasses override.
    """
    raise NotImplementedError(
        "get_optimizer_params() is not implemented for {}".format(
            self.__class__.__name__
        )
    )

get_optimizer

get_optimizer(lrs=[0.0001], optim_mat=False)

Build an Adam optimizer over the surface's differentiable parameters.

Parameters:

Name	Type	Description	Default
lrs	list[float]	Learning rates passed to get_optimizer_params. Defaults to [1e-4].	[0.0001]
optim_mat	bool	Whether to optimize the material. Defaults to False.	False

Returns:

Name	Type	Description
optimizer	Adam	Adam optimizer for the surface.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def get_optimizer(self, lrs=[1e-4], optim_mat=False):
    """Build an Adam optimizer over the surface's differentiable parameters.

    Args:
        lrs (list[float], optional): Learning rates passed to
            `get_optimizer_params`. Defaults to [1e-4].
        optim_mat (bool, optional): Whether to optimize the material.
            Defaults to False.

    Returns:
        optimizer (torch.optim.Adam): Adam optimizer for the surface.
    """
    params = self.get_optimizer_params(lrs, optim_mat=optim_mat)
    return torch.optim.Adam(params)

update_r

update_r(r)

Update the aperture radius, clamped to max_height.

Parameters:

Name	Type	Description	Default
r	float	Requested aperture radius [mm].	required

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def update_r(self, r):
    """Update the aperture radius, clamped to `max_height`.

    Args:
        r (float): Requested aperture radius [mm].
    """
    r_max = self.max_height()
    self.r = min(r, r_max)

draw_r

draw_r()

Return the effective drawing radius [mm], clamped to max_height.

Returns:

Name	Type	Description
r_eff	float	Effective drawing radius [mm].

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def draw_r(self):
    """Return the effective drawing radius [mm], clamped to `max_height`.

    Returns:
        r_eff (float): Effective drawing radius [mm].
    """
    return min(self.r, self.max_height())

draw_widget

draw_widget(ax, color='black', linestyle='solid', d=0.0)

Draw the surface profile as a 2D line on a Matplotlib axis.

Plots the meridional (y-z) cross section sampled across the aperture.

Parameters:

Name	Type	Description	Default
ax	Axes	Axis to draw on.	required
color	str	Line color. Defaults to "black".	'black'
linestyle	str	Matplotlib line style. Defaults to "solid".	'solid'

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def draw_widget(self, ax, color="black", linestyle="solid", d=0.0):
    """Draw the surface profile as a 2D line on a Matplotlib axis.

    Plots the meridional (y-z) cross section sampled across the aperture.

    Args:
        ax (matplotlib.axes.Axes): Axis to draw on.
        color (str, optional): Line color. Defaults to "black".
        linestyle (str, optional): Matplotlib line style. Defaults to "solid".
    """
    r_eff = self.draw_r()
    r = torch.linspace(-r_eff, r_eff, 128, device=self.device, dtype=self.dtype)
    z = self.surface_with_offset(
        r,
        torch.zeros(len(r), device=self.device, dtype=self.dtype),
        valid_check=False,
        d=d,
    )
    ax.plot(
        z.cpu().detach().numpy(),
        r.cpu().detach().numpy(),
        color=color,
        linestyle=linestyle,
        linewidth=0.75,
    )

create_mesh

create_mesh(n_rings=32, n_arms=128, color=[0.06, 0.3, 0.6], d=0.0)

Create a triangulated mesh of the surface for 3D visualization.

Populates self.vertices, self.faces, self.rim, and self.mesh_color.

Parameters:

Name	Type	Description	Default
n_rings	int	Number of concentric rings for radial sampling. Defaults to 32.	32
n_arms	int	Number of angular divisions. Defaults to 128.	128
color	list[float]	RGB mesh color in [0, 1]. Defaults to [0.06, 0.3, 0.6].	[0.06, 0.3, 0.6]

Returns:

Name	Type	Description
self	Surface	The surface with mesh data (for chaining).

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def create_mesh(self, n_rings=32, n_arms=128, color=[0.06, 0.3, 0.6], d=0.0):
    """Create a triangulated mesh of the surface for 3D visualization.

    Populates `self.vertices`, `self.faces`, `self.rim`, and `self.mesh_color`.

    Args:
        n_rings (int, optional): Number of concentric rings for radial
            sampling. Defaults to 32.
        n_arms (int, optional): Number of angular divisions. Defaults to 128.
        color (list[float], optional): RGB mesh color in [0, 1]. Defaults to
            [0.06, 0.3, 0.6].

    Returns:
        self (Surface): The surface with mesh data (for chaining).
    """
    self.vertices = self._create_vertices(n_rings, n_arms, d=d)
    self.faces = self._create_faces(n_rings, n_arms)
    self.rim = self._create_rim(n_rings, n_arms)
    self.mesh_color = color
    return self

get_polydata

get_polydata()

Build a PyVista PolyData object from the cached vertices and faces.

Requires create_mesh to have been called first. The PolyData is used to draw the surface and export it as an .obj file.

Returns:

Name	Type	Description
polydata	PolyData	Mesh as a PyVista PolyData object.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def get_polydata(self):
    """Build a PyVista PolyData object from the cached vertices and faces.

    Requires `create_mesh` to have been called first. The PolyData is used to
    draw the surface and export it as an .obj file.

    Returns:
        polydata (pyvista.PolyData): Mesh as a PyVista PolyData object.
    """
    from pyvista import PolyData

    face_vertex_n = 3  # vertices per triangle
    formatted_faces = np.hstack(
        [
            face_vertex_n * np.ones((self.faces.shape[0], 1), dtype=np.uint32),
            self.faces,
        ]
    )
    return PolyData(self.vertices, formatted_faces)

surf_dict

surf_dict()

Serialize the surface's common parameters to a dict.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters (type, r, d_next, pos_xy, vec_local, is_square, mat2, plus informational (mat2_n)/(mat2_V)), with numeric values rounded to 4 decimals.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def surf_dict(self):
    """Serialize the surface's common parameters to a dict.

    Returns:
        surf_dict (dict): Surface parameters (type, `r`, `d_next`, `pos_xy`,
            `vec_local`, `is_square`, `mat2`, plus informational
            `(mat2_n)`/`(mat2_V)`), with numeric values rounded to 4 decimals.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "d_next": self.d_next.item(),
        "pos_xy": (self.pos_x.item(), self.pos_y.item()),
        "vec_local": tuple(self.vec_local.tolist()),
        "is_square": self.is_square,
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Return the Zemax (.zmx) text block describing this surface.

Parameters:

Name	Type	Description	Default
surf_idx	int	Index of this surface in the Zemax surface list.	required
d_next	float	Axial distance [mm] to the next surface (thickness).	required

Returns:

Name	Type	Description
zmx_str	str	Zemax-formatted surface definition string.

Raises:

Type	Description
NotImplementedError	Always, on the base class; subclasses override.

Source code in deeplens-src/deeplens/geometric_surface/base_surface.py

def zmx_str(self, surf_idx, d_next):
    """Return the Zemax (.zmx) text block describing this surface.

    Args:
        surf_idx (int): Index of this surface in the Zemax surface list.
        d_next (float): Axial distance [mm] to the next surface (thickness).

    Returns:
        zmx_str (str): Zemax-formatted surface definition string.

    Raises:
        NotImplementedError: Always, on the base class; subclasses override.
    """
    raise NotImplementedError(
        "zmx_str() is not implemented for {}".format(self.__class__.__name__)
    )

Spherical surface defined by curvature c = 1/R.

deeplens.geometric_surface.Spheric

Spheric(
    c,
    r,
    d_next,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
)

Bases: Surface

Spherical refractive surface parameterized by curvature.

A sphere of radius \(R = 1/c\) whose vertex sits at the optical axis. The sag (surface height along \(z\)) is:

\[ z(x, y) = \frac{c \rho^2}{1 + \sqrt{1 - c^2 \rho^2}}, \quad \rho^2 = x^2 + y^2 \]

Attributes:

Name	Type	Description
c	Tensor	Surface curvature \(1/R\) [1/mm], scalar tensor. Gradients are enabled by get_optimizer_params for optimization.

Initialize a spherical surface.

Parameters:

Name	Type	Description	Default
c	float	Surface curvature \(1/R\) [1/mm]. Use 0 for a flat surface (treated as a plane).	required
r	float	Aperture radius [mm].	required
d	float	Axial vertex position [mm].	required
mat2	str or Material	Material on the transmission side.	required
pos_xy	list[float]	Lateral offset [x, y] [mm]. Defaults to [0.0, 0.0].	[0.0, 0.0]
vec_local	list[float]	Local surface normal direction. Defaults to [0.0, 0.0, 1.0].	[0.0, 0.0, 1.0]
is_square	bool	Use a square aperture instead of a circular one. Defaults to False.	False
device	str	Compute device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def __init__(
    self,
    c,
    r,
    d_next,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize a spherical surface.

    Args:
        c (float): Surface curvature $1/R$ [1/mm]. Use 0 for a flat
            surface (treated as a plane).
        r (float): Aperture radius [mm].
        d (float): Axial vertex position [mm].
        mat2 (str or Material): Material on the transmission side.
        pos_xy (list[float], optional): Lateral offset `[x, y]` [mm].
            Defaults to `[0.0, 0.0]`.
        vec_local (list[float], optional): Local surface normal direction.
            Defaults to `[0.0, 0.0, 1.0]`.
        is_square (bool, optional): Use a square aperture instead of a
            circular one. Defaults to False.
        device (str, optional): Compute device. Defaults to `"cpu"`.
    """
    super(Spheric, self).__init__(
        r=r,
        d_next=d_next,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.c = torch.as_tensor(c, dtype=self.d_next.dtype, device=device)
    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Construct a Spheric surface from a parameter dictionary.

Accepts either a radius of curvature roc [mm] (converted to curvature c, with roc == 0 mapped to c = 0) or a curvature c [1/mm] directly. The aperture radius r [mm], vertex position d [mm], and transmission material mat2 are read from the dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Must contain r, d, mat2 and either roc or c.	required

Returns:

Name	Type	Description
surface	Spheric	The constructed spherical surface.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Construct a `Spheric` surface from a parameter dictionary.

    Accepts either a radius of curvature `roc` [mm] (converted to curvature
    `c`, with `roc == 0` mapped to `c = 0`) or a curvature `c` [1/mm]
    directly. The aperture radius `r` [mm], vertex position `d` [mm], and
    transmission material `mat2` are read from the dictionary.

    Args:
        surf_dict (dict): Surface parameters. Must contain `r`, `d`, `mat2`
            and either `roc` or `c`.

    Returns:
        surface (Spheric): The constructed spherical surface.
    """
    if "roc" in surf_dict:
        if surf_dict["roc"] != 0:
            c = 1 / surf_dict["roc"]
        else:
            c = 0.0
    else:
        c = surf_dict["c"]

    return cls(
        c=c,
        r=surf_dict["r"],
        d_next=surf_dict["d_next"],
        mat2=surf_dict["mat2"],
        pos_xy=surf_dict.get("pos_xy", [0.0, 0.0]),
        vec_local=surf_dict.get("vec_local", [0.0, 0.0, 1.0]),
        is_square=surf_dict.get("is_square", False),
        device=surf_dict.get("device", "cpu"),
    )

paraxial_power

paraxial_power(n1, n2)

Return the paraxial optical power of this surface [1/mm].

Only the vertex curvature contributes; higher-order geometry does not affect first-order properties.

Parameters:

Name	Type	Description	Default
n1	Tensor	Refractive index of the incident medium.	required
n2	Tensor	Refractive index of the transmission medium.	required

Returns:

Name	Type	Description
power	Tensor	Surface power (n2 - n1) * c [1/mm], scalar.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def paraxial_power(self, n1, n2):
    """Return the paraxial optical power of this surface [1/mm].

    Only the vertex curvature contributes; higher-order geometry does
    not affect first-order properties.

    Args:
        n1 (torch.Tensor): Refractive index of the incident medium.
        n2 (torch.Tensor): Refractive index of the transmission medium.

    Returns:
        power (torch.Tensor): Surface power `(n2 - n1) * c` [1/mm], scalar.
    """
    return (n2 - n1) * self.c

intersect

intersect(ray, n=1.0)

Solve the ray-surface intersection analytically in local coordinates.

Substitutes the ray \(p(t) = o + t\,d\) into the sphere \(x^2 + y^2 + (z - R)^2 = R^2\) (with \(R = 1/c\)) and solves the resulting quadratic for \(t\), picking the root whose intersection lies closest to the surface vertex at \(z = 0\). A flat surface ($|c| < $ EPSILON) is handled as a plane. Rays falling outside the aperture or with no real root are flagged invalid. Updates ray.o, ray.is_valid, and, for coherent rays, ray.opl (adding \(n\,t\)).

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray, modified in place.	required
n	float	Refractive index of the incident medium, used for the optical path length update. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
ray	Ray	The same ray with updated position, validity, and opl.

Raises:

Type	Description
Exception	If a coherent ray travels more than 100 mm under float32, where OPL accumulation loses precision.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def intersect(self, ray, n=1.0):
    """Solve the ray-surface intersection analytically in local coordinates.

    Substitutes the ray $p(t) = o + t\\,d$ into the sphere
    $x^2 + y^2 + (z - R)^2 = R^2$ (with $R = 1/c$) and solves the resulting
    quadratic for $t$, picking the root whose intersection lies closest to
    the surface vertex at $z = 0$. A flat surface ($|c| < $ `EPSILON`) is
    handled as a plane. Rays falling outside the aperture or with no real
    root are flagged invalid. Updates `ray.o`, `ray.is_valid`, and, for
    coherent rays, `ray.opl` (adding $n\\,t$).

    Args:
        ray (Ray): Input ray, modified in place.
        n (float, optional): Refractive index of the incident medium, used
            for the optical path length update. Defaults to 1.0.

    Returns:
        ray (Ray): The same ray with updated position, validity, and opl.

    Raises:
        Exception: If a coherent ray travels more than 100 mm under
            float32, where OPL accumulation loses precision.
    """
    c = self.c
    original_o = ray.o

    # Use the vertex-anchored sphere equation. The surface passes through
    # the local origin, so
    #
    #   x² + y² + z² - 2 z R = 0,  R = 1 / c.
    #
    # Multiplying by c before substituting o + t*d gives a quadratic with
    # no R² subtraction. Flat surfaces are selected with tensor operations
    # so tracing does not extract curvature to the host at every surface.
    is_flat = torch.abs(c) < EPSILON
    c_solve = torch.where(is_flat, torch.ones_like(c), c)
    od = torch.sum(original_o * ray.d, dim=-1)
    dd = torch.sum(ray.d * ray.d, dim=-1)
    oo = torch.sum(original_o * original_o, dim=-1)

    a = c_solve * dd
    b = 2.0 * (c_solve * od - ray.d[..., 2])
    c_coeff = c_solve * oo - 2.0 * original_o[..., 2]

    discriminant = b * b - 4 * a * c_coeff
    sphere_valid = discriminant >= 0
    sqrt_discriminant = torch.sqrt(torch.clamp(discriminant, min=EPSILON))

    # Stable root pair. q avoids subtracting nearly equal values; q/a and
    # c_coeff/q are the two roots by Vieta's formula.
    q = torch.where(
        b >= 0,
        -(b + sqrt_discriminant) / 2.0,
        (sqrt_discriminant - b) / 2.0,
    )
    q_safe = torch.where(
        q.abs() < EPSILON,
        torch.where(
            q < 0,
            -torch.full_like(q, EPSILON),
            torch.full_like(q, EPSILON),
        ),
        q,
    )
    t1 = q / a
    t2 = c_coeff / q_safe

    z1 = original_o[..., 2] + t1 * ray.d[..., 2]
    z2 = original_o[..., 2] + t2 * ray.d[..., 2]
    sphere_t = torch.where(torch.abs(z1) < torch.abs(z2), t1, t2)

    dz_safe = torch.where(
        ray.d[..., 2].abs() < EPSILON,
        torch.where(
            ray.d[..., 2] < 0,
            -torch.full_like(ray.d[..., 2], EPSILON),
            torch.full_like(ray.d[..., 2], EPSILON),
        ),
        ray.d[..., 2],
    )
    plane_t = -original_o[..., 2] / dz_safe
    t = torch.where(is_flat, plane_t, sphere_t)
    plane_valid = (ray.d[..., 2].abs() >= EPSILON) & torch.isfinite(plane_t)
    valid_intersect = torch.where(is_flat, plane_valid, sphere_valid)

    new_o = original_o + t.unsqueeze(-1) * ray.d
    within_aperture = self.is_within_boundary(new_o[..., 0], new_o[..., 1])
    valid = (
        valid_intersect & torch.isfinite(t) & within_aperture & (ray.is_valid > 0)
    )

    # Update ray position
    ray.o = torch.where(valid.unsqueeze(-1), new_o, original_o)
    ray.is_valid = ray.is_valid * valid

    if ray.is_coherent:
        total_t = t
        if ray.o.dtype == torch.float32 and total_t.abs().max() > 100:
            raise ValueError(
                "Coherent ray tracing over long paths requires float64 rays."
            )
        new_opl = ray.opl + n * total_t.unsqueeze(-1)
        ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    return ray

is_within_data_range

is_within_data_range(x, y)

Check whether points lie within the sag-defined region.

Points are valid only where \(x^2 + y^2 < 1/c^2\), i.e. inside the radius where the sphere's sag is real-valued.

Parameters:

Name	Type	Description	Default
x	Tensor	Local x coordinate [mm].	required
y	Tensor	Local y coordinate [mm], same shape as x.	required

Returns:

Name	Type	Description
valid	Tensor	Boolean mask, same shape as x.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def is_within_data_range(self, x, y):
    """Check whether points lie within the sag-defined region.

    Points are valid only where $x^2 + y^2 < 1/c^2$, i.e. inside the radius
    where the sphere's sag is real-valued.

    Args:
        x (torch.Tensor): Local x coordinate [mm].
        y (torch.Tensor): Local y coordinate [mm], same shape as `x`.

    Returns:
        valid (torch.Tensor): Boolean mask, same shape as `x`.
    """
    c = self.c
    valid = (x**2 + y**2) < 1 / c**2
    return valid

max_height

max_height()

Return the maximum valid radial height of the surface.

Equal to \(|R| = 1/|c|\) minus a small 0.001 mm margin to stay inside the region where the sag is well-defined.

Returns:

Name	Type	Description
max_height	float	Maximum radial height [mm].

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def max_height(self):
    """Return the maximum valid radial height of the surface.

    Equal to $|R| = 1/|c|$ minus a small 0.001 mm margin to stay inside the
    region where the sag is well-defined.

    Returns:
        max_height (float): Maximum radial height [mm].
    """
    c = self.c
    max_height = torch.sqrt(1 / c**2).item() - 0.001
    return max_height

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001], optim_mat=False)

Enable gradients on c and d and build optimizer parameter groups.

Parameters:

Name	Type	Description	Default
lrs	list[float]	Learning rates [lr_d, lr_c] for the vertex position and curvature. Defaults to [1e-4, 1e-4].	[0.0001, 0.0001]
optim_mat	bool	Also optimize the transmission material parameters (skipped when the material is air). Defaults to False.	False

Returns:

Name	Type	Description
params	list[dict]	Optimizer parameter groups, each with params and lr keys.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def get_optimizer_params(self, lrs=[1e-4, 1e-4], optim_mat=False):
    """Enable gradients on `c` and `d` and build optimizer parameter groups.

    Args:
        lrs (list[float], optional): Learning rates `[lr_d, lr_c]` for the
            vertex position and curvature. Defaults to `[1e-4, 1e-4]`.
        optim_mat (bool, optional): Also optimize the transmission material
            parameters (skipped when the material is air). Defaults to False.

    Returns:
        params (list[dict]): Optimizer parameter groups, each with `params`
            and `lr` keys.
    """
    self.c.requires_grad_(True)
    self.d_next.requires_grad_(True)

    params = []
    params.append({"params": [self.d_next], "lr": lrs[0]})
    params.append({"params": [self.c], "lr": lrs[1]})

    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

surf_dict

surf_dict()

Serialize the surface to a parameter dictionary.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters with keys type, r, (c), roc, (d), mat2, plus informational (mat2_n)/(mat2_V). Lengths are in [mm], curvature in [1/mm], rounded to 4 decimals.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

def surf_dict(self):
    """Serialize the surface to a parameter dictionary.

    Returns:
        surf_dict (dict): Surface parameters with keys `type`, `r`, `(c)`,
            `roc`, `(d)`, `mat2`, plus informational `(mat2_n)`/`(mat2_V)`.
            Lengths are in [mm], curvature in
            [1/mm], rounded to 4 decimals.
    """
    c_value = self.c.item()
    roc = 1 / c_value if c_value != 0 else 0.0
    surf_dict = {
        "type": "Spheric",
        "r": self.r,
        "(c)": c_value,
        "roc": roc,
        "d_next": self.d_next.item(),
        "pos_xy": [self.pos_x.item(), self.pos_y.item()],
        "vec_local": self.vec_local.tolist(),
        "is_square": self.is_square,
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Format the surface as a Zemax STANDARD surface block.

Parameters:

Name	Type	Description	Default
surf_idx	int	Surface index in the Zemax file.	required
d_next	Tensor	Axial distance to the next surface [mm], scalar tensor.	required

Returns:

Name	Type	Description
zmx_str	str	Multi-line Zemax surface description.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py

    def zmx_str(self, surf_idx, d_next):
        """Format the surface as a Zemax STANDARD surface block.

        Args:
            surf_idx (int): Surface index in the Zemax file.
            d_next (torch.Tensor): Axial distance to the next surface [mm],
                scalar tensor.

        Returns:
            zmx_str (str): Multi-line Zemax surface description.
        """
        if self.mat2.get_name() == "air":
            zmx_str = f"""SURF {surf_idx} 
    TYPE STANDARD 
    CURV {self.c.item()} 
    DISZ {d_next.item()} 
    DIAM {self.r} 1 0 0 1 ""
"""
        else:
            zmx_str = f"""SURF {surf_idx} 
    TYPE STANDARD 
    CURV {self.c.item()} 
    DISZ {d_next.item()} 
    GLAS ___BLANK 1 0 {self.mat2.n} {self.mat2.V}
    DIAM {self.r} 1 0 0 1 ""
"""
        return zmx_str

Even-asphere surface: spherical base with polynomial corrections.

deeplens.geometric_surface.Aspheric

Aspheric(
    r,
    d_next,
    c,
    k,
    ai,
    mat2,
    ai2=None,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
)

Bases: Surface

Even-order aspheric surface.

The sag function is:

\[ z(\rho) = \frac{c\,\rho^2}{1 + \sqrt{1-(1+k)c^2\rho^2}} + \sum_{i=2}^{n} a_{2i}\,\rho^{2i}, \quad \rho^2 = x^2 + y^2 \]

The polynomial starts at the 4th-order term (a4) because the 2nd-order term competes with the base curvature c.

All coefficients c, k, and ai are differentiable torch tensors so they can be optimised with gradient descent.

Attributes:

Name	Type	Description
c	Tensor	Base curvature [1/mm].
k	Tensor	Conic constant.
ai2	Tensor or None	2nd-order aspheric coefficient (legacy).
ai	Tensor	Even-order aspheric coefficients [a4, a6, a8, ...].

Initialize an aspheric surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius [mm].	required
d	float	Axial vertex position [mm].	required
c	float	Base curvature 1/R [1/mm].	required
k	float	Conic constant (0 = sphere, -1 = paraboloid).	required
ai	list[float] or None	Even-order aspheric coefficients starting from the 4th-order term: [a4, a6, a8, ...]. Pass None or an empty list for a pure conic.	required
mat2	str or Material	Material on the transmission side.	required
ai2	float or None	2nd-order aspheric coefficient from legacy data. Included in sag but not optimised. Defaults to None.	None
pos_xy	list[float]	Lateral offset [x, y] [mm]. Defaults to [0.0, 0.0].	[0.0, 0.0]
vec_local	list[float]	Local normal direction. Defaults to [0.0, 0.0, 1.0].	[0.0, 0.0, 1.0]
is_square	bool	Square aperture flag. Defaults to False.	False
device	str	Compute device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def __init__(
    self,
    r,
    d_next,
    c,
    k,
    ai,
    mat2,
    ai2=None,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize an aspheric surface.

    Args:
        r (float): Aperture radius [mm].
        d (float): Axial vertex position [mm].
        c (float): Base curvature `1/R` [1/mm].
        k (float): Conic constant (`0` = sphere, `-1` = paraboloid).
        ai (list[float] or None): Even-order aspheric coefficients
            starting from the 4th-order term: `[a4, a6, a8, ...]`.
            Pass `None` or an empty list for a pure conic.
        mat2 (str or Material): Material on the transmission side.
        ai2 (float or None, optional): 2nd-order aspheric coefficient
            from legacy data. Included in sag but not optimised.
            Defaults to None.
        pos_xy (list[float], optional): Lateral offset `[x, y]` [mm].
            Defaults to `[0.0, 0.0]`.
        vec_local (list[float], optional): Local normal direction.
            Defaults to `[0.0, 0.0, 1.0]`.
        is_square (bool, optional): Square aperture flag.
            Defaults to False.
        device (str, optional): Compute device. Defaults to `"cpu"`.
    """
    Surface.__init__(
        self,
        r=r,
        d_next=d_next,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    tensor_kwargs = {"dtype": self.d_next.dtype, "device": device}
    self.c = torch.as_tensor(c, **tensor_kwargs)
    self.k = torch.as_tensor(k, **tensor_kwargs)

    # 2nd-order coefficient (legacy, not optimised)
    if ai2 is not None:
        self.ai2 = torch.as_tensor(ai2, **tensor_kwargs)
    else:
        self.ai2 = None

    if ai is not None and len(ai) > 0:
        self.ai = torch.as_tensor(ai, **tensor_kwargs)
        self.ai_degree = len(ai)
        # ai[0] -> ai4, ai[1] -> ai6, ai[2] -> ai8, ...
        for i, a in enumerate(ai):
            setattr(self, f"ai{2 * (i + 2)}", torch.as_tensor(a, **tensor_kwargs))
    else:
        self.ai = None
        self.ai_degree = 0

    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Create an aspheric surface from a serialized dict.

The base curvature is read from roc (radius of curvature [mm], converted to c = 1/roc) when present, otherwise from c [1/mm]. For legacy data where use_ai2 is True (or absent), the first element of ai is interpreted as the 2nd-order coefficient a2 and the rest as [a4, a6, a8, ...].

Parameters:

Name	Type	Description	Default
surf_dict	dict	Serialized surface with keys r, d, k, mat2, either roc or c, and optionally ai and use_ai2.	required

Returns:

Name	Type	Description
surface	Aspheric	The reconstructed aspheric surface.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Create an aspheric surface from a serialized dict.

    The base curvature is read from `roc` (radius of curvature [mm],
    converted to `c = 1/roc`) when present, otherwise from `c` [1/mm].
    For legacy data where `use_ai2` is True (or absent), the first
    element of `ai` is interpreted as the 2nd-order coefficient `a2`
    and the rest as `[a4, a6, a8, ...]`.

    Args:
        surf_dict (dict): Serialized surface with keys `r`, `d`, `k`,
            `mat2`, either `roc` or `c`, and optionally `ai` and
            `use_ai2`.

    Returns:
        surface (Aspheric): The reconstructed aspheric surface.
    """
    if "roc" in surf_dict:
        if surf_dict["roc"] != 0:
            c = 1 / surf_dict["roc"]
        else:
            c = 0.0
    else:
        c = surf_dict["c"]

    ai = surf_dict.get("ai", [])
    ai2_val = None

    # Backward compatibility: old format includes a2 as first element.
    # New files written by this code set use_ai2 explicitly.
    if surf_dict.get("use_ai2", True) and len(ai) > 0:
        if "use_ai2" not in surf_dict:
            print(
                f"Surface dict lacks 'use_ai2'; assuming ai[0]={ai[0]:.4g} is the "
                "2nd-order coefficient (legacy format)."
            )
        ai2_val = ai[0]  # Extract the a2 coefficient
        ai = ai[1:]  # Remaining: [a4, a6, a8, ...]

    return cls(
        r=surf_dict["r"],
        d_next=surf_dict["d_next"],
        c=c,
        k=surf_dict["k"],
        ai=ai,
        ai2=ai2_val,
        mat2=surf_dict["mat2"],
        pos_xy=surf_dict.get("pos_xy", [0.0, 0.0]),
        vec_local=surf_dict.get("vec_local", [0.0, 0.0, 1.0]),
        is_square=surf_dict.get("is_square", False),
        device=surf_dict.get("device", "cpu"),
    )

paraxial_power

paraxial_power(n1, n2)

Return the paraxial optical power of this surface [1/mm].

Only the vertex curvature contributes, so the conic constant and the 4th- and higher-order coefficients do not affect first-order properties. The legacy \(a_2\rho^2\) term is the exception: near the vertex the sag is \((c/2 + a_2)\rho^2\), so it shifts the vertex curvature to \(c + 2 a_2\) and must be included.

Parameters:

Name	Type	Description	Default
n1	Tensor	Refractive index of the incident medium.	required
n2	Tensor	Refractive index of the transmission medium.	required

Returns:

Name	Type	Description
power	Tensor	Surface power (n2 - n1) * (c + 2 * a2) [1/mm], scalar.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def paraxial_power(self, n1, n2):
    """Return the paraxial optical power of this surface [1/mm].

    Only the vertex curvature contributes, so the conic constant and the
    4th- and higher-order coefficients do not affect first-order
    properties. The legacy $a_2\\rho^2$ term is the exception: near the
    vertex the sag is $(c/2 + a_2)\\rho^2$, so it shifts the vertex
    curvature to $c + 2 a_2$ and must be included.

    Args:
        n1 (torch.Tensor): Refractive index of the incident medium.
        n2 (torch.Tensor): Refractive index of the transmission medium.

    Returns:
        power (torch.Tensor): Surface power `(n2 - n1) * (c + 2 * a2)`
            [1/mm], scalar.
    """
    c = self.c if self.ai2 is None else self.c + 2.0 * self.ai2
    return (n2 - n1) * c

is_within_data_range

is_within_data_range(x, y)

Return a mask of points where the conic sag is real-valued.

A point is valid when \((1+k)c^2\rho^2 < 1\), i.e. inside the conic's real boundary. Fully tensorized (no Python branch on the tensor value of k) so the function is safe to trace through torch.compile. When \(k \le -1\) the conic has no real boundary, so every point is treated as valid.

Parameters:

Name	Type	Description	Default
x	Tensor	x coordinate(s) [mm], any shape.	required
y	Tensor	y coordinate(s) [mm], broadcastable with x.	required

Returns:

Name	Type	Description
valid	Tensor	Boolean mask, same shape as the broadcast of x and y.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def is_within_data_range(self, x, y):
    """Return a mask of points where the conic sag is real-valued.

    A point is valid when $(1+k)c^2\\rho^2 < 1$, i.e. inside the conic's
    real boundary. Fully tensorized (no Python branch on the tensor value
    of `k`) so the function is safe to trace through `torch.compile`.
    When $k \\le -1$ the conic has no real boundary, so every point is
    treated as valid.

    Args:
        x (torch.Tensor): x coordinate(s) [mm], any shape.
        y (torch.Tensor): y coordinate(s) [mm], broadcastable with `x`.

    Returns:
        valid (torch.Tensor): Boolean mask, same shape as the broadcast
            of `x` and `y`.
    """
    c, k = self._get_curvature_params()
    one_plus_k = 1 + k
    # Avoid division by zero / negative when computing the limit; the
    # bogus value is masked out by the where below.
    safe = torch.where(one_plus_k > 0, one_plus_k, torch.ones_like(one_plus_k))
    limit_sq = 1.0 / (c * c * safe)
    inside = (x * x + y * y) < limit_sq
    return torch.where(one_plus_k > 0, inside, torch.ones_like(inside))

max_height

max_height()

Return the maximum valid radial height of the surface.

For an oblate/ellipsoidal conic (\(k > -1\)) the sag is real only up to \(\rho_{max} = \sqrt{1/((k+1)c^2)}\); a small margin (0.001 mm) is subtracted. For \(k \le -1\) there is no boundary and a large value (10000 mm) is returned.

Returns:

Name	Type	Description
max_height	float	Maximum valid radial height [mm].

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def max_height(self):
    """Return the maximum valid radial height of the surface.

    For an oblate/ellipsoidal conic ($k > -1$) the sag is real only up to
    $\\rho_{max} = \\sqrt{1/((k+1)c^2)}$; a small margin (0.001 mm) is
    subtracted. For $k \\le -1$ there is no boundary and a large value
    (10000 mm) is returned.

    Returns:
        max_height (float): Maximum valid radial height [mm].
    """
    c, k = self._get_curvature_params()
    if k > -1:
        return torch.sqrt(1 / (k + 1) / (c**2)).item() - 0.001
    return 10e3

newton_initial_t

newton_initial_t(ray)

Seed Newton's method with the base-sphere intersection.

The base sphere is defined by the asphere curvature c with conic constant and polynomial coefficients set to zero. Its vertex is at the local origin and its center is at (0, 0, 1/c). Of the two analytic intersections, this method selects the point closest to the vertex, matching the sag branch represented by _sag.

A repeated root is handled explicitly. A flat base, a ray that misses the sphere, or a non-finite analytic result falls back safely to the vertex-plane approximation supplied by Surface.

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray bundle in local surface coordinates.	required

Returns:

Name	Type	Description
t	Tensor	Initial intersection parameter [mm], shape [...] matching the ray batch.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def newton_initial_t(self, ray):
    """Seed Newton's method with the base-sphere intersection.

    The base sphere is defined by the asphere curvature ``c`` with
    conic constant and polynomial coefficients set to zero. Its vertex is
    at the local origin and its center is at ``(0, 0, 1/c)``. Of the two
    analytic intersections, this method selects the point closest to the
    vertex, matching the sag branch represented by :meth:`_sag`.

    A repeated root is handled explicitly. A flat base, a ray that misses
    the sphere, or a non-finite analytic result falls back safely to the
    vertex-plane approximation supplied by :class:`Surface`.

    Args:
        ray (Ray): Input ray bundle in local surface coordinates.

    Returns:
        t (torch.Tensor): Initial intersection parameter [mm], shape [...]
            matching the ray batch.
    """
    t_plane = super().newton_initial_t(ray)
    c = self.c

    if c.abs() < EPSILON:
        return t_plane

    # Vertex-anchored base-sphere equation:
    #
    #   c * (x^2 + y^2 + z^2) - 2z = 0.
    #
    # This form avoids the R^2 subtraction in the center-anchored
    # equation and remains accurate for shallow curvatures in float32.
    od = torch.sum(ray.o * ray.d, dim=-1)
    dd = torch.sum(ray.d * ray.d, dim=-1)
    oo = torch.sum(ray.o * ray.o, dim=-1)

    a = c * dd
    b = 2.0 * (c * od - ray.d[..., 2])
    c_coeff = c * oo - 2.0 * ray.o[..., 2]
    discriminant = b * b - 4.0 * a * c_coeff
    sqrt_discriminant = torch.sqrt(torch.clamp(discriminant, min=0.0))

    # Stable quadratic roots. For a repeated root q is zero, so use the
    # standard repeated-root expression for the second candidate.
    q = torch.where(
        b >= 0,
        -(b + sqrt_discriminant) / 2.0,
        (sqrt_discriminant - b) / 2.0,
    )
    t1 = q / a
    repeated_root = -b / (2.0 * a)
    q_is_nonzero = q.abs() > EPSILON
    q_safe = torch.where(q_is_nonzero, q, torch.ones_like(q))
    t2 = torch.where(q_is_nonzero, c_coeff / q_safe, repeated_root)

    z1 = ray.o[..., 2] + t1 * ray.d[..., 2]
    z2 = ray.o[..., 2] + t2 * ray.d[..., 2]
    t_sphere = torch.where(z1.abs() < z2.abs(), t1, t2)

    has_real_finite_root = (discriminant >= 0) & torch.isfinite(t_sphere)
    return torch.where(has_real_finite_root, t_sphere, t_plane)

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001, 0.01, 0.0001], optim_mat=False)

Get optimizer parameters for different parameters.

The learning rate for each aspheric coefficient \(a_{2n}\) is scaled by \(1 / \max(r, 1)^{2n}\) so that the effective sag perturbation per Adam step is approximately constant (~lr_base mm) regardless of surface semi-diameter. Without this normalisation, gradients scale as \(O(r^{2n})\) and can reach \(10^5\) for camera-sized surfaces, causing NaN within a few dozen iterations.

Parameters:

Name	Type	Description	Default
lrs	list[float]	Learning rates for [d, c, k, ai]. Defaults to [1e-4, 1e-4, 1e-2, 1e-4].	[0.0001, 0.0001, 0.01, 0.0001]
optim_mat	bool	Whether to also optimize the material parameters. Defaults to False.	False

Returns:

Name	Type	Description
params	list[dict]	Parameter groups (each a dict with params and lr) ready to pass to a torch optimizer.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def get_optimizer_params(self, lrs=[1e-4, 1e-4, 1e-2, 1e-4], optim_mat=False):
    """Get optimizer parameters for different parameters.

    The learning rate for each aspheric coefficient $a_{2n}$ is scaled
    by $1 / \\max(r, 1)^{2n}$ so that the effective sag perturbation per
    Adam step is approximately constant (~lr_base mm) regardless of
    surface semi-diameter. Without this normalisation, gradients scale
    as $O(r^{2n})$ and can reach $10^5$ for camera-sized surfaces,
    causing NaN within a few dozen iterations.

    Args:
        lrs (list[float], optional): Learning rates for `[d, c, k, ai]`.
            Defaults to `[1e-4, 1e-4, 1e-2, 1e-4]`.
        optim_mat (bool, optional): Whether to also optimize the
            material parameters. Defaults to False.

    Returns:
        params (list[dict]): Parameter groups (each a dict with `params`
            and `lr`) ready to pass to a torch optimizer.
    """
    params = []

    # Optimize distance
    self.d_next.requires_grad_(True)
    params.append({"params": [self.d_next], "lr": lrs[0]})

    # Optimize curvature
    self.c.requires_grad_(True)
    params.append({"params": [self.c], "lr": lrs[1]})

    # Optimize conic constant
    self.k.requires_grad_(True)
    params.append({"params": [self.k], "lr": lrs[2]})

    # Optimize aspheric coefficients with r-normalised learning rates.
    # Gradient of sag w.r.t. a_{2n} scales as r^{2n}.  Dividing the lr
    # by r^{2n} keeps the effective sag change per step ≈ lr_base,
    # so every order contributes equally to surface shape evolution.
    if self.ai is not None:
        if self.ai_degree > 0:
            r_norm = max(self.r, 1.0)
            lr_base = lrs[3] if len(lrs) > 3 else 1e-4
            for i in range(self.ai_degree):
                p_name = f"ai{2 * (i + 2)}"
                p = getattr(self, p_name)
                p.requires_grad_(True)
                order = 2 * (i + 2)  # 4, 6, 8, 10, ...
                lr_ai = lr_base / r_norm**order
                params.append({"params": [p], "lr": lr_ai})

    # Optimize material parameters
    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

surf_dict

surf_dict()

Serialize the surface to a dict.

The aspheric coefficients are written into the ai list as [a4, a6, a8, ...]; when the legacy ai2 coefficient is present it is prepended so ai[0] = a2 and use_ai2 is set True.

Returns:

Name	Type	Description
surf_dict	dict	Serialized surface (keys type, r, roc, d, k, ai, use_ai2, mat2, plus informational (c)/(ai*)/(mat2_n)/(mat2_V) entries). Lengths in [mm], c in [1/mm].

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

def surf_dict(self):
    """Serialize the surface to a dict.

    The aspheric coefficients are written into the `ai` list as
    `[a4, a6, a8, ...]`; when the legacy `ai2` coefficient is present it
    is prepended so `ai[0] = a2` and `use_ai2` is set True.

    Returns:
        surf_dict (dict): Serialized surface (keys `type`, `r`, `roc`,
            `d`, `k`, `ai`, `use_ai2`, `mat2`, plus informational
            `(c)`/`(ai*)`/`(mat2_n)`/`(mat2_V)` entries). Lengths in [mm], `c` in [1/mm].
    """
    has_ai2 = self.ai2 is not None
    c_value = self.c.item()
    surf_dict = {
        "type": "Aspheric",
        "r": self.r,
        "(c)": c_value,
        "roc": 0.0 if c_value == 0.0 else 1.0 / c_value,
        "d_next": self.d_next.item(),
        "k": self.k.item(),
        "ai": [],
        "use_ai2": has_ai2,
        "pos_xy": [self.pos_x.item(), self.pos_y.item()],
        "vec_local": self.vec_local.tolist(),
        "is_square": self.is_square,
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }

    # Prepend a2 to ai list if present (ai2 key is informational;
    # deserialization reads ai[0] when use_ai2=True)
    if has_ai2:
        surf_dict["ai2"] = self.ai2.item()
        surf_dict["ai"].append(self.ai2.item())

    for i in range(self.ai_degree):
        order = i + 2
        coeff = getattr(self, f"ai{2 * order}")
        surf_dict[f"(ai{2 * order})"] = coeff.item()
        surf_dict["ai"].append(coeff.item())

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Return the Zemax (.zmx) text block for this surface.

Emits an EVENASPH surface. PARM 1 holds the legacy 2nd-order coefficient a2, and PARM 2 onward hold a4, a6, a8, ..., padded with zeros up to PARM 8 (a16).

Parameters:

Name	Type	Description	Default
surf_idx	int	Surface index in the Zemax file.	required
d_next	Tensor	Axial distance [mm] to the next surface.	required

Returns:

Name	Type	Description
zmx_str	str	Multi-line Zemax surface description.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py

    def zmx_str(self, surf_idx, d_next):
        """Return the Zemax (.zmx) text block for this surface.

        Emits an `EVENASPH` surface. PARM 1 holds the legacy 2nd-order
        coefficient `a2`, and PARM 2 onward hold `a4, a6, a8, ...`, padded
        with zeros up to PARM 8 (a16).

        Args:
            surf_idx (int): Surface index in the Zemax file.
            d_next (torch.Tensor): Axial distance [mm] to the next surface.

        Returns:
            zmx_str (str): Multi-line Zemax surface description.
        """
        assert self.c.item() != 0, (
            "Aperture surface is re-implemented in Aperture class."
        )
        assert self.ai is not None or self.k != 0, (
            "Spheric surface is re-implemented in Spheric class."
        )

        # Collect absolute ai values, PARM 1 = a2, PARM 2+ = a4, a6, ...
        abs_ai = [self.ai2.item() if self.ai2 is not None else 0.0]
        for i in range(self.ai_degree):
            abs_ai.append(getattr(self, f"ai{2 * (i + 2)}").item())

        # Pad with zeros for Zemax PARM format (needs 8 PARMs for a2–a16)
        while len(abs_ai) < 8:
            abs_ai.append(0.0)

        if self.mat2.get_name() == "air":
            zmx_str = f"""SURF {surf_idx}
    TYPE EVENASPH
    CURV {self.c.item()}
    DISZ {d_next.item()}
    DIAM {self.r} 1 0 0 1 ""
    CONI {self.k}
    PARM 1 {abs_ai[0]}
    PARM 2 {abs_ai[1]}
    PARM 3 {abs_ai[2]}
    PARM 4 {abs_ai[3]}
    PARM 5 {abs_ai[4]}
    PARM 6 {abs_ai[5]}
    PARM 7 {abs_ai[6]}
    PARM 8 {abs_ai[7]}
"""
        else:
            zmx_str = f"""SURF {surf_idx}
    TYPE EVENASPH
    CURV {self.c.item()}
    DISZ {d_next.item()}
    GLAS ___BLANK 1 0 {self.mat2.n} {self.mat2.V}
    DIAM {self.r} 1 0 0 1 ""
    CONI {self.k}
    PARM 1 {abs_ai[0]}
    PARM 2 {abs_ai[1]}
    PARM 3 {abs_ai[2]}
    PARM 4 {abs_ai[3]}
    PARM 5 {abs_ai[4]}
    PARM 6 {abs_ai[5]}
    PARM 7 {abs_ai[6]}
    PARM 8 {abs_ai[7]}
"""
        return zmx_str

deeplens.geometric_surface.Aperture

Aperture(
    r,
    d_next,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
)

Bases: Plane

Aperture stop surface.

A flat circular (or square) opening that blocks rays falling outside its clear aperture. Inherits the planar intersection logic from Plane and always sits in air (no refraction).

Attributes:

Name	Type	Description
r	float	Aperture radius (clear half-diameter) in [mm].
d_next	Tensor	Thickness to the next vertex in [mm].
is_square	bool	If True, the aperture is square instead of circular.
tolerancing	bool	Whether tolerancing perturbations are enabled.

Initialize an aperture surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius (clear half-diameter) in [mm].	required
d_next	float	Thickness to the next vertex in [mm].	required
pos_xy	list	Lateral (x, y) offset of the surface in [mm]. Defaults to [0.0, 0.0].	[0.0, 0.0]
vec_local	list	Local surface normal (z-axis) direction. Defaults to [0.0, 0.0, 1.0].	[0.0, 0.0, 1.0]
is_square	bool	If True, use a square aperture. Defaults to False.	False
device	str	Torch device for surface tensors. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def __init__(
    self,
    r,
    d_next,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize an aperture surface.

    Args:
        r (float): Aperture radius (clear half-diameter) in [mm].
        d_next (float): Thickness to the next vertex in [mm].
        pos_xy (list, optional): Lateral (x, y) offset of the surface in [mm]. Defaults to [0.0, 0.0].
        vec_local (list, optional): Local surface normal (z-axis) direction. Defaults to [0.0, 0.0, 1.0].
        is_square (bool, optional): If True, use a square aperture. Defaults to False.
        device (str, optional): Torch device for surface tensors. Defaults to "cpu".
    """
    Plane.__init__(
        self,
        r=r,
        d_next=d_next,
        mat2="air",
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.tolerancing = False
    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Construct an Aperture from a surface dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Requires "r" and "d_next"; optional keys "is_square", "pos_xy", "vec_local", and "device".	required

Returns:

Name	Type	Description
aperture	Aperture	The constructed aperture surface.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Construct an Aperture from a surface dictionary.

    Args:
        surf_dict (dict): Surface parameters. Requires "r" and "d_next"; optional
            keys "is_square", "pos_xy", "vec_local", and "device".

    Returns:
        aperture (Aperture): The constructed aperture surface.
    """
    return cls(
        r=surf_dict["r"],
        d_next=surf_dict["d_next"],
        is_square=surf_dict["is_square"] if "is_square" in surf_dict else False,
        pos_xy=surf_dict["pos_xy"] if "pos_xy" in surf_dict else [0.0, 0.0],
        vec_local=surf_dict["vec_local"]
        if "vec_local" in surf_dict
        else [0.0, 0.0, 1.0],
        device=surf_dict["device"] if "device" in surf_dict else "cpu",
    )

ray_reaction

ray_reaction(ray, n1=1.0, n2=1.0, refraction=False)

Trace a ray through the aperture.

Transforms the ray into local coordinates, intersects it with the aperture plane (rays outside the clear aperture are marked invalid), then transforms back to global coordinates. The aperture does not refract, so n1, n2, and refraction are ignored.

Parameters:

Name	Type	Description	Default
ray	Ray	Input ray batch in global coordinates.	required
n1	float	Refractive index before the surface (unused). Defaults to 1.0.	1.0
n2	float	Refractive index after the surface (unused). Defaults to 1.0.	1.0
refraction	bool	Ignored for an aperture. Defaults to False.	False

Returns:

Name	Type	Description
ray	Ray	Ray after intersection, in global coordinates, with is_valid updated.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def ray_reaction(self, ray, n1=1.0, n2=1.0, refraction=False):
    """Trace a ray through the aperture.

    Transforms the ray into local coordinates, intersects it with the
    aperture plane (rays outside the clear aperture are marked invalid),
    then transforms back to global coordinates. The aperture does not
    refract, so `n1`, `n2`, and `refraction` are ignored.

    Args:
        ray (Ray): Input ray batch in global coordinates.
        n1 (float, optional): Refractive index before the surface (unused). Defaults to 1.0.
        n2 (float, optional): Refractive index after the surface (unused). Defaults to 1.0.
        refraction (bool, optional): Ignored for an aperture. Defaults to False.

    Returns:
        ray (Ray): Ray after intersection, in global coordinates, with `is_valid` updated.
    """
    ray = self.to_local_coord(ray)
    ray = self.intersect(ray)
    ray = self.to_global_coord(ray)
    return ray

draw_widget

draw_widget(ax, color='orange', linestyle='solid', d=0.0)

Draw the aperture as wedge marks on a 2D cross-section plot.

Parameters:

Name	Type	Description	Default
ax	Axes	Axes to draw on (z-x cross-section).	required
color	str	Line color. Defaults to "orange".	'orange'
linestyle	str	Matplotlib line style. Defaults to "solid".	'solid'

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def draw_widget(self, ax, color="orange", linestyle="solid", d=0.0):
    """Draw the aperture as wedge marks on a 2D cross-section plot.

    Args:
        ax (matplotlib.axes.Axes): Axes to draw on (z-x cross-section).
        color (str, optional): Line color. Defaults to "orange".
        linestyle (str, optional): Matplotlib line style. Defaults to "solid".
    """
    d = float(d)
    aper_wedge_l = 0.05 * self.r  # [mm]
    aper_wedge_h = 0.15 * self.r  # [mm]

    # Parallel edges
    z = np.linspace(d - aper_wedge_l, d + aper_wedge_l, 3)
    x = -self.r * np.ones(3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)
    x = self.r * np.ones(3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)

    # Vertical edges
    z = d * np.ones(3)
    x = np.linspace(self.r, self.r + aper_wedge_h, 3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)
    x = np.linspace(-self.r - aper_wedge_h, -self.r, 3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)

draw_widget3D

draw_widget3D(ax, color='black', d=0.0)

Draw the aperture as an edge circle in a 3D plot.

Parameters:

Name	Type	Description	Default
ax	Axes3D	3D axes to draw on.	required
color	str	Line color. Defaults to "black".	'black'

Returns:

Name	Type	Description
line	list	The Line3D objects returned by ax.plot.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def draw_widget3D(self, ax, color="black", d=0.0):
    """Draw the aperture as an edge circle in a 3D plot.

    Args:
        ax (mpl_toolkits.mplot3d.axes3d.Axes3D): 3D axes to draw on.
        color (str, optional): Line color. Defaults to "black".

    Returns:
        line (list): The Line3D objects returned by `ax.plot`.
    """
    # Draw the edge circle
    theta = np.linspace(0, 2 * np.pi, 100)
    edge_x = self.r * np.cos(theta)
    edge_y = self.r * np.sin(theta)
    edge_z = np.full_like(edge_x, float(d))

    # Plot the edge circle
    line = ax.plot(edge_z, edge_x, edge_y, color=color, linewidth=1.5)

    return line

create_mesh

create_mesh(n_rings=32, n_arms=128, color=[0.0, 0.0, 0.0], d=0.0)

Create a triangulated surface mesh for the aperture.

Builds vertices, faces, and rim, then stores them on the surface.

Parameters:

Name	Type	Description	Default
n_rings	int	Number of concentric rings for sampling. Defaults to 32.	32
n_arms	int	Number of angular divisions. Defaults to 128.	128
color	list	RGB color of the mesh. Defaults to [0.0, 0.0, 0.0].	[0.0, 0.0, 0.0]

Returns:

Name	Type	Description
self	Aperture	The aperture with vertices, faces, rim, and mesh_color set (for chaining).

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def create_mesh(self, n_rings=32, n_arms=128, color=[0.0, 0.0, 0.0], d=0.0):
    """Create a triangulated surface mesh for the aperture.

    Builds vertices, faces, and rim, then stores them on the surface.

    Args:
        n_rings (int, optional): Number of concentric rings for sampling. Defaults to 32.
        n_arms (int, optional): Number of angular divisions. Defaults to 128.
        color (list, optional): RGB color of the mesh. Defaults to [0.0, 0.0, 0.0].

    Returns:
        self (Aperture): The aperture with `vertices`, `faces`, `rim`, and `mesh_color` set (for chaining).
    """
    self.vertices = self._create_vertices(n_rings, n_arms, d=d)
    self.faces = self._create_faces(n_rings, n_arms)
    self.rim = self._create_rim(n_rings, n_arms)
    self.mesh_color = color
    return self

get_optimizer_params

get_optimizer_params(lrs=[0.0001])

Enable gradients on the axial position and build optimizer param groups.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates; lrs[0] is applied to d. Defaults to [1e-4].	[0.0001]

Returns:

Name	Type	Description
params	list	List with one optimizer param group dict for d.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def get_optimizer_params(self, lrs=[1e-4]):
    """Enable gradients on the axial position and build optimizer param groups.

    Args:
        lrs (list, optional): Learning rates; `lrs[0]` is applied to `d`. Defaults to [1e-4].

    Returns:
        params (list): List with one optimizer param group dict for `d`.
    """
    self.d_next.requires_grad_(True)

    params = []
    params.append({"params": [self.d_next], "lr": lrs[0]})

    return params

surf_dict

surf_dict()

Serialize the aperture parameters to a dictionary.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters with keys "type", "r", "d_next", "mat2", and "is_square". Radius and position are in [mm].

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

def surf_dict(self):
    """Serialize the aperture parameters to a dictionary.

    Returns:
        surf_dict (dict): Surface parameters with keys "type", "r", "d_next",
            "mat2", and "is_square". Radius and position are in [mm].
    """
    surf_dict = {
        "type": "Aperture",
        "r": self.r,
        "d_next": self.d_next.item(),
        "mat2": "air",
        "pos_xy": [self.pos_x.item(), self.pos_y.item()],
        "vec_local": self.vec_local.tolist(),
        "is_square": self.is_square,
    }
    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Format the aperture as a Zemax (.zmx) STOP surface block.

Parameters:

Name	Type	Description	Default
surf_idx	int	Surface index in the Zemax file.	required
d_next	Tensor	Distance to the next surface in [mm].	required

Returns:

Name	Type	Description
zmx_str	str	Zemax surface definition string for this aperture.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py

    def zmx_str(self, surf_idx, d_next):
        """Format the aperture as a Zemax (.zmx) STOP surface block.

        Args:
            surf_idx (int): Surface index in the Zemax file.
            d_next (torch.Tensor): Distance to the next surface in [mm].

        Returns:
            zmx_str (str): Zemax surface definition string for this aperture.
        """
        zmx_str = f"""SURF {surf_idx}
    STOP
    TYPE STANDARD
    CURV 0.0
    DISZ {d_next.item()}
    DIAM {self.r} 1 0 0 1 ""
"""
        return zmx_str

Phase Surfaces

Phase surfaces model flat diffractive optical elements (DOEs) and metasurfaces via a wavelength-scaled phase profile. All classes inherit from Phase and implement phi() (phase map) and dphi_dxy() (phase gradient) for the generalized Snell's law deflection in diffract().

Note: Phase.diffract() treats the phase profile as wavelength-independent. Only the λ scaling in the generalized Snell's law and the OPL accumulation vary with wavelength. For a physical DOE whose phase profile changes with wavelength via the height–index relation (n(λ)−1)·h, use DiffractiveSurface instead.

deeplens.phase_surface.Phase

Phase(
    r,
    d_next,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: DeepObj

Base phase profile for diffractive surfaces (metasurface or DOE).

Represents a flat (zero-sag) substrate carrying a phase pattern \(\phi(x, y)\), using sequential geometry. It owns only d_next, the thickness to the next vertex; the containing GeoLens derives its absolute position. Provides the common ray-tracing machinery (intersection, refraction, generalized-Snell diffraction, local/global transforms); the phase profile \(\phi\) and its gradient are defined by subclasses.

Attributes:

Name	Type	Description
vec_global	Tensor	Global axis direction \([0, 0, 1]\), shape [3].
d_next	Tensor	Thickness to the next vertex in [mm], scalar.
pos_x	Tensor	Surface x-offset in [mm], scalar.
pos_y	Tensor	Surface y-offset in [mm], scalar.
vec_local	Tensor	Unit surface normal in global coordinates, shape [3].
mat2	Material	Material on the exit side of the surface.
r	float	Surface radius / half-aperture in [mm].
is_square	bool	If True the aperture is a square of side \(r\sqrt{2}\); otherwise a circle of radius \(r\).
w	float	Square aperture width \(r\sqrt{2}\) in [mm].
h	float	Square aperture height \(r\sqrt{2}\) in [mm].
diffraction_order	int	Diffraction order \(m\) used in the generalized Snell's law. Defaults to 1.
norm_radii	float	Radius in [mm] used to normalize coordinates for the phase polynomial. Defaults to r.
device	str or device	Device holding the tensor state.

Reference

[1] https://support.zemax.com/hc/en-us/articles/1500005489061-How-diffractive-surfaces-are-modeled-in-OpticStudio [2] https://optics.ansys.com/hc/en-us/articles/360042097313-Small-Scale-Metalens-Field-Propagation [3] https://optics.ansys.com/hc/en-us/articles/18254409091987-Large-Scale-Metalens-Ray-Propagation

Initialize a flat phase substrate.

Parameters:

Name	Type	Description	Default
r	float	Surface radius / half-aperture in [mm].	required
d_next	float	Axial thickness to the next vertex in [mm].	required
norm_radii	float or None	Radius in [mm] used to normalize coordinates for the phase polynomial. Defaults to None, which uses r.	None
mat2	str	Material on the exit side of the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral (x, y) offset of the surface center in [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Surface normal direction (not necessarily normalized) in global coordinates. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True the aperture is a square of side \(r\sqrt{2}\); otherwise a circle of radius \(r\). Defaults to True.	True
device	str	Device for the tensor state. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def __init__(
    self,
    r,
    d_next,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a flat phase substrate.

    Args:
        r (float): Surface radius / half-aperture in [mm].
        d_next (float): Axial thickness to the next vertex in [mm].
        norm_radii (float or None, optional): Radius in [mm] used to normalize
            coordinates for the phase polynomial. Defaults to None, which uses `r`.
        mat2 (str, optional): Material on the exit side of the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral (x, y) offset of the surface center in [mm].
            Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Surface normal direction (not necessarily
            normalized) in global coordinates. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True the aperture is a square of side
            $r\\sqrt{2}$; otherwise a circle of radius $r$. Defaults to True.
        device (str, optional): Device for the tensor state. Defaults to "cpu".
    """
    super().__init__()

    self.d_next = (
        d_next.detach().clone()
        if torch.is_tensor(d_next)
        else torch.tensor(d_next, dtype=torch.get_default_dtype())
    )
    if not self.d_next.is_floating_point():
        self.d_next = self.d_next.to(torch.get_default_dtype())

    state_dtype = self.d_next.dtype
    state_device = self.d_next.device
    self.vec_global = torch.tensor(
        [0.0, 0.0, 1.0], dtype=state_dtype, device=state_device
    )
    self.pos_x = torch.as_tensor(pos_xy[0], dtype=state_dtype, device=state_device)
    self.pos_y = torch.as_tensor(pos_xy[1], dtype=state_dtype, device=state_device)
    self.vec_local = F.normalize(
        torch.as_tensor(vec_local, dtype=state_dtype, device=state_device),
        p=2,
        dim=-1,
    )

    self.mat2 = Material(mat2)
    self.r = float(r)
    self.is_square = is_square
    if is_square:
        self.w = self.r * float(np.sqrt(2))
        self.h = self.r * float(np.sqrt(2))

    self.device = device if device is not None else torch.device("cpu")
    self.to(self.device)
    self._cache_rotation_matrices()

    # Phase surfaces retain rectangular extents even for circular apertures
    # because phase-map visualization and fabrication helpers use them.
    self.w = self.r * float(np.sqrt(2))
    self.h = self.r * float(np.sqrt(2))

    self.diffraction_order = 1
    self.norm_radii = self.r if norm_radii is None else norm_radii

phi

phi(x, y)

Reference phase map at design wavelength. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def phi(self, x, y):
    """Reference phase map at design wavelength. Must be implemented by subclasses."""
    raise NotImplementedError("phi() must be implemented by subclasses")

dphi_dxy

dphi_dxy(x, y)

Calculate phase derivatives. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def dphi_dxy(self, x, y):
    """Calculate phase derivatives. Must be implemented by subclasses."""
    raise NotImplementedError("dphi_dxy() must be implemented by subclasses")

ray_reaction

ray_reaction(ray, n1, n2)

Trace a ray through the phase surface.

Transforms the ray to local coordinates, intersects it with the plane, applies refraction then diffraction, and transforms back to global coordinates.

Parameters:

Name	Type	Description	Default
ray	Ray	Incident ray in global coordinates.	required
n1	float	Refractive index of the medium before the surface.	required
n2	float	Refractive index of the medium after the surface.	required

Returns:

Name	Type	Description
ray	Ray	Updated ray in global coordinates.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def ray_reaction(self, ray, n1, n2):
    """Trace a ray through the phase surface.

    Transforms the ray to local coordinates, intersects it with the plane,
    applies refraction then diffraction, and transforms back to global
    coordinates.

    Args:
        ray (Ray): Incident ray in global coordinates.
        n1 (float): Refractive index of the medium before the surface.
        n2 (float): Refractive index of the medium after the surface.

    Returns:
        ray (Ray): Updated ray in global coordinates.
    """
    ray = self.to_local_coord(ray)
    ray = self.intersect(ray, n1)
    ray = self.refract(ray, n1 / n2)
    ray = self.diffract(ray, n2=n2)
    ray = self.to_global_coord(ray)
    return ray

intersect

intersect(ray, n=1.0)

Solve ray-plane intersection in local coordinates and update the ray.

Advances each ray to the \(z = 0\) plane, marks rays falling outside the aperture as invalid, and (for coherent rays) accumulates optical path length. Rays nearly parallel to the plane are guarded against division by a near-zero z-direction.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray in local coordinates.	required
n	float	Refractive index of the medium the ray travels through, used for OPL accumulation. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
ray	Ray	Ray advanced to the surface plane with updated o, is_valid, and opl.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def intersect(self, ray, n=1.0):
    """Solve ray-plane intersection in local coordinates and update the ray.

    Advances each ray to the $z = 0$ plane, marks rays falling outside the
    aperture as invalid, and (for coherent rays) accumulates optical path
    length. Rays nearly parallel to the plane are guarded against division
    by a near-zero z-direction.

    Args:
        ray (Ray): Ray in local coordinates.
        n (float, optional): Refractive index of the medium the ray travels
            through, used for OPL accumulation. Defaults to 1.0.

    Returns:
        ray (Ray): Ray advanced to the surface plane with updated `o`,
            `is_valid`, and `opl`.
    """
    # Solve intersection. Guard against a near-zero z-direction (rays
    # parallel to the plane) before dividing, matching ray.py prop_to.
    dz = ray.d[..., 2]
    dz = torch.where(dz.abs() < EPSILON, torch.full_like(dz, EPSILON), dz)
    t = (0.0 - ray.o[..., 2]) / dz
    new_o = ray.o + t.unsqueeze(-1) * ray.d
    if self.is_square:
        valid = (
            (torch.abs(new_o[..., 0]) < self.w / 2)
            & (torch.abs(new_o[..., 1]) < self.h / 2)
            & (ray.is_valid > 0)
        )
    else:
        valid = (new_o[..., 0] ** 2 + new_o[..., 1] ** 2 < self.r**2) & (
            ray.is_valid > 0
        )

    # Update rays
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    ray.o = torch.where(valid.unsqueeze(-1), new_o, ray.o)
    ray.is_valid = ray.is_valid * valid

    if ray.is_coherent:
        ray.opl = torch.where(
            valid.unsqueeze(-1), ray.opl + n * t.unsqueeze(-1), ray.opl
        )

    return ray

diffract

diffract(ray, n2=1.0)

Apply phase-surface diffraction to a ray.

Two effects are applied:

1. The phase \(\phi\) (in radians) adds to the optical path length as \(\phi \cdot \lambda / (2\pi)\), where \(\lambda\) is converted from [µm] to [mm] internally.
2. The phase gradient bends the ray via the generalized Snell's law \(n_2 \sin\theta_2 = n_1 \sin\theta_1 + m\,\lambda / (2\pi)\,\partial\phi/\partial x\). Since standard refraction is already applied, the remaining deflection added to the unit direction is \(\Delta l = m\,\lambda / (2\pi n_2)\,\partial\phi/\partial x\). The deflection sign is flipped for backward-propagating rays.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray with position, direction, and wavelength in [µm].	required
n2	float	Refractive index of the medium after the surface; the deflection scales as \(1/n_2\). Defaults to 1.0.	1.0

Returns:

Name	Type	Description
ray	Ray	Ray with updated direction d and (for coherent rays) opl.

Note

Material dispersion is not modelled here. The phase profile \(\phi(x, y)\) is treated as wavelength-independent; only the \(\lambda\) scaling in the generalized Snell's law and the OPL accumulation vary with wavelength. For a physical DOE whose phase profile itself changes with wavelength (via \((n(\lambda) - 1)\,h\)), use DiffractiveSurface instead.

Reference

[1] https://support.zemax.com/hc/en-us/articles/1500005489061-How-diffractive-surfaces-are-modeled-in-OpticStudio [2] Light propagation with phase discontinuities: generalized laws of reflection and refraction. Science 2011.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def diffract(self, ray, n2=1.0):
    """Apply phase-surface diffraction to a ray.

    Two effects are applied:

    1. The phase $\\phi$ (in radians) adds to the optical path length as
       $\\phi \\cdot \\lambda / (2\\pi)$, where $\\lambda$ is converted from [µm]
       to [mm] internally.
    2. The phase gradient bends the ray via the generalized Snell's law
       $n_2 \\sin\\theta_2 = n_1 \\sin\\theta_1 + m\\,\\lambda / (2\\pi)\\,\\partial\\phi/\\partial x$.
       Since standard refraction is already applied, the remaining deflection
       added to the unit direction is $\\Delta l = m\\,\\lambda / (2\\pi n_2)\\,\\partial\\phi/\\partial x$.
       The deflection sign is flipped for backward-propagating rays.

    Args:
        ray (Ray): Ray with position, direction, and wavelength in [µm].
        n2 (float, optional): Refractive index of the medium after the surface;
            the deflection scales as $1/n_2$. Defaults to 1.0.

    Returns:
        ray (Ray): Ray with updated direction `d` and (for coherent rays) `opl`.

    Note:
        Material dispersion is not modelled here. The phase profile $\\phi(x, y)$
        is treated as wavelength-independent; only the $\\lambda$ scaling in the
        generalized Snell's law and the OPL accumulation vary with wavelength.
        For a physical DOE whose phase profile itself changes with wavelength
        (via $(n(\\lambda) - 1)\\,h$), use `DiffractiveSurface` instead.

    Reference:
        [1] https://support.zemax.com/hc/en-us/articles/1500005489061-How-diffractive-surfaces-are-modeled-in-OpticStudio
        [2] Light propagation with phase discontinuities: generalized laws of reflection and refraction. Science 2011.
    """
    forward = (ray.d * ray.is_valid.unsqueeze(-1))[..., 2].sum() > 0
    valid = ray.is_valid > 0

    # Step 1: DOE phase modulation
    if ray.is_coherent:
        phi = self.phi(ray.o[..., 0], ray.o[..., 1])
        new_opl = ray.opl + phi.unsqueeze(-1) * (ray.wvln * 1e-3) / (2 * torch.pi)
        ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    # Step 2: bend rays via generalized Snell's law
    # n₂·l₂ = n₁·l₁ + M·λ/(2π)·dφ/dx
    # After refraction: l₂ = l_refracted + M·λ/(2π·n₂)·dφ/dx
    dphidx, dphidy = self.dphi_dxy(ray.o[..., 0], ray.o[..., 1])

    wvln_mm = ray.wvln * 1e-3
    order = self.diffraction_order
    phase_deflection_scale = wvln_mm / (2 * torch.pi * n2)
    if forward:
        new_d_x = ray.d[..., 0] + phase_deflection_scale * dphidx * order
        new_d_y = ray.d[..., 1] + phase_deflection_scale * dphidy * order
    else:
        new_d_x = ray.d[..., 0] - phase_deflection_scale * dphidx * order
        new_d_y = ray.d[..., 1] - phase_deflection_scale * dphidy * order

    new_d = torch.stack([new_d_x, new_d_y, ray.d[..., 2]], dim=-1)
    new_d = F.normalize(new_d, p=2, dim=-1)
    ray.d = torch.where(valid.unsqueeze(-1), new_d, ray.d)

    return ray

refract

refract(ray, eta)

Refract a ray with vector Snell's law in local coordinates.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def refract(self, ray, eta):
    """Refract a ray with vector Snell's law in local coordinates."""
    normal_vec = self.normal_vec(ray)
    dot_product = (-normal_vec * ray.d).sum(-1).unsqueeze(-1)
    k = 1 - eta**2 * (1 - dot_product**2)

    valid = (k >= 0).squeeze(-1) & (ray.is_valid > 0)
    k = k * valid.unsqueeze(-1)

    new_d = eta * ray.d + (eta * dot_product - torch.sqrt(k + EPSILON)) * normal_vec
    ray.d = torch.where(valid.unsqueeze(-1), new_d, ray.d)
    ray.is_valid = ray.is_valid * valid
    return ray

to_local_coord

to_local_coord(ray)

Transform a ray from the surface reference frame to local coordinates.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def to_local_coord(self, ray):
    """Transform a ray from the surface reference frame to local coordinates."""
    offset = torch.stack(
        [self.pos_x, self.pos_y, torch.zeros_like(self.pos_x)]
    ).expand_as(ray.o)
    ray.o = ray.o - offset

    if self._R_to_local is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_local)
        ray.d = self._apply_rotation(ray.d, self._R_to_local)
        ray.d = F.normalize(ray.d, p=2, dim=-1)
    return ray

to_global_coord

to_global_coord(ray)

Transform a ray from local coordinates to the surface reference frame.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def to_global_coord(self, ray):
    """Transform a ray from local coordinates to the surface reference frame."""
    if self._R_to_global is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_global)
        ray.d = self._apply_rotation(ray.d, self._R_to_global)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    offset = torch.stack(
        [self.pos_x, self.pos_y, torch.zeros_like(self.pos_x)]
    ).expand_as(ray.o)
    ray.o = ray.o + offset
    return ray

normal_vec

normal_vec(ray)

Calculate the surface normal vector at intersection points.

The normal points from the surface toward the side the light comes from (i.e. it is flipped to oppose the ray's z-direction).

Parameters:

Name	Type	Description	Default
ray	Ray	Ray providing the propagation direction.	required

Returns:

Name	Type	Description
normal_vec	Tensor	Unit normal vectors, same shape as ray.d.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def normal_vec(self, ray):
    """Calculate the surface normal vector at intersection points.

    The normal points from the surface toward the side the light comes from
    (i.e. it is flipped to oppose the ray's z-direction).

    Args:
        ray (Ray): Ray providing the propagation direction.

    Returns:
        normal_vec (torch.Tensor): Unit normal vectors, same shape as `ray.d`.
    """
    normal_vec = torch.zeros_like(ray.d)
    normal_vec[..., 2] = -1
    is_forward = ray.d[..., 2].unsqueeze(-1) > 0
    normal_vec = torch.where(is_forward, normal_vec, -normal_vec)
    return normal_vec

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.01], optim_mat=False)

Generate optimizer parameters. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def get_optimizer_params(self, lrs=[1e-4, 1e-2], optim_mat=False):
    """Generate optimizer parameters. Must be implemented by subclasses."""
    raise NotImplementedError(
        "get_optimizer_params() must be implemented by subclasses"
    )

get_optimizer

get_optimizer(lrs)

Build an Adam optimizer over the surface's learnable parameters.

Parameters:

Name	Type	Description	Default
lrs	list or float	Learning rate(s) for the parameter groups. A single float is wrapped into a one-element list.	required

Returns:

Name	Type	Description
optimizer	Adam	Adam optimizer over the parameters returned by get_optimizer_params.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def get_optimizer(self, lrs):
    """Build an Adam optimizer over the surface's learnable parameters.

    Args:
        lrs (list or float): Learning rate(s) for the parameter groups. A
            single float is wrapped into a one-element list.

    Returns:
        optimizer (torch.optim.Adam): Adam optimizer over the parameters
            returned by `get_optimizer_params`.
    """
    if isinstance(lrs, float):
        lrs = [lrs]
    params = self.get_optimizer_params(lrs)
    optimizer = torch.optim.Adam(params)
    return optimizer

update_r

update_r(r)

Update surface radius / half-aperture and the square aperture extents.

A flat phase surface has no geometric height constraint, and because the polynomial is normalized by a fixed norm_radii, phase coefficients do not need rescaling.

Parameters:

Name	Type	Description	Default
r	float	New surface radius / half-aperture in [mm].	required

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def update_r(self, r):
    """Update surface radius / half-aperture and the square aperture extents.

    A flat phase surface has no geometric height constraint, and because the
    polynomial is normalized by a fixed `norm_radii`, phase coefficients do
    not need rescaling.

    Args:
        r (float): New surface radius / half-aperture in [mm].
    """
    self.r = float(r)
    self.w = self.r * float(np.sqrt(2))
    self.h = self.r * float(np.sqrt(2))

phase2height_map

phase2height_map(design_wvln, refractive_idx=1.5, res=512)

Convert the phase map to a physical height map for DOE fabrication.

Derived from the phase-height relation of a transmissive DOE in air, \(\phi = (2\pi/\lambda)(n - 1)h\), giving \(h = \phi\lambda / (2\pi(n - 1))\).

Parameters:

Name	Type	Description	Default
design_wvln	float	Design wavelength in [µm].	required
refractive_idx	float	Refractive index of the DOE material at design_wvln. Defaults to 1.5.	1.5
res	int	Pixel resolution of the returned square height map. Defaults to 512.	512

Returns:

Name	Type	Description
height_map	Tensor	Height map of shape [res, res] in the same units as design_wvln ([µm]).

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def phase2height_map(self, design_wvln, refractive_idx=1.5, res=512):
    """Convert the phase map to a physical height map for DOE fabrication.

    Derived from the phase-height relation of a transmissive DOE in air,
    $\\phi = (2\\pi/\\lambda)(n - 1)h$, giving $h = \\phi\\lambda / (2\\pi(n - 1))$.

    Args:
        design_wvln (float): Design wavelength in [µm].
        refractive_idx (float, optional): Refractive index of the DOE material
            at `design_wvln`. Defaults to 1.5.
        res (int, optional): Pixel resolution of the returned square height map.
            Defaults to 512.

    Returns:
        height_map (torch.Tensor): Height map of shape [res, res] in the same
            units as `design_wvln` ([µm]).
    """
    x, y = torch.meshgrid(
        torch.linspace(-self.w / 2, self.w / 2, res),
        torch.linspace(self.h / 2, -self.h / 2, res),
        indexing="xy",
    )
    x, y = x.to(self.device), y.to(self.device)
    phi = self.phi(x, y)  # [0, 2π], shape [res, res]
    height_map = phi * design_wvln / (2 * torch.pi * (refractive_idx - 1))
    return height_map

draw_r

draw_r()

Effective drawing radius for 2D layout drawing.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def draw_r(self):
    """Effective drawing radius for 2D layout drawing."""
    return self.r

surface_with_offset

surface_with_offset(*args, d=0.0, **kwargs)

Return the caller-provided vertex position for layout drawing.

The phase surface is flat and does not own an absolute position.

Returns:

Name	Type	Description
d	Tensor	Axial plane position in [mm], scalar.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def surface_with_offset(self, *args, d=0.0, **kwargs):
    """Return the caller-provided vertex position for layout drawing.

    The phase surface is flat and does not own an absolute position.

    Returns:
        d (torch.Tensor): Axial plane position in [mm], scalar.
    """
    if torch.is_tensor(d):
        return d
    return torch.tensor(float(d), device=self.device)

draw_phase_map

draw_phase_map(save_name='./DOE_phase_map.png')

Draw the phase map (clipped to \([0, 2\pi]\)) and save it to a file.

Parameters:

Name	Type	Description	Default
save_name	str	Output image path. Defaults to "./DOE_phase_map.png".	'./DOE_phase_map.png'

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def draw_phase_map(self, save_name="./DOE_phase_map.png"):
    """Draw the phase map (clipped to $[0, 2\\pi]$) and save it to a file.

    Args:
        save_name (str, optional): Output image path. Defaults to
            "./DOE_phase_map.png".
    """
    x, y = torch.meshgrid(
        torch.linspace(-self.w / 2, self.w / 2, 2000),
        torch.linspace(self.h / 2, -self.h / 2, 2000),
        indexing="xy",
    )
    x, y = x.to(self.device), y.to(self.device)
    pmap = self.phi(x, y)

    fig, ax = plt.subplots(1, 1, figsize=(6, 5))
    im = ax.imshow(pmap.cpu().numpy(), vmin=0, vmax=2 * torch.pi)
    ax.set_title("Phase map 0.55um", fontsize=10)
    ax.grid(False)
    fig.colorbar(im)
    fig.savefig(save_name, dpi=600, bbox_inches="tight")
    plt.close(fig)

draw_widget

draw_widget(ax, color='black', linestyle='-', d=0.0)

Draw the DOE as a sawtooth (blazed) profile on a 2D layout axis.

Parameters:

Name	Type	Description	Default
ax	Axes	Axis to draw on.	required
color	str	Accepted for API consistency but ignored; the profile is always drawn in orange. Defaults to "black".	'black'
linestyle	str	Matplotlib line style for the profile. Defaults to "-".	'-'

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def draw_widget(self, ax, color="black", linestyle="-", d=0.0):
    """Draw the DOE as a sawtooth (blazed) profile on a 2D layout axis.

    Args:
        ax (matplotlib.axes.Axes): Axis to draw on.
        color (str, optional): Accepted for API consistency but ignored; the
            profile is always drawn in orange. Defaults to "black".
        linestyle (str, optional): Matplotlib line style for the profile.
            Defaults to "-".
    """
    max_offset = self.r / 100
    d = float(d)

    # Draw DOE
    roc = self.r * 2
    x = np.linspace(-self.r, self.r, 128)
    y = np.zeros_like(x)
    r = np.sqrt(x**2 + y**2 + EPSILON)
    sag = roc * (1 - np.sqrt(1 - r**2 / roc**2))
    sag = max_offset - np.fmod(sag, max_offset)
    ax.plot(d + sag, x, color="orange", linestyle=linestyle, linewidth=0.75)

save_ckpt

save_ckpt(save_path='./doe.pth')

Save DOE parameters. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def save_ckpt(self, save_path="./doe.pth"):
    """Save DOE parameters. Must be implemented by subclasses."""
    raise NotImplementedError("save_ckpt() must be implemented by subclasses")

load_ckpt

load_ckpt(load_path='./doe.pth')

Load DOE parameters. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def load_ckpt(self, load_path="./doe.pth"):
    """Load DOE parameters. Must be implemented by subclasses."""
    raise NotImplementedError("load_ckpt() must be implemented by subclasses")

surf_dict

surf_dict()

Return surface parameters. Must be implemented by subclasses.

Source code in deeplens-src/deeplens/phase_surface/base_phase.py

def surf_dict(self):
    """Return surface parameters. Must be implemented by subclasses."""
    raise NotImplementedError("surf_dict() must be implemented by subclasses")

deeplens.phase_surface.Binary2Phase

Binary2Phase(
    r,
    d_next,
    order2=0.0,
    order4=0.0,
    order6=0.0,
    order8=0.0,
    order10=0.0,
    order12=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Zemax BINARY_2 phase profile on a flat substrate.

Parameterizes the diffractive phase as an even radial polynomial in the normalized radius \(\rho = r / r_\text{norm}\):

\[\phi(\rho) = \sum_{i=1}^{6} a_{2i}\,\rho^{2i}\]

where the coefficients order2..order12 are stored in radians [rad]. The phase is evaluated with Horner's method and wrapped to \([0, 2\pi)\).

Attributes:

Name	Type	Description
order2	Tensor	Coefficient of \(\rho^2\), scalar [rad].
order4	Tensor	Coefficient of \(\rho^4\), scalar [rad].
order6	Tensor	Coefficient of \(\rho^6\), scalar [rad].
order8	Tensor	Coefficient of \(\rho^8\), scalar [rad].
order10	Tensor	Coefficient of \(\rho^{10}\), scalar [rad].
order12	Tensor	Coefficient of \(\rho^{12}\), scalar [rad].
param_model	str	Parameterization tag, always "binary2".
norm_radii	float	Normalization radius \(r_\text{norm}\) [mm].

Initialize a Binary2 phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius (half-diameter) [mm].	required
d	float	Axial position of the surface in global coordinates [mm].	required
order2	float	Coefficient of \(\rho^2\) [rad]. Defaults to 0.0.	0.0
order4	float	Coefficient of \(\rho^4\) [rad]. Defaults to 0.0.	0.0
order6	float	Coefficient of \(\rho^6\) [rad]. Defaults to 0.0.	0.0
order8	float	Coefficient of \(\rho^8\) [rad]. Defaults to 0.0.	0.0
order10	float	Coefficient of \(\rho^{10}\) [rad]. Defaults to 0.0.	0.0
order12	float	Coefficient of \(\rho^{12}\) [rad]. Defaults to 0.0.	0.0
norm_radii	float or None	Normalization radius for the polynomial [mm]. Defaults to None, in which case r is used.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral (x, y) offset of the surface center [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True, use a square aperture; otherwise circular. Defaults to True.	True
device	str	Torch device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def __init__(
    self,
    r,
    d_next,
    order2=0.0,
    order4=0.0,
    order6=0.0,
    order8=0.0,
    order10=0.0,
    order12=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a Binary2 phase surface.

    Args:
        r (float): Aperture radius (half-diameter) [mm].
        d (float): Axial position of the surface in global coordinates [mm].
        order2 (float, optional): Coefficient of $\\rho^2$ [rad]. Defaults to 0.0.
        order4 (float, optional): Coefficient of $\\rho^4$ [rad]. Defaults to 0.0.
        order6 (float, optional): Coefficient of $\\rho^6$ [rad]. Defaults to 0.0.
        order8 (float, optional): Coefficient of $\\rho^8$ [rad]. Defaults to 0.0.
        order10 (float, optional): Coefficient of $\\rho^{10}$ [rad]. Defaults to 0.0.
        order12 (float, optional): Coefficient of $\\rho^{12}$ [rad]. Defaults to 0.0.
        norm_radii (float or None, optional): Normalization radius for the polynomial [mm].
            Defaults to None, in which case `r` is used.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral (x, y) offset of the surface center [mm]. Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True, use a square aperture; otherwise circular. Defaults to True.
        device (str, optional): Torch device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    # Initialize polynomial coefficients
    self.order2 = torch.tensor(order2)
    self.order4 = torch.tensor(order4)
    self.order6 = torch.tensor(order6)
    self.order8 = torch.tensor(order8)
    self.order10 = torch.tensor(order10)
    self.order12 = torch.tensor(order12)

    self.param_model = "binary2"
    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Construct a Binary2 phase surface from a parameter dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Requires keys "r" and "d_next"; optionally "order2".."order12", "norm_radii", "mat2", and "is_square".	required

Returns:

Name	Type	Description
obj	Binary2Phase	The constructed phase surface.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Construct a Binary2 phase surface from a parameter dictionary.

    Args:
        surf_dict (dict): Surface parameters. Requires keys "r" and "d_next";
            optionally "order2".."order12", "norm_radii", "mat2", and
            "is_square".

    Returns:
        obj (Binary2Phase): The constructed phase surface.
    """
    mat2 = surf_dict.get("mat2", "air")
    norm_radii = surf_dict.get("norm_radii", None)
    is_square = surf_dict.get("is_square", True)
    obj = cls(
        surf_dict["r"],
        surf_dict["d_next"],
        surf_dict.get("order2", 0.0),
        surf_dict.get("order4", 0.0),
        surf_dict.get("order6", 0.0),
        surf_dict.get("order8", 0.0),
        surf_dict.get("order10", 0.0),
        surf_dict.get("order12", 0.0),
        norm_radii,
        mat2,
        is_square=is_square,
    )
    return obj

phi

phi(x, y)

Compute the reference phase at the design wavelength.

Evaluates the even radial polynomial in normalized radius via Horner's method and wraps the result to \([0, 2\pi)\) with torch.remainder.

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase values [rad] wrapped to \([0, 2\pi)\), same shape as x.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def phi(self, x, y):
    """Compute the reference phase at the design wavelength.

    Evaluates the even radial polynomial in normalized radius via Horner's
    method and wraps the result to $[0, 2\\pi)$ with `torch.remainder`.

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], same shape as `x`.

    Returns:
        phi (torch.Tensor): Phase values [rad] wrapped to $[0, 2\\pi)$, same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    r2 = x_norm * x_norm + y_norm * y_norm + EPSILON

    # Horner's method: r2*(o2 + r2*(o4 + r2*(o6 + r2*(o8 + r2*(o10 + r2*o12)))))
    phi = r2 * (
        self.order2
        + r2
        * (
            self.order4
            + r2
            * (
                self.order6
                + r2 * (self.order8 + r2 * (self.order10 + r2 * self.order12))
            )
        )
    )

    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the lateral phase gradient at the given points.

Differentiates the (unwrapped) phase polynomial and applies the chain rule through the normalized radius.

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Partial derivative \(\partial\phi/\partial x\) [rad/mm], same shape as x.
dphidy	Tensor	Partial derivative \(\partial\phi/\partial y\) [rad/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def dphi_dxy(self, x, y):
    """Compute the lateral phase gradient at the given points.

    Differentiates the (unwrapped) phase polynomial and applies the chain
    rule through the normalized radius.

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], same shape as `x`.

    Returns:
        dphidx (torch.Tensor): Partial derivative $\\partial\\phi/\\partial x$ [rad/mm], same shape as `x`.
        dphidy (torch.Tensor): Partial derivative $\\partial\\phi/\\partial y$ [rad/mm], same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    r2 = x_norm * x_norm + y_norm * y_norm + EPSILON

    # d/dr2 of polynomial, then chain rule: dphi/dx = dphi/dr2 * 2*x_norm / norm_radii
    # Horner's: o2 + r2*(2*o4 + r2*(3*o6 + r2*(4*o8 + r2*(5*o10 + r2*6*o12))))
    dphidr2 = self.order2 + r2 * (
        2 * self.order4
        + r2
        * (
            3 * self.order6
            + r2
            * (4 * self.order8 + r2 * (5 * self.order10 + r2 * 6 * self.order12))
        )
    )
    dphidx = dphidr2 * 2 * x_norm / self.norm_radii
    dphidy = dphidr2 * 2 * y_norm / self.norm_radii

    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.01], optim_mat=False)

Build optimizer parameter groups for the phase surface.

Enables gradients on d_next and the six polynomial coefficients, grouping d_next with the first learning rate and all coefficients with the second.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates [lr_position, lr_coeffs]. Defaults to [1e-4, 1e-2].	[0.0001, 0.01]
optim_mat	bool	Must be False; materials are not optimized for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List of parameter-group dicts for a torch optimizer.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def get_optimizer_params(self, lrs=[1e-4, 1e-2], optim_mat=False):
    """Build optimizer parameter groups for the phase surface.

    Enables gradients on `d_next` and the six polynomial coefficients,
    grouping `d_next` with the first learning rate and all
    coefficients with the second.

    Args:
        lrs (list, optional): Learning rates ``[lr_position, lr_coeffs]``. Defaults to [1e-4, 1e-2].
        optim_mat (bool, optional): Must be False; materials are not optimized for phase surfaces. Defaults to False.

    Returns:
        params (list): List of parameter-group dicts for a torch optimizer.

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []

    # Optimize sequential thickness.
    self.d_next.requires_grad = True
    params.append({"params": [self.d_next], "lr": lrs[0]})

    # Optimize polynomial coefficients
    self.order2.requires_grad = True
    self.order4.requires_grad = True
    self.order6.requires_grad = True
    self.order8.requires_grad = True
    self.order10.requires_grad = True
    self.order12.requires_grad = True
    params.append({"params": [self.order2], "lr": lrs[1]})
    params.append({"params": [self.order4], "lr": lrs[1]})
    params.append({"params": [self.order6], "lr": lrs[1]})
    params.append({"params": [self.order8], "lr": lrs[1]})
    params.append({"params": [self.order10], "lr": lrs[1]})
    params.append({"params": [self.order12], "lr": lrs[1]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./binary2_doe.pth')

Save the Binary2 phase coefficients to disk.

Parameters:

Name	Type	Description	Default
save_path	str	Output checkpoint path. Defaults to "./binary2_doe.pth".	'./binary2_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def save_ckpt(self, save_path="./binary2_doe.pth"):
    """Save the Binary2 phase coefficients to disk.

    Args:
        save_path (str, optional): Output checkpoint path. Defaults to "./binary2_doe.pth".
    """
    torch.save(
        {
            "param_model": self.param_model,
            "order2": self.order2.clone().detach().cpu(),
            "order4": self.order4.clone().detach().cpu(),
            "order6": self.order6.clone().detach().cpu(),
            "order8": self.order8.clone().detach().cpu(),
            "order10": self.order10.clone().detach().cpu(),
            "order12": self.order12.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./binary2_doe.pth')

Load Binary2 phase coefficients from disk onto the surface device.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path to load. Defaults to "./binary2_doe.pth".	'./binary2_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def load_ckpt(self, load_path="./binary2_doe.pth"):
    """Load Binary2 phase coefficients from disk onto the surface device.

    Args:
        load_path (str, optional): Checkpoint path to load. Defaults to "./binary2_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.order2 = ckpt["order2"].to(self.device)
    self.order4 = ckpt["order4"].to(self.device)
    self.order6 = ckpt["order6"].to(self.device)
    self.order8 = ckpt["order8"].to(self.device)
    self.order10 = ckpt["order10"].to(self.device)
    self.order12 = ckpt["order12"].to(self.device)

zmx_str

zmx_str(surf_idx, d_next)

Return the Zemax BINARY_2 surface block as a string.

PARM 1-8 are set to zero (flat substrate, no aspheric sag) so that Zemax interprets the XDAT entries purely as phase polynomial coefficients.

Parameters:

Name	Type	Description	Default
surf_idx	int	Surface index used in the SURF header.	required
d_next	Tensor	Distance to the next surface [mm], scalar tensor (read via .item()).	required

Returns:

Name	Type	Description
zmx_str	str	Multi-line Zemax surface description.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

    def zmx_str(self, surf_idx, d_next):
        """Return the Zemax BINARY_2 surface block as a string.

        PARM 1-8 are set to zero (flat substrate, no aspheric sag) so that
        Zemax interprets the XDAT entries purely as phase polynomial
        coefficients.

        Args:
            surf_idx (int): Surface index used in the SURF header.
            d_next (torch.Tensor): Distance to the next surface [mm], scalar tensor (read via `.item()`).

        Returns:
            zmx_str (str): Multi-line Zemax surface description.
        """
        coeffs = [
            self.order2.item(),
            self.order4.item(),
            self.order6.item(),
            self.order8.item(),
            self.order10.item(),
            self.order12.item(),
        ]
        n_terms = len(coeffs)

        # Build XDAT block: term count, norm radius, then coefficients
        xdat_str = f"    XDAT 1 {n_terms} 0 0\n"
        xdat_str += f"    XDAT 2 {self.norm_radii} 0 0\n"
        for j, coeff in enumerate(coeffs, start=3):
            xdat_str += f"    XDAT {j} {coeff} 0 0\n"

        zmx_str = f"""SURF {surf_idx}
    TYPE BINARY_2
    CURV 0.0
    DISZ {d_next.item()}
    DIAM {self.r} 1 0 0 1 ""
    PARM 1 0
    PARM 2 0
    PARM 3 0
    PARM 4 0
    PARM 5 0
    PARM 6 0
    PARM 7 0
    PARM 8 0
{xdat_str}"""
        return zmx_str

surf_dict

surf_dict()

Return a serializable dictionary of the surface parameters.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters including type, radius r [mm], polynomial coefficients (rounded), norm_radii [mm], position d [mm], and material name.

Source code in deeplens-src/deeplens/phase_surface/binary2.py

def surf_dict(self):
    """Return a serializable dictionary of the surface parameters.

    Returns:
        surf_dict (dict): Surface parameters including type, radius `r` [mm],
            polynomial coefficients (rounded), `norm_radii` [mm], position `d` [mm],
            and material name.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "order2": round(self.order2.item(), 4),
        "order4": round(self.order4.item(), 4),
        "order6": round(self.order6.item(), 4),
        "order8": round(self.order8.item(), 4),
        "order10": round(self.order10.item(), 4),
        "order12": round(self.order12.item(), 4),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.FresnelPhase

FresnelPhase(
    r,
    d_next,
    f0=100.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Ideal Fresnel-lens phase profile on a plane substrate.

Implements the quadratic phase of a thin lens, \(\phi = -\pi (x^2 + y^2) / (\lambda_0 f_0)\), wrapped to \([0, 2\pi)\), where the design wavelength is fixed at \(\lambda_0 = 0.55\,\mu m\). The single optimizable parameter is the focal length f0 [mm].

Attributes:

Name	Type	Description
f0	Tensor	Scalar tensor, focal length at the design wavelength (550 nm) [mm].
param_model	str	Parameterization identifier, always "fresnel".

Initialize a Fresnel-lens phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius (half-diameter) of the surface [mm].	required
d	float	Axial position of the surface along the optical axis [mm].	required
f0	float	Focal length at the design wavelength (550 nm) [mm]. Defaults to 100.0.	100.0
norm_radii	float or None	Normalization radius for the phase profile [mm]. Defaults to None, in which case r is used.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral (x, y) position of the surface center [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction in the global frame. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True the aperture is a square of full width \(r\sqrt{2}\); otherwise it is a circle of radius r. Defaults to True.	True
device	str	Torch device for tensors. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def __init__(
    self,
    r,
    d_next,
    f0=100.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a Fresnel-lens phase surface.

    Args:
        r (float): Aperture radius (half-diameter) of the surface [mm].
        d (float): Axial position of the surface along the optical axis [mm].
        f0 (float, optional): Focal length at the design wavelength (550 nm) [mm]. Defaults to 100.0.
        norm_radii (float or None, optional): Normalization radius for the phase profile [mm].
            Defaults to None, in which case `r` is used.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral (x, y) position of the surface center [mm].
            Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction in the global frame.
            Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True the aperture is a square of full
            width $r\\sqrt{2}$; otherwise it is a circle of radius `r`. Defaults to True.
        device (str, optional): Torch device for tensors. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    # Focal length at 550nm
    self.f0 = torch.tensor(f0)
    self.param_model = "fresnel"
    self.to(device)

init_from_dict classmethod

init_from_dict(param_dict)

Initialize a FresnelPhase from a dictionary of parameters.

Parameters:

Name	Type	Description	Default
param_dict	dict	Surface parameters. Recognized keys are "r", "d_next", "f0", "norm_radii", "mat2", "pos_xy", "vec_local", "is_square", and "device", matching the __init__ arguments. Missing optional keys fall back to defaults.	required

Returns:

Name	Type	Description
surf	FresnelPhase	The constructed Fresnel phase surface.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

@classmethod
def init_from_dict(cls, param_dict):
    """Initialize a FresnelPhase from a dictionary of parameters.

    Args:
        param_dict (dict): Surface parameters. Recognized keys are "r", "d_next", "f0",
            "norm_radii", "mat2", "pos_xy", "vec_local", "is_square", and "device",
            matching the `__init__` arguments. Missing optional keys fall back to defaults.

    Returns:
        surf (FresnelPhase): The constructed Fresnel phase surface.
    """
    r = param_dict.get("r")
    d_next = param_dict.get("d_next")
    f0 = param_dict.get("f0", 100.0)
    norm_radii = param_dict.get("norm_radii", None)
    mat2 = param_dict.get("mat2", "air")
    pos_xy = param_dict.get("pos_xy", [0.0, 0.0])
    vec_local = param_dict.get("vec_local", [0.0, 0.0, 1.0])
    is_square = param_dict.get("is_square", True)
    device = param_dict.get("device", "cpu")
    return cls(
        r=r,
        d_next=d_next,
        f0=f0,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

phi

phi(x, y)

Compute the wrapped Fresnel-lens phase at the design wavelength.

Evaluates the ideal thin-lens quadratic phase \(\phi = -\pi (x^2 + y^2) / (\lambda_0 f_0)\) with \(\lambda_0 = 0.55\,\mu m\) (0.55e-3 mm), wrapped to \([0, 2\pi)\).

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates on the surface [mm], any broadcastable shape.	required
y	Tensor	Y coordinates on the surface [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase in radians, wrapped to \([0, 2\pi)\), same shape as x.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def phi(self, x, y):
    """Compute the wrapped Fresnel-lens phase at the design wavelength.

    Evaluates the ideal thin-lens quadratic phase
    $\\phi = -\\pi (x^2 + y^2) / (\\lambda_0 f_0)$ with $\\lambda_0 = 0.55\\,\\mu m$
    (0.55e-3 mm), wrapped to $[0, 2\\pi)$.

    Args:
        x (torch.Tensor): X coordinates on the surface [mm], any broadcastable shape.
        y (torch.Tensor): Y coordinates on the surface [mm], same shape as `x`.

    Returns:
        phi (torch.Tensor): Phase in radians, wrapped to $[0, 2\\pi)$, same shape as `x`.
    """
    phi = (
        -2 * torch.pi * torch.fmod((x**2 + y**2) / (2 * 0.55e-3 * self.f0), 1)
    )  # unit [mm]
    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the phase gradient (dphi/dx, dphi/dy) of the unwrapped phase.

Differentiates the unwrapped quadratic phase, giving \(\partial\phi/\partial x = -2\pi x / (\lambda_0 f_0)\) and likewise for \(y\), with \(\lambda_0 = 0.55\,\mu m\) (0.55e-3 mm).

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates on the surface [mm], any broadcastable shape.	required
y	Tensor	Y coordinates on the surface [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase derivative along x [rad/mm], same shape as x.
dphidy	Tensor	Phase derivative along y [rad/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def dphi_dxy(self, x, y):
    """Compute the phase gradient (dphi/dx, dphi/dy) of the unwrapped phase.

    Differentiates the unwrapped quadratic phase, giving
    $\\partial\\phi/\\partial x = -2\\pi x / (\\lambda_0 f_0)$ and likewise for $y$,
    with $\\lambda_0 = 0.55\\,\\mu m$ (0.55e-3 mm).

    Args:
        x (torch.Tensor): X coordinates on the surface [mm], any broadcastable shape.
        y (torch.Tensor): Y coordinates on the surface [mm], same shape as `x`.

    Returns:
        dphidx (torch.Tensor): Phase derivative along x [rad/mm], same shape as `x`.
        dphidy (torch.Tensor): Phase derivative along y [rad/mm], same shape as `x`.
    """
    dphidx = -2 * torch.pi * x / (0.55e-3 * self.f0)  # unit [mm]
    dphidy = -2 * torch.pi * y / (0.55e-3 * self.f0)
    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Build optimizer parameter groups for the focal length.

Enables gradients on f0 and returns a single parameter group using lrs[0] as its learning rate. Material parameters are not optimized for a phase surface.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates; only lrs[0] is used (for f0). Defaults to [1e-4].	[0.0001]
optim_mat	bool	Must be False; material optimization is unsupported. Defaults to False.	False

Returns:

Name	Type	Description
params	list	A list with one parameter group dict {"params": [f0], "lr": lrs[0]}.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Build optimizer parameter groups for the focal length.

    Enables gradients on `f0` and returns a single parameter group using `lrs[0]`
    as its learning rate. Material parameters are not optimized for a phase surface.

    Args:
        lrs (list, optional): Learning rates; only `lrs[0]` is used (for `f0`).
            Defaults to [1e-4].
        optim_mat (bool, optional): Must be False; material optimization is unsupported.
            Defaults to False.

    Returns:
        params (list): A list with one parameter group dict {"params": [f0], "lr": lrs[0]}.

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []

    # Optimize focal length
    self.f0.requires_grad = True
    params.append({"params": [self.f0], "lr": lrs[0]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./fresnel_doe.pth')

Save Fresnel DOE parameters (param_model and f0) to a checkpoint file.

Parameters:

Name	Type	Description	Default
save_path	str	Output checkpoint path. Defaults to "./fresnel_doe.pth".	'./fresnel_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def save_ckpt(self, save_path="./fresnel_doe.pth"):
    """Save Fresnel DOE parameters (param_model and f0) to a checkpoint file.

    Args:
        save_path (str, optional): Output checkpoint path. Defaults to "./fresnel_doe.pth".
    """
    torch.save(
        {
            "param_model": self.param_model,
            "f0": self.f0.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./fresnel_doe.pth')

Load Fresnel DOE parameters (param_model and f0) from a checkpoint file.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path to load. Defaults to "./fresnel_doe.pth".	'./fresnel_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def load_ckpt(self, load_path="./fresnel_doe.pth"):
    """Load Fresnel DOE parameters (param_model and f0) from a checkpoint file.

    Args:
        load_path (str, optional): Checkpoint path to load. Defaults to "./fresnel_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.f0 = ckpt["f0"].to(self.device)

surf_dict

surf_dict()

Return a serializable dict of surface parameters.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters including "type", "r", "is_square", "param_model", "f0", "norm_radii", "d", "mat2", plus informational "(mat2_n)"/"(mat2_V)", suitable for reconstruction via init_from_dict.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py

def surf_dict(self):
    """Return a serializable dict of surface parameters.

    Returns:
        surf_dict (dict): Surface parameters including "type", "r", "is_square",
            "param_model", "f0", "norm_radii", "d", "mat2", plus informational
            "(mat2_n)"/"(mat2_V)", suitable for
            reconstruction via `init_from_dict`.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "f0": self.f0.item(),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.ZernikePhase

ZernikePhase(
    r,
    d_next,
    zernike_order=37,
    zernike_coeff=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=False,
    device="cpu",
)

Bases: Phase

Diffractive phase surface parameterized by Zernike polynomials.

Implements a flat (plane-substrate) diffractive surface whose phase profile is the weighted sum of the first 37 standard (normalized) Zernike polynomials evaluated over a circular pupil. Inherits ray-tracing, refraction, and diffraction from Phase.

Attributes:

Name	Type	Description
param_model	str	Parameterization tag, always "zernike".
zernike_order	int	Number of Zernike terms (up to 37).
z_coeff	Tensor	Zernike coefficients, shape [zernike_order].
norm_radii	float	Radius [mm] used to normalize pupil coordinates.

Reference

[1] https://support.zemax.com/hc/en-us/articles/1500005489061-How-diffractive-surfaces-are-modeled-in-OpticStudio [2] https://optics.ansys.com/hc/en-us/articles/360042097313-Small-Scale-Metalens-Field-Propagation [3] https://optics.ansys.com/hc/en-us/articles/18254409091987-Large-Scale-Metalens-Ray-Propagation

Initialize a Zernike phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius of the surface [mm].	required
d	float	Axial position (distance to next surface) [mm].	required
zernike_order	int	Number of Zernike terms. Defaults to 37.	37
zernike_coeff	list or Tensor or None	Initial Zernike coefficients of length zernike_order. If None, they are randomly initialized as randn * 1e-3. Defaults to None.	None
norm_radii	float or None	Radius [mm] used to normalize pupil coordinates. Defaults to None, meaning r is used.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral (x, y) position of the surface center [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction in global coordinates. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	Whether the aperture is square. Defaults to False.	False
device	str	Computation device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def __init__(
    self,
    r,
    d_next,
    zernike_order=37,
    zernike_coeff=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=False,
    device="cpu",
):
    """Initialize a Zernike phase surface.

    Args:
        r (float): Aperture radius of the surface [mm].
        d (float): Axial position (distance to next surface) [mm].
        zernike_order (int, optional): Number of Zernike terms. Defaults to 37.
        zernike_coeff (list or torch.Tensor or None, optional): Initial Zernike
            coefficients of length `zernike_order`. If None, they are randomly
            initialized as `randn * 1e-3`. Defaults to None.
        norm_radii (float or None, optional): Radius [mm] used to normalize pupil
            coordinates. Defaults to None, meaning `r` is used.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral (x, y) position of the surface center [mm].
            Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction in global
            coordinates. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): Whether the aperture is square. Defaults to False.
        device (str, optional): Computation device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.param_model = "zernike"

    # Zernike polynomial parameterization
    self.zernike_order = zernike_order
    if zernike_coeff is None:
        self.z_coeff = torch.randn(self.zernike_order) * 1e-3
    else:
        self.z_coeff = torch.tensor(zernike_coeff)

    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Construct a Zernike phase surface from a serialized dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Requires keys "r" and "d_next"; optionally reads "mat2", "norm_radii", "zernike_order", and "z_coeff" (the Zernike coefficient list/tensor).	required

Returns:

Name	Type	Description
obj	ZernikePhase	The reconstructed phase surface.

Source code in deeplens-src/deeplens/phase_surface/zernike.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Construct a Zernike phase surface from a serialized dictionary.

    Args:
        surf_dict (dict): Surface parameters. Requires keys "r" and "d_next";
            optionally reads "mat2", "norm_radii", "zernike_order", and
            "z_coeff" (the Zernike coefficient list/tensor).

    Returns:
        obj (ZernikePhase): The reconstructed phase surface.
    """
    mat2 = surf_dict.get("mat2", "air")
    norm_radii = surf_dict.get("norm_radii", None)
    zernike_order = surf_dict.get("zernike_order", 37)

    obj = cls(
        surf_dict["r"],
        surf_dict["d_next"],
        zernike_order=zernike_order,
        norm_radii=norm_radii,
        mat2=mat2,
    )

    # Load Zernike coefficients
    z_coeff = surf_dict.get("z_coeff", None)
    if z_coeff is not None:
        obj.z_coeff = (
            torch.tensor(z_coeff, device=obj.device)
            if not isinstance(z_coeff, torch.Tensor)
            else z_coeff.to(obj.device)
        )

    return obj

phi

phi(x, y)

Compute the reference phase map from the Zernike polynomials.

Overrides Phase.phi. Coordinates are normalized by norm_radii before evaluation, and the result is wrapped into \([0, 2\pi)\).

Parameters:

Name	Type	Description	Default
x	Tensor	Lateral x coordinates [mm], any shape.	required
y	Tensor	Lateral y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase in radians, wrapped to \([0, 2\pi)\), same shape as x.

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def phi(self, x, y):
    """Compute the reference phase map from the Zernike polynomials.

    Overrides `Phase.phi`. Coordinates are normalized by `norm_radii` before
    evaluation, and the result is wrapped into $[0, 2\\pi)$.

    Args:
        x (torch.Tensor): Lateral x coordinates [mm], any shape.
        y (torch.Tensor): Lateral y coordinates [mm], same shape as `x`.

    Returns:
        phi (torch.Tensor): Phase in radians, wrapped to $[0, 2\\pi)$, same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    phi = self._calculate_zernike_phase(x_norm, y_norm)
    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the phase gradient (dphi/dx, dphi/dy) at given points.

Overrides Phase.dphi_dxy. Coordinates are normalized by norm_radii before evaluation; the returned derivatives are with respect to the physical x, y in [mm].

Parameters:

Name	Type	Description	Default
x	Tensor	Lateral x coordinates [mm], any shape.	required
y	Tensor	Lateral y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase derivative dphi/dx [1/mm], same shape as x.
dphidy	Tensor	Phase derivative dphi/dy [1/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def dphi_dxy(self, x, y):
    """Compute the phase gradient (dphi/dx, dphi/dy) at given points.

    Overrides `Phase.dphi_dxy`. Coordinates are normalized by `norm_radii`
    before evaluation; the returned derivatives are with respect to the
    physical x, y in [mm].

    Args:
        x (torch.Tensor): Lateral x coordinates [mm], any shape.
        y (torch.Tensor): Lateral y coordinates [mm], same shape as `x`.

    Returns:
        dphidx (torch.Tensor): Phase derivative dphi/dx [1/mm], same shape as `x`.
        dphidy (torch.Tensor): Phase derivative dphi/dy [1/mm], same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    dphidx, dphidy = self._calculate_zernike_derivatives(x_norm, y_norm)
    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Build optimizer parameter groups for the Zernike coefficients.

Enables gradients on z_coeff and assigns it the first learning rate.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates; lrs[0] is used for z_coeff. Defaults to [1e-4].	[0.0001]
optim_mat	bool	Must be False; material optimization is unsupported for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List with one parameter-group dict for the optimizer.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Build optimizer parameter groups for the Zernike coefficients.

    Enables gradients on `z_coeff` and assigns it the first learning rate.

    Args:
        lrs (list, optional): Learning rates; `lrs[0]` is used for `z_coeff`.
            Defaults to [1e-4].
        optim_mat (bool, optional): Must be False; material optimization is
            unsupported for phase surfaces. Defaults to False.

    Returns:
        params (list): List with one parameter-group dict for the optimizer.

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []
    self.z_coeff.requires_grad = True
    params.append({"params": [self.z_coeff], "lr": lrs[0]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./zernike_doe.pth')

Save the Zernike coefficients and order to a checkpoint file.

Parameters:

Name	Type	Description	Default
save_path	str	Output path. Defaults to "./zernike_doe.pth".	'./zernike_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def save_ckpt(self, save_path="./zernike_doe.pth"):
    """Save the Zernike coefficients and order to a checkpoint file.

    Args:
        save_path (str, optional): Output path. Defaults to "./zernike_doe.pth".
    """
    torch.save(
        {
            "param_model": "zernike",
            "z_coeff": self.z_coeff.clone().detach().cpu(),
            "zernike_order": self.zernike_order,
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./zernike_doe.pth')

Load Zernike coefficients and order from a checkpoint file.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path. Defaults to "./zernike_doe.pth".	'./zernike_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def load_ckpt(self, load_path="./zernike_doe.pth"):
    """Load Zernike coefficients and order from a checkpoint file.

    Args:
        load_path (str, optional): Checkpoint path. Defaults to "./zernike_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.z_coeff = ckpt["z_coeff"].to(self.device)
    self.zernike_order = ckpt["zernike_order"]

surf_dict

surf_dict()

Serialize the surface parameters to a JSON-friendly dictionary.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters including type, radius r [mm], is_square, param_model, the Zernike coefficient list z_coeff, zernike_order, norm_radii [mm], axial position d [mm], and the right-side material name.

Source code in deeplens-src/deeplens/phase_surface/zernike.py

def surf_dict(self):
    """Serialize the surface parameters to a JSON-friendly dictionary.

    Returns:
        surf_dict (dict): Surface parameters including type, radius `r` [mm],
            `is_square`, `param_model`, the Zernike coefficient list
            `z_coeff`, `zernike_order`, `norm_radii` [mm], axial position
            `d` [mm], and the right-side material name.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "z_coeff": self.z_coeff.clone().detach().cpu().tolist(),
        "zernike_order": self.zernike_order,
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.PolyPhase

PolyPhase(
    r,
    d_next,
    order2=0.0,
    order3=0.0,
    order4=0.0,
    order5=0.0,
    order6=0.0,
    order7=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Polynomial phase profile on a plane substrate.

Models a diffractive (DOE/metasurface) phase pattern as a 1D radial polynomial with even terms in the normalized radius plus odd terms in the normalized \(x\)/\(y\) coordinates. Coordinates are normalized by norm_radii before evaluation. The reference phase is

\[ \phi(x, y) = \sum_{k \in \{2,4,6\}} a_k\, r_n^{k} + \sum_{k \in \{3,5,7\}} a_k\, (x_n^{k} + y_n^{k}), \]

where \(r_n = \sqrt{x_n^2 + y_n^2}\) and \(x_n, y_n\) are the normalized coordinates. Phase is in radians and wrapped to \([0, 2\pi)\).

Attributes:

Name	Type	Description
order2	Tensor	Coefficient \(a_2\) of the \(r_n^2\) term (scalar tensor).
order3	Tensor	Coefficient \(a_3\) of the \((x_n^3 + y_n^3)\) term (scalar tensor).
order4	Tensor	Coefficient \(a_4\) of the \(r_n^4\) term (scalar tensor).
order5	Tensor	Coefficient \(a_5\) of the \((x_n^5 + y_n^5)\) term (scalar tensor).
order6	Tensor	Coefficient \(a_6\) of the \(r_n^6\) term (scalar tensor).
order7	Tensor	Coefficient \(a_7\) of the \((x_n^7 + y_n^7)\) term (scalar tensor).
param_model	str	Parameterization identifier, always "poly1d".

Initialize a polynomial phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius of the substrate [mm].	required
d	float	Axial position of the surface along the optical axis [mm].	required
order2	float	Coefficient of the \(r_n^2\) term. Defaults to 0.0.	0.0
order3	float	Coefficient of the \((x_n^3 + y_n^3)\) term. Defaults to 0.0.	0.0
order4	float	Coefficient of the \(r_n^4\) term. Defaults to 0.0.	0.0
order5	float	Coefficient of the \((x_n^5 + y_n^5)\) term. Defaults to 0.0.	0.0
order6	float	Coefficient of the \(r_n^6\) term. Defaults to 0.0.	0.0
order7	float	Coefficient of the \((x_n^7 + y_n^7)\) term. Defaults to 0.0.	0.0
norm_radii	float or None	Normalization radius [mm] for the coordinates. Defaults to the aperture radius r when None.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral \((x, y)\) offset of the surface [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local direction vector (surface normal) in global coordinates; normalized internally. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True, the aperture is square; otherwise circular. Defaults to True.	True
device	str	Torch device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/poly.py

def __init__(
    self,
    r,
    d_next,
    order2=0.0,
    order3=0.0,
    order4=0.0,
    order5=0.0,
    order6=0.0,
    order7=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a polynomial phase surface.

    Args:
        r (float): Aperture radius of the substrate [mm].
        d (float): Axial position of the surface along the optical axis [mm].
        order2 (float, optional): Coefficient of the $r_n^2$ term. Defaults to 0.0.
        order3 (float, optional): Coefficient of the $(x_n^3 + y_n^3)$ term. Defaults to 0.0.
        order4 (float, optional): Coefficient of the $r_n^4$ term. Defaults to 0.0.
        order5 (float, optional): Coefficient of the $(x_n^5 + y_n^5)$ term. Defaults to 0.0.
        order6 (float, optional): Coefficient of the $r_n^6$ term. Defaults to 0.0.
        order7 (float, optional): Coefficient of the $(x_n^7 + y_n^7)$ term. Defaults to 0.0.
        norm_radii (float or None, optional): Normalization radius [mm] for the
            coordinates. Defaults to the aperture radius `r` when None.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral $(x, y)$ offset of the surface [mm].
            Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local direction vector (surface normal) in
            global coordinates; normalized internally. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True, the aperture is square; otherwise circular.
            Defaults to True.
        device (str, optional): Torch device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.order2 = torch.tensor(order2)
    self.order3 = torch.tensor(order3)
    self.order4 = torch.tensor(order4)
    self.order5 = torch.tensor(order5)
    self.order6 = torch.tensor(order6)
    self.order7 = torch.tensor(order7)

    self.param_model = "poly1d"
    self.to(device)

phi

phi(x, y)

Evaluate the reference phase map at design wavelength.

Parameters:

Name	Type	Description	Default
x	Tensor	\(x\) coordinates [mm], any shape.	required
y	Tensor	\(y\) coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase in radians wrapped to \([0, 2\pi)\), same shape as x.

Source code in deeplens-src/deeplens/phase_surface/poly.py

def phi(self, x, y):
    """Evaluate the reference phase map at design wavelength.

    Args:
        x (torch.Tensor): $x$ coordinates [mm], any shape.
        y (torch.Tensor): $y$ coordinates [mm], same shape as `x`.

    Returns:
        phi (torch.Tensor): Phase in radians wrapped to $[0, 2\\pi)$, same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    r_norm = torch.sqrt(x_norm**2 + y_norm**2 + EPSILON)

    phi_even = (
        self.order2 * r_norm**2 + self.order4 * r_norm**4 + self.order6 * r_norm**6
    )
    phi_odd = (
        self.order3 * (x_norm**3 + y_norm**3)
        + self.order5 * (x_norm**5 + y_norm**5)
        + self.order7 * (x_norm**7 + y_norm**7)
    )
    phi = phi_even + phi_odd

    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the spatial phase gradient at the given points.

Returns the partial derivatives of the unwrapped phase polynomial, in units of radians per millimetre [rad/mm].

Parameters:

Name	Type	Description	Default
x	Tensor	\(x\) coordinates [mm], any shape.	required
y	Tensor	\(y\) coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	\(\partial\phi/\partial x\) [rad/mm], same shape as x.
dphidy	Tensor	\(\partial\phi/\partial y\) [rad/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/poly.py

def dphi_dxy(self, x, y):
    """Compute the spatial phase gradient at the given points.

    Returns the partial derivatives of the unwrapped phase polynomial,
    in units of radians per millimetre [rad/mm].

    Args:
        x (torch.Tensor): $x$ coordinates [mm], any shape.
        y (torch.Tensor): $y$ coordinates [mm], same shape as `x`.

    Returns:
        dphidx (torch.Tensor): $\\partial\\phi/\\partial x$ [rad/mm], same shape as `x`.
        dphidy (torch.Tensor): $\\partial\\phi/\\partial y$ [rad/mm], same shape as `x`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii
    r_norm = torch.sqrt(x_norm**2 + y_norm**2 + EPSILON)

    dphi_even_dr = (
        2 * self.order2 * r_norm
        + 4 * self.order4 * r_norm**3
        + 6 * self.order6 * r_norm**5
    )
    dphi_even_dx = dphi_even_dr * x_norm / r_norm / self.norm_radii
    dphi_even_dy = dphi_even_dr * y_norm / r_norm / self.norm_radii

    dphi_odd_dx = (
        3 * self.order3 * x_norm**2
        + 5 * self.order5 * x_norm**4
        + 7 * self.order7 * x_norm**6
    ) / self.norm_radii
    dphi_odd_dy = (
        3 * self.order3 * y_norm**2
        + 5 * self.order5 * y_norm**4
        + 7 * self.order7 * y_norm**6
    ) / self.norm_radii

    dphidx = dphi_even_dx + dphi_odd_dx
    dphidy = dphi_even_dy + dphi_odd_dy

    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.001, 0.0001, 1e-05, 1e-06, 1e-07], optim_mat=False)

Build per-parameter optimizer groups for the polynomial coefficients.

Marks the six polynomial coefficients (order2 through order7) as requiring gradients and assigns each its own learning rate.

Parameters:

Name	Type	Description	Default
lrs	list	Six learning rates applied to order2...order7 respectively. Defaults to [1e-4, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7].	[0.0001, 0.001, 0.0001, 1e-05, 1e-06, 1e-07]
optim_mat	bool	Must be False; material parameters are not optimized for a phase surface. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List of parameter-group dicts, one per coefficient, each with keys "params" and "lr".

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/poly.py

def get_optimizer_params(
    self, lrs=[1e-4, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7], optim_mat=False
):
    """Build per-parameter optimizer groups for the polynomial coefficients.

    Marks the six polynomial coefficients (`order2` through `order7`) as
    requiring gradients and assigns each its own learning rate.

    Args:
        lrs (list, optional): Six learning rates applied to `order2`...`order7`
            respectively. Defaults to [1e-4, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7].
        optim_mat (bool, optional): Must be False; material parameters are not
            optimized for a phase surface. Defaults to False.

    Returns:
        params (list): List of parameter-group dicts, one per coefficient,
            each with keys "params" and "lr".

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []

    # Optimize polynomial coefficients with different learning rates
    self.order2.requires_grad = True
    self.order3.requires_grad = True
    self.order4.requires_grad = True
    self.order5.requires_grad = True
    self.order6.requires_grad = True
    self.order7.requires_grad = True

    params.append({"params": [self.order2], "lr": lrs[0]})
    params.append({"params": [self.order3], "lr": lrs[1]})
    params.append({"params": [self.order4], "lr": lrs[2]})
    params.append({"params": [self.order5], "lr": lrs[3]})
    params.append({"params": [self.order6], "lr": lrs[4]})
    params.append({"params": [self.order7], "lr": lrs[5]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./poly1d_doe.pth')

Save the Poly1D DOE parameters to a checkpoint file.

Parameters:

Name	Type	Description	Default
save_path	str	Output checkpoint path. Defaults to "./poly1d_doe.pth".	'./poly1d_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/poly.py

def save_ckpt(self, save_path="./poly1d_doe.pth"):
    """Save the Poly1D DOE parameters to a checkpoint file.

    Args:
        save_path (str, optional): Output checkpoint path. Defaults to "./poly1d_doe.pth".
    """
    torch.save(
        {
            "param_model": self.param_model,
            "order2": self.order2.clone().detach().cpu(),
            "order3": self.order3.clone().detach().cpu(),
            "order4": self.order4.clone().detach().cpu(),
            "order5": self.order5.clone().detach().cpu(),
            "order6": self.order6.clone().detach().cpu(),
            "order7": self.order7.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./poly1d_doe.pth')

Load the Poly1D DOE parameters from a checkpoint file.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path to load. Defaults to "./poly1d_doe.pth".	'./poly1d_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/poly.py

def load_ckpt(self, load_path="./poly1d_doe.pth"):
    """Load the Poly1D DOE parameters from a checkpoint file.

    Args:
        load_path (str, optional): Checkpoint path to load. Defaults to "./poly1d_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.order2 = ckpt["order2"].to(self.device)
    self.order3 = ckpt["order3"].to(self.device)
    self.order4 = ckpt["order4"].to(self.device)
    self.order5 = ckpt["order5"].to(self.device)
    self.order6 = ckpt["order6"].to(self.device)
    self.order7 = ckpt["order7"].to(self.device)

surf_dict

surf_dict()

Return a serializable dict of the surface parameters.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters, including the surface type, aperture radius r [mm], is_square, param_model, the six polynomial coefficients (rounded to 4 decimals), norm_radii [mm], sequential thickness d_next [mm], and the material name.

Source code in deeplens-src/deeplens/phase_surface/poly.py

def surf_dict(self):
    """Return a serializable dict of the surface parameters.

    Returns:
        surf_dict (dict): Surface parameters, including the surface type,
            aperture radius `r` [mm], `is_square`, `param_model`, the six
            polynomial coefficients (rounded to 4 decimals), `norm_radii` [mm],
            sequential thickness `d_next` [mm], and the material name.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "order2": round(self.order2.item(), 4),
        "order3": round(self.order3.item(), 4),
        "order4": round(self.order4.item(), 4),
        "order5": round(self.order5.item(), 4),
        "order6": round(self.order6.item(), 4),
        "order7": round(self.order7.item(), 4),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.CubicPhase

CubicPhase(
    r,
    d_next,
    coeff_x3=0.0,
    coeff_y3=0.0,
    coeff_x2y=0.0,
    coeff_xy2=0.0,
    coeff_x3y=0.0,
    coeff_xy3=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Cubic phase profile on a plane substrate.

Diffractive surface whose phase is a cubic 2D polynomial in normalized coordinates \((x/r_n,\, y/r_n)\), where \(r_n\) is norm_radii. Used to model cubic-phase (wavefront-coding) plates that extend depth of field.

Attributes:

Name	Type	Description
coeff_x3	Tensor	Scalar coefficient of \(x^3\) [rad].
coeff_y3	Tensor	Scalar coefficient of \(y^3\) [rad].
coeff_x2y	Tensor	Scalar coefficient of \(x^2 y\) [rad].
coeff_xy2	Tensor	Scalar coefficient of \(x y^2\) [rad].
coeff_x3y	Tensor	Scalar coefficient of \(x^3 y\) [rad].
coeff_xy3	Tensor	Scalar coefficient of \(x y^3\) [rad].
norm_radii	float	Normalization radius [mm] used to scale \((x, y)\).
param_model	str	Parameterization name, always "cubic".

Initialize a cubic phase surface.

Parameters:

Name	Type	Description	Default
r	float	Surface aperture radius [mm].	required
d	float	Axial position of the surface in the global frame [mm].	required
coeff_x3	float	Coefficient of \(x^3\) [rad]. Defaults to 0.0.	0.0
coeff_y3	float	Coefficient of \(y^3\) [rad]. Defaults to 0.0.	0.0
coeff_x2y	float	Coefficient of \(x^2 y\) [rad]. Defaults to 0.0.	0.0
coeff_xy2	float	Coefficient of \(x y^2\) [rad]. Defaults to 0.0.	0.0
coeff_x3y	float	Coefficient of \(x^3 y\) [rad]. Defaults to 0.0.	0.0
coeff_xy3	float	Coefficient of \(x y^3\) [rad]. Defaults to 0.0.	0.0
norm_radii	float or None	Normalization radius [mm] for the polynomial coordinates. Defaults to None, in which case the base class sets it to r.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral \((x, y)\) offset of the surface center [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True, use a square aperture; otherwise circular. Defaults to True.	True
device	str	Torch device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def __init__(
    self,
    r,
    d_next,
    coeff_x3=0.0,
    coeff_y3=0.0,
    coeff_x2y=0.0,
    coeff_xy2=0.0,
    coeff_x3y=0.0,
    coeff_xy3=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a cubic phase surface.

    Args:
        r (float): Surface aperture radius [mm].
        d (float): Axial position of the surface in the global frame [mm].
        coeff_x3 (float, optional): Coefficient of $x^3$ [rad]. Defaults to 0.0.
        coeff_y3 (float, optional): Coefficient of $y^3$ [rad]. Defaults to 0.0.
        coeff_x2y (float, optional): Coefficient of $x^2 y$ [rad]. Defaults to 0.0.
        coeff_xy2 (float, optional): Coefficient of $x y^2$ [rad]. Defaults to 0.0.
        coeff_x3y (float, optional): Coefficient of $x^3 y$ [rad]. Defaults to 0.0.
        coeff_xy3 (float, optional): Coefficient of $x y^3$ [rad]. Defaults to 0.0.
        norm_radii (float or None, optional): Normalization radius [mm] for the
            polynomial coordinates. Defaults to None, in which case the base
            class sets it to `r`.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral $(x, y)$ offset of the surface center
            [mm]. Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction. Defaults to
            (0.0, 0.0, 1.0).
        is_square (bool, optional): If True, use a square aperture; otherwise
            circular. Defaults to True.
        device (str, optional): Torch device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.coeff_x3 = torch.tensor(coeff_x3)
    self.coeff_y3 = torch.tensor(coeff_y3)
    self.coeff_x2y = torch.tensor(coeff_x2y)
    self.coeff_xy2 = torch.tensor(coeff_xy2)
    self.coeff_x3y = torch.tensor(coeff_x3y)
    self.coeff_xy3 = torch.tensor(coeff_xy3)

    self.param_model = "cubic"
    self.to(device)

phi

phi(x, y)

Compute the reference phase map at the design wavelength.

Evaluates the cubic polynomial in normalized coordinates \((x_n, y_n) = (x/r_n,\, y/r_n)\) and wraps the result into \([0, 2\pi)\) via torch.remainder:

\[ \phi = c_{x3} x_n^3 + c_{y3} y_n^3 + c_{x2y} x_n^2 y_n + c_{xy2} x_n y_n^2 + c_{x3y} x_n^3 y_n + c_{xy3} x_n y_n^3 \]

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], broadcastable with x.	required

Returns:

Name	Type	Description
phi	Tensor	Wrapped phase [rad] in \([0, 2\pi)\), same shape as the broadcast of x and y.

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def phi(self, x, y):
    """Compute the reference phase map at the design wavelength.

    Evaluates the cubic polynomial in normalized coordinates
    $(x_n, y_n) = (x/r_n,\\, y/r_n)$ and wraps the result into $[0, 2\\pi)$
    via `torch.remainder`:

    $$
    \\phi = c_{x3} x_n^3 + c_{y3} y_n^3 + c_{x2y} x_n^2 y_n
          + c_{xy2} x_n y_n^2 + c_{x3y} x_n^3 y_n + c_{xy3} x_n y_n^3
    $$

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], broadcastable with `x`.

    Returns:
        phi (torch.Tensor): Wrapped phase [rad] in $[0, 2\\pi)$, same shape
            as the broadcast of `x` and `y`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii

    phi = (
        self.coeff_x3 * x_norm**3
        + self.coeff_y3 * y_norm**3
        + self.coeff_x2y * x_norm**2 * y_norm
        + self.coeff_xy2 * x_norm * y_norm**2
        + self.coeff_x3y * x_norm**3 * y_norm
        + self.coeff_xy3 * x_norm * y_norm**3
    )

    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the spatial phase derivatives at given points.

Returns the analytic gradient of the unwrapped cubic polynomial (not the wrapped phi). Derivatives are taken w.r.t. normalized coordinates and converted to physical coordinates by dividing by norm_radii.

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], broadcastable with x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase derivative \(\partial\phi/\partial x\) [rad/mm], same shape as the broadcast of x and y.
dphidy	Tensor	Phase derivative \(\partial\phi/\partial y\) [rad/mm], same shape as the broadcast of x and y.

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def dphi_dxy(self, x, y):
    """Compute the spatial phase derivatives at given points.

    Returns the analytic gradient of the unwrapped cubic polynomial (not the
    wrapped `phi`). Derivatives are taken w.r.t. normalized coordinates and
    converted to physical coordinates by dividing by `norm_radii`.

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], broadcastable with `x`.

    Returns:
        dphidx (torch.Tensor): Phase derivative $\\partial\\phi/\\partial x$
            [rad/mm], same shape as the broadcast of `x` and `y`.
        dphidy (torch.Tensor): Phase derivative $\\partial\\phi/\\partial y$
            [rad/mm], same shape as the broadcast of `x` and `y`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii

    # Derivatives with respect to normalized coordinates
    dphi_dx_norm = (
        3 * self.coeff_x3 * x_norm**2
        + 2 * self.coeff_x2y * x_norm * y_norm
        + 3 * self.coeff_x3y * x_norm**2 * y_norm
        + self.coeff_xy2 * y_norm**2
        + self.coeff_xy3 * y_norm**3
    )

    dphi_dy_norm = (
        3 * self.coeff_y3 * y_norm**2
        + self.coeff_x2y * x_norm**2
        + 2 * self.coeff_xy2 * x_norm * y_norm
        + self.coeff_x3y * x_norm**3
        + 3 * self.coeff_xy3 * x_norm * y_norm**2
    )

    # Convert back to physical coordinates
    dphidx = dphi_dx_norm / self.norm_radii
    dphidy = dphi_dy_norm / self.norm_radii

    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(
    lrs=[0.0001, 0.0001, 0.0001, 0.0001, 1e-05, 1e-05], optim_mat=False
)

Build optimizer parameter groups for the cubic coefficients.

Enables gradients on the six polynomial coefficients and assigns each a per-group learning rate, in the order [coeff_x3, coeff_y3, coeff_x2y, coeff_xy2, coeff_x3y, coeff_xy3].

Parameters:

Name	Type	Description	Default
lrs	list	Six learning rates, one per coefficient. Defaults to [1e-4, 1e-4, 1e-4, 1e-4, 1e-5, 1e-5].	[0.0001, 0.0001, 0.0001, 0.0001, 1e-05, 1e-05]
optim_mat	bool	Must be False; material parameters are not optimized for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List of six optimizer parameter-group dicts, each with "params" and "lr" keys.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def get_optimizer_params(
    self, lrs=[1e-4, 1e-4, 1e-4, 1e-4, 1e-5, 1e-5], optim_mat=False
):
    """Build optimizer parameter groups for the cubic coefficients.

    Enables gradients on the six polynomial coefficients and assigns each a
    per-group learning rate, in the order
    `[coeff_x3, coeff_y3, coeff_x2y, coeff_xy2, coeff_x3y, coeff_xy3]`.

    Args:
        lrs (list, optional): Six learning rates, one per coefficient.
            Defaults to [1e-4, 1e-4, 1e-4, 1e-4, 1e-5, 1e-5].
        optim_mat (bool, optional): Must be False; material parameters are
            not optimized for phase surfaces. Defaults to False.

    Returns:
        params (list): List of six optimizer parameter-group dicts, each with
            "params" and "lr" keys.

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []

    # Optimize cubic polynomial coefficients with different learning rates
    self.coeff_x3.requires_grad = True
    self.coeff_y3.requires_grad = True
    self.coeff_x2y.requires_grad = True
    self.coeff_xy2.requires_grad = True
    self.coeff_x3y.requires_grad = True
    self.coeff_xy3.requires_grad = True

    params.append({"params": [self.coeff_x3], "lr": lrs[0]})
    params.append({"params": [self.coeff_y3], "lr": lrs[1]})
    params.append({"params": [self.coeff_x2y], "lr": lrs[2]})
    params.append({"params": [self.coeff_xy2], "lr": lrs[3]})
    params.append({"params": [self.coeff_x3y], "lr": lrs[4]})
    params.append({"params": [self.coeff_xy3], "lr": lrs[5]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./cubic_doe.pth')

Save the cubic DOE parameters to a checkpoint file.

Saves param_model and the six cubic coefficients (detached, on CPU).

Parameters:

Name	Type	Description	Default
save_path	str	Output checkpoint path. Defaults to "./cubic_doe.pth".	'./cubic_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def save_ckpt(self, save_path="./cubic_doe.pth"):
    """Save the cubic DOE parameters to a checkpoint file.

    Saves `param_model` and the six cubic coefficients (detached, on CPU).

    Args:
        save_path (str, optional): Output checkpoint path. Defaults to
            "./cubic_doe.pth".
    """
    torch.save(
        {
            "param_model": self.param_model,
            "coeff_x3": self.coeff_x3.clone().detach().cpu(),
            "coeff_y3": self.coeff_y3.clone().detach().cpu(),
            "coeff_x2y": self.coeff_x2y.clone().detach().cpu(),
            "coeff_xy2": self.coeff_xy2.clone().detach().cpu(),
            "coeff_x3y": self.coeff_x3y.clone().detach().cpu(),
            "coeff_xy3": self.coeff_xy3.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./cubic_doe.pth')

Load the cubic DOE parameters from a checkpoint file.

Restores param_model and the six cubic coefficients onto self.device.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path to load. Defaults to "./cubic_doe.pth".	'./cubic_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def load_ckpt(self, load_path="./cubic_doe.pth"):
    """Load the cubic DOE parameters from a checkpoint file.

    Restores `param_model` and the six cubic coefficients onto `self.device`.

    Args:
        load_path (str, optional): Checkpoint path to load. Defaults to
            "./cubic_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.coeff_x3 = ckpt["coeff_x3"].to(self.device)
    self.coeff_y3 = ckpt["coeff_y3"].to(self.device)
    self.coeff_x2y = ckpt["coeff_x2y"].to(self.device)
    self.coeff_xy2 = ckpt["coeff_xy2"].to(self.device)
    self.coeff_x3y = ckpt["coeff_x3y"].to(self.device)
    self.coeff_xy3 = ckpt["coeff_xy3"].to(self.device)

surf_dict

surf_dict()

Return a serializable dict of the surface parameters.

Coefficients, norm_radii, and d are rounded to 4 decimals. Used for exporting the surface to a lens file.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters including type, r, is_square, param_model, the six cubic coefficients, norm_radii [mm], d [mm], and mat2 name.

Source code in deeplens-src/deeplens/phase_surface/cubic.py

def surf_dict(self):
    """Return a serializable dict of the surface parameters.

    Coefficients, `norm_radii`, and `d` are rounded to 4 decimals. Used for
    exporting the surface to a lens file.

    Returns:
        surf_dict (dict): Surface parameters including type, `r`, `is_square`,
            `param_model`, the six cubic coefficients, `norm_radii` [mm],
            `d` [mm], and `mat2` name.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "coeff_x3": round(self.coeff_x3.item(), 4),
        "coeff_y3": round(self.coeff_y3.item(), 4),
        "coeff_x2y": round(self.coeff_x2y.item(), 4),
        "coeff_xy2": round(self.coeff_xy2.item(), 4),
        "coeff_x3y": round(self.coeff_x3y.item(), 4),
        "coeff_xy3": round(self.coeff_xy3.item(), 4),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.GratingPhase

GratingPhase(
    r,
    d_next,
    theta=0.0,
    alpha=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Linear (blazed) grating phase on a plane substrate.

The phase profile is a linear ramp along a direction set by theta:

\[ \phi(x, y) = \alpha \left( \frac{x}{R} \sin\theta + \frac{y}{R} \cos\theta \right) \]

where \(R\) is norm_radii, \(\alpha\) (alpha) is the ramp magnitude, and \(\theta\) (theta) is the grating-vector angle measured from the x-axis [rad]. The returned phase is wrapped to \([0, 2\pi)\). Inherits ray tracing, diffraction, and coordinate transforms from Phase.

Attributes:

Name	Type	Description
theta	Tensor	Scalar grating-vector angle from the x-axis [rad].
alpha	Tensor	Scalar phase-ramp magnitude [rad] over norm_radii.
param_model	str	Parameterization tag, "grating".

Initialize a linear grating phase surface.

Parameters:

Name	Type	Description	Default
r	float	Surface aperture radius [mm].	required
d	float	Axial position of the surface in global coordinates [mm].	required
theta	float	Grating-vector angle from the x-axis [rad]. Defaults to 0.0.	0.0
alpha	float	Phase-ramp magnitude [rad] across norm_radii. Defaults to 0.0.	0.0
norm_radii	float or None	Normalization radius [mm] for the phase profile. Defaults to None, which uses r.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Lateral (x, y) surface offset [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True, use a square aperture; otherwise circular. Defaults to True.	True
device	str	Torch device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/grating.py

def __init__(
    self,
    r,
    d_next,
    theta=0.0,
    alpha=0.0,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a linear grating phase surface.

    Args:
        r (float): Surface aperture radius [mm].
        d (float): Axial position of the surface in global coordinates [mm].
        theta (float, optional): Grating-vector angle from the x-axis [rad]. Defaults to 0.0.
        alpha (float, optional): Phase-ramp magnitude [rad] across `norm_radii`. Defaults to 0.0.
        norm_radii (float or None, optional): Normalization radius [mm] for the phase
            profile. Defaults to None, which uses `r`.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Lateral (x, y) surface offset [mm]. Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True, use a square aperture; otherwise circular.
            Defaults to True.
        device (str, optional): Torch device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    # Grating parameters
    self.theta = torch.tensor(theta)  # angle from x-axis to grating vector
    self.alpha = torch.tensor(alpha)  # slope of the grating

    self.param_model = "grating"
    self.to(device)

init_from_dict classmethod

init_from_dict(param_dict)

Initialize a GratingPhase from a parameter dictionary.

Parameters:

Name	Type	Description	Default
param_dict	dict	Parameter dictionary. Recognized keys match the __init__ arguments ("r", "d", "theta", "alpha", "norm_radii", "mat2", "pos_xy", "vec_local", "is_square", "device"); missing keys fall back to the __init__ defaults.	required

Returns:

Name	Type	Description
grating	GratingPhase	The constructed grating phase surface.

Source code in deeplens-src/deeplens/phase_surface/grating.py

@classmethod
def init_from_dict(cls, param_dict):
    """Initialize a GratingPhase from a parameter dictionary.

    Args:
        param_dict (dict): Parameter dictionary. Recognized keys match the
            `__init__` arguments ("r", "d", "theta", "alpha", "norm_radii",
            "mat2", "pos_xy", "vec_local", "is_square", "device"); missing
            keys fall back to the `__init__` defaults.

    Returns:
        grating (GratingPhase): The constructed grating phase surface.
    """
    # Extract parameters with defaults matching __init__ signature
    r = param_dict.get("r")
    d_next = param_dict.get("d_next")
    theta = param_dict.get("theta", 0.0)
    alpha = param_dict.get("alpha", 0.0)
    norm_radii = param_dict.get("norm_radii", None)
    mat2 = param_dict.get("mat2", "air")
    pos_xy = param_dict.get("pos_xy", [0.0, 0.0])
    vec_local = param_dict.get("vec_local", [0.0, 0.0, 1.0])
    is_square = param_dict.get("is_square", True)
    device = param_dict.get("device", "cpu")
    return cls(
        r=r,
        d_next=d_next,
        theta=theta,
        alpha=alpha,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

phi

phi(x, y)

Compute the grating phase at given points, wrapped to \([0, 2\pi)\).

Evaluates \(\phi = \alpha\,(x/R \sin\theta + y/R \cos\theta)\) where \(R\) is norm_radii, then wraps the result into \([0, 2\pi)\).

Parameters:

Name	Type	Description	Default
x	Tensor	Lateral x-coordinates [mm], any shape.	required
y	Tensor	Lateral y-coordinates [mm], broadcastable with x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase values [rad] in \([0, 2\pi)\), same shape as the broadcast of x and y.

Source code in deeplens-src/deeplens/phase_surface/grating.py

def phi(self, x, y):
    """Compute the grating phase at given points, wrapped to $[0, 2\\pi)$.

    Evaluates $\\phi = \\alpha\\,(x/R \\sin\\theta + y/R \\cos\\theta)$ where
    $R$ is `norm_radii`, then wraps the result into $[0, 2\\pi)$.

    Args:
        x (torch.Tensor): Lateral x-coordinates [mm], any shape.
        y (torch.Tensor): Lateral y-coordinates [mm], broadcastable with `x`.

    Returns:
        phi (torch.Tensor): Phase values [rad] in $[0, 2\\pi)$, same shape as the
            broadcast of `x` and `y`.
    """
    x_norm = x / self.norm_radii
    y_norm = y / self.norm_radii

    phi = self.alpha * (
        x_norm * torch.sin(self.theta) + y_norm * torch.cos(self.theta)
    )

    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the phase gradient (dphi/dx, dphi/dy) at given points.

For this linear grating the gradient is spatially constant: \(\partial\phi/\partial x = \alpha\sin\theta / R\) and \(\partial\phi/\partial y = \alpha\cos\theta / R\), with \(R\) = norm_radii.

Parameters:

Name	Type	Description	Default
x	Tensor	Lateral x-coordinates [mm]; only its shape is used.	required
y	Tensor	Lateral y-coordinates [mm]; only its shape is used.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase x-derivative [rad/mm], broadcast to the shape of x.
dphidy	Tensor	Phase y-derivative [rad/mm], broadcast to the shape of y.

Source code in deeplens-src/deeplens/phase_surface/grating.py

def dphi_dxy(self, x, y):
    """Compute the phase gradient (dphi/dx, dphi/dy) at given points.

    For this linear grating the gradient is spatially constant:
    $\\partial\\phi/\\partial x = \\alpha\\sin\\theta / R$ and
    $\\partial\\phi/\\partial y = \\alpha\\cos\\theta / R$, with $R$ = `norm_radii`.

    Args:
        x (torch.Tensor): Lateral x-coordinates [mm]; only its shape is used.
        y (torch.Tensor): Lateral y-coordinates [mm]; only its shape is used.

    Returns:
        dphidx (torch.Tensor): Phase x-derivative [rad/mm], broadcast to the shape of `x`.
        dphidy (torch.Tensor): Phase y-derivative [rad/mm], broadcast to the shape of `y`.
    """
    # Scalar derivatives broadcast to input tensor shape without allocation
    dphidx = (self.alpha * torch.sin(self.theta) / self.norm_radii).expand_as(x)
    dphidy = (self.alpha * torch.cos(self.theta) / self.norm_radii).expand_as(y)
    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.001], optim_mat=False)

Build Adam parameter groups for the grating phase parameters.

Enables gradients on theta and alpha and returns one parameter group each, using lrs[0] for theta and lrs[1] for alpha.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates [lr_theta, lr_alpha]. Defaults to [1e-4, 1e-3].	[0.0001, 0.001]
optim_mat	bool	Must be False; material parameters are not optimized for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List of two parameter-group dicts for theta and alpha.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/grating.py

def get_optimizer_params(self, lrs=[1e-4, 1e-3], optim_mat=False):
    """Build Adam parameter groups for the grating phase parameters.

    Enables gradients on `theta` and `alpha` and returns one parameter group
    each, using `lrs[0]` for `theta` and `lrs[1]` for `alpha`.

    Args:
        lrs (list, optional): Learning rates `[lr_theta, lr_alpha]`.
            Defaults to [1e-4, 1e-3].
        optim_mat (bool, optional): Must be False; material parameters are not
            optimized for phase surfaces. Defaults to False.

    Returns:
        params (list): List of two parameter-group dicts for `theta` and `alpha`.

    Raises:
        AssertionError: If `optim_mat` is True.
    """
    params = []

    # Optimize grating parameters
    self.theta.requires_grad = True
    self.alpha.requires_grad = True
    params.append({"params": [self.theta], "lr": lrs[0]})
    params.append({"params": [self.alpha], "lr": lrs[1]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./grating_doe.pth')

Save grating parameters (param_model, theta, alpha) to a checkpoint.

Parameters:

Name	Type	Description	Default
save_path	str	Output checkpoint path. Defaults to "./grating_doe.pth".	'./grating_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/grating.py

def save_ckpt(self, save_path="./grating_doe.pth"):
    """Save grating parameters (`param_model`, `theta`, `alpha`) to a checkpoint.

    Args:
        save_path (str, optional): Output checkpoint path. Defaults to "./grating_doe.pth".
    """
    torch.save(
        {
            "param_model": self.param_model,
            "theta": self.theta.clone().detach().cpu(),
            "alpha": self.alpha.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./grating_doe.pth')

Load grating parameters (param_model, theta, alpha) from a checkpoint.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path to load. Defaults to "./grating_doe.pth".	'./grating_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/grating.py

def load_ckpt(self, load_path="./grating_doe.pth"):
    """Load grating parameters (`param_model`, `theta`, `alpha`) from a checkpoint.

    Args:
        load_path (str, optional): Checkpoint path to load. Defaults to "./grating_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.theta = ckpt["theta"].to(self.device)
    self.alpha = ckpt["alpha"].to(self.device)

surf_dict

surf_dict()

Return a serializable dict of the grating surface parameters.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters including "type", "r", "is_square", "param_model", "theta", "alpha", "norm_radii", "d", "mat2", plus informational "(mat2_n)"/"(mat2_V)". Numeric "theta", "alpha", "norm_radii", and "d" are rounded to 4 decimals.

Source code in deeplens-src/deeplens/phase_surface/grating.py

def surf_dict(self):
    """Return a serializable dict of the grating surface parameters.

    Returns:
        surf_dict (dict): Surface parameters including "type", "r", "is_square",
            "param_model", "theta", "alpha", "norm_radii", "d", "mat2", plus
            informational "(mat2_n)"/"(mat2_V)".
            Numeric "theta", "alpha", "norm_radii", and "d" are rounded to
            4 decimals.
    """
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "theta": round(self.theta.item(), 4),
        "alpha": round(self.alpha.item(), 4),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.NURBSPhase

NURBSPhase(
    r,
    d_next,
    control_points_u=8,
    control_points_v=8,
    degree_u=3,
    degree_v=3,
    control_points=None,
    weights=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Diffractive phase surface parameterized by a NURBS surface.

The phase profile is the z-component of a NURBS (Non-Uniform Rational B-Spline) surface defined by a 2D grid of control points and clamped knot vectors in the u and v directions. The surface is evaluated with B-spline basis functions via the Cox-de Boor recursion. The (x, y) ray coordinates are normalized to the NURBS parameter domain [0, 1]; the returned phase is in radians and wrapped to [0, 2π).

Attributes:

Name	Type	Description
control_points	Tensor	Control point coordinates (x, y, z) of shape (control_points_u, control_points_v, 3); z is the phase [rad].
weights	Tensor	Rational B-spline weights of shape (control_points_u, control_points_v).
knots_u	Tensor	Clamped knot vector in u of shape (control_points_u + degree_u + 1,).
knots_v	Tensor	Clamped knot vector in v of shape (control_points_v + degree_v + 1,).
control_points_u	int	Number of control points in u.
control_points_v	int	Number of control points in v.
degree_u	int	B-spline degree in u.
degree_v	int	B-spline degree in v.
param_model	str	Parameterization tag, "nurbs".

Reference

[1] The NURBS Book by Piegl and Tiller. [2] https://en.wikipedia.org/wiki/Non-uniform_rational_B-spline

Initialize a NURBS phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius of the surface [mm].	required
d	float	Axial distance to the next surface [mm].	required
control_points_u	int	Number of control points in u. Defaults to 8.	8
control_points_v	int	Number of control points in v. Defaults to 8.	8
degree_u	int	B-spline degree in u. Defaults to 3.	3
degree_v	int	B-spline degree in v. Defaults to 3.	3
control_points	Tensor or None	Control point coordinates of shape (control_points_u, control_points_v, 3) holding (x, y, z) where z is the phase [rad]. If None, x, y are placed on an even grid in [-1, 1] and z is initialized with small random values (std 1e-3 [rad]). Defaults to None.	None
weights	Tensor or None	Rational B-spline weights of shape (control_points_u, control_points_v). If None, all weights are 1. Defaults to None.	None
norm_radii	float or None	Radius [mm] used to normalize (x, y) into the parameter domain. If None, defaults to r. Defaults to None.	None
mat2	str	Material after the surface. Defaults to "air".	'air'
pos_xy	tuple	Surface (x, y) position [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction. Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	Whether the aperture is square. Defaults to True.	True
device	str	Computation device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def __init__(
    self,
    r,
    d_next,
    control_points_u=8,
    control_points_v=8,
    degree_u=3,
    degree_v=3,
    control_points=None,
    weights=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a NURBS phase surface.

    Args:
        r (float): Aperture radius of the surface [mm].
        d (float): Axial distance to the next surface [mm].
        control_points_u (int, optional): Number of control points in u. Defaults to 8.
        control_points_v (int, optional): Number of control points in v. Defaults to 8.
        degree_u (int, optional): B-spline degree in u. Defaults to 3.
        degree_v (int, optional): B-spline degree in v. Defaults to 3.
        control_points (torch.Tensor or None, optional): Control point coordinates
            of shape (control_points_u, control_points_v, 3) holding (x, y, z) where
            z is the phase [rad]. If None, x, y are placed on an even grid in [-1, 1]
            and z is initialized with small random values (std 1e-3 [rad]). Defaults to None.
        weights (torch.Tensor or None, optional): Rational B-spline weights of shape
            (control_points_u, control_points_v). If None, all weights are 1. Defaults to None.
        norm_radii (float or None, optional): Radius [mm] used to normalize (x, y) into
            the parameter domain. If None, defaults to r. Defaults to None.
        mat2 (str, optional): Material after the surface. Defaults to "air".
        pos_xy (tuple, optional): Surface (x, y) position [mm]. Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction. Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): Whether the aperture is square. Defaults to True.
        device (str, optional): Computation device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    # NURBS surface parameters
    self.control_points_u = control_points_u
    self.control_points_v = control_points_v
    self.degree_u = degree_u
    self.degree_v = degree_v

    # Generate knot vectors (clamped B-splines)
    self.knots_u = self._generate_clamped_knots(control_points_u, degree_u)
    self.knots_v = self._generate_clamped_knots(control_points_v, degree_v)

    # Initialize control points (x, y, z) where z represents phase.
    # Use the default dtype (not a hardcoded float32) so float64 runs stay
    # double precision.
    if control_points is None:
        # Initialize with small random phase values
        cp = (
            torch.randn(control_points_u, control_points_v, 3, device=device) * 1e-3
        )
        # Set x,y coordinates to be evenly spaced in [-1, 1] range
        u_coords = torch.linspace(0, 1, control_points_u, device=device)
        v_coords = torch.linspace(0, 1, control_points_v, device=device)
        u_grid, v_grid = torch.meshgrid(u_coords, v_coords, indexing="ij")
        cp[..., 0] = u_grid * 2 - 1  # x coordinates
        cp[..., 1] = v_grid * 2 - 1  # y coordinates
    else:
        cp = torch.as_tensor(
            control_points, dtype=torch.get_default_dtype(), device=device
        )
        assert cp.shape == (control_points_u, control_points_v, 3), (
            f"control_points must have shape ({control_points_u}, {control_points_v}, 3)"
        )
    self.control_points = cp

    # Initialize weights for rational B-splines
    if weights is None:
        w = torch.ones(control_points_u, control_points_v, device=device)
    else:
        w = torch.as_tensor(weights, dtype=torch.get_default_dtype(), device=device)
        assert w.shape == (control_points_u, control_points_v), (
            f"weights must have shape ({control_points_u}, {control_points_v})"
        )
    self.weights = w

    self.param_model = "nurbs"
    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Initialize a NURBS phase surface from a parameter dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Requires "r" and "d"; optional keys include "mat2", "norm_radii", "control_points_u", "control_points_v", "degree_u", "degree_v", "control_points", and "weights".	required

Returns:

Name	Type	Description
obj	NURBSPhase	The constructed NURBS phase surface.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize a NURBS phase surface from a parameter dictionary.

    Args:
        surf_dict (dict): Surface parameters. Requires "r" and "d"; optional
            keys include "mat2", "norm_radii", "control_points_u",
            "control_points_v", "degree_u", "degree_v", "control_points",
            and "weights".

    Returns:
        obj (NURBSPhase): The constructed NURBS phase surface.
    """
    mat2 = surf_dict.get("mat2", "air")
    norm_radii = surf_dict.get("norm_radii", None)
    control_points_u = surf_dict.get("control_points_u", 8)
    control_points_v = surf_dict.get("control_points_v", 8)
    degree_u = surf_dict.get("degree_u", 3)
    degree_v = surf_dict.get("degree_v", 3)

    obj = cls(
        surf_dict["r"],
        surf_dict["d_next"],
        control_points_u=control_points_u,
        control_points_v=control_points_v,
        degree_u=degree_u,
        degree_v=degree_v,
        norm_radii=norm_radii,
        mat2=mat2,
    )

    # Load control points and weights
    control_points = surf_dict.get("control_points", None)
    if control_points is not None:
        obj.control_points = torch.as_tensor(control_points, device=obj.device)

    weights = surf_dict.get("weights", None)
    if weights is not None:
        obj.weights = torch.as_tensor(weights, device=obj.device)

    return obj

phi

phi(x, y)

Compute the reference phase map at the design wavelength.

Coordinates are normalized by norm_radii and mapped to the NURBS parameter domain [0, 1]. Points outside the unit circle (normalized radius greater than 1) are set to 0, and the result is wrapped to [0, 2π).

Parameters:

Name	Type	Description	Default
x	Tensor	x coordinates [mm].	required
y	Tensor	y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Phase [rad] in [0, 2π), same shape as x.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def phi(self, x, y):
    """Compute the reference phase map at the design wavelength.

    Coordinates are normalized by `norm_radii` and mapped to the NURBS
    parameter domain [0, 1]. Points outside the unit circle (normalized
    radius greater than 1) are set to 0, and the result is wrapped to [0, 2π).

    Args:
        x (torch.Tensor): x coordinates [mm].
        y (torch.Tensor): y coordinates [mm], same shape as x.

    Returns:
        phi (torch.Tensor): Phase [rad] in [0, 2π), same shape as x.
    """
    # Normalize coordinates to [0, 1] range for NURBS parameter space
    x_norm = (x / self.norm_radii + 1.0) / 2.0  # Map [-1, 1] to [0, 1]
    y_norm = (y / self.norm_radii + 1.0) / 2.0  # Map [-1, 1] to [0, 1]

    # Vectorized NURBS evaluation over all points (z-component is the phase).
    phi = self._evaluate_z_batch(x_norm.flatten(), y_norm.flatten()).reshape(
        x_norm.shape
    )

    # Apply circular aperture mask (set phase to 0 outside unit circle)
    r_squared = (x / self.norm_radii) ** 2 + (y / self.norm_radii) ** 2
    mask = r_squared > 1
    phi = torch.where(mask, torch.zeros_like(phi), phi)

    # Ensure phase is in [0, 2π) range
    phi = torch.remainder(phi, 2 * torch.pi)

    return phi

dphi_dxy

dphi_dxy(x, y)

Compute phase derivatives (dphi/dx, dphi/dy) by central differences.

Uses a finite-difference step of 1e-6 [mm] on phi. Points outside the unit circle (normalized radius greater than 1) are set to 0.

Parameters:

Name	Type	Description	Default
x	Tensor	x coordinates [mm].	required
y	Tensor	y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase derivative along x [rad/mm], same shape as x.
dphidy	Tensor	Phase derivative along y [rad/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def dphi_dxy(self, x, y):
    """Compute phase derivatives (dphi/dx, dphi/dy) by central differences.

    Uses a finite-difference step of 1e-6 [mm] on `phi`. Points outside the
    unit circle (normalized radius greater than 1) are set to 0.

    Args:
        x (torch.Tensor): x coordinates [mm].
        y (torch.Tensor): y coordinates [mm], same shape as x.

    Returns:
        dphidx (torch.Tensor): Phase derivative along x [rad/mm], same shape as x.
        dphidy (torch.Tensor): Phase derivative along y [rad/mm], same shape as x.
    """
    # For numerical differentiation, compute phi at slightly offset positions
    eps = 1e-6

    # Compute dphi/dx
    phi_x_plus = self.phi(x + eps, y)
    phi_x_minus = self.phi(x - eps, y)
    dphidx = (phi_x_plus - phi_x_minus) / (2 * eps)

    # Compute dphi/dy
    phi_y_plus = self.phi(x, y + eps)
    phi_y_minus = self.phi(x, y - eps)
    dphidy = (phi_y_plus - phi_y_minus) / (2 * eps)

    # Apply circular mask
    r_squared = (x / self.norm_radii) ** 2 + (y / self.norm_radii) ** 2
    mask = r_squared > 1
    dphidx = torch.where(mask, torch.zeros_like(dphidx), dphidx)
    dphidy = torch.where(mask, torch.zeros_like(dphidy), dphidy)

    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.01], optim_mat=False)

Build optimizer parameter groups for the NURBS control points.

Control points are always optimized at lrs[0]. If a second learning rate is given, the weights are also optimized at lrs[1].

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates [control_points_lr, weights_lr]. Defaults to [1e-4, 1e-2].	[0.0001, 0.01]
optim_mat	bool	Must be False; material parameters are not optimized for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	List of parameter-group dicts for torch.optim.

Raises:

Type	Description
AssertionError	If optim_mat is True.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def get_optimizer_params(self, lrs=[1e-4, 1e-2], optim_mat=False):
    """Build optimizer parameter groups for the NURBS control points.

    Control points are always optimized at `lrs[0]`. If a second learning
    rate is given, the weights are also optimized at `lrs[1]`.

    Args:
        lrs (list, optional): Learning rates [control_points_lr, weights_lr].
            Defaults to [1e-4, 1e-2].
        optim_mat (bool, optional): Must be False; material parameters are not
            optimized for phase surfaces. Defaults to False.

    Returns:
        params (list): List of parameter-group dicts for `torch.optim`.

    Raises:
        AssertionError: If optim_mat is True.
    """
    params = []

    # Enable gradients for control points (only z-coordinate for phase)
    self.control_points.requires_grad = True
    params.append({"params": [self.control_points], "lr": lrs[0]})

    # Optionally optimize weights
    if len(lrs) > 1:
        self.weights.requires_grad = True
        params.append({"params": [self.weights], "lr": lrs[1]})

    # We do not optimize material parameters for phase surface.
    assert optim_mat is False, (
        "Material parameters are not optimized for phase surface."
    )

    return params

save_ckpt

save_ckpt(save_path='./nurbs_doe.pth')

Save NURBS DOE parameters to a checkpoint file.

Parameters:

Name	Type	Description	Default
save_path	str	Output path. Defaults to "./nurbs_doe.pth".	'./nurbs_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def save_ckpt(self, save_path="./nurbs_doe.pth"):
    """Save NURBS DOE parameters to a checkpoint file.

    Args:
        save_path (str, optional): Output path. Defaults to "./nurbs_doe.pth".
    """
    torch.save(
        {
            "param_model": "nurbs",
            "control_points": self.control_points.clone().detach().cpu(),
            "weights": self.weights.clone().detach().cpu(),
            "control_points_u": self.control_points_u,
            "control_points_v": self.control_points_v,
            "degree_u": self.degree_u,
            "degree_v": self.degree_v,
            "knots_u": self.knots_u.clone().detach().cpu(),
            "knots_v": self.knots_v.clone().detach().cpu(),
        },
        save_path,
    )

load_ckpt

load_ckpt(load_path='./nurbs_doe.pth')

Load NURBS DOE parameters from a checkpoint file.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path. Defaults to "./nurbs_doe.pth".	'./nurbs_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def load_ckpt(self, load_path="./nurbs_doe.pth"):
    """Load NURBS DOE parameters from a checkpoint file.

    Args:
        load_path (str, optional): Checkpoint path. Defaults to "./nurbs_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.control_points_u = ckpt["control_points_u"]
    self.control_points_v = ckpt["control_points_v"]
    self.control_points = ckpt["control_points"].to(self.device)
    self.weights = ckpt["weights"].to(self.device)
    self.degree_u = ckpt["degree_u"]
    self.degree_v = ckpt["degree_v"]
    self.knots_u = ckpt["knots_u"].to(self.device)
    self.knots_v = ckpt["knots_v"].to(self.device)

surf_dict

surf_dict()

Return surface parameters as a serializable dictionary.

Returns:

Name	Type	Description
surf_dict	dict	Surface parameters (control points, weights, knot vectors, degrees, radii, distance, and material) suitable for JSON export.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py

def surf_dict(self):
    """Return surface parameters as a serializable dictionary.

    Returns:
        surf_dict (dict): Surface parameters (control points, weights, knot
            vectors, degrees, radii, distance, and material) suitable for
            JSON export.
    """
    surf_dict = {
        "type": "Phase",
        "r": self.r,
        "is_square": self.is_square,
        "param_model": "nurbs",
        "control_points": self.control_points.clone().detach().cpu().tolist(),
        "weights": self.weights.clone().detach().cpu().tolist(),
        "control_points_u": self.control_points_u,
        "control_points_v": self.control_points_v,
        "degree_u": self.degree_u,
        "degree_v": self.degree_v,
        "knots_u": self.knots_u.clone().detach().cpu().tolist(),
        "knots_v": self.knots_v.clone().detach().cpu().tolist(),
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    return surf_dict

deeplens.phase_surface.VortexPhase

VortexPhase(
    r,
    d_next,
    charge=1,
    f0=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
)

Bases: Phase

Vortex phase surface combining a spiral phase and an optional Fresnel lens.

The phase profile is

\[\phi(x, y) = \text{charge} \cdot \operatorname{atan2}(y, x) - \frac{\pi (x^2 + y^2)}{f_0}\]

where the first term imparts orbital angular momentum (topological charge charge) and the second term is a Fresnel focusing phase. f0 directly sets the phase curvature [mm²]; no design wavelength is assumed, so chromatic behaviour emerges naturally from the ray wavelength during diffraction. Setting f0=None disables the Fresnel term, leaving a pure vortex.

Attributes:

Name	Type	Description
charge	int	Topological charge of the spiral phase term.
f0	Tensor or None	Scalar Fresnel phase-curvature parameter [mm²], or None if the Fresnel term is disabled.
param_model	str	Parameterization tag, always "vortex".

Reference

Yu et al., "Light Propagation with Phase Discontinuities," Science 2011.

Initialize a vortex phase surface.

Parameters:

Name	Type	Description	Default
r	float	Aperture radius [mm].	required
d	float	Axial position [mm].	required
charge	int	Topological charge of the spiral term. Positive for a left-handed helix, negative for right-handed. Defaults to 1.	1
f0	float or None	Fresnel phase-curvature parameter [mm²], setting the quadratic term \(-\pi r^2 / f_0\). No design wavelength is assumed, so the implied focal length is \(f_0 / \lambda\) and varies with the ray wavelength \(\lambda\). None disables the Fresnel term. Defaults to None.	None
norm_radii	float or None	Normalization radius [mm] for phase-map display. Defaults to None, which falls back to r.	None
mat2	str	Material name after the surface. Defaults to "air".	'air'
pos_xy	tuple	(x, y) position offset in the global frame [mm]. Defaults to (0.0, 0.0).	(0.0, 0.0)
vec_local	tuple	Local surface normal direction (x, y, z). Defaults to (0.0, 0.0, 1.0).	(0.0, 0.0, 1.0)
is_square	bool	If True use a square aperture, else circular. Defaults to True.	True
device	str	Torch device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def __init__(
    self,
    r,
    d_next,
    charge=1,
    f0=None,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    """Initialize a vortex phase surface.

    Args:
        r (float): Aperture radius [mm].
        d (float): Axial position [mm].
        charge (int, optional): Topological charge of the spiral term. Positive for a
            left-handed helix, negative for right-handed. Defaults to 1.
        f0 (float or None, optional): Fresnel phase-curvature parameter [mm²], setting
            the quadratic term $-\\pi r^2 / f_0$. No design wavelength is assumed, so the
            implied focal length is $f_0 / \\lambda$ and varies with the ray wavelength
            $\\lambda$. None disables the Fresnel term. Defaults to None.
        norm_radii (float or None, optional): Normalization radius [mm] for phase-map
            display. Defaults to None, which falls back to `r`.
        mat2 (str, optional): Material name after the surface. Defaults to "air".
        pos_xy (tuple, optional): (x, y) position offset in the global frame [mm].
            Defaults to (0.0, 0.0).
        vec_local (tuple, optional): Local surface normal direction (x, y, z).
            Defaults to (0.0, 0.0, 1.0).
        is_square (bool, optional): If True use a square aperture, else circular.
            Defaults to True.
        device (str, optional): Torch device. Defaults to "cpu".
    """
    super().__init__(
        r=r,
        d_next=d_next,
        norm_radii=norm_radii,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.charge = int(charge)
    self.f0 = torch.tensor(float(f0)) if f0 is not None else None
    self.param_model = "vortex"
    self.to(device)

init_from_dict classmethod

init_from_dict(surf_dict)

Initialize a VortexPhase from a parameter dictionary.

Parameters:

Name	Type	Description	Default
surf_dict	dict	Surface parameters. Recognized keys: "r", "d" (required), and optionally "charge", "f0", "norm_radii", "mat2", "pos_xy", "vec_local", "is_square", "device".	required

Returns:

Name	Type	Description
surf	VortexPhase	The constructed vortex phase surface.

Source code in deeplens-src/deeplens/phase_surface/vortex.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize a VortexPhase from a parameter dictionary.

    Args:
        surf_dict (dict): Surface parameters. Recognized keys: "r", "d" (required),
            and optionally "charge", "f0", "norm_radii", "mat2", "pos_xy",
            "vec_local", "is_square", "device".

    Returns:
        surf (VortexPhase): The constructed vortex phase surface.
    """
    f0_raw = surf_dict.get("f0", None)
    return cls(
        r=surf_dict["r"],
        d_next=surf_dict["d_next"],
        charge=surf_dict.get("charge", 1),
        f0=f0_raw,
        norm_radii=surf_dict.get("norm_radii", None),
        mat2=surf_dict.get("mat2", "air"),
        pos_xy=surf_dict.get("pos_xy", [0.0, 0.0]),
        vec_local=surf_dict.get("vec_local", [0.0, 0.0, 1.0]),
        is_square=surf_dict.get("is_square", True),
        device=surf_dict.get("device", "cpu"),
    )

phi

phi(x, y)

Compute the phase map, wrapped to [0, 2π) [rad].

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
phi	Tensor	Wrapped phase [rad], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def phi(self, x, y):
    """Compute the phase map, wrapped to [0, 2π) [rad].

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], same shape as x.

    Returns:
        phi (torch.Tensor): Wrapped phase [rad], same shape as x.
    """
    phi = self.charge * torch.atan2(y, x)  # spiral term, in (-charge·π, charge·π]
    if self.f0 is not None:
        r2 = x * x + y * y
        phi = phi - torch.pi * r2 / self.f0
    phi = torch.remainder(phi, 2 * torch.pi)
    return phi

dphi_dxy

dphi_dxy(x, y)

Compute the analytical (unwrapped) phase gradient for generalized Snell's law.

Parameters:

Name	Type	Description	Default
x	Tensor	X coordinates [mm], any shape.	required
y	Tensor	Y coordinates [mm], same shape as x.	required

Returns:

Name	Type	Description
dphidx	Tensor	Phase derivative along x [rad/mm], same shape as x.
dphidy	Tensor	Phase derivative along y [rad/mm], same shape as x.

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def dphi_dxy(self, x, y):
    """Compute the analytical (unwrapped) phase gradient for generalized Snell's law.

    Args:
        x (torch.Tensor): X coordinates [mm], any shape.
        y (torch.Tensor): Y coordinates [mm], same shape as x.

    Returns:
        dphidx (torch.Tensor): Phase derivative along x [rad/mm], same shape as x.
        dphidy (torch.Tensor): Phase derivative along y [rad/mm], same shape as x.
    """
    r2 = x * x + y * y + EPSILON
    # d/dx [charge·atan2(y,x)] = -charge·y / r²
    # d/dy [charge·atan2(y,x)] =  charge·x / r²
    dphidx = self.charge * (-y / r2)
    dphidy = self.charge * (x / r2)
    if self.f0 is not None:
        scale = torch.pi / self.f0
        dphidx = dphidx - 2.0 * scale * x
        dphidy = dphidy - 2.0 * scale * y
    return dphidx, dphidy

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Return optimizer parameter groups.

Only f0 is differentiable; the topological charge charge is discrete and therefore never optimized. When f0 is None the returned list is empty.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates; only lrs[0] (for f0) is used. Defaults to [1e-4].	[0.0001]
optim_mat	bool	Must be False; material parameters are not optimized for phase surfaces. Defaults to False.	False

Returns:

Name	Type	Description
params	list	Optimizer parameter groups, one dict per optimized tensor.

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Return optimizer parameter groups.

    Only `f0` is differentiable; the topological charge `charge` is discrete and
    therefore never optimized. When `f0` is None the returned list is empty.

    Args:
        lrs (list, optional): Learning rates; only `lrs[0]` (for `f0`) is used.
            Defaults to [1e-4].
        optim_mat (bool, optional): Must be False; material parameters are not
            optimized for phase surfaces. Defaults to False.

    Returns:
        params (list): Optimizer parameter groups, one dict per optimized tensor.
    """
    assert not optim_mat, (
        "Material parameters are not optimized for phase surfaces."
    )
    params = []
    if self.f0 is not None:
        self.f0.requires_grad_(True)
        params.append({"params": [self.f0], "lr": lrs[0]})
    return params

save_ckpt

save_ckpt(save_path='./vortex_doe.pth')

Save VortexPhase parameters to a checkpoint file.

Parameters:

Name	Type	Description	Default
save_path	str	Output path. Defaults to "./vortex_doe.pth".	'./vortex_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def save_ckpt(self, save_path="./vortex_doe.pth"):
    """Save VortexPhase parameters to a checkpoint file.

    Args:
        save_path (str, optional): Output path. Defaults to "./vortex_doe.pth".
    """
    ckpt = {
        "param_model": self.param_model,
        "charge": self.charge,
        "f0": self.f0.clone().detach().cpu() if self.f0 is not None else None,
    }
    torch.save(ckpt, save_path)

load_ckpt

load_ckpt(load_path='./vortex_doe.pth')

Load VortexPhase parameters from a checkpoint file.

Parameters:

Name	Type	Description	Default
load_path	str	Checkpoint path. Defaults to "./vortex_doe.pth".	'./vortex_doe.pth'

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def load_ckpt(self, load_path="./vortex_doe.pth"):
    """Load VortexPhase parameters from a checkpoint file.

    Args:
        load_path (str, optional): Checkpoint path. Defaults to "./vortex_doe.pth".
    """
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.charge = int(ckpt["charge"])
    f0 = ckpt.get("f0")
    self.f0 = f0.to(self.device) if f0 is not None else None

surf_dict

surf_dict()

Return surface parameters as a serializable dictionary.

The "f0" key is included only when the Fresnel term is enabled.

Returns:

Name	Type	Description
d	dict	Surface parameters (type, geometry, charge, f0, material, etc.).

Source code in deeplens-src/deeplens/phase_surface/vortex.py

def surf_dict(self):
    """Return surface parameters as a serializable dictionary.

    The "f0" key is included only when the Fresnel term is enabled.

    Returns:
        d (dict): Surface parameters (type, geometry, charge, f0, material, etc.).
    """
    d = {
        "type": self.__class__.__name__,
        "r": self.r,
        "is_square": self.is_square,
        "param_model": self.param_model,
        "charge": self.charge,
        "norm_radii": round(self.norm_radii, 4),
        "d_next": round(self.d_next.item(), 4),
        "mat2": self.mat2.get_name(),
        "(mat2_n)": round(float(self.mat2.n), 4),
        "(mat2_V)": round(float(self.mat2.V), 4),
    }
    if self.f0 is not None:
        d["f0"] = round(self.f0.item(), 4)
    return d

Light

Geometric ray representation carrying origin, direction, wavelength, validity mask, energy, and optical path length (OPL).

deeplens.Ray

Ray(o, d, wvln, is_coherent=False, device='cpu')

Bases: DeepObj

Batched ray bundle for optical simulation.

Stores ray origins, directions, wavelength, validity mask, energy, bend penalty, and optical path length. Per-ray tensors share a batch shape, conventionally (..., num_rays), with trailing feature axes where noted below. Wavelength is a scalar. With no remaining batch axes, a single ray has origin and direction shapes (3,).

Attributes:

Name	Type	Description
o	Tensor	Ray origins, shape (..., num_rays, 3) [mm].
d	Tensor	Unit ray directions, shape (..., num_rays, 3).
wvln	Tensor	Wavelength scalar [µm].
shape	Size	Batch shape (..., num_rays) shared by the ray tensors.
is_valid	Tensor	Binary validity mask, shape (..., num_rays).
en	Tensor	Energy weight, shape (..., num_rays, 1).
stop_dist	Tensor	Per-ray distance from the physical aperture-stop centre in units of the stop radius, shape (..., num_rays). Initialized to inf; tracing records finite distances for rays valid at the stop and inf for invalid rays.
bend_penalty	Tensor	Accumulated per-surface bend penalty, shape (..., num_rays, 1).
opl	Tensor	Optical path length, shape (..., num_rays, 1) [mm]. Only accumulated when is_coherent is True.
is_coherent	bool	Whether optical path length tracking is enabled.
device	str | device	Compute device holding the ray tensors.

Initialize a ray object.

The direction d is normalized to unit length on construction. Auxiliary tensors (is_valid, en, bend_penalty, opl, stop_dist) are allocated over the batch shape using the origin's dtype and device.

Parameters:

Name	Type	Description	Default
o	Tensor	Ray origin, shape (..., num_rays, 3) [mm].	required
d	Tensor	Ray direction, shape (..., num_rays, 3). Normalized to unit length internally.	required
wvln	float	Ray wavelength [µm], must satisfy 0.1 < wvln < 10.0. Required and passed explicitly (the Lens carries primary_wvln/ wvln_rgb, not the Ray).	required
is_coherent	bool	Enable optical path length tracking for coherent tracing. Defaults to False.	False
device	str | device	Compute device. Defaults to "cpu".	'cpu'

Source code in deeplens-src/deeplens/light/ray.py

def __init__(self, o, d, wvln, is_coherent=False, device="cpu"):
    """Initialize a ray object.

    The direction `d` is normalized to unit length on construction. Auxiliary
    tensors (`is_valid`, `en`, `bend_penalty`, `opl`, `stop_dist`) are
    allocated over the batch shape using the origin's dtype and device.

    Args:
        o (torch.Tensor): Ray origin, shape `(..., num_rays, 3)` [mm].
        d (torch.Tensor): Ray direction, shape `(..., num_rays, 3)`.
            Normalized to unit length internally.
        wvln (float): Ray wavelength [µm], must satisfy 0.1 < wvln < 10.0.
            Required and passed explicitly (the Lens carries `primary_wvln`/
            `wvln_rgb`, not the Ray).
        is_coherent (bool, optional): Enable optical path length tracking for
            coherent tracing. Defaults to False.
        device (str | torch.device, optional): Compute device. Defaults to "cpu".
    """
    # Basic ray parameters. Directions and all auxiliary tensors inherit the
    # origin dtype so ray state does not depend on the process-wide default
    # after construction.
    self.o = torch.as_tensor(o, device=device)
    if not self.o.is_floating_point():
        self.o = self.o.to(torch.get_default_dtype())
    self.d = torch.as_tensor(d, device=device, dtype=self.o.dtype)
    super().__init__(dtype=self.o.dtype)
    self.shape = self.o.shape[:-1]

    # Wavelength
    assert wvln > 0.1 and wvln < 10.0, "Ray wavelength unit should be [um]"
    self.wvln = torch.as_tensor(wvln, device=device, dtype=self.o.dtype)

    # Auxiliary ray parameters - create directly on device
    self.is_valid = torch.ones(self.shape, device=device, dtype=self.o.dtype)
    self.en = torch.ones((*self.shape, 1), device=device, dtype=self.o.dtype)
    self.bend_penalty = torch.zeros(
        (*self.shape, 1), device=device, dtype=self.o.dtype
    )
    self.stop_dist = torch.full_like(self.is_valid, float("inf"))

    # Coherent ray tracing
    self.is_coherent = is_coherent  # bool
    self.opl = torch.zeros((*self.shape, 1), device=device, dtype=self.o.dtype)

    self.device = device
    self.d = F.normalize(self.d, p=2, dim=-1)

prop_to

prop_to(z, n=1.0)

Propagate the ray to a given depth plane in place.

Moves each valid ray origin along its direction toward the depth plane at axial coordinate \(z\). The denominator for nearly parallel rays is clamped to magnitude EPSILON, preserving its sign; their displacement is approximate. Invalid rays retain their origins and optical paths. In coherent mode, float64 rays are required and the optical path length is incremented by \(n \cdot t\), where \(t\) is the signed propagation distance.

Parameters:

Name	Type	Description	Default
z	float | Tensor	Target axial coordinate [mm], scalar or broadcastable to the ray batch shape.	required
n	float | Tensor	Refractive index, scalar or broadcastable to the ray batch shape. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
self	Ray	The updated ray (for chaining).

Raises:

Type	Description
ValueError	If coherent propagation is requested for non-float64 rays.

Source code in deeplens-src/deeplens/light/ray.py

def prop_to(self, z, n=1.0):
    """Propagate the ray to a given depth plane in place.

    Moves each valid ray origin along its direction toward the depth plane
    at axial coordinate $z$. The denominator for nearly parallel rays is
    clamped to magnitude `EPSILON`, preserving its sign; their displacement
    is approximate. Invalid rays retain their origins and optical paths.
    In coherent mode, float64 rays are required and the optical path length
    is incremented by $n \\cdot t$, where $t$ is the signed propagation
    distance.

    Args:
        z (float | torch.Tensor): Target axial coordinate [mm], scalar or
            broadcastable to the ray batch shape.
        n (float | torch.Tensor, optional): Refractive index, scalar or
            broadcastable to the ray batch shape. Defaults to 1.0.

    Returns:
        self (Ray): The updated ray (for chaining).

    Raises:
        ValueError: If coherent propagation is requested for non-float64 rays.
    """
    if self.is_coherent and self.o.dtype != torch.float64:
        raise ValueError("Coherent ray tracing requires float64 rays.")

    valid = self.is_valid > 0
    valid_mask = valid.unsqueeze(-1)
    direction = torch.where(valid_mask, self.d, 0.0)
    origin_z = torch.where(valid, self.o[..., 2], z)

    # Guard against rays (nearly) parallel to the target plane: d_z ~ 0 would
    # make t = inf/NaN and contaminate gradients through the torch.where below.
    dz = direction[..., 2]
    dz_safe = torch.where(dz < 0, dz.clamp(max=-EPSILON), dz.clamp(min=EPSILON))
    t = (z - origin_z) / dz_safe
    new_o = torch.where(valid_mask, self.o + direction * t.unsqueeze(-1), self.o)

    if self.is_coherent:
        new_opl = self.opl + (n * t).unsqueeze(-1)
        self.opl = torch.where(valid_mask, new_opl, self.opl)

    self.o = new_o
    return self

centroid

centroid(mode: Literal['geometric', 'chief_ray'] = 'geometric') -> torch.Tensor

Compute a geometric or chief-ray reference position for each field.

In chief-ray mode, selection precedes the final validity check. If the selected ray was clipped after the stop, or no finite stop distance was recorded, that field falls back to its geometric centroid. The sample selection is discrete; gradients through ray positions are preserved.

Parameters:

Name	Type	Description	Default
mode	str	"geometric" returns the energy-unweighted mean of valid ray origins. "chief_ray" returns the origin of the sampled ray with the smallest recorded stop_dist. Defaults to "geometric".	'geometric'

Returns:

Name	Type	Description
centroid	Tensor	Centroid position, shape (..., 3) [mm]. Fields with no valid rays return zero. A squeezed single ray returns its position if valid, or zero otherwise.

Raises:

Type	Description
ValueError	If mode is not "geometric" or "chief_ray".

Source code in deeplens-src/deeplens/light/ray.py

def centroid(
    self, mode: Literal["geometric", "chief_ray"] = "geometric"
) -> torch.Tensor:
    """Compute a geometric or chief-ray reference position for each field.

    In chief-ray mode, selection precedes the final validity check. If the
    selected ray was clipped after the stop, or no finite stop distance was
    recorded, that field falls back to its geometric centroid. The sample
    selection is discrete; gradients through ray positions are preserved.

    Args:
        mode (str): ``"geometric"`` returns the energy-unweighted mean of
            valid ray origins. ``"chief_ray"`` returns the origin of the
            sampled ray with the smallest recorded ``stop_dist``. Defaults
            to ``"geometric"``.

    Returns:
        centroid (torch.Tensor): Centroid position, shape `(..., 3)` [mm].
            Fields with no valid rays return zero. A squeezed single ray
            returns its position if valid, or zero otherwise.

    Raises:
        ValueError: If ``mode`` is not ``"geometric"`` or ``"chief_ray"``.
    """
    if mode not in ("geometric", "chief_ray"):
        raise ValueError(f"Unsupported centroid mode: {mode}.")

    valid_o = torch.where((self.is_valid > 0).unsqueeze(-1), self.o, 0.0)
    if self.o.ndim == 1:
        return valid_o

    geometric_centroid = (valid_o * self.is_valid.unsqueeze(-1)).sum(
        -2
    ) / self.is_valid.sum(-1).add(EPSILON).unsqueeze(-1)
    if mode == "geometric" or self.o.shape[-2] == 0:
        return geometric_centroid

    stop_dist = self.stop_dist.masked_fill(self.stop_dist.isnan(), float("inf"))
    min_dist, sample_index = stop_dist.min(dim=-1, keepdim=True)
    index = sample_index.unsqueeze(-1).expand(*sample_index.shape, 3)
    chief_centroid = self.o.gather(-2, index).squeeze(-2)
    chief_valid = torch.isfinite(min_dist) & (
        self.is_valid.gather(-1, sample_index) > 0
    )
    return torch.where(chief_valid, chief_centroid, geometric_centroid)

rms_error

rms_error(center_ref=None)

Compute the mean RMS spot radius over valid rays.

For each field, the RMS radius is computed from the in-plane (x, y) deviation of valid ray origins about center_ref, then averaged across fields. A shifted square root keeps zero-radius gradients finite. Fields with no valid rays contribute inf.

Parameters:

Name	Type	Description	Default
center_ref	Tensor	Reference center, shape (..., 3) [mm]. If None, the per-field geometric centroid is used without differentiating through its calculation. Defaults to None.	None

Returns:

Name	Type	Description
rms_error	Tensor	Scalar mean RMS spot radius [mm].

Source code in deeplens-src/deeplens/light/ray.py

def rms_error(self, center_ref=None):
    """Compute the mean RMS spot radius over valid rays.

    For each field, the RMS radius is computed from the in-plane (x, y)
    deviation of valid ray origins about `center_ref`, then averaged across
    fields. A shifted square root keeps zero-radius gradients finite.
    Fields with no valid rays contribute `inf`.

    Args:
        center_ref (torch.Tensor, optional): Reference center, shape `(..., 3)`
            [mm]. If None, the per-field geometric centroid is used without
            differentiating through its calculation. Defaults to None.

    Returns:
        rms_error (torch.Tensor): Scalar mean RMS spot radius [mm].
    """
    # Calculate the centroid of the ray as reference
    if center_ref is None:
        with torch.no_grad():
            center_ref = self.centroid()

    center_ref = center_ref.unsqueeze(-2)

    # Calculate RMS error for each region
    offset = torch.where(
        (self.is_valid > 0).unsqueeze(-1),
        self.o[..., :2] - center_ref[..., :2],
        0.0,
    )
    squared_radius = offset.square().sum(-1)
    valid_count = self.is_valid.sum(-1)
    mean_squared_radius = (squared_radius * self.is_valid).sum(
        -1
    ) / valid_count.clamp_min(1)

    # ``sqrt(0)`` has an infinite derivative and yielded NaN gradients for
    # coincident or all-invalid bundles. The shifted safe square root keeps
    # an exact zero value with a finite zero gradient. A bundle with no valid
    # rays is an invalid optical result, not a perfect zero-RMS spot.
    epsilon = torch.as_tensor(EPSILON, device=self.o.device, dtype=self.o.dtype)
    rms_error = torch.sqrt(mean_squared_radius + epsilon) - torch.sqrt(epsilon)
    rms_error = torch.where(
        valid_count > 0,
        rms_error,
        torch.full_like(rms_error, float("inf")),
    )

    # Average RMS error
    return rms_error.mean()

flip_xy

flip_xy()

Negate the x and y components of ray origins and directions in place.

Used when computing the point spread function and wavefront distribution.

Returns:

Name	Type	Description
self	Ray	The updated ray (for chaining).

Source code in deeplens-src/deeplens/light/ray.py

def flip_xy(self):
    """Negate the x and y components of ray origins and directions in place.

    Used when computing the point spread function and wavefront distribution.

    Returns:
        self (Ray): The updated ray (for chaining).
    """
    self.o = torch.cat([-self.o[..., :2], self.o[..., 2:]], dim=-1)
    self.d = torch.cat([-self.d[..., :2], self.d[..., 2:]], dim=-1)
    return self

clone

clone(device=None)

Copy ray tensors and tracing state, optionally to another device.

Tensor storage is independent; autograd connections are preserved. This copies the defined ray state rather than arbitrary attributes added by callers.

Parameters:

Name	Type	Description	Default
device	str | device | None	Target device for the clone. If None, the source device is used. Defaults to None.	None

Returns:

Name	Type	Description
ray	Ray	A new ray with cloned tensors on the target device.

Source code in deeplens-src/deeplens/light/ray.py

def clone(self, device=None):
    """Copy ray tensors and tracing state, optionally to another device.

    Tensor storage is independent; autograd connections are preserved.
    This copies the defined ray state rather than arbitrary attributes
    added by callers.

    Args:
        device (str | torch.device | None, optional): Target device for the
            clone. If None, the source device is used. Defaults to None.

    Returns:
        ray (Ray): A new ray with cloned tensors on the target device.
    """
    target_device = self.device if device is None else device

    ray = Ray.__new__(Ray)
    ray.o = self.o.clone().to(target_device)
    ray.d = self.d.clone().to(target_device)
    ray.wvln = self.wvln.clone().to(target_device)
    ray.is_valid = self.is_valid.clone().to(target_device)
    ray.en = self.en.clone().to(target_device)
    ray.bend_penalty = self.bend_penalty.clone().to(target_device)
    ray.opl = self.opl.clone().to(target_device)
    ray.stop_dist = self.stop_dist.clone().to(target_device)

    ray.is_coherent = self.is_coherent
    ray.device = torch.device(target_device)
    ray.dtype = ray.o.dtype
    ray.shape = ray.o.shape[:-1]
    return ray

squeeze

squeeze(dim=0)

Squeeze the leading batch dimension of all ray tensors in place.

Only the leading batch axis is supported: it is the one axis shared by every ray tensor regardless of its trailing feature axis, so trailing feature axes and the scalar wavelength are preserved. shape is updated. Squeezing a non-singleton axis is a no-op.

Parameters:

Name	Type	Description	Default
dim	int	Batch dimension to squeeze. Only 0 is supported. Defaults to 0.	0

Returns:

Name	Type	Description
self	Ray	The updated ray (for chaining).

Raises:

Type	Description
ValueError	If dim is not 0.

Source code in deeplens-src/deeplens/light/ray.py

def squeeze(self, dim=0):
    """Squeeze the leading batch dimension of all ray tensors in place.

    Only the leading batch axis is supported: it is the one axis shared by
    every ray tensor regardless of its trailing feature axis, so trailing
    feature axes and the scalar wavelength are preserved. `shape` is
    updated. Squeezing a non-singleton axis is a no-op.

    Args:
        dim (int, optional): Batch dimension to squeeze. Only 0 is
            supported. Defaults to 0.

    Returns:
        self (Ray): The updated ray (for chaining).

    Raises:
        ValueError: If `dim` is not 0.
    """
    if dim != 0:
        raise ValueError(f"Unsupported squeeze dimension: {dim}, expected 0.")

    self.o = self.o.squeeze(0)
    self.d = self.d.squeeze(0)
    self.is_valid = self.is_valid.squeeze(0)
    self.en = self.en.squeeze(0)
    self.opl = self.opl.squeeze(0)
    self.bend_penalty = self.bend_penalty.squeeze(0)
    self.stop_dist = self.stop_dist.squeeze(0)
    self.shape = self.o.shape[:-1]
    return self

unsqueeze

unsqueeze(dim=0)

Insert a leading size-1 batch dimension into all ray tensors in place.

Only the leading batch axis is supported: it is the one axis shared by every ray tensor regardless of its trailing feature axis, so trailing feature axes and the scalar wavelength are preserved. shape is updated.

Parameters:

Name	Type	Description	Default
dim	int	Position at which to insert the batch dimension. Only 0 is supported. Defaults to 0.	0

Returns:

Name	Type	Description
self	Ray	The updated ray (for chaining).

Raises:

Type	Description
ValueError	If dim is not 0.

Source code in deeplens-src/deeplens/light/ray.py

def unsqueeze(self, dim=0):
    """Insert a leading size-1 batch dimension into all ray tensors in place.

    Only the leading batch axis is supported: it is the one axis shared by
    every ray tensor regardless of its trailing feature axis, so trailing
    feature axes and the scalar wavelength are preserved. `shape` is
    updated.

    Args:
        dim (int, optional): Position at which to insert the batch
            dimension. Only 0 is supported. Defaults to 0.

    Returns:
        self (Ray): The updated ray (for chaining).

    Raises:
        ValueError: If `dim` is not 0.
    """
    if dim != 0:
        raise ValueError(f"Unsupported unsqueeze dimension: {dim}, expected 0.")

    self.o = self.o.unsqueeze(0)
    self.d = self.d.unsqueeze(0)
    self.is_valid = self.is_valid.unsqueeze(0)
    self.en = self.en.unsqueeze(0)
    self.opl = self.opl.unsqueeze(0)
    self.bend_penalty = self.bend_penalty.unsqueeze(0)
    self.stop_dist = self.stop_dist.unsqueeze(0)
    self.shape = self.o.shape[:-1]
    return self

Complex electromagnetic field with Angular Spectrum Method (ASM), Fresnel, and Fraunhofer propagation via torch.fft.

deeplens.ComplexWave

ComplexWave(u=None, wvln=0.55, z=0.0, phy_size=(4.0, 4.0), res=(2000, 2000))

Bases: DeepObj

Complex scalar wave field for diffraction simulation.

Represents a monochromatic, coherent complex amplitude on a uniform rectangular grid. Propagation methods (band-limited ASM, Fresnel) are implemented as member functions and use torch.fft for efficiency.

Attributes:

Name	Type	Description
u	Tensor	Complex amplitude, shape [1, 1, H, W].
wvln	float	Wavelength [µm].
k	float	Wave number \(2\pi / (\lambda \times 10^{-3})\) [mm⁻¹].
phy_size	tuple	Physical aperture size (W, H) [mm].
ps	float	Pixel pitch [mm] (square pixels).
res	tuple	Grid resolution (H, W) in pixels.
x	Tensor	x coordinate grid, shape [H, W][mm].
y	Tensor	y coordinate grid, shape [H, W][mm].
z	Tensor	Axial position grid, shape [H, W][mm].
plain_asm_dist_max	float	Nyquist limit of plain ASM [mm] (reference only).
fresnel_dist_min	float	Distance above which single-FFT Fresnel is well-sampled [mm].

Initialize a complex wave field.

Parameters:

Name	Type	Description	Default
u	Tensor or None	Initial complex amplitude. Accepted shapes: [H, W], [1, H, W], or [1, 1, H, W]. If None, a zero field is created with the given res. Defaults to None.	None
wvln	float	Wavelength [µm]. Defaults to 0.55.	0.55
z	float	Initial axial position [mm]. Defaults to 0.0.	0.0
phy_size	tuple	Physical aperture (W, H) [mm]. Defaults to (4.0, 4.0).	(4.0, 4.0)
res	tuple	Grid resolution (H, W) [pixels]. Only used when u is None. Defaults to (2000, 2000).	(2000, 2000)

Raises:

Type	Description
AssertionError	If the pixel pitch is not square or the wavelength is outside the range (0.1, 10) µm.

Source code in deeplens-src/deeplens/light/wave.py

def __init__(
    self,
    u=None,
    wvln=0.55,
    z=0.0,
    phy_size=(4.0, 4.0),
    res=(2000, 2000),
):
    """Initialize a complex wave field.

    Args:
        u (torch.Tensor or None, optional): Initial complex amplitude.
            Accepted shapes: [H, W], [1, H, W], or [1, 1, H, W]. If None,
            a zero field is created with the given res. Defaults to None.
        wvln (float, optional): Wavelength [µm]. Defaults to 0.55.
        z (float, optional): Initial axial position [mm]. Defaults to 0.0.
        phy_size (tuple, optional): Physical aperture (W, H) [mm].
            Defaults to (4.0, 4.0).
        res (tuple, optional): Grid resolution (H, W) [pixels]. Only used
            when u is None. Defaults to (2000, 2000).

    Raises:
        AssertionError: If the pixel pitch is not square or the wavelength
            is outside the range (0.1, 10) µm.
    """
    if u is not None:
        if not u.dtype == torch.complex128:
            print(
                "A complex wave field is created with single precision. "
                "In the future, we want to always use double precision."
            )

        self.u = u if torch.is_tensor(u) else torch.from_numpy(u)
        if not self.u.is_complex():
            self.u = self.u.to(torch.complex64)

        # [H, W] or [1, H, W] to [1, 1, H, W]
        if len(u.shape) == 2:
            self.u = u.unsqueeze(0).unsqueeze(0)
        elif len(self.u.shape) == 3:
            self.u = self.u.unsqueeze(0)

        self.res = self.u.shape[-2:]

    else:
        # Initialize a zero complex wave field
        amp = torch.zeros(res).unsqueeze(0).unsqueeze(0)
        phi = torch.zeros(res).unsqueeze(0).unsqueeze(0)
        self.u = amp + 1j * phi
        self.res = res

    # Wave field parameters
    assert wvln > 0.1 and wvln < 10.0, "Wavelength should be in [um]."
    self.wvln = wvln  # [um], wavelength
    self.k = 2 * torch.pi / (self.wvln * 1e-3)  # [mm^-1], wave number

    # Physical size and pixel size
    self.phy_size = phy_size  # [mm], physical size
    px = phy_size[0] / self.res[0]
    py = phy_size[1] / self.res[1]
    assert abs(px - py) <= 1e-9 * max(abs(px), abs(py)) + 1e-12, (
        "Pixel size is not square."
    )
    self.ps = phy_size[0] / self.res[0]  # [mm], pixel size

    # Wave field grid
    self.x, self.y = self.gen_xy_grid()  # x, y grid
    self.z = torch.full_like(self.x, z)  # z grid

    # Cached reference distances (depend only on wvln, ps, phy_size).
    # plain_asm_dist_max: Nyquist limit of plain ASM. prop() uses band-limited
    #   ASM, which stays valid past this, so it is kept only for reference.
    # fresnel_dist_min: distance above which single-FFT Fresnel is well-sampled.
    self.plain_asm_dist_max = Nyquist_ASM_zmax(
        wvln=self.wvln, ps=self.ps, side_length=self.phy_size[0]
    )
    self.fresnel_dist_min = Fresnel_zmin(
        wvln=self.wvln, ps=self.ps, side_length=self.phy_size[0]
    )

point_wave classmethod

point_wave(
    point=(0.0, 0.0, -1000.0),
    wvln=0.55,
    z=0.0,
    phy_size=(4.0, 4.0),
    res=(2000, 2000),
    valid_r=None,
)

Create a spherical wave field on the x0y plane from a point source.

The phase is \(\pm k r\) where \(r\) is the distance from the source to each grid point; the sign is positive for a diverging wave (source behind the plane, \(z_{src} < z\)) and negative otherwise. The amplitude is normalized to \(r_{min} / r\).

Parameters:

Name	Type	Description	Default
point	tuple	Point source position (x, y, z) in object space [mm]. Defaults to (0.0, 0.0, -1000.0).	(0.0, 0.0, -1000.0)
wvln	float	Wavelength [µm]. Defaults to 0.55.	0.55
z	float	Field z position [mm]. Defaults to 0.0.	0.0
phy_size	tuple	Physical size (W, H) of the x0y plane [mm]. Defaults to (4.0, 4.0).	(4.0, 4.0)
res	tuple	Grid resolution (H, W) [pixels]. Defaults to (2000, 2000).	(2000, 2000)
valid_r	float or None	If set, zero the field outside a circle of this radius [mm], e.g. a lens aperture. Defaults to None.	None

Returns:

Name	Type	Description
field	ComplexWave	Complex field on the x0y plane.

Source code in deeplens-src/deeplens/light/wave.py

@classmethod
def point_wave(
    cls,
    point=(0.0, 0.0, -1000.0),
    wvln=0.55,
    z=0.0,
    phy_size=(4.0, 4.0),
    res=(2000, 2000),
    valid_r=None,
):
    """Create a spherical wave field on the x0y plane from a point source.

    The phase is $\\pm k r$ where $r$ is the distance from the source to
    each grid point; the sign is positive for a diverging wave (source
    behind the plane, $z_{src} < z$) and negative otherwise. The amplitude
    is normalized to $r_{min} / r$.

    Args:
        point (tuple, optional): Point source position (x, y, z) in object
            space [mm]. Defaults to (0.0, 0.0, -1000.0).
        wvln (float, optional): Wavelength [µm]. Defaults to 0.55.
        z (float, optional): Field z position [mm]. Defaults to 0.0.
        phy_size (tuple, optional): Physical size (W, H) of the x0y plane
            [mm]. Defaults to (4.0, 4.0).
        res (tuple, optional): Grid resolution (H, W) [pixels]. Defaults to
            (2000, 2000).
        valid_r (float or None, optional): If set, zero the field outside a
            circle of this radius [mm], e.g. a lens aperture. Defaults to None.

    Returns:
        field (ComplexWave): Complex field on the x0y plane.
    """
    assert wvln > 0.1 and wvln < 10.0, "Wavelength should be in [um]."
    k = 2 * torch.pi / (wvln * 1e-3)  # [mm^-1], wave number

    # Create meshgrid on target plane
    x, y = torch.meshgrid(
        torch.linspace(
            -0.5 * phy_size[0], 0.5 * phy_size[0], res[0], dtype=torch.float64
        ),
        torch.linspace(
            0.5 * phy_size[1], -0.5 * phy_size[1], res[1], dtype=torch.float64
        ),
        indexing="xy",
    )

    # Calculate distance to point source, and calculate spherical wave phase
    # Add EPSILON inside the sqrt so r is never exactly 0 (avoids 1/r blow-up
    # and r.min()->0 when the source lies on the plane and on a grid node).
    r = torch.sqrt(
        (x - point[0]) ** 2 + (y - point[1]) ** 2 + (z - point[2]) ** 2 + EPSILON
    )
    if point[2] < z:
        phi = k * r
    else:
        phi = -k * r
    u = (r.min() / r) * torch.exp(1j * phi)

    # Apply valid circle if provided, e.g., the aperture of a lens
    if valid_r is not None:
        mask = (x - point[0]) ** 2 + (y - point[1]) ** 2 < valid_r**2
        u = u * mask

    # Create wave field
    return cls(u=u, wvln=wvln, phy_size=phy_size, res=res, z=z)

plane_wave classmethod

plane_wave(
    wvln=0.55,
    z=0.0,
    phy_size=(4.0, 4.0),
    res=(2000, 2000),
    theta_x=0.0,
    theta_y=0.0,
    valid_r=None,
)

Create a planar wave field on the x0y plane.

With theta_x = theta_y = 0 the result is a uniform unit-amplitude plane wave travelling along +z. Non-zero angles produce a tilted (obliquely incident / off-axis) plane wave whose wavevector makes the given angles with the optical axis; this adds the linear phase ramp \(\exp(i k (x \sin\theta_x + y \sin\theta_y))\) while the amplitude stays uniform.

Parameters:

Name	Type	Description	Default
wvln	float	Wavelength [µm]. Defaults to 0.55.	0.55
z	float	Field z position [mm]. Defaults to 0.0.	0.0
phy_size	tuple	Physical size (W, H) of the field [mm]. Defaults to (4.0, 4.0).	(4.0, 4.0)
res	tuple	Grid resolution (H, W) [pixels]. Defaults to (2000, 2000).	(2000, 2000)
theta_x	float	Tilt angle of the wavevector in the x-z plane [rad]. Defaults to 0.0.	0.0
theta_y	float	Tilt angle of the wavevector in the y-z plane [rad]. Defaults to 0.0.	0.0
valid_r	float or None	If set, zero the field outside a circle of this radius [mm]. Defaults to None.	None

Returns:

Name	Type	Description
field	ComplexWave	Complex field.

Source code in deeplens-src/deeplens/light/wave.py

@classmethod
def plane_wave(
    cls,
    wvln=0.55,
    z=0.0,
    phy_size=(4.0, 4.0),
    res=(2000, 2000),
    theta_x=0.0,
    theta_y=0.0,
    valid_r=None,
):
    """Create a planar wave field on the x0y plane.

    With theta_x = theta_y = 0 the result is a uniform unit-amplitude plane
    wave travelling along +z. Non-zero angles produce a tilted (obliquely
    incident / off-axis) plane wave whose wavevector makes the given angles
    with the optical axis; this adds the linear phase ramp
    $\\exp(i k (x \\sin\\theta_x + y \\sin\\theta_y))$ while the amplitude
    stays uniform.

    Args:
        wvln (float, optional): Wavelength [µm]. Defaults to 0.55.
        z (float, optional): Field z position [mm]. Defaults to 0.0.
        phy_size (tuple, optional): Physical size (W, H) of the field [mm].
            Defaults to (4.0, 4.0).
        res (tuple, optional): Grid resolution (H, W) [pixels]. Defaults to
            (2000, 2000).
        theta_x (float, optional): Tilt angle of the wavevector in the x-z
            plane [rad]. Defaults to 0.0.
        theta_y (float, optional): Tilt angle of the wavevector in the y-z
            plane [rad]. Defaults to 0.0.
        valid_r (float or None, optional): If set, zero the field outside a
            circle of this radius [mm]. Defaults to None.

    Returns:
        field (ComplexWave): Complex field.
    """
    assert wvln > 0.1 and wvln < 10.0, "Wavelength should be in [um]."

    # Create a plane wave field
    if theta_x == 0.0 and theta_y == 0.0:
        # On-axis: uniform unit-amplitude field.
        u = torch.ones(res, dtype=torch.float64) + 0j
    else:
        # Off-axis: tilted plane wave, i.e. a linear phase ramp.
        k = 2 * torch.pi / (wvln * 1e-3)  # [mm^-1], wave number
        x, y = torch.meshgrid(
            torch.linspace(
                -0.5 * phy_size[0], 0.5 * phy_size[0], res[0], dtype=torch.float64
            ),
            torch.linspace(
                0.5 * phy_size[1], -0.5 * phy_size[1], res[1], dtype=torch.float64
            ),
            indexing="xy",
        )
        u = torch.exp(1j * k * (x * math.sin(theta_x) + y * math.sin(theta_y)))

    # Apply valid circle if provided
    if valid_r is not None:
        x, y = torch.meshgrid(
            torch.linspace(-0.5 * phy_size[0], 0.5 * phy_size[0], res[0]),
            torch.linspace(-0.5 * phy_size[1], 0.5 * phy_size[1], res[1]),
            indexing="xy",
        )
        mask = (x**2 + y**2) < valid_r**2
        u = u * mask

    # Create wave field
    return cls(u=u, phy_size=phy_size, wvln=wvln, res=res, z=z)

image_wave classmethod

image_wave(img, wvln=0.55, z=0.0, phy_size=(4.0, 4.0))

Initialize a complex wave field from an image.

The image is interpreted as intensity in range [0, 1]; the field amplitude is its square root and the phase is zero.

Parameters:

Name	Type	Description	Default
img	Tensor	Input image, shape [H, W] or [B, C, H, W], data range [0, 1], dtype float32.	required
wvln	float	Wavelength [µm]. Defaults to 0.55.	0.55
z	float	Field z position [mm]. Defaults to 0.0.	0.0
phy_size	tuple	Physical size (W, H) of the field [mm]. Defaults to (4.0, 4.0).	(4.0, 4.0)

Returns:

Name	Type	Description
field	ComplexWave	Complex field.

Source code in deeplens-src/deeplens/light/wave.py

@classmethod
def image_wave(cls, img, wvln=0.55, z=0.0, phy_size=(4.0, 4.0)):
    """Initialize a complex wave field from an image.

    The image is interpreted as intensity in range [0, 1]; the field
    amplitude is its square root and the phase is zero.

    Args:
        img (torch.Tensor): Input image, shape [H, W] or [B, C, H, W], data
            range [0, 1], dtype float32.
        wvln (float, optional): Wavelength [µm]. Defaults to 0.55.
        z (float, optional): Field z position [mm]. Defaults to 0.0.
        phy_size (tuple, optional): Physical size (W, H) of the field [mm].
            Defaults to (4.0, 4.0).

    Returns:
        field (ComplexWave): Complex field.
    """
    assert img.dtype == torch.float32, "Image must be float32."

    amp = torch.sqrt(img)
    phi = torch.zeros_like(amp)
    u = amp + 1j * phi

    return cls(u=u, wvln=wvln, phy_size=phy_size, res=u.shape[-2:], z=z)

prop

prop(prop_dist, n=1.0)

Propagate the field forward by prop_dist and update self.

Selects the diffraction method from the propagation distance: a near-zero distance is a no-op; sub-wavelength distances raise (not implemented); distances up to fresnel_dist_min use band-limited ASM; larger distances use single-FFT Fresnel diffraction. The axial grid z is advanced by prop_dist.

Parameters:

Name	Type	Description	Default
prop_dist	float	Propagation distance [mm].	required
n	float	Refractive index of the medium. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
self	ComplexWave	The propagated wave field (for chaining).

Raises:

Type	Description
Exception	If the propagation distance is in the sub-wavelength range (full-wave methods such as FDTD are not implemented).

Reference

[1] Modeling and propagation of near-field diffraction patterns: A more complete approach. Table 1. [2] https://github.com/kaanaksit/odak/blob/master/odak/wave/classical.py [3] https://spie.org/samples/PM103.pdf [4] "Non-approximated Rayleigh Sommerfeld diffraction integral: advantages and disadvantages in the propagation of complex wave fields"

Source code in deeplens-src/deeplens/light/wave.py

def prop(self, prop_dist, n=1.0):
    """Propagate the field forward by `prop_dist` and update `self`.

    Selects the diffraction method from the propagation distance: a
    near-zero distance is a no-op; sub-wavelength distances raise (not
    implemented); distances up to `fresnel_dist_min` use band-limited ASM;
    larger distances use single-FFT Fresnel diffraction. The axial grid `z`
    is advanced by `prop_dist`.

    Args:
        prop_dist (float): Propagation distance [mm].
        n (float, optional): Refractive index of the medium. Defaults to 1.0.

    Returns:
        self (ComplexWave): The propagated wave field (for chaining).

    Raises:
        Exception: If the propagation distance is in the sub-wavelength
            range (full-wave methods such as FDTD are not implemented).

    Reference:
        [1] Modeling and propagation of near-field diffraction patterns: A more complete approach. Table 1.
        [2] https://github.com/kaanaksit/odak/blob/master/odak/wave/classical.py
        [3] https://spie.org/samples/PM103.pdf
        [4] "Non-approximated Rayleigh Sommerfeld diffraction integral: advantages and disadvantages in the propagation of complex wave fields"
    """
    # Determine propagation method using cached boundaries
    wvln_mm = self.wvln * 1e-3  # [um] to [mm]

    # Wave propagation methods
    if prop_dist < DELTA:
        # Zero distance: do nothing
        pass

    elif prop_dist < wvln_mm:
        # Sub-wavelength distance: full wave method (e.g., FDTD)
        raise Exception(
            "The propagation distance in sub-wavelength range is not implemented yet. "
            "Have to use full wave method (e.g., FDTD)."
        )

    elif prop_dist <= self.fresnel_dist_min:
        # Band-limited ASM (Matsushima & Shimobaba 2009): rigorous angular
        # spectrum with a band-limit that suppresses aliasing. Valid across
        # the near and intermediate fields, so it covers the former gap
        # between the Nyquist-ASM and Fresnel regimes.
        self.u = BandLimitedASM(
            self.u, z=prop_dist, wvln=self.wvln, ps=self.ps, n=n
        )

    else:
        # Fresnel diffraction (far field)
        self.u = FresnelDiffraction(
            self.u, z=prop_dist, wvln=self.wvln, ps=self.ps, n=n
        )

    # Update z grid
    self.z += prop_dist
    return self

prop_to

prop_to(z, n=1)

Propagate the field to the absolute plane z and update self.

Computes the relative distance from the current axial position and delegates to prop.

Parameters:

Name	Type	Description	Default
z	float	Destination plane z coordinate [mm].	required
n	float	Refractive index of the medium. Defaults to 1.	1

Returns:

Name	Type	Description
self	ComplexWave	The propagated wave field (for chaining).

Source code in deeplens-src/deeplens/light/wave.py

def prop_to(self, z, n=1):
    """Propagate the field to the absolute plane `z` and update `self`.

    Computes the relative distance from the current axial position and
    delegates to `prop`.

    Args:
        z (float): Destination plane z coordinate [mm].
        n (float, optional): Refractive index of the medium. Defaults to 1.

    Returns:
        self (ComplexWave): The propagated wave field (for chaining).
    """
    # Use float() instead of .item() to avoid GPU-CPU sync on CUDA tensors
    # (self.z is a full grid but all values are identical; [0,0] is representative)
    prop_dist = float(z) - float(self.z[0, 0])
    self.prop(prop_dist, n=n)
    return self

gen_xy_grid

gen_xy_grid()

Generate the x and y coordinate grids, shape [H, W].

x runs along the width (res[1] columns, extent phy_size[0]) and y along the height (res[0] rows, extent phy_size[1]), consistent with point_wave / plane_wave. With indexing="xy" the outputs have shape (len(y_1d), len(x_1d)) = (H, W).

Returns:

Name	Type	Description
x	Tensor	x coordinate grid, shape [H, W][mm].
y	Tensor	y coordinate grid, shape [H, W][mm].

Source code in deeplens-src/deeplens/light/wave.py

def gen_xy_grid(self):
    """Generate the x and y coordinate grids, shape [H, W].

    x runs along the width (res[1] columns, extent phy_size[0]) and y along
    the height (res[0] rows, extent phy_size[1]), consistent with
    `point_wave` / `plane_wave`. With indexing="xy" the outputs have shape
    (len(y_1d), len(x_1d)) = (H, W).

    Returns:
        x (torch.Tensor): x coordinate grid, shape [H, W] [mm].
        y (torch.Tensor): y coordinate grid, shape [H, W] [mm].
    """
    x, y = torch.meshgrid(
        torch.linspace(
            -0.5 * self.phy_size[0], 0.5 * self.phy_size[0], self.res[1]
        ),
        torch.linspace(
            0.5 * self.phy_size[1], -0.5 * self.phy_size[1], self.res[0]
        ),
        indexing="xy",
    )
    return x, y

gen_freq_grid

gen_freq_grid()

Generate the spatial-frequency grids, shape [H, W].

Returns:

Name	Type	Description
fx	Tensor	x-frequency grid, shape [H, W][mm⁻¹].
fy	Tensor	y-frequency grid, shape [H, W][mm⁻¹].

Source code in deeplens-src/deeplens/light/wave.py

def gen_freq_grid(self):
    """Generate the spatial-frequency grids, shape [H, W].

    Returns:
        fx (torch.Tensor): x-frequency grid, shape [H, W] [mm⁻¹].
        fy (torch.Tensor): y-frequency grid, shape [H, W] [mm⁻¹].
    """
    x, y = self.gen_xy_grid()
    fx = x / (self.ps * self.phy_size[0])
    fy = y / (self.ps * self.phy_size[1])
    return fx, fy

load

load(filepath)

Load a wave field from file (only .npz is supported).

Parameters:

Name	Type	Description	Default
filepath	str	Path to the file to load.	required

Raises:

Type	Description
Exception	If the file format is not supported.

Source code in deeplens-src/deeplens/light/wave.py

def load(self, filepath):
    """Load a wave field from file (only `.npz` is supported).

    Args:
        filepath (str): Path to the file to load.

    Raises:
        Exception: If the file format is not supported.
    """
    if filepath.endswith(".npz"):
        self.load_npz(filepath)
    else:
        raise Exception("Unimplemented file format.")

load_npz

load_npz(filepath)

Load the complex wave field and grids from a .npz file.

Parameters:

Name	Type	Description	Default
filepath	str	Path to the .npz file.	required

Source code in deeplens-src/deeplens/light/wave.py

def load_npz(self, filepath):
    """Load the complex wave field and grids from a `.npz` file.

    Args:
        filepath (str): Path to the `.npz` file.
    """
    data = np.load(filepath)
    self.u = torch.from_numpy(data["u"])
    self.x = torch.from_numpy(data["x"])
    self.y = torch.from_numpy(data["y"])
    self.wvln = data["wvln"].item()
    self.phy_size = data["phy_size"].tolist()
    self.res = self.u.shape[-2:]

save

save(filepath='./wavefield.npz')

Save the complex wave field to file (only .npz is supported).

Parameters:

Name	Type	Description	Default
filepath	str	Output path. Defaults to "./wavefield.npz".	'./wavefield.npz'

Raises:

Type	Description
Exception	If the file format is not supported.

Source code in deeplens-src/deeplens/light/wave.py

def save(self, filepath="./wavefield.npz"):
    """Save the complex wave field to file (only `.npz` is supported).

    Args:
        filepath (str, optional): Output path. Defaults to "./wavefield.npz".

    Raises:
        Exception: If the file format is not supported.
    """
    if filepath.endswith(".npz"):
        self.save_npz(filepath)
    else:
        raise Exception("Unimplemented file format.")

save_npz

save_npz(filepath='./wavefield.npz')

Save the field to a .npz file plus intensity/amplitude/phase PNGs.

Writes u, x, y, wvln, and phy_size to the .npz archive, and additionally saves normalized intensity, amplitude, and phase images next to it.

Parameters:

Name	Type	Description	Default
filepath	str	Output .npz path. Defaults to "./wavefield.npz".	'./wavefield.npz'

Source code in deeplens-src/deeplens/light/wave.py

def save_npz(self, filepath="./wavefield.npz"):
    """Save the field to a `.npz` file plus intensity/amplitude/phase PNGs.

    Writes `u`, `x`, `y`, `wvln`, and `phy_size` to the `.npz` archive, and
    additionally saves normalized intensity, amplitude, and phase images
    next to it.

    Args:
        filepath (str, optional): Output `.npz` path. Defaults to
            "./wavefield.npz".
    """
    from torchvision.utils import save_image

    # Save data
    np.savez_compressed(
        filepath,
        u=self.u.cpu().numpy(),
        x=self.x.cpu().numpy(),
        y=self.y.cpu().numpy(),
        wvln=np.array(self.wvln),
        phy_size=np.array(self.phy_size),
    )

    # Save intensity, amplitude, and phase images
    u = self.u.cpu()
    save_image(u.abs() ** 2, f"{filepath[:-4]}_intensity.png", normalize=True)
    save_image(u.abs(), f"{filepath[:-4]}_amp.png", normalize=True)
    save_image(u.angle(), f"{filepath[:-4]}_phase.png", normalize=True)

save_image

save_image(save_name=None, data='irr')

Render the field to an image (alias of show).

Parameters:

Name	Type	Description	Default
save_name	str or None	Output image path; if None, plot with matplotlib instead. Defaults to None.	None
data	str	Quantity to visualize, one of "irr", "amp", "phi"/"phase", "real", "imag". Defaults to "irr".	'irr'

Source code in deeplens-src/deeplens/light/wave.py

def save_image(self, save_name=None, data="irr"):
    """Render the field to an image (alias of `show`).

    Args:
        save_name (str or None, optional): Output image path; if None, plot
            with matplotlib instead. Defaults to None.
        data (str, optional): Quantity to visualize, one of "irr", "amp",
            "phi"/"phase", "real", "imag". Defaults to "irr".
    """
    return self.show(save_name=save_name, data=data)

show

show(save_name=None, data='irr')

Render the field as an image, either saved to disk or plotted.

Parameters:

Name	Type	Description	Default
save_name	str or None	Output image path; if None, the field is shown with matplotlib. Defaults to None.	None
data	str	Quantity to visualize: "irr" (intensity), "amp" (amplitude), "phi"/"phase", "real", or "imag". Defaults to "irr".	'irr'

Raises:

Type	Description
Exception	If data is unrecognized or the field shape is unsupported.

Source code in deeplens-src/deeplens/light/wave.py

def show(self, save_name=None, data="irr"):
    """Render the field as an image, either saved to disk or plotted.

    Args:
        save_name (str or None, optional): Output image path; if None, the
            field is shown with matplotlib. Defaults to None.
        data (str, optional): Quantity to visualize: "irr" (intensity),
            "amp" (amplitude), "phi"/"phase", "real", or "imag". Defaults
            to "irr".

    Raises:
        Exception: If `data` is unrecognized or the field shape is
            unsupported.
    """
    from torchvision.utils import save_image

    cmap = "gray"
    if data == "irr":
        value = self.u.detach().abs() ** 2
    elif data == "amp":
        value = self.u.detach().abs()
    elif data == "phi" or data == "phase":
        value = torch.angle(self.u).detach()
        cmap = "hsv"
    elif data == "real":
        value = self.u.real.detach()
    elif data == "imag":
        value = self.u.imag.detach()
    else:
        raise Exception(f"Unimplemented visualization: {data}.")

    if len(self.u.shape) == 2:
        raise Exception("Deprecated.")
        if save_name is not None:
            save_image(value, save_name, normalize=True)
        else:
            value = value.cpu().numpy()
            plt.imshow(
                value,
                cmap=cmap,
                extent=[
                    -self.phy_size[0] / 2,
                    self.phy_size[0] / 2,
                    -self.phy_size[1] / 2,
                    self.phy_size[1] / 2,
                ],
            )

    elif len(self.u.shape) == 4:
        B, C, H, W = self.u.shape
        if B == 1:
            if save_name is not None:
                save_image(value, save_name, normalize=True)
            else:
                value = value.cpu().numpy()
                plt.imshow(
                    value[0, 0, :, :],
                    cmap=cmap,
                    extent=[
                        -self.phy_size[0] / 2,
                        self.phy_size[0] / 2,
                        -self.phy_size[1] / 2,
                        self.phy_size[1] / 2,
                    ],
                )
        else:
            if save_name is not None:
                plt.savefig(save_name)
            else:
                value = value.cpu().numpy()
                fig, axs = plt.subplots(1, B)
                for i in range(B):
                    axs[i].imshow(
                        value[i, 0, :, :],
                        cmap=cmap,
                        extent=[
                            -self.phy_size[0] / 2,
                            self.phy_size[0] / 2,
                            -self.phy_size[1] / 2,
                            self.phy_size[1] / 2,
                        ],
                    )
                fig.show()
    else:
        raise Exception("Unsupported complex field shape.")

pad

pad(Hpad, Wpad)

Zero-pad the field and expand its physical size accordingly.

Pads Hpad pixels on the top and bottom and Wpad pixels on the left and right, then updates res, phy_size, and the coordinate grids so the pixel pitch stays constant. Modifies self in place.

Parameters:

Name	Type	Description	Default
Hpad	int	Number of pixels to pad on the top and bottom.	required
Wpad	int	Number of pixels to pad on the left and right.	required

Source code in deeplens-src/deeplens/light/wave.py

def pad(self, Hpad, Wpad):
    """Zero-pad the field and expand its physical size accordingly.

    Pads `Hpad` pixels on the top and bottom and `Wpad` pixels on the left
    and right, then updates `res`, `phy_size`, and the coordinate grids so
    the pixel pitch stays constant. Modifies `self` in place.

    Args:
        Hpad (int): Number of pixels to pad on the top and bottom.
        Wpad (int): Number of pixels to pad on the left and right.
    """
    self.u = F.pad(self.u, (Hpad, Hpad, Wpad, Wpad), mode="constant", value=0)

    Horg, Worg = self.res
    self.res = [Horg + 2 * Hpad, Worg + 2 * Wpad]
    self.phy_size = [
        self.phy_size[0] * self.res[0] / Horg,
        self.phy_size[1] * self.res[1] / Worg,
    ]
    self.x, self.y = self.gen_xy_grid()
    self.z = torch.full_like(self.x, float(self.z[0, 0]))

flip

flip()

Flip the field and its grids horizontally and vertically.

Returns:

Name	Type	Description
self	ComplexWave	The flipped wave field (for chaining).

Source code in deeplens-src/deeplens/light/wave.py

def flip(self):
    """Flip the field and its grids horizontally and vertically.

    Returns:
        self (ComplexWave): The flipped wave field (for chaining).
    """
    self.u = torch.flip(self.u, [-1, -2])
    self.x = torch.flip(self.x, [-1, -2])
    self.y = torch.flip(self.y, [-1, -2])
    self.z = torch.flip(self.z, [-1, -2])
    return self

Image Simulation

PSF Rendering

Functions for rendering images with point spread functions.

deeplens.imgsim.psf.conv_psf

conv_psf(img, psf, method='conv')

Render an image batch with one spatially invariant PSF.

Applies a per-channel, size-preserving ("same") 2-D convolution using reflect padding so that the output keeps the input spatial dimensions. The "conv" and "fft" backends apply the identical reflect-padded convolution and agree to FFT round-off; they differ only in compute cost.

Parameters:

Name	Type	Description	Default
img	Tensor	Input image batch, shape [B, C, H, W].	required
psf	Tensor	Shared PSF kernel with shape [C, ks, ks], or per-image kernels with shape [B, C, ks, ks]. ks may be odd or even.	required
method	str	Convolution backend. "conv" uses a direct F.conv2d (cost ~O(ks^2) per pixel); "fft" uses an FFT linear convolution (cost roughly independent of ks). Prefer "fft" for the large chromatic PSFs of a diffractive lens (ks ~= 512-768), where direct convolution is impractical. Defaults to "conv".	'conv'

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Raises:

Type	Description
ValueError	If method is not "conv" or "fft".

Examples:

psf = lens.psf_rgb(points=torch.tensor([0.0, 0.0, -10000.0]))
img_blur = conv_psf(img, psf)

Source code in deeplens-src/deeplens/imgsim/psf.py

def conv_psf(img, psf, method="conv"):
    """Render an image batch with one spatially invariant PSF.

    Applies a per-channel, size-preserving ("same") 2-D convolution using
    reflect padding so that the output keeps the input spatial dimensions. The
    ``"conv"`` and ``"fft"`` backends apply the identical reflect-padded
    convolution and agree to FFT round-off; they differ only in compute cost.

    Args:
        img (torch.Tensor): Input image batch, shape ``[B, C, H, W]``.
        psf (torch.Tensor): Shared PSF kernel with shape ``[C, ks, ks]``, or
            per-image kernels with shape ``[B, C, ks, ks]``. ``ks`` may be odd
            or even.
        method (str, optional): Convolution backend.  ``"conv"`` uses a direct
            ``F.conv2d`` (cost ``~O(ks^2)`` per pixel); ``"fft"`` uses an FFT
            linear convolution (cost roughly independent of ``ks``).  Prefer
            ``"fft"`` for the large chromatic PSFs of a diffractive lens
            (``ks ~= 512-768``), where direct convolution is impractical.
            Defaults to ``"conv"``.

    Returns:
        img_render (torch.Tensor): Rendered image, shape ``[B, C, H, W]``.

    Raises:
        ValueError: If ``method`` is not ``"conv"`` or ``"fft"``.

    Example:
        ```python
        psf = lens.psf_rgb(points=torch.tensor([0.0, 0.0, -10000.0]))
        img_blur = conv_psf(img, psf)
        ```
    """
    if img.ndim != 4:
        raise ValueError(f"img must have shape [B, C, H, W], got {tuple(img.shape)}.")
    B, C, H, W = img.shape
    if psf.ndim == 3:
        C_psf, ks, kw = psf.shape
        batched_psf = False
    elif psf.ndim == 4:
        B_psf, C_psf, ks, kw = psf.shape
        batched_psf = True
        if B_psf != B:
            raise ValueError(
                f"psf batch size ({B_psf}) must match image batch size ({B})."
            )
    else:
        raise ValueError(
            f"psf must have shape [C, K, K] or [B, C, K, K], got {tuple(psf.shape)}."
        )
    if C_psf != C:
        raise ValueError(f"psf channels ({C_psf}) must match image channels ({C}).")
    if ks != kw:
        raise ValueError("psf kernels must be square.")

    # Size-preserving ("same") padding that totals ks - 1, split so both odd and
    # even kernels keep the output shape equal to the input (symmetric ks // 2
    # only preserves size for odd ks; even ks would output N + 1).
    pad_top = (ks - 1) // 2
    pad_bottom = ks // 2
    pad_left = (ks - 1) // 2
    pad_right = ks // 2
    img_pad = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom), mode="reflect")

    if method == "conv":
        # Flip the PSF because F.conv2d computes cross-correlation, not convolution.
        if not batched_psf:
            psf_k = torch.flip(psf, [-2, -1]).unsqueeze(1)
            return F.conv2d(img_pad, psf_k, groups=C)

        # Group over both batch and channel so each image receives its own PSF.
        psf_k = torch.flip(psf, [-2, -1]).reshape(B * C, 1, ks, ks)
        img_grouped = img_pad.reshape(1, B * C, *img_pad.shape[-2:])
        rendered = F.conv2d(img_grouped, psf_k, groups=B * C)
        return rendered.reshape(B, C, H, W)

    if method == "fft":
        Hp, Wp = img_pad.shape[-2:]
        # Linear (not circular) convolution: zero-pad both operands to at least
        # ``Hp + ks - 1`` so the FFT product equals the linear convolution, then
        # keep the "valid" window (length Hp - ks + 1 == H) at offset ks - 1.
        fh, fw = Hp + ks - 1, Wp + ks - 1
        fimg = torch.fft.rfft2(img_pad, s=(fh, fw))
        fpsf = torch.fft.rfft2(psf, s=(fh, fw))
        if not batched_psf:
            fpsf = fpsf.unsqueeze(0)
        conv_full = torch.fft.irfft2(fimg * fpsf, s=(fh, fw))  # [B, C, fh, fw]
        return conv_full[..., ks - 1 : ks - 1 + H, ks - 1 : ks - 1 + W]

    raise ValueError(f"Unknown conv_psf method: {method!r} (expected 'conv' or 'fft').")

deeplens.imgsim.psf.conv_psf_map

conv_psf_map(img, psf_map)

Render an image batch with a spatially varying PSF map.

Divides the image into grid_h × grid_w non-overlapping patches and convolves each patch with its corresponding PSF kernel. The full image is padded before patch extraction to avoid artificial seams from independent per-patch padding.

Parameters:

Name	Type	Description	Default
img	Tensor	Input image batch, shape [B, C, H, W].	required
psf_map	Tensor	PSF map, shape [grid_h, grid_w, C, ks, ks].	required

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Source code in deeplens-src/deeplens/imgsim/psf.py

def conv_psf_map(img, psf_map):
    """Render an image batch with a spatially varying PSF map.

    Divides the image into ``grid_h × grid_w`` non-overlapping patches and
    convolves each patch with its corresponding PSF kernel. The full image is
    padded before patch extraction to avoid artificial seams from independent
    per-patch padding.

    Args:
        img (torch.Tensor): Input image batch, shape ``[B, C, H, W]``.
        psf_map (torch.Tensor): PSF map, shape ``[grid_h, grid_w, C, ks, ks]``.

    Returns:
        img_render (torch.Tensor): Rendered image, shape ``[B, C, H, W]``.
    """
    B, C, H, W = img.shape
    grid_h, grid_w, C_psf, ks, _ = psf_map.shape
    assert C_psf == C, f"PSF map channels ({C_psf}) must match image channels ({C})."

    # Padding
    pad_top = (ks - 1) // 2
    pad_bottom = ks // 2
    pad_left = (ks - 1) // 2
    pad_right = ks // 2
    img_pad = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom), mode="reflect")

    # Pre-flip entire PSF map once (instead of flipping each PSF inside the loop)
    psf_map_flipped = torch.flip(psf_map, dims=(-2, -1))

    # Render image patch by patch
    img_render = torch.zeros_like(img)
    for i in range(grid_h):
        h_low = (i * H) // grid_h
        h_high = ((i + 1) * H) // grid_h

        for j in range(grid_w):
            w_low = (j * W) // grid_w
            w_high = ((j + 1) * W) // grid_w

            # PSF, [C, 1, ks, ks]
            psf = psf_map_flipped[i, j].unsqueeze(1)

            # Consider overlap to avoid boundary artifacts
            img_pad_patch = img_pad[
                :,
                :,
                h_low : h_high + pad_top + pad_bottom,
                w_low : w_high + pad_left + pad_right,
            ]

            # Convolution, [B, C, h_high-h_low, w_high-w_low]
            render_patch = F.conv2d(img_pad_patch, psf, groups=C)
            img_render[:, :, h_low:h_high, w_low:w_high] = render_patch

    return img_render

deeplens.imgsim.psf.conv_psf_depth_interp

conv_psf_depth_interp(
    img, depth, psf_kernels, psf_depths, interp_mode="depth", padding_mode="reflect"
)

Depth-interpolated PSF convolution for a spatially-uniform but depth-varying blur.

Pre-convolves the image with PSFs at each reference depth, then blends the results using per-pixel linear interpolation weights derived from depth. This approximates defocus blur for a single field position across a depth range without computing a separate PSF per pixel.

Parameters:

Name	Type	Description	Default
img	Tensor	Image batch, shape [B, C, H, W], values in [0, 1].	required
depth	Tensor	Depth map, shape [B, 1, H, W], values in (-∞, 0) mm (negative convention).	required
psf_kernels	Tensor	PSF stack at reference depths, shape [num_depth, C, ks, ks].	required
psf_depths	Tensor	Depth of each PSF layer, shape [num_depth], values in (-∞, 0) mm. Must be monotone.	required
interp_mode	str	Interpolation space. "depth" interpolates linearly in depth; "disparity" interpolates linearly in 1/depth. Defaults to "depth".	'depth'
padding_mode	str or None	Padding mode passed to F.pad before convolution. If None, assumes img is already padded and applies no additional padding. Defaults to "reflect".	'reflect'

Returns:

Name	Type	Description
img_render	Tensor	Blurred image, shape [B, C, H, W].

Raises:

Type	Description
AssertionError	If depth or psf_depths contain non-negative values, or if interp_mode is not "depth" or "disparity".

Source code in deeplens-src/deeplens/imgsim/psf.py

def conv_psf_depth_interp(
    img, depth, psf_kernels, psf_depths, interp_mode="depth", padding_mode="reflect"
):
    """Depth-interpolated PSF convolution for a spatially-uniform but depth-varying blur.

    Pre-convolves the image with PSFs at each reference depth, then blends the
    results using per-pixel linear interpolation weights derived from `depth`.
    This approximates defocus blur for a single field position across a depth
    range without computing a separate PSF per pixel.

    Args:
        img (torch.Tensor): Image batch, shape ``[B, C, H, W]``, values in
            ``[0, 1]``.
        depth (torch.Tensor): Depth map, shape ``[B, 1, H, W]``, values in
            ``(-∞, 0)`` mm (negative convention).
        psf_kernels (torch.Tensor): PSF stack at reference depths, shape
            ``[num_depth, C, ks, ks]``.
        psf_depths (torch.Tensor): Depth of each PSF layer, shape
            ``[num_depth]``, values in ``(-∞, 0)`` mm.  Must be monotone.
        interp_mode (str, optional): Interpolation space.  ``"depth"``
            interpolates linearly in depth; ``"disparity"`` interpolates
            linearly in 1/depth.  Defaults to ``"depth"``.
        padding_mode (str or None, optional): Padding mode passed to
            `F.pad` before convolution. If ``None``, assumes ``img`` is already
            padded and applies no additional padding. Defaults to "reflect".

    Returns:
        img_render (torch.Tensor): Blurred image, shape ``[B, C, H, W]``.

    Raises:
        AssertionError: If `depth` or `psf_depths` contain non-negative values,
            or if `interp_mode` is not ``"depth"`` or ``"disparity"``.
    """
    assert interp_mode in ["depth", "disparity"], (
        f"interp_mode must be 'depth' or 'disparity', got {interp_mode}"
    )
    assert depth.min() < 0 and depth.max() < 0, (
        f"depth must be negative, got {depth.min()} and {depth.max()}"
    )
    assert psf_depths.min() < 0 and psf_depths.max() < 0, (
        f"psf_depths must be negative, got {psf_depths.min()} and {psf_depths.max()}"
    )

    num_depths, C_psf, ks, _ = psf_kernels.shape
    psf_depths = psf_depths.to(device=depth.device, dtype=depth.dtype)

    # =================================
    # PSF convolution for all depths
    # =================================
    B, C, _, _ = img.shape
    assert C_psf == C, f"PSF channels ({C_psf}) must match image channels ({C})."
    assert psf_depths.numel() == num_depths, (
        f"psf_depths length ({psf_depths.numel()}) must match PSF depth count ({num_depths})."
    )

    # Prepare PSF kernel: [num_depths, C, ks, ks] -> [num_depths*C, 1, ks, ks]
    # Flip the PSF because F.conv2d uses cross-correlation
    psf_stacked = torch.flip(psf_kernels, [-2, -1]).reshape(num_depths * C, 1, ks, ks)

    if padding_mode is None:
        img_padded_small = img
    else:
        # Pad before expand: pad [B, C, H, W] first (C channels), then expand to num_depths*C
        # This reduces padding work by a factor of num_depths
        pad_top = (ks - 1) // 2
        pad_bottom = ks // 2
        pad_left = (ks - 1) // 2
        pad_right = ks // 2
        img_padded_small = F.pad(
            img, (pad_left, pad_right, pad_top, pad_bottom), mode=padding_mode
        )

    # Expand padded img: [B, C, Hpad, Wpad] -> [B, num_depths*C, Hpad, Wpad]
    img_padded = img_padded_small.repeat(1, num_depths, 1, 1)

    # Grouped convolution: each of the num_depths*C channels is convolved with its own kernel
    imgs_blur = F.conv2d(
        img_padded, psf_stacked, groups=num_depths * C
    )  # [B, num_depths*C, Hout, Wout]
    H, W = imgs_blur.shape[-2:]

    # Reshape to [num_depths, B, C, H, W]
    imgs_blur = imgs_blur.reshape(B, num_depths, C, H, W).permute(1, 0, 2, 3, 4)

    # =================================
    # Depth/Disparity interpolation
    # =================================
    B_depth, _, H_depth, W_depth = depth.shape
    assert B_depth == B, (
        f"Depth batch size ({B_depth}) must match image batch size ({B})."
    )
    assert H_depth == H and W_depth == W, (
        f"Depth shape ({H_depth}, {W_depth}) must match rendered shape ({H}, {W})."
    )
    depth_flat = depth.flatten(1)  # shape [B, H*W]
    depth_flat = depth_flat.clamp(psf_depths[0], psf_depths[-1])
    indices = torch.searchsorted(psf_depths, depth_flat, right=True)  # shape [B, H*W]
    indices = indices.clamp(1, num_depths - 1)
    idx0 = indices - 1
    idx1 = indices

    # Calculate weights for depth interpolation
    d0 = psf_depths[idx0]  # shape [B, H*W]
    d1 = psf_depths[idx1]

    if interp_mode == "depth":
        # Interpolate in depth space
        denom = d1 - d0
        denom[denom == 0] = 1e-6  # Avoid division by zero
        w1 = (depth_flat - d0) / denom  # shape [B, H*W]
    else:
        # Interpolate in disparity space (disparity = 1/depth)
        disp_flat = 1.0 / depth_flat
        disp0 = 1.0 / d0
        disp1 = 1.0 / d1
        denom = disp1 - disp0
        denom[denom == 0] = 1e-6  # Avoid division by zero
        w1 = (disp_flat - disp0) / denom  # shape [B, H*W]

    w0 = 1 - w1

    # Create a weight tensor
    weights = torch.zeros(num_depths, B, H * W, device=img.device, dtype=img.dtype)
    weights.scatter_add_(0, idx0.unsqueeze(0).long(), w0.unsqueeze(0))
    weights.scatter_add_(0, idx1.unsqueeze(0).long(), w1.unsqueeze(0))
    weights = weights.view(num_depths, B, 1, H, W)

    # Apply weights to the blurred images
    img_render = torch.sum(imgs_blur * weights, dim=0)
    return img_render

deeplens.imgsim.psf.conv_psf_map_depth_interp

conv_psf_map_depth_interp(img, depth, psf_map, psf_depths, interp_mode='depth')

Render with a spatially varying, depth-interpolated PSF map.

The image is divided into PSF-map grid cells. For each cell, the image patch is convolved with all reference-depth PSFs for that cell, then the convolved results are blended per pixel using interpolation weights from the depth map.

Parameters:

Name	Type	Description	Default
img	Tensor	Image batch, shape [B, C, H, W], values in [0, 1].	required
depth	Tensor	Depth map, shape [B, 1, H, W], values in (-inf, 0) using the negative-depth convention.	required
psf_map	Tensor	PSF map, shape [grid_h, grid_w, num_depth, C, ks, ks].	required
psf_depths	Tensor	Reference depths, shape [num_depth], values in (-inf, 0). Used to interpolate psf_map.	required
interp_mode	str	"depth" for linear depth interpolation or "disparity" for linear interpolation in \(1/ ext{depth}\). Defaults to "depth".	'depth'

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Source code in deeplens-src/deeplens/imgsim/psf.py

def conv_psf_map_depth_interp(img, depth, psf_map, psf_depths, interp_mode="depth"):
    """Render with a spatially varying, depth-interpolated PSF map.

    The image is divided into PSF-map grid cells. For each cell, the image
    patch is convolved with all reference-depth PSFs for that cell, then the
    convolved results are blended per pixel using interpolation weights from
    the depth map.

    Args:
        img (torch.Tensor): Image batch, shape ``[B, C, H, W]``, values in
            ``[0, 1]``.
        depth (torch.Tensor): Depth map, shape ``[B, 1, H, W]``, values in
            ``(-inf, 0)`` using the negative-depth convention.
        psf_map (torch.Tensor): PSF map, shape
            ``[grid_h, grid_w, num_depth, C, ks, ks]``.
        psf_depths (torch.Tensor): Reference depths, shape ``[num_depth]``,
            values in ``(-inf, 0)``. Used to interpolate ``psf_map``.
        interp_mode (str, optional): ``"depth"`` for linear depth interpolation
            or ``"disparity"`` for linear interpolation in $1/\text{depth}$.
            Defaults to "depth".

    Returns:
        img_render (torch.Tensor): Rendered image, shape ``[B, C, H, W]``.
    """
    _, _, H, W = img.shape
    grid_h, grid_w, _, _, ks, _ = psf_map.shape

    # Pad the full image once to avoid boundary artifacts at patch seams.
    # Without this, each patch would be padded independently (reflecting within
    # its own boundary), producing visible seams at grid boundaries.
    pad_top = (ks - 1) // 2
    pad_bottom = ks // 2
    pad_left = (ks - 1) // 2
    pad_right = ks // 2
    img_pad = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom), mode="reflect")

    # Render image patch by patch
    img_render = torch.zeros_like(img)
    for i in range(grid_h):
        h_low = (i * H) // grid_h
        h_high = ((i + 1) * H) // grid_h

        for j in range(grid_w):
            w_low = (j * W) // grid_w
            w_high = ((j + 1) * W) // grid_w

            # Extract overlapping patch from pre-padded image (no per-patch padding needed)
            img_pad_patch = img_pad[
                :,
                :,
                h_low : h_high + pad_top + pad_bottom,
                w_low : w_high + pad_left + pad_right,
            ]
            depth_patch = depth[:, :, h_low:h_high, w_low:w_high]
            render_patch = conv_psf_depth_interp(
                img_pad_patch,
                depth_patch,
                psf_map[i, j],
                psf_depths,
                interp_mode=interp_mode,
                padding_mode=None,
            )
            img_render[:, :, h_low:h_high, w_low:w_high] = render_patch

    return img_render

deeplens.imgsim.psf.conv_psf_occlusion

conv_psf_occlusion(img, depth, psf_kernels, psf_depths)

Occlusion-aware bokeh rendering using back-to-front layered compositing.

Discretizes the scene into depth layers and composites them from back (far) to front (near). Each layer is blurred independently with its depth-specific PSF, and composited using the over-operator. This prevents color bleeding at depth discontinuities.

Reference

[1] "Dr.Bokeh: DiffeRentiable Occlusion-aware Bokeh Rendering", CVPR 2024.

Parameters:

Name	Type	Description	Default
img	Tensor	Input image, shape [B, C, H, W], values in [0, 1].	required
depth	Tensor	Depth map, shape [B, 1, H, W], values in (-inf, 0) mm (negative convention).	required
psf_kernels	Tensor	PSF at each depth layer, shape [num_layers, C, ks, ks].	required
psf_depths	Tensor	Depth value for each layer, shape [num_layers] mm. Must be negative and sorted ascending (far to near, e.g. -5000 ... -200).	required

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Source code in deeplens-src/deeplens/imgsim/psf.py

def conv_psf_occlusion(img, depth, psf_kernels, psf_depths):
    """Occlusion-aware bokeh rendering using back-to-front layered compositing.

    Discretizes the scene into depth layers and composites them from back (far)
    to front (near). Each layer is blurred independently with its depth-specific
    PSF, and composited using the over-operator. This prevents color bleeding at
    depth discontinuities.

    Reference:
        [1] "Dr.Bokeh: DiffeRentiable Occlusion-aware Bokeh Rendering", CVPR 2024.

    Args:
        img (torch.Tensor): Input image, shape ``[B, C, H, W]``, values in
            ``[0, 1]``.
        depth (torch.Tensor): Depth map, shape ``[B, 1, H, W]``, values in
            ``(-inf, 0)`` mm (negative convention).
        psf_kernels (torch.Tensor): PSF at each depth layer, shape
            ``[num_layers, C, ks, ks]``.
        psf_depths (torch.Tensor): Depth value for each layer, shape
            ``[num_layers]`` mm. Must be negative and sorted ascending (far to
            near, e.g. -5000 ... -200).

    Returns:
        img_render (torch.Tensor): Rendered image, shape ``[B, C, H, W]``.
    """
    assert depth.min() < 0 and depth.max() < 0, (
        f"depth must be negative, got min={depth.min()} max={depth.max()}"
    )
    assert psf_depths.min() < 0 and psf_depths.max() < 0, (
        f"psf_depths must be negative, got min={psf_depths.min()} max={psf_depths.max()}"
    )

    num_layers, C, ks, _ = psf_kernels.shape
    B, C_img, H, W = img.shape
    assert C == C_img, f"PSF channels ({C}) must match image channels ({C_img})"

    device = img.device
    dtype = img.dtype
    psf_depths = psf_depths.to(device=device, dtype=dtype)

    # Assign each pixel to its nearest depth layer
    # depth and psf_depths are both negative; depth_map shape [B, 1, H, W]
    depth_expanded = depth.view(B, 1, H, W).expand(B, num_layers, H, W)
    psf_depths_view = psf_depths.view(1, num_layers, 1, 1)
    dist = torch.abs(depth_expanded - psf_depths_view)  # [B, num_layers, H, W]
    layer_assignment = dist.argmin(dim=1, keepdim=True)  # [B, 1, H, W]

    # Pre-compute flipped PSFs and padding for convolution
    psf_flipped = torch.flip(psf_kernels, [-2, -1])  # [num_layers, C, ks, ks]
    pad_top = (ks - 1) // 2
    pad_bottom = ks // 2
    pad_left = (ks - 1) // 2
    pad_right = ks // 2

    # Back-to-front compositing (layer 0 is farthest, layer num_layers-1 is nearest)
    result = torch.zeros(B, C, H, W, device=device, dtype=dtype)

    for i in range(num_layers):
        # Create soft mask for this layer: 1 where pixels belong to this layer
        mask = (layer_assignment == i).to(dtype=dtype)  # [B, 1, H, W]

        # Run unconditionally: an all-zero mask convolves to zero and composites
        # to a no-op (result = 0 + result * (1 - 0)). Skipping it via
        # `if mask.sum() == 0` forced a GPU->CPU sync every iteration.

        # Layer RGB: pixels in this layer, zero elsewhere
        layer_rgb = img * mask  # [B, C, H, W]

        # Convolve layer RGB with this layer's PSF
        psf_i = psf_flipped[i].unsqueeze(1)  # [C, 1, ks, ks]
        layer_rgb_pad = F.pad(
            layer_rgb,
            (pad_left, pad_right, pad_top, pad_bottom),
            mode="constant",
            value=0,
        )
        blurred_rgb = F.conv2d(layer_rgb_pad, psf_i, groups=C)  # [B, C, H, W]

        # Convolve mask with the same PSF (use one channel of PSF, since PSF sums to 1 per channel)
        # Average across channels for mask blurring (PSF is same across channels for paraxial)
        psf_i_mono = psf_flipped[i, 0:1].unsqueeze(1)  # [1, 1, ks, ks]
        mask_pad = F.pad(
            mask, (pad_left, pad_right, pad_top, pad_bottom), mode="constant", value=0
        )
        blurred_mask = F.conv2d(mask_pad, psf_i_mono, groups=1)  # [B, 1, H, W]
        blurred_mask = blurred_mask.clamp(0, 1)

        # Over-compositing (back-to-front):
        # result = blurred_rgb + result * (1 - blurred_mask)
        result = blurred_rgb + result * (1 - blurred_mask)

    return result

deeplens.imgsim.psf.splat_psf_per_pixel

splat_psf_per_pixel(img, psf, chunk_size=None)

Render an image batch by splatting each pixel through its own PSF.

Uses a different PSF kernel for each source pixel and accumulates the scattered contributions with F.fold. When chunk_size is set, source pixels are processed tile by tile to reduce peak memory while preserving PSF contributions that cross tile boundaries.

Parameters:

Name	Type	Description	Default
img	Tensor	Image batch to be blurred, shape [B, C, H, W].	required
psf	Tensor	Per-pixel local PSFs, shape [H, W, C, ks, ks]. ks may be odd or even.	required
chunk_size	int or None	Source tile size for memory-efficient rendering. If None, render the whole image at once. Defaults to None.	None

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Source code in deeplens-src/deeplens/imgsim/psf.py

def splat_psf_per_pixel(img, psf, chunk_size=None):
    """Render an image batch by splatting each pixel through its own PSF.

    Uses a different PSF kernel for each source pixel and accumulates the
    scattered contributions with ``F.fold``. When ``chunk_size`` is set, source
    pixels are processed tile by tile to reduce peak memory while preserving
    PSF contributions that cross tile boundaries.

    Args:
        img (torch.Tensor): Image batch to be blurred, shape ``[B, C, H, W]``.
        psf (torch.Tensor): Per-pixel local PSFs, shape ``[H, W, C, ks, ks]``.
            ``ks`` may be odd or even.
        chunk_size (int or None, optional): Source tile size for
            memory-efficient rendering. If ``None``, render the whole image at
            once. Defaults to None.

    Returns:
        img_render (torch.Tensor): Rendered image, shape ``[B, C, H, W]``.
    """
    B, C, H, W = img.shape
    H_psf, W_psf, C_psf, ks, _ = psf.shape
    assert C == C_psf, "Image and PSF channels mismatch."
    assert H == H_psf and W == W_psf, "Image and PSF size mismatch."

    pad_top = (ks - 1) // 2
    pad_bottom = ks // 2
    pad_left = (ks - 1) // 2
    pad_right = ks // 2

    if chunk_size is None:
        img_expand = img.unsqueeze(-1).unsqueeze(-1)  # [B, C, H, W, 1, 1]
        kernels = psf.permute(2, 0, 1, 3, 4).unsqueeze(0)  # [1, C, H, W, ks, ks]
        img_render = img_expand * kernels  # [B, C, H, W, ks, ks]
        img_render = img_render.permute(0, 1, 4, 5, 2, 3).reshape(B, C * ks * ks, H * W)
        img_render = F.fold(img_render, (H + ks - 1, W + ks - 1), (ks, ks), padding=0)
    else:
        assert chunk_size > 0, "chunk_size must be positive."

        img_render = img.new_zeros(
            B,
            C,
            H + pad_top + pad_bottom,
            W + pad_left + pad_right,
        )

        for y0 in range(0, H, chunk_size):
            y1 = min(y0 + chunk_size, H)
            for x0 in range(0, W, chunk_size):
                x1 = min(x0 + chunk_size, W)
                img_patch = img[:, :, y0:y1, x0:x1]
                psf_patch = psf[y0:y1, x0:x1, :, :, :]

                patch_h, patch_w = y1 - y0, x1 - x0
                img_patch = img_patch.unsqueeze(-1).unsqueeze(-1)
                kernels = psf_patch.permute(2, 0, 1, 3, 4).unsqueeze(0)
                render_patch = img_patch * kernels
                render_patch = render_patch.permute(0, 1, 4, 5, 2, 3).reshape(
                    B, C * ks * ks, patch_h * patch_w
                )
                img_render[:, :, y0 : y1 + ks - 1, x0 : x1 + ks - 1] += F.fold(
                    render_patch,
                    (patch_h + ks - 1, patch_w + ks - 1),
                    (ks, ks),
                    padding=0,
                )

    return img_render[
        :,
        :,
        pad_top : pad_top + H,
        pad_left : pad_left + W,
    ]

deeplens.imgsim.psf.interp_psf_map

interp_psf_map(psf_map, grid_old, grid_new)

Resample a PSF map to a different spatial grid size.

Supports either a packed map [C, grid_old*ks, grid_old*ks] or an unpacked map [grid_old, grid_old, C, ks, ks]. Each kernel sample location is bilinearly interpolated across the PSF grid, and the result is returned in packed-map layout.

Parameters:

Name	Type	Description	Default
psf_map	Tensor	Packed or unpacked PSF map.	required
grid_old	int	Input grid size. Ignored for unpacked input, where the grid size is read from psf_map.	required
grid_new	int	Output grid size.	required

Returns:

Name	Type	Description
psf_map_interp	Tensor	Interpolated packed PSF map, shape [C, grid_new*ks, grid_new*ks].

Source code in deeplens-src/deeplens/imgsim/psf.py

def interp_psf_map(psf_map, grid_old, grid_new):
    """Resample a PSF map to a different spatial grid size.

    Supports either a packed map ``[C, grid_old*ks, grid_old*ks]`` or an
    unpacked map ``[grid_old, grid_old, C, ks, ks]``. Each kernel sample
    location is bilinearly interpolated across the PSF grid, and the result is
    returned in packed-map layout.

    Args:
        psf_map (torch.Tensor): Packed or unpacked PSF map.
        grid_old (int): Input grid size. Ignored for unpacked input, where the
            grid size is read from ``psf_map``.
        grid_new (int): Output grid size.

    Returns:
        psf_map_interp (torch.Tensor): Interpolated packed PSF map, shape
            ``[C, grid_new*ks, grid_new*ks]``.
    """
    if len(psf_map.shape) == 3:
        # [C, grid_old*ks, grid_old*ks]
        C, H, W = psf_map.shape
        assert H % grid_old == 0 and W % grid_old == 0, (
            "PSF map size should be divisible by grid"
        )
        ks = int(H / grid_old)
        assert ks % 2 == 1, "PSF kernel size should be odd"

        # Reshape from [C, grid*ks, grid*ks] to [grid_old, grid_old, C, ks, ks]
        psf_map_interp = psf_map.reshape(C, grid_old, ks, grid_old, ks).permute(
            1, 3, 0, 2, 4
        )  # .reshape(grid_old, grid_old, C, ks, ks)
    elif len(psf_map.shape) == 5:
        # [grid_old, grid_old, C, ks, ks]
        grid_h, grid_w, C, ks_h, ks_w = psf_map.shape
        assert grid_h == grid_w, f"PSF map grid must be square, got {grid_h}x{grid_w}"
        assert ks_h == ks_w, f"PSF kernel must be square, got {ks_h}x{ks_w}"
        grid_old = grid_h
        ks = ks_h
        psf_map_interp = psf_map
    else:
        raise ValueError(
            "PSF map should be [C, grid_old*ks, grid_old*ks] or [grid_old, grid_old, C, ks, ks]"
        )

    # Reshape from [grid_old, grid_old, C, ks, ks] to [ks*ks, C, grid_old, grid_old]
    psf_map_interp = psf_map_interp.permute(3, 4, 2, 0, 1).reshape(
        ks * ks, C, grid_old, grid_old
    )

    # Interpolate from [ks*ks, C, grid_old, grid_old] to [ks*ks, C, grid_new, grid_new]
    psf_map_interp = F.interpolate(
        psf_map_interp, size=(grid_new, grid_new), mode="bilinear", align_corners=True
    )

    # Reshape from [ks*ks, C, grid_new, grid_new] to [C, grid_new*ks, grid_new*ks]
    psf_map_interp = (
        psf_map_interp.reshape(ks, ks, C, grid_new, grid_new)
        .permute(2, 3, 0, 4, 1)
        .reshape(C, grid_new * ks, grid_new * ks)
    )

    return psf_map_interp

deeplens.imgsim.psf.rotate_psf

rotate_psf(psf, theta)

Rotate a batch of RGB PSF kernels counter-clockwise.

Rotation is performed around the center of each square PSF kernel using F.grid_sample.

Parameters:

Name	Type	Description	Default
psf	Tensor	PSF batch, shape [N, 3, ks, ks].	required
theta	Tensor	Rotation angles in radians, shape [N].	required

Returns:

Name	Type	Description
rotated_psf	Tensor	Rotated PSFs, shape [N, 3, ks, ks].

Source code in deeplens-src/deeplens/imgsim/psf.py

def rotate_psf(psf, theta):
    """Rotate a batch of RGB PSF kernels counter-clockwise.

    Rotation is performed around the center of each square PSF kernel using
    ``F.grid_sample``.

    Args:
        psf (torch.Tensor): PSF batch, shape ``[N, 3, ks, ks]``.
        theta (torch.Tensor): Rotation angles in radians, shape ``[N]``.

    Returns:
        rotated_psf (torch.Tensor): Rotated PSFs, shape ``[N, 3, ks, ks]``.
    """
    assert len(psf.shape) == 4, "PSF should be [N, 3, ks, ks]"

    N, _, ks, _ = psf.shape
    assert ks == psf.shape[3], "PSF kernel should be square"

    # To rotate the image counter-clockwise, the sampling grid must be rotated clockwise.
    # The matrix for a clockwise rotation by theta is:
    # [ cos(theta)  sin(theta) ]
    # [ -sin(theta) cos(theta) ]
    rotation_matrices = torch.zeros(N, 2, 3, device=psf.device, dtype=psf.dtype)
    rotation_matrices[:, 0, 0] = torch.cos(theta)
    rotation_matrices[:, 0, 1] = torch.sin(theta)
    rotation_matrices[:, 1, 0] = -torch.sin(theta)
    rotation_matrices[:, 1, 1] = torch.cos(theta)

    # Rotate PSFs
    grid = F.affine_grid(rotation_matrices, psf.shape, align_corners=True)
    rotated_psf = F.grid_sample(psf, grid, align_corners=True)

    return rotated_psf

Monte Carlo Integrals

Utilities for ray-bundle accumulation and ray-traced image sampling.

deeplens.imgsim.monte_carlo.forward_integral

forward_integral(ray, ps, ks, pointc=None, interpolate=True)

Differentiable Monte-Carlo integral over a ray bundle onto a pixel grid.

Bins ray hit positions into a ks x ks grid centred on pointc (or the valid-ray centroid when pointc is None). In coherent mode the complex amplitude sqrt(|dz|) * exp(i * phase) is accumulated; in incoherent mode unit intensity is accumulated. All N field points scatter into their own output slices in batched index_put_(accumulate=True) calls.

Parameters:

Name	Type	Description	Default
ray	Ray	Traced ray bundle with origin ray.o of shape [N, spp, 3] (or [spp, 3] for a single field point).	required
ps	float	Pixel size [mm].	required
ks	int	Output grid size in pixels (square).	required
pointc	Tensor or None	Reference centre [mm] for each field point, shape [N, 2]. If None, the valid-ray centroid is used. Defaults to None.	None
interpolate	bool	If True, each ray splits its contribution across the four surrounding pixels via bilinear weights. If False, each ray is hard-binned into the floor pixel (faster, no gradient w.r.t. in-pixel position). Defaults to True.	True

Returns:

Name	Type	Description
grid	Tensor	Accumulated field, shape [N, ks, ks] (or [ks, ks] for a single input point). Dtype is complex when ray.is_coherent is True, otherwise float.

Source code in deeplens-src/deeplens/imgsim/monte_carlo.py

def forward_integral(ray, ps, ks, pointc=None, interpolate=True):
    """Differentiable Monte-Carlo integral over a ray bundle onto a pixel grid.

    Bins ray hit positions into a `ks` x `ks` grid centred on `pointc` (or the
    valid-ray centroid when `pointc` is None). In coherent mode the complex
    amplitude `sqrt(|dz|) * exp(i * phase)` is accumulated; in incoherent mode
    unit intensity is accumulated. All `N` field points scatter into their own
    output slices in batched `index_put_(accumulate=True)` calls.

    Args:
        ray (Ray): Traced ray bundle with origin `ray.o` of shape
            [N, spp, 3] (or [spp, 3] for a single field point).
        ps (float): Pixel size [mm].
        ks (int): Output grid size in pixels (square).
        pointc (torch.Tensor or None, optional): Reference centre [mm] for each
            field point, shape [N, 2]. If None, the valid-ray centroid is used.
            Defaults to None.
        interpolate (bool, optional): If True, each ray splits its contribution
            across the four surrounding pixels via bilinear weights. If False,
            each ray is hard-binned into the floor pixel (faster, no gradient
            w.r.t. in-pixel position). Defaults to True.

    Returns:
        grid (torch.Tensor): Accumulated field, shape [N, ks, ks] (or [ks, ks]
            for a single input point). Dtype is complex when `ray.is_coherent`
            is True, otherwise float.
    """
    if ray.o.ndim == 2:
        single_point = True
        ray = ray.unsqueeze(0)
    else:
        single_point = False

    points = ray.o[..., :2]  # [N, spp, 2]
    valid = ray.is_valid  # [N, spp]
    N, spp = valid.shape
    device = valid.device

    # Centre the grid on pointc (or the valid-ray centroid).
    if pointc is None:
        pointc = (points * valid.unsqueeze(-1)).sum(-2) / valid.unsqueeze(-1).sum(
            -2
        ).add(EPSILON)
    points_shift = points - pointc.unsqueeze(-2)  # [N, spp, 2]

    # Reject points that fall outside the grid window.
    field_max = (ks / 2 - 0.5) * ps
    in_window = (points_shift[..., 0].abs() < (field_max - 0.001 * ps)) & (
        points_shift[..., 1].abs() < (field_max - 0.001 * ps)
    )
    valid = valid * in_window.to(valid.dtype)

    # Per-ray intensity (real) or complex amplitude.
    if ray.is_coherent:
        # Add EPSILON: sqrt'(0) is infinite -> NaN gradient for steep rays (dz~0).
        amp = torch.sqrt(ray.d[..., 2].abs() + EPSILON)  # sqrt(|dz|)
        opl = ray.opl.squeeze(-1)  # [N, spp]
        opl_min = opl.min(dim=-1, keepdim=True).values
        wvln_mm = ray.wvln * 1e-3
        phase = torch.fmod((opl - opl_min) / wvln_mm, 1) * (2 * torch.pi)
        value = amp * torch.exp(1j * phase)
    else:
        value = torch.ones_like(valid)

    # Fractional pixel indices: y up -> row down, x right -> col right.
    # Pixel centres lie on an integer grid in [0, ks-1].
    norm_row = (field_max - points_shift[..., 1]) / (2 * field_max)
    norm_col = (points_shift[..., 0] + field_max) / (2 * field_max)
    pix_row = norm_row * (ks - 1)
    pix_col = norm_col * (ks - 1)
    r_floor = pix_row.floor()
    c_floor = pix_col.floor()

    r0 = r_floor.long().clamp(0, ks - 1)
    c0 = c_floor.long().clamp(0, ks - 1)

    masked_value = valid * value

    # Batched scatter: all N field points accumulate simultaneously via a
    # batch-aware ``index_put_``, which does support per-batch accumulation
    # when the index tuple carries a batch dimension.
    batch_idx = torch.arange(N, device=device).unsqueeze(-1).expand(N, spp)
    grid = torch.zeros(N, ks, ks, dtype=value.dtype, device=device)
    if interpolate:
        w_r = pix_row - r_floor
        w_c = pix_col - c_floor
        r1 = (r0 + 1).clamp(0, ks - 1)
        c1 = (c0 + 1).clamp(0, ks - 1)
        grid.index_put_(
            (batch_idx, r0, c0), (1 - w_r) * (1 - w_c) * masked_value, accumulate=True
        )
        grid.index_put_(
            (batch_idx, r0, c1), (1 - w_r) * w_c * masked_value, accumulate=True
        )
        grid.index_put_(
            (batch_idx, r1, c0), w_r * (1 - w_c) * masked_value, accumulate=True
        )
        grid.index_put_((batch_idx, r1, c1), w_r * w_c * masked_value, accumulate=True)
    else:
        grid.index_put_((batch_idx, r0, c0), masked_value, accumulate=True)

    if single_point:
        grid = grid.squeeze(0)
        ray = ray.squeeze(0)  # restore caller's ray shape (unsqueeze mutates)

    return grid

deeplens.imgsim.monte_carlo.backward_integral

backward_integral(
    ray, img_obj, ps, interpolate=True, energy_correction=None, vignetting=False
)

Backward Monte-Carlo integration for ray-tracing-based rendering.

Sample an input image at each ray's hit position and average over the samples-per-pixel (spp) axis to render the output. The input image is always replicate-padded by one pixel on each side so that rays landing within half a pixel of the edge can still be bilinearly sampled without silently truncating.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object. Shape of ray.o is [h, w, spp, 3], with positions in [mm].	required
img_obj	Tensor	Source image, shape [B, C, H, W]. Spatial size H, W is read from this tensor.	required
ps	float	Pixel size [mm].	required
interpolate	bool	If True, bilinearly sample the four surrounding pixels; if False, nearest-pixel sampling. Defaults to True.	True
energy_correction	Tensor or None	Per-ray weight tensor of shape [h, w, spp, 1] (e.g. ray.en). When supplied, it is used as an importance weight; under the default (non-vignetting) mode it enters both numerator and denominator, yielding a proper weighted Monte-Carlo mean. Under vignetting it scales only the numerator (fixed denominator). Defaults to None (uniform weights).	None
vignetting	bool	If True, divide by a fixed denominator (torch.numel(ray.is_valid)) instead of the sum of weights; pixels hit by few or attenuated rays therefore appear dimmer (mechanical vignetting). Defaults to False.	False

Returns:

Name	Type	Description
output	Tensor	Rendered image, shape [B, C, h, w].

Raises:

Type	Description
Exception	If ray.is_coherent is True (coherent backward integral is not supported).

Source code in deeplens-src/deeplens/imgsim/monte_carlo.py

def backward_integral(
    ray,
    img_obj,
    ps,
    interpolate=True,
    energy_correction=None,
    vignetting=False,
):
    """Backward Monte-Carlo integration for ray-tracing-based rendering.

    Sample an input image at each ray's hit position and average over the
    samples-per-pixel (spp) axis to render the output. The input image is
    always replicate-padded by one pixel on each side so that rays landing
    within half a pixel of the edge can still be bilinearly sampled without
    silently truncating.

    Args:
        ray (Ray): Ray object. Shape of `ray.o` is [h, w, spp, 3], with
            positions in [mm].
        img_obj (torch.Tensor): Source image, shape [B, C, H, W]. Spatial size
            H, W is read from this tensor.
        ps (float): Pixel size [mm].
        interpolate (bool, optional): If True, bilinearly sample the four
            surrounding pixels; if False, nearest-pixel sampling. Defaults to
            True.
        energy_correction (torch.Tensor or None, optional): Per-ray weight
            tensor of shape [h, w, spp, 1] (e.g. `ray.en`). When supplied, it
            is used as an importance weight; under the default (non-vignetting)
            mode it enters both numerator and denominator, yielding a proper
            weighted Monte-Carlo mean. Under vignetting it scales only the
            numerator (fixed denominator). Defaults to None (uniform weights).
        vignetting (bool, optional): If True, divide by a fixed denominator
            (`torch.numel(ray.is_valid)`) instead of the sum of weights; pixels
            hit by few or attenuated rays therefore appear dimmer (mechanical
            vignetting). Defaults to False.

    Returns:
        output (torch.Tensor): Rendered image, shape [B, C, h, w].

    Raises:
        Exception: If `ray.is_coherent` is True (coherent backward integral is
            not supported).
    """
    assert len(img_obj.shape) == 4
    H, W = img_obj.shape[-2:]
    p = ray.o[..., :2]  # shape [h, w, spp, 2]
    img_obj = F.pad(img_obj, (1, 1, 1, 1), "replicate")

    # Convert ray positions to uv coordinates
    u = torch.clamp(W / 2 + p[..., 0] / ps, min=-0.99, max=W - 0.01)
    v = torch.clamp(H / 2 + p[..., 1] / ps, min=0.01, max=H + 0.99)

    # (idx_i, idx_j) denotes left-top pixel (reference); indices don't carry gradients.
    # (idx + 1 because we did padding)
    idx_i = H - v.ceil().long() + 1
    idx_j = u.floor().long() + 1

    # Gradients are stored in interpolation weight parameters
    w_i = v - v.floor().long()
    w_j = u.ceil().long() - u

    if ray.is_coherent:
        raise Exception("Backward coherent integral needs to be checked.")

    # Monte-Carlo integration over the spp axis (last dim).
    if interpolate:
        # Bilinear splatting
        # img_obj [B, C, H+2, W+2], idx_i/idx_j [h, w, spp] -> out_img [B, C, h, w, spp]
        out_img = img_obj[..., idx_i, idx_j] * w_i * w_j
        out_img += img_obj[..., idx_i + 1, idx_j] * (1 - w_i) * w_j
        out_img += img_obj[..., idx_i, idx_j + 1] * w_i * (1 - w_j)
        out_img += img_obj[..., idx_i + 1, idx_j + 1] * (1 - w_i) * (1 - w_j)
    else:
        out_img = img_obj[..., idx_i, idx_j]

    # Extra per-ray energy correction factor (e.g. for non-uniform ray sampling).
    weight = ray.is_valid
    if energy_correction is not None:
        weight = weight * energy_correction.squeeze(-1)

    # Normalize by the sum of weights (or fixed denominator if vignetting) to get the Monte-Carlo mean.
    if vignetting:
        output = torch.sum(out_img * weight, -1) / torch.numel(ray.is_valid)
    else:
        output = torch.sum(out_img * weight, -1) / (torch.sum(weight, -1) + EPSILON)

    return output

deeplens.imgsim.monte_carlo.assign_points_to_pixels

assign_points_to_pixels(points, mask, ks, x_range, y_range, value, interpolate=True)

Scatter point samples onto a ks x ks pixel grid.

Bins each point's value (intensity or complex amplitude) into the grid spanned by x_range x y_range using index_put_(accumulate=True). Supports both incoherent and coherent ray tracing. Handles a single point source only, constrained by the advanced-indexing scatter.

Parameters:

Name	Type	Description	Default
points	Tensor	Sample positions [mm], shape [spp, 2] as (x, y).	required
mask	Tensor	Validity mask, shape [spp].	required
ks	int	Output grid size in pixels (square).	required
x_range	tuple	Grid extent (x_min, x_max) [mm].	required
y_range	tuple	Grid extent (y_min, y_max) [mm].	required
value	Tensor	Per-point value to accumulate (intensity or complex amplitude), shape [spp].	required
interpolate	bool	If True, split each point across the four surrounding pixels via bilinear weights; if False, hard-bin into the floor pixel. Defaults to True.	True

Returns:

Name	Type	Description
grid	Tensor	Accumulated intensity or complex amplitude, shape [ks, ks]. Dtype matches value.

Source code in deeplens-src/deeplens/imgsim/monte_carlo.py

def assign_points_to_pixels(
    points,
    mask,
    ks,
    x_range,
    y_range,
    value,
    interpolate=True,
):
    """Scatter point samples onto a `ks` x `ks` pixel grid.

    Bins each point's `value` (intensity or complex amplitude) into the grid
    spanned by `x_range` x `y_range` using `index_put_(accumulate=True)`.
    Supports both incoherent and coherent ray tracing. Handles a single point
    source only, constrained by the advanced-indexing scatter.

    Args:
        points (torch.Tensor): Sample positions [mm], shape [spp, 2] as (x, y).
        mask (torch.Tensor): Validity mask, shape [spp].
        ks (int): Output grid size in pixels (square).
        x_range (tuple): Grid extent (x_min, x_max) [mm].
        y_range (tuple): Grid extent (y_min, y_max) [mm].
        value (torch.Tensor): Per-point value to accumulate (intensity or
            complex amplitude), shape [spp].
        interpolate (bool, optional): If True, split each point across the four
            surrounding pixels via bilinear weights; if False, hard-bin into the
            floor pixel. Defaults to True.

    Returns:
        grid (torch.Tensor): Accumulated intensity or complex amplitude, shape
            [ks, ks]. Dtype matches `value`.
    """
    # Parameters
    device = points.device
    x_min, x_max = x_range
    y_min, y_max = y_range

    # Normalize points to the range [0, 1] (direct computation, no intermediate allocation)
    norm_0 = (points[:, 1] - y_max) / (y_min - y_max)
    norm_1 = (points[:, 0] - x_min) / (x_max - x_min)

    # Check if points are within valid range
    valid_points = (norm_0 >= 0) & (norm_0 <= 1) & (norm_1 >= 0) & (norm_1 <= 1)
    mask = mask * valid_points

    if interpolate:
        # Compute float pixel indices
        pix_0 = norm_0 * (ks - 1)
        pix_1 = norm_1 * (ks - 1)
        pix_0_floor = pix_0.floor()
        pix_1_floor = pix_1.floor()

        # Bilinear weights
        w_b = pix_0 - pix_0_floor
        w_r = pix_1 - pix_1_floor
        w_b_1 = 1 - w_b
        w_r_1 = 1 - w_r

        # Pixel indices for 4 corners (clamped)
        r0 = pix_0_floor.long().clamp(0, ks - 1)
        c0 = pix_1_floor.long().clamp(0, ks - 1)
        r1 = (r0 + 1).clamp(0, ks - 1)
        c1 = (c0 + 1).clamp(0, ks - 1)

        # Pre-compute masked value once
        masked_value = mask * value

        # Use advanced indexing to increment the count for each corresponding pixel
        grid = torch.zeros(ks, ks, dtype=value.dtype, device=device)
        grid.index_put_((r0, c0), w_b_1 * w_r_1 * masked_value, accumulate=True)
        grid.index_put_((r0, c1), w_b_1 * w_r * masked_value, accumulate=True)
        grid.index_put_((r1, c0), w_b * w_r_1 * masked_value, accumulate=True)
        grid.index_put_((r1, c1), w_b * w_r * masked_value, accumulate=True)

    else:
        pix_0 = (norm_0 * (ks - 1)).floor().long().clamp(0, ks - 1)
        pix_1 = (norm_1 * (ks - 1)).floor().long().clamp(0, ks - 1)

        grid = torch.zeros(ks, ks, dtype=value.dtype, device=device)
        grid.index_put_((pix_0, pix_1), mask * value, accumulate=True)

    return grid