Skip to content

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

Current floating-point dtype of all owned tensors.
device dtype

Current compute device (set by to).
Source code in deeplens-src/deeplens/base.py 
22
23
def __init__(self, dtype=None):
    self.dtype = torch.get_default_dtype() if dtype is None else dtype

__str__ 

__str__()

Called when using print() and str()

Source code in deeplens-src/deeplens/base.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
def __str__(self):
    """Called when using print() and str()"""
    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)

Call the forward function.

Source code in deeplens-src/deeplens/base.py 
39
40
41
def __call__(self, inp):
    """Call the forward function."""
    return self.forward(inp)

clone 

clone()

Clone a DeepObj object.

Source code in deeplens-src/deeplens/base.py 
43
44
45
def clone(self):
    """Clone a DeepObj object."""
    return copy.deepcopy(self)

to 

to(device)

Move all tensors and nested objects to device.

Recursively walks over every instance attribute and moves tensors, nn.Module sub-objects, and nested DeepObj objects to the requested device.

Parameters:

Name Type Description Default
device

Target device, e.g. "cuda", "cpu", or a torch.device instance.

 required

Returns:

Name Type Description
DeepObj

self (for chaining).
Example

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

Source code in deeplens-src/deeplens/base.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def to(self, device):
    """Move all tensors and nested objects to *device*.

    Recursively walks over every instance attribute and moves tensors,
    ``nn.Module`` sub-objects, and nested ``DeepObj`` objects to the
    requested device.

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

    Returns:
        DeepObj: ``self`` (for chaining).

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

    for key, val in vars(self).items():
        if torch.is_tensor(val):
            setattr(self, key, val.to(device))
        elif isinstance(val, nn.Module):
            val.to(device)
        elif issubclass(type(val), DeepObj):
            val.to(device)
        elif val.__class__.__name__ in ("list", "tuple"):
            for i, v in enumerate(val):
                if torch.is_tensor(v):
                    val[i] = v.to(device)
                elif issubclass(type(v), DeepObj):
                    v.to(device)
    return self

astype 

astype(dtype)

Convert all floating-point tensors to dtype.

Also calls torch.set_default_dtype(dtype) so that subsequent tensor creation uses the same precision.

Parameters:

Name Type Description Default
dtype dtype

Target floating-point dtype. Must be one of torch.float16, torch.float32, or torch.float64. Pass None to be a no-op.

 required

Returns:

Name Type Description
DeepObj

self (for chaining).

Raises:

Type Description
AssertionError

If dtype is not a recognised floating-point dtype.
Example

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

Source code in deeplens-src/deeplens/base.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def astype(self, dtype):
    """Convert all floating-point tensors to *dtype*.

    Also calls ``torch.set_default_dtype(dtype)`` so that subsequent
    tensor creation uses the same precision.

    Args:
        dtype (torch.dtype): Target floating-point dtype.  Must be one of
            ``torch.float16``, ``torch.float32``, or ``torch.float64``.
            Pass ``None`` to be a no-op.

    Returns:
        DeepObj: ``self`` (for chaining).

    Raises:
        AssertionError: If *dtype* is not a recognised floating-point dtype.

    Example:
        >>> 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."

    if torch.get_default_dtype() != dtype:
        torch.set_default_dtype(dtype)
        print(f"Set {dtype} as default torch dtype.")

    self.dtype = dtype
    for key, val in vars(self).items():
        if torch.is_tensor(val) and val.dtype in dtype_ls:
            setattr(self, key, val.to(dtype))
        elif issubclass(type(val), DeepObj):
            val.astype(dtype)
        elif issubclass(type(val), list):
            for i, v in enumerate(val):
                if torch.is_tensor(v) and v.dtype in dtype_ls:
                    val[i] = v.to(dtype)
                elif issubclass(type(v), DeepObj):
                    v.astype(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, or specified inline as "n/V" (Cauchy approximation from Abbe number V).

Supported dispersion models: "sellmeier", "cauchy", "schott", and "interp" (lookup table).

Attributes:

Name Type Description
name str

Lowercase material name.
dispersion str

Dispersion model used ("sellmeier", "cauchy", "schott", or "interp").
n float

Refractive index at the d-line (587 nm).
V float

Abbe number.

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, e.g. "1.5168/64.17"
  • Custom name registered in materials_data.json

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.
Example

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 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
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, e.g. ``"1.5168/64.17"``
            * Custom name registered in ``materials_data.json``

            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:
        >>> 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.load_dispersion()
    self.device = device

load_dispersion 

load_dispersion()

Load material dispersion equation.

Source code in deeplens-src/deeplens/material/materials.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def load_dispersion(self):
    """Load material dispersion equation."""
    # 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.dispersion = "interp"
        mat_data = CUSTOM_data["INTERP_TABLE"][self.name]
        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 = float(np.interp(0.5893, 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

    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)

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

set_material_param_agf 

set_material_param_agf(material_data, material_name)

Set the material parameters and dispersion equation from AGF file.

Source code in deeplens-src/deeplens/material/materials.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def set_material_param_agf(self, material_data, material_name):
    """Set the material parameters and dispersion equation from AGF file."""
    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"]
    else:
        print(f"error: not {material_name}")

set_sellmeier_param 

set_sellmeier_param(params=None)

Manually set sellmeier parameters k1, l1, k2, l2, k3, l3.

This function is used when we want to manually set the sellmeier parameters for a custom material.

Source code in deeplens-src/deeplens/material/materials.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def set_sellmeier_param(self, params=None):
    """Manually set sellmeier parameters k1, l1, k2, l2, k3, l3.

    This function is used when we want to manually set the sellmeier parameters for a custom material.
    """
    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 given wvln.

Source code in deeplens-src/deeplens/material/materials.py 
260
261
262
263
264
265
266
def refractive_index(self, wvln):
    """Compute the refractive index at given wvln."""
    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 at given wvln.

Source code in deeplens-src/deeplens/material/materials.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
def ior(self, wvln):
    """Compute the refractive index at given wvln."""
    assert wvln.min() > 0.1 and wvln.max() < 10, "Wavelength should be in [um]."

    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 == "optimizable":
        # Cauchy's equation, calculate (A, B) on the fly
        B = (self.n - 1) / self.V / (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) paramters to (A, B) parameters to find the material.

Source code in deeplens-src/deeplens/material/materials.py 
345
346
347
348
349
350
351
352
353
354
355
@staticmethod
def nV_to_AB(n, V):
    """Convert (n ,V) paramters to (A, B) parameters to find the material."""

    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)

Find the closest material in the CDGM common glasses database.

Source code in deeplens-src/deeplens/material/materials.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def match_material(self, mat_table=None):
    """Find the closest material in the CDGM common glasses database."""
    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])

Optimize the material parameters (n, V).

Optimizing refractive index is more important than optimizing Abbe number.

Parameters:

Name Type Description Default
lrs list

learning rates for n and V. Defaults to [1e-4, 1e-4].

 [0.0001, 0.01]
Source code in deeplens-src/deeplens/material/materials.py 
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
def get_optimizer_params(self, lrs=[1e-4, 1e-2]):
    """Optimize the material parameters (n, V). 

    Optimizing refractive index is more important than optimizing Abbe number.

    Args:
        lrs (list): learning rates for n and V. Defaults to [1e-4, 1e-4].
    """
    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

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, 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.

A surface sits at axial position d (mm) in the global coordinate system, 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 by three stages, implemented in ray_reaction:

  1. Coordinate transform – 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 Tensor

Axial position of the surface vertex [mm].
r float

Aperture radius [mm].
mat2 Material

Optical material on the transmission side.
is_square bool

If True the aperture is square; otherwise circular.

Initialize a generic optical surface.

Parameters:

Name Type Description Default
r float

Aperture radius [mm].

 required
d float

Axial position of the surface 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 normal direction. Defaults to [0.0, 0.0, 1.0] (on-axis).

 [0.0, 0.0, 1.0]
is_square bool

Use a square aperture. Defaults to False.

 False
device str

Compute device. Defaults to "cpu".

 'cpu'
Source code in deeplens-src/deeplens/geometric_surface/base.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def __init__(
    self,
    r,
    d,
    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].
        d (float): Axial position of the surface 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 normal direction.
            Defaults to ``[0.0, 0.0, 1.0]`` (on-axis).
        is_square (bool, optional): Use a square aperture.
            Defaults to ``False``.
        device (str, optional): Compute device. Defaults to ``"cpu"``.
    """
    super(Surface, self).__init__()

    # Global direction vector, always pointing to the positive z-axis
    self.vec_global = torch.tensor([0.0, 0.0, 1.0])

    # Surface position in global coordinate system
    self.d = torch.tensor(d)
    self.pos_x = torch.tensor(pos_xy[0])
    self.pos_y = torch.tensor(pos_xy[1])

    # Surface direction vector in global coordinate system
    self.vec_local = F.normalize(torch.tensor(vec_local), p=2, dim=-1)

    # Material after the surface
    self.mat2 = Material(mat2)

    # Surface aperture radius (non-differentiable).
    # For a square aperture, r is the circumscribed-circle radius
    # (i.e. the half-diagonal), so the side length is r * sqrt(2).
    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))

    # 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

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

    # Pre-compute rotation matrices (depends only on static vec_local/vec_global)
    self._cache_rotation_matrices()

init_from_dict classmethod 

init_from_dict(surf_dict)

Initialize surface from a dict.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
123
124
125
126
127
128
@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize surface from a dict."""
    raise NotImplementedError(
        f"init_from_dict() is not implemented for {cls.__name__}."
    )

ray_reaction 

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

Compute the output ray after intersection and refraction/reflection.

Transforms the ray to the local surface frame, solves the intersection via Newton's method, applies vector Snell's law (or specular reflection), then transforms back to global coordinates.

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 (default) refract the ray; if False reflect it.

 True

Returns:

Name Type Description
Ray

Updated ray bundle after the surface interaction.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def ray_reaction(self, ray, n1, n2, refraction=True):
    """Compute the output ray after intersection and refraction/reflection.

    Transforms the ray to the local surface frame, solves the intersection
    via Newton's method, applies vector Snell's law (or specular reflection),
    then transforms back to global coordinates.

    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`` (default) refract the ray;
            if ``False`` reflect it.

    Returns:
        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 to global coordinate system
    ray = self.to_global_coord(ray)

    return ray

intersect 

intersect(ray, n=1.0)

Solve ray-surface intersection in local coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray.

 required
n float

refractive index. Defaults to 1.0.

 1.0
Source code in deeplens-src/deeplens/geometric_surface/base.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in local coordinate system.

    Args:
        ray (Ray): input ray.
        n (float, optional): refractive index. Defaults to 1.0.
    """
    # 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:
        if t.abs().max() > 100 and torch.get_default_dtype() == torch.float32:
            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

newtons_method 

newtons_method(ray)

Solve intersection by Newton's method in local coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray.

 required

Returns:

Name Type Description
t tensor

intersection time.
valid tensor

valid mask.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
def newtons_method(self, ray):
    """Solve intersection by Newton's method in local coordinate system.

    Args:
        ray (Ray): input ray.

    Returns:
        t (tensor): intersection time.
        valid (tensor): valid mask.
    """
    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]

    # Initial guess of t (can also use spherical surface for initial guess)
    t = -ray.o[..., 2] / dzdt

    # 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. Determine valid solutions — reuse ft and valid from the diff step
    with torch.no_grad():
        valid = valid & (ft.abs() < newton_convergence)

    return t, valid

refract 

refract(ray, eta)

Calculate refracted ray according to Snell's law in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from. d is already normalized if both n and ray.d are normalized.

Parameters:

Name Type Description Default
ray Ray

incident ray.

 required
eta float

ratio of indices of refraction, eta = n_i / n_t

 required

Returns:

Name Type Description
ray Ray

refracted ray.
References

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

Source code in deeplens-src/deeplens/geometric_surface/base.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def refract(self, ray, eta):
    """Calculate refracted ray according to Snell's law in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from. d is already normalized if both n and ray.d are normalized.

    Args:
        ray (Ray): incident ray.
        eta (float): ratio of indices of refraction, eta = n_i / n_t

    Returns:
        ray (Ray): refracted ray.

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

    # Compute refraction according to Snell's law, normal_vec * ray_d
    dot_product = (-normal_vec * ray.d).sum(-1).unsqueeze(-1)
    k = 1 - eta**2 * (1 - dot_product**2)

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

    # Update ray direction
    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)

    # Update ray valid mask
    ray.is_valid = ray.is_valid * valid

    return ray

bend_penalty 

bend_penalty(ray, old_d)

Accumulate soft per-surface bend penalty onto ray.

Penalty rises smoothly once the bend angle between old_d and the refracted ray.d exceeds bend_angle_max and stays near zero for mild refractions. Normalization by bend_scale = 1 - cos_bend_min keeps the gradient magnitude consistent with the CRA loss in optim.py.

Parameters:

Name Type Description Default
ray Ray

ray after refraction (ray.d is the new direction).

 required
old_d Tensor

pre-refraction ray directions, same shape as ray.d.

 required

Returns:

Name Type Description
Ray

ray with updated bend_penalty.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
def bend_penalty(self, ray, old_d):
    """Accumulate soft per-surface bend penalty onto ray.

    Penalty rises smoothly once the bend angle between ``old_d`` and the
    refracted ``ray.d`` exceeds ``bend_angle_max`` and stays near zero for
    mild refractions.  Normalization by ``bend_scale = 1 - cos_bend_min`` keeps
    the gradient magnitude consistent with the CRA loss in optim.py.

    Args:
        ray (Ray): ray after refraction (ray.d is the new direction).
        old_d (Tensor): pre-refraction ray directions, same shape as ray.d.

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

reflect 

reflect(ray)

Calculate reflected ray in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from.

Parameters:

Name Type Description Default
ray Ray

incident ray.

 required

Returns:

Name Type Description
ray Ray

reflected ray.
References

[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.py 
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def reflect(self, ray):
    """Calculate reflected ray in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from.

    Args:
        ray (Ray): incident ray.

    Returns:
        ray (Ray): reflected ray.

    References:
        [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)

Calculate surface normal vector at the intersection point in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from.

Parameters:

Name Type Description Default
ray Ray

input ray.

 required

Returns:

Name Type Description
n_vec tensor

surface normal vector.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
def normal_vec(self, ray):
    """Calculate surface normal vector at the intersection point in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from.

    Args:
        ray (Ray): input ray.

    Returns:
        n_vec (tensor): surface normal vector.
    """
    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

to_local_coord 

to_local_coord(ray)

Transform ray to local coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray in global coordinate system.

 required

Returns:

Name Type Description
ray Ray

transformed ray in local coordinate system.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
def to_local_coord(self, ray):
    """Transform ray to local coordinate system.

    Args:
        ray (Ray): input ray in global coordinate system.

    Returns:
        ray (Ray): transformed ray in local coordinate system.
    """
    # Shift ray origin to surface origin
    offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o - offset

    # Rotate ray origin and direction
    if torch.abs(torch.dot(self.vec_local, self.vec_global) - 1.0) > EPSILON:
        R = self._get_rotation_matrix(self.vec_local, self.vec_global)
        ray.o = self._apply_rotation(ray.o, R)
        ray.d = self._apply_rotation(ray.d, R)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    return ray

to_global_coord 

to_global_coord(ray)

Transform ray to global coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray in local coordinate system.

 required

Returns:

Name Type Description
ray Ray

transformed ray in global coordinate system.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
def to_global_coord(self, ray):
    """Transform ray to global coordinate system.

    Args:
        ray (Ray): input ray in local coordinate system.

    Returns:
        ray (Ray): transformed ray in global coordinate system.
    """
    # Rotate ray origin and direction
    if torch.abs(torch.dot(self.vec_local, self.vec_global) - 1.0) > EPSILON:
        R = self._get_rotation_matrix(self.vec_global, self.vec_local)
        ray.o = self._apply_rotation(ray.o, R)
        ray.d = self._apply_rotation(ray.d, R)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    # Shift ray origin back to global coordinates
    offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o + offset

    return ray

sag 

sag(x, y, valid=None)

Calculate sag (z) of the surface: z = f(x, y).

Valid term is used to avoid NaN when x, y exceed the data range, which happens in spherical and aspherical surfaces.

Calculating r = sqrt(x2, y2) may cause an NaN error during back-propagation. Because dr/dx = x / sqrt(x2 + y2), NaN will occur when x=y=0.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
483
484
485
486
487
488
489
490
491
492
493
494
def sag(self, x, y, valid=None):
    """Calculate sag (z) of the surface: z = f(x, y).

    Valid term is used to avoid NaN when x, y exceed the data range, which happens in spherical and aspherical surfaces.

    Calculating r = sqrt(x**2, y**2) may cause an NaN error during back-propagation. Because dr/dx = x / sqrt(x**2 + y**2), NaN will occur when x=y=0.
    """
    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 derivatives of surface function. Surface function: f(x, y, z): sag(x, y) - z = 0. This function is used in Newton's method and normal vector calculation.

There are several methods to compute derivatives of surfaces

[1] Analytical derivatives: The current implementation is based on this method. But the implementation only works for surfaces which can be written as z = sag(x, y). For implicit surfaces, we need to compute derivatives (df/dx, df/dy, df/dz). [2] Numerical derivatives: Use finite difference method to compute derivatives. This can be used for those very complex surfaces, for example, NURBS. But it may suffer from numerical instability when the surface is very steep. [3] Automatic differentiation: Use torch.autograd to compute derivatives. This can work for almost all the surfaces and is accurate, but it requires an extra backward pass to compute the derivatives of the surface function.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
511
512
513
514
515
516
517
518
519
520
521
522
523
524
def dfdxyz(self, x, y, valid=None):
    """Compute derivatives of surface function. Surface function: f(x, y, z): sag(x, y) - z = 0. This function is used in Newton's method and normal vector calculation.

    There are several methods to compute derivatives of surfaces:
        [1] Analytical derivatives: The current implementation is based on this method. But the implementation only works for surfaces which can be written as z = sag(x, y). For implicit surfaces, we need to compute derivatives (df/dx, df/dy, df/dz).
        [2] Numerical derivatives: Use finite difference method to compute derivatives. This can be used for those very complex surfaces, for example, NURBS. But it may suffer from numerical instability when the surface is very steep.
        [3] Automatic differentiation: Use torch.autograd to compute derivatives. This can work for almost all the surfaces and is accurate, but it requires an extra backward pass to compute the derivatives of the surface function.
    """
    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 surface function f(x, y, z): sag(x, y) - z = 0. This function is currently only used for surfaces constraints.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
def d2fdxyz2(self, x, y, valid=None):
    """Compute second-order partial derivatives of the surface function f(x, y, z): sag(x, y) - z = 0. This function is currently only used for surfaces constraints."""
    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)

Valid points within the data range and boundary of the surface.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
586
587
588
def is_valid(self, x, y):
    """Valid points within the data range and boundary of the surface."""
    return self.is_within_data_range(x, y) & self.is_within_boundary(x, y)

is_within_boundary 

is_within_boundary(x, y)

Valid points within the boundary of the surface.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
590
591
592
593
594
595
596
597
598
599
600
def is_within_boundary(self, x, y):
    """Valid points within the boundary of the surface."""
    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)

Valid points inside the data region of the sag function.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
602
603
604
def is_within_data_range(self, x, y):
    """Valid points inside the data region of the sag function."""
    return torch.ones_like(x, dtype=torch.bool)

max_height 

max_height()

Maximum valid height.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
606
607
608
def max_height(self):
    """Maximum valid height."""
    return 10e3

surface_with_offset 

surface_with_offset(x, y, valid_check=True)

Calculate z coordinate of the surface at (x, y).

This function is used in lens setup plotting and lens self-intersection detection.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
610
611
612
613
614
615
616
617
618
619
620
def surface_with_offset(self, x, y, valid_check=True):
    """Calculate z coordinate of the surface at (x, y).

    This function is used in lens setup plotting and lens self-intersection detection.
    """
    x = x if torch.is_tensor(x) else torch.tensor(x, device=self.device)
    y = y if torch.is_tensor(y) else torch.tensor(y, device=self.device)
    if valid_check:
        return self.sag(x, y) + self.d
    else:
        return self._sag(x, y) + self.d

surface_sag 

surface_sag(x, y)

Calculate sag of the surface at (x, y).

This function is currently not used.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
622
623
624
625
626
627
628
629
def surface_sag(self, x, y):
    """Calculate sag of the surface at (x, y).

    This function is currently not used.
    """
    x = x if torch.is_tensor(x) else torch.tensor(x, device=self.device)
    y = y if torch.is_tensor(y) else torch.tensor(y, device=self.device)
    return self.sag(x, y).item()

get_optimizer_params 

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

Get optimizer parameters for different parameters.

Parameters:

Name Type Description Default
lrs list

learning rates for different parameters.

 [0.0001]
optim_mat bool

whether to optimize material. Defaults to False.

 False
Source code in deeplens-src/deeplens/geometric_surface/base.py 
635
636
637
638
639
640
641
642
643
644
645
646
def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Get optimizer parameters for different parameters.

    Args:
        lrs (list): learning rates for different parameters.
        optim_mat (bool): whether to optimize material. Defaults to False.
    """
    raise NotImplementedError(
        "get_optimizer_params() is not implemented for {}".format(
            self.__class__.__name__
        )
    )

get_optimizer 

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

Get optimizer for the surface.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
648
649
650
651
def get_optimizer(self, lrs=[1e-4], optim_mat=False):
    """Get 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 surface radius.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
653
654
655
656
def update_r(self, r):
    """Update surface radius."""
    r_max = self.max_height()
    self.r = min(r, r_max)

draw_r 

draw_r()

Effective drawing radius, clamped to the valid data range.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
661
662
663
def draw_r(self):
    """Effective drawing radius, clamped to the valid data range."""
    return min(self.r, self.max_height())

draw_widget 

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

Draw widget for the surface on the 2D plot.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
665
666
667
668
669
670
671
672
673
674
675
676
677
678
def draw_widget(self, ax, color="black", linestyle="solid"):
    """Draw widget for the surface on the 2D plot."""
    r_eff = self.draw_r()
    r = torch.linspace(-r_eff, r_eff, 128, device=self.device)
    z = self.surface_with_offset(
        r, torch.zeros(len(r), device=self.device), valid_check=False
    )
    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])

Create triangulated surface mesh.

Parameters:

Name Type Description Default
n_rings int

Number of concentric rings for sampling.

 32
n_arms int

Number of angular divisions.

 128
color List[float]

The color of the mesh.

 [0.06, 0.3, 0.6]

Returns:

Name Type Description
self

The surface with mesh data.
Source code in deeplens-src/deeplens/geometric_surface/base.py 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
def create_mesh(self, n_rings=32, n_arms=128, color=[0.06, 0.3, 0.6]):
    """Create triangulated surface mesh.

    Args:
        n_rings (int): Number of concentric rings for sampling.
        n_arms (int): Number of angular divisions.
        color (List[float]): The color of the mesh.

    Returns:
        self: The surface with mesh data.
    """
    self.vertices = self._create_vertices(n_rings, n_arms)
    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()

Get PyVista PolyData object from previously generated vertices and faces.

PolyData object will be used to draw the surface and export as .obj file.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
def get_polydata(self):
    """Get PyVista PolyData object from previously generated vertices and faces.

    PolyData object will be used to draw the surface and export as .obj file.
    """
    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)

zmx_str 

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in deeplens-src/deeplens/geometric_surface/base.py 
811
812
813
814
815
def zmx_str(self, surf_idx, d_next):
    """Return Zemax surface string."""
    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, 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.

The sag function is:

.. math::

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]. Differentiable with respect to gradient-based optimization.

Initialize a spherical surface.

Parameters:

Name Type Description Default
c float

Surface curvature 1/R [1/mm]. Use 0 for a flat surface (equivalent to 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 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/spheric.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def __init__(
    self,
    c,
    r,
    d,
    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 (equivalent to ``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 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"``.
    """
    super(Spheric, self).__init__(
        r=r,
        d=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.c = torch.tensor(c)
    self.to(device)

intersect 

intersect(ray, n=1.0)

Solve ray-surface intersection in local coordinate system using analytical method.

Sphere equation: (x)^2 + (y)^2 + (z - R)^2 = R^2, where R = 1/c Ray equation: p(t) = o + t*d Solve quadratic equation for intersection parameter t.

Parameters:

Name Type Description Default
ray Ray

input ray.

 required
n float

refractive index. Defaults to 1.0.

 1.0

Returns:

Name Type Description
ray Ray

ray with updated position and opl.
Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in local coordinate system using analytical method.

    Sphere equation: (x)^2 + (y)^2 + (z - R)^2 = R^2, where R = 1/c
    Ray equation: p(t) = o + t*d
    Solve quadratic equation for intersection parameter t.

    Args:
        ray (Ray): input ray.
        n (float, optional): refractive index. Defaults to 1.0.

    Returns:
        ray (Ray): ray with updated position and opl.
    """
    c = self.c

    if torch.abs(c) < EPSILON:
        # Handle flat surface as a plane
        t = (0.0 - ray.o[..., 2]) / ray.d[..., 2]
        new_o = ray.o + t.unsqueeze(-1) * ray.d
        valid = (new_o[..., 0] ** 2 + new_o[..., 1] ** 2 < self.r**2) & (
            ray.is_valid > 0
        )
    else:
        R = 1.0 / c

        # Vector from ray origin to sphere center at (0, 0, R)
        oc = ray.o.clone()
        oc[..., 2] = oc[..., 2] - R

        # Quadratic equation: a*t^2 + b*t + c = 0
        # a = d·d = 1 (since ray direction is normalized)
        # b = 2*(o-center)·d
        # c = (o-center)·(o-center) - R^2

        a = torch.sum(ray.d * ray.d, dim=-1)  # Should be 1 for normalized rays
        b = 2.0 * torch.sum(oc * ray.d, dim=-1)
        c_coeff = torch.sum(oc * oc, dim=-1) - R * R

        discriminant = b * b - 4 * a * c_coeff
        valid_intersect = discriminant >= 0

        sqrt_discriminant = torch.sqrt(torch.clamp(discriminant, min=EPSILON))
        t1 = (-b - sqrt_discriminant) / (2 * a + EPSILON)
        t2 = (-b + sqrt_discriminant) / (2 * a + EPSILON)

        # Choose intersection closest to z=0 (surface vertex)
        z1 = ray.o[..., 2] + t1 * ray.d[..., 2]
        z2 = ray.o[..., 2] + t2 * ray.d[..., 2]
        use_t1 = torch.abs(z1) < torch.abs(z2)
        t = torch.where(use_t1, t1, t2)

        new_o = ray.o + t.unsqueeze(-1) * ray.d

        # Check aperture
        r_squared = new_o[..., 0] ** 2 + new_o[..., 1] ** 2
        within_aperture = r_squared <= (self.r**2 + EPSILON)

        valid = valid_intersect & within_aperture & (ray.is_valid > 0)

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

    if ray.is_coherent:
        if t.abs().max() > 100 and torch.get_default_dtype() == torch.float32:
            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

is_within_data_range 

is_within_data_range(x, y)

Invalid when shape is non-defined.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
213
214
215
216
217
def is_within_data_range(self, x, y):
    """Invalid when shape is non-defined."""
    c = self.c
    valid = (x**2 + y**2) < 1 / c**2
    return valid

max_height 

max_height()

Maximum valid height.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
219
220
221
222
223
def max_height(self):
    """Maximum valid height."""
    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)

Activate gradient computation for c and d and return optimizer parameters.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
def get_optimizer_params(self, lrs=[1e-4, 1e-4], optim_mat=False):
    """Activate gradient computation for c and d and return optimizer parameters."""
    self.c.requires_grad_(True)
    self.d.requires_grad_(True)

    params = []
    params.append({"params": [self.d], "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()

Return surface parameters.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
245
246
247
248
249
250
251
252
253
254
255
256
257
def surf_dict(self):
    """Return surface parameters."""
    roc = 1 / self.c.item() if self.c.item() != 0 else 0.0
    surf_dict = {
        "type": "Spheric",
        "r": round(self.r, 4),
        "(c)": round(self.c.item(), 4),
        "roc": round(roc, 4),
        "(d)": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }

    return surf_dict

zmx_str 

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in deeplens-src/deeplens/geometric_surface/spheric.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
    def zmx_str(self, surf_idx, d_next):
        """Return Zemax surface string."""
        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, 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:

.. math::

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 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def __init__(
    self,
    r,
    d,
    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=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.c = torch.tensor(c)
    self.k = torch.tensor(k)

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

    if ai is not None and len(ai) > 0:
        self.ai = torch.tensor(ai)
        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.tensor(a))
    else:
        self.ai = None
        self.ai_degree = 0

    self.to(device)

is_within_data_range 

is_within_data_range(x, y)

Invalid when shape is non-defined.

Fully tensorized (no Python branch on the tensor value of k) so the function is safe to trace through torch.compile. When k <= -1 the conic has no real boundary, so every point is treated as valid.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def is_within_data_range(self, x, y):
    """Invalid when shape is non-defined.

    Fully tensorized (no Python branch on the tensor value of ``k``) so
    the function is safe to trace through ``torch.compile``. When
    ``k <= -1`` the conic has no real boundary, so every point is
    treated as valid.
    """
    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()

Maximum valid height.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py 
227
228
229
230
231
232
def max_height(self):
    """Maximum valid height."""
    c, k = self._get_curvature_params()
    if k > -1:
        return torch.sqrt(1 / (k + 1) / (c**2)).item() - 0.001
    return 10e3

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

learning rates for [d, c, k, ai].

 [0.0001, 0.0001, 0.01, 0.0001]
optim_mat bool

whether to optimize material. Defaults to False.

 False
Source code in deeplens-src/deeplens/geometric_surface/aspheric.py 
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
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, optional): learning rates for ``[d, c, k, ai]``.
        optim_mat (bool, optional): whether to optimize material.
            Defaults to False.
    """
    params = []

    # Optimize distance
    self.d.requires_grad_(True)
    params.append({"params": [self.d], "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()

Return a dict of surface.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
def surf_dict(self):
    """Return a dict of surface."""
    has_ai2 = self.ai2 is not None
    surf_dict = {
        "type": "Aspheric",
        "r": round(self.r, 4),
        "(c)": round(self.c.item(), 4),
        "roc": round(1 / self.c.item(), 4),
        "d": round(self.d.item(), 4),
        "k": round(self.k.item(), 4),
        "ai": [],
        "use_ai2": has_ai2,
        "mat2": self.mat2.get_name(),
    }

    # 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"] = float(format(self.ai2.item(), ".6e"))
        surf_dict["ai"].append(float(format(self.ai2.item(), ".6e")))

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

    return surf_dict

zmx_str 

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in deeplens-src/deeplens/geometric_surface/aspheric.py 
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
    def zmx_str(self, surf_idx, d_next):
        """Return Zemax surface string."""
        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, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Plane

Aperture surface.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(
    self,
    r,
    d,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Aperture surface."""
    Plane.__init__(
        self,
        r=r,
        d=d,
        mat2="air",
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.tolerancing = False
    self.to(device)

ray_reaction 

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

Compute output ray after intersection and refraction.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
43
44
45
46
47
48
def ray_reaction(self, ray, n1=1.0, n2=1.0, refraction=False):
    """Compute output ray after intersection and refraction."""
    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')

Draw aperture wedge on the figure.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def draw_widget(self, ax, color="orange", linestyle="solid"):
    """Draw aperture wedge on the figure."""
    d = self.d.item()
    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')

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

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
73
74
75
76
77
78
79
80
81
82
83
84
def draw_widget3D(self, ax, color="black"):
    """Draw the aperture as a circle in a 3D 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, self.d.item())  # Constant z at aperture position

    # 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])

Create triangulated surface mesh.

Parameters:

Name Type Description Default
n_rings int

Number of concentric rings for sampling.

 32
n_arms int

Number of angular divisions.

 128
color List[float]

The color of the mesh.

 [0.0, 0.0, 0.0]

Returns:

Name Type Description
self

The surface with mesh data.
Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def create_mesh(self, n_rings=32, n_arms=128, color=[0.0, 0.0, 0.0]):
    """Create triangulated surface mesh.

    Args:
        n_rings (int): Number of concentric rings for sampling.
        n_arms (int): Number of angular divisions.
        color (List[float]): The color of the mesh.

    Returns:
        self: The surface with mesh data.
    """
    self.vertices = self._create_vertices(n_rings, n_arms)
    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])

Activate gradient computation for d and return optimizer parameters.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
166
167
168
169
170
171
172
173
def get_optimizer_params(self, lrs=[1e-4]):
    """Activate gradient computation for d and return optimizer parameters."""
    self.d.requires_grad_(True)

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

    return params

surf_dict 

surf_dict()

Dict of surface parameters.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
178
179
180
181
182
183
184
185
186
187
def surf_dict(self):
    """Dict of surface parameters."""
    surf_dict = {
        "type": "Aperture",
        "r": round(self.r, 4),
        "(d)": round(self.d.item(), 4),
        "mat2": "air",
        "is_square": self.is_square,
    }
    return surf_dict

zmx_str 

zmx_str(surf_idx, d_next)

Zemax surface string.

Source code in deeplens-src/deeplens/geometric_surface/aperture.py 
189
190
191
192
193
194
195
196
197
    def zmx_str(self, surf_idx, d_next):
        """Zemax surface string."""
        zmx_str = f"""SURF {surf_idx}
    STOP
    TYPE STANDARD
    CURV 0.0
    DISZ {d_next.item()}
"""
        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, 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).

This is the base class that provides common functionality for all phase parameterizations. Specific parameterizations should inherit from this class.

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

Source code in deeplens-src/deeplens/phase_surface/phase.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    r,
    d,
    norm_radii=None,
    mat2="air",
    pos_xy=(0.0, 0.0),
    vec_local=(0.0, 0.0, 1.0),
    is_square=True,
    device="cpu",
):
    super().__init__()

    # Global direction vector, always pointing to the positive z-axis
    self.vec_global = torch.tensor([0.0, 0.0, 1.0])

    # Surface position in global coordinate system
    self.d = torch.tensor(d)
    self.pos_x = torch.tensor(pos_xy[0])
    self.pos_y = torch.tensor(pos_xy[1])

    # Surface direction vector in global coordinate system
    self.vec_local = F.normalize(torch.tensor(vec_local), p=2, dim=-1)

    # Material after the surface
    self.mat2 = Material(mat2)

    # DOE geometry
    self.r = float(r)
    self.is_square = is_square
    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

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

    # Pre-compute rotation matrices (depends only on static vec_local/vec_global)
    self._cache_rotation_matrices()

phi 

phi(x, y)

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

Source code in deeplens-src/deeplens/phase_surface/phase.py 
86
87
88
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/phase.py 
90
91
92
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)

Ray reaction on DOE surface.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
 97
 98
 99
100
101
102
103
104
def ray_reaction(self, ray, n1, n2):
    """Ray reaction on DOE surface."""
    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 coordinate system and update ray data.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def intersect(self, ray, n=1.0):
    """Solve ray-plane intersection in local coordinate system and update ray data."""
    # Solve intersection
    t = (0.0 - ray.o[..., 2]) / ray.d[..., 2]
    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)

Diffraction of a phase surface.

Step 1. The phase φ in radians adds to the optical path length of the ray. Step 2. The gradient of the phase profile (phase slope) changes the direction of rays via the generalized Snell's law: n₂·sin(θ₂) = n₁·sin(θ₁) + m · λ / (2π) · dφ/dx

After standard refraction has already been applied, the remaining phase deflection on the unit direction vector is: Δl = m · λ / (2π · n₂) · dφ/dx

Parameters:

Name Type Description Default
ray

Ray object with position, direction, and wavelength.

 required
n2

Refractive index of the medium after the surface. Required for correct generalized Snell's law: deflection scales as 1/n₂.

 1.0
Note

Material dispersion is not modelled here. The phase profile φ(x,y) is treated 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 itself changes with wavelength (via (n(λ)-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/phase.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def diffract(self, ray, n2=1.0):
    """Diffraction of a phase surface.

    Step 1. The phase φ in radians adds to the optical path length of the ray.
    Step 2. The gradient of the phase profile (phase slope) changes the direction
       of rays via the generalized Snell's law:
       ``n₂·sin(θ₂) = n₁·sin(θ₁) + m · λ / (2π) · dφ/dx``

       After standard refraction has already been applied, the remaining
       phase deflection on the unit direction vector is:
       ``Δl = m · λ / (2π · n₂) · dφ/dx``

    Args:
        ray: Ray object with position, direction, and wavelength.
        n2: Refractive index of the medium after the surface. Required for
            correct generalized Snell's law: deflection scales as 1/n₂.

    Note:
        Material dispersion is not modelled here. The phase profile ``φ(x,y)``
        is treated 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 itself changes with wavelength
        (via ``(n(λ)-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)

Calculate refracted ray according to Snell's law in local coordinate system.

Parameters:

Name Type Description Default
ray Ray

incident ray.

 required
eta float

ratio of indices of refraction, eta = n_i / n_t

 required

Returns:

Name Type Description
ray Ray

refracted ray.
Source code in deeplens-src/deeplens/phase_surface/phase.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def refract(self, ray, eta):
    """Calculate refracted ray according to Snell's law in local coordinate system.

    Args:
        ray (Ray): incident ray.
        eta (float): ratio of indices of refraction, eta = n_i / n_t

    Returns:
        ray (Ray): refracted ray.
    """
    # Compute normal vectors
    normal_vec = self.normal_vec(ray)

    # Compute refraction according to Snell's law
    dot_product = (-normal_vec * ray.d).sum(-1).unsqueeze(-1)
    k = 1 - eta**2 * (1 - dot_product**2)

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

    # Update ray direction
    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)

    # Update ray valid mask
    ray.is_valid = ray.is_valid * valid

    return ray

normal_vec 

normal_vec(ray)

Calculate surface normal vector at intersection points.

Normal vector points from the surface toward the side where the light is coming from.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
222
223
224
225
226
227
228
229
230
231
def normal_vec(self, ray):
    """Calculate surface normal vector at intersection points.

    Normal vector points from the surface toward the side where the light is coming from.
    """
    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

to_local_coord 

to_local_coord(ray)

Transform ray to local coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray in global coordinate system.

 required

Returns:

Name Type Description
ray Ray

transformed ray in local coordinate system.
Source code in deeplens-src/deeplens/phase_surface/phase.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
def to_local_coord(self, ray):
    """Transform ray to local coordinate system.

    Args:
        ray (Ray): input ray in global coordinate system.

    Returns:
        ray (Ray): transformed ray in local coordinate system.
    """
    # Shift ray origin to surface origin
    offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o - offset

    # Rotate ray origin and direction
    if torch.abs(torch.dot(self.vec_local, self.vec_global) - 1.0) > EPSILON:
        R = self._get_rotation_matrix(self.vec_local, self.vec_global)
        ray.o = self._apply_rotation(ray.o, R)
        ray.d = self._apply_rotation(ray.d, R)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    return ray

to_global_coord 

to_global_coord(ray)

Transform ray to global coordinate system.

Parameters:

Name Type Description Default
ray Ray

input ray in local coordinate system.

 required

Returns:

Name Type Description
ray Ray

transformed ray in global coordinate system.
Source code in deeplens-src/deeplens/phase_surface/phase.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
def to_global_coord(self, ray):
    """Transform ray to global coordinate system.

    Args:
        ray (Ray): input ray in local coordinate system.

    Returns:
        ray (Ray): transformed ray in global coordinate system.
    """
    # Rotate ray origin and direction
    if torch.abs(torch.dot(self.vec_local, self.vec_global) - 1.0) > EPSILON:
        R = self._get_rotation_matrix(self.vec_global, self.vec_local)
        ray.o = self._apply_rotation(ray.o, R)
        ray.d = self._apply_rotation(ray.d, R)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    # Shift ray origin back to global coordinates
    offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o + offset

    return ray

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/phase.py 
323
324
325
326
327
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)

Generate optimizer.

Parameters:

Name Type Description Default
lrs list or float

Learning rates for different parameters.

 required
Source code in deeplens-src/deeplens/phase_surface/phase.py 
329
330
331
332
333
334
335
336
337
338
339
def get_optimizer(self, lrs):
    """Generate optimizer.

    Args:
        lrs (list or float): Learning rates for different parameters.
    """
    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. 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.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
341
342
343
344
345
346
347
348
def update_r(self, r):
    """Update surface radius. 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.
    """
    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

φ = 2π/λ · (n−1) · h → h = φ · λ / (2π · (n−1))

Parameters:

Name Type Description Default
design_wvln

Design wavelength [um].

 required
refractive_idx

Refractive index of the DOE material at design_wvln.

 1.5
res

Pixel resolution of the returned height map (square grid).

 512

Returns:

Type Description

Tensor of shape [res, res] with height values in the same units

as design_wvln (um).
Source code in deeplens-src/deeplens/phase_surface/phase.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
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:
        φ = 2π/λ · (n−1) · h  →  h = φ · λ / (2π · (n−1))

    Args:
        design_wvln: Design wavelength [um].
        refractive_idx: Refractive index of the DOE material at ``design_wvln``.
        res: Pixel resolution of the returned height map (square grid).

    Returns:
        Tensor of shape ``[res, res]`` with height values in the same units
        as ``design_wvln`` (um).
    """
    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/phase.py 
378
379
380
def draw_r(self):
    """Effective drawing radius for 2D layout drawing."""
    return self.r

surface_with_offset 

surface_with_offset(*args, **kwargs)

Surface sag with offset, only used in layout drawing.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
382
383
384
def surface_with_offset(self, *args, **kwargs):
    """Surface sag with offset, only used in layout drawing."""
    return self.d

draw_phase_map 

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

Draw phase map. Range from [0, 2*pi].

Source code in deeplens-src/deeplens/phase_surface/phase.py 
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
def draw_phase_map(self, save_name="./DOE_phase_map.png"):
    """Draw phase map. Range from [0, 2*pi]."""
    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='-')

Draw DOE as a two-side surface.

Source code in deeplens-src/deeplens/phase_surface/phase.py 
404
405
406
407
408
409
410
411
412
413
414
415
416
def draw_widget(self, ax, color="black", linestyle="-"):
    """Draw DOE as a two-side surface."""
    max_offset = self.d.item() / 100
    d = self.d.item()

    # 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/phase.py 
421
422
423
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/phase.py 
425
426
427
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/phase.py 
429
430
431
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, 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

Binary2 phase on a plane substrate.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def __init__(
    self,
    r,
    d,
    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",
):
    super().__init__(
        r=r,
        d=d,
        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)

Initialize Binary2 phase surface from dictionary.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize Binary2 phase surface from dictionary."""
    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"],
        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)

Reference phase map at design wavelength.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def phi(self, x, y):
    """Reference phase map at design wavelength."""
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points."""
    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)

Generate optimizer parameters.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def get_optimizer_params(self, lrs=[1e-4, 1e-2], optim_mat=False):
    """Generate optimizer parameters."""
    params = []

    # Optimize position
    self.d.requires_grad = True
    params.append({"params": [self.d], "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 Binary2 DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def save_ckpt(self, save_path="./binary2_doe.pth"):
    """Save Binary2 DOE parameters."""
    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 DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
148
149
150
151
152
153
154
155
156
157
def load_ckpt(self, load_path="./binary2_doe.pth"):
    """Load Binary2 DOE parameters."""
    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 Zemax BINARY_2 surface 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.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
    def zmx_str(self, surf_idx, d_next):
        """Return Zemax BINARY_2 surface 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.
        """
        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 surface parameters.

Source code in deeplens-src/deeplens/phase_surface/binary2.py 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.FresnelPhase 

FresnelPhase(r, d, 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

Fresnel phase on a plane substrate.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def __init__(
    self,
    r,
    d,
    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",
):
    super().__init__(
        r=r,
        d=d,
        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 FresnelPhase from a dictionary of parameters.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@classmethod
def init_from_dict(cls, param_dict):
    """Initialize FresnelPhase from a dictionary of parameters."""
    r = param_dict.get("r")
    d = param_dict.get("d")
    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=d,
        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)

Reference phase map at design wavelength.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
63
64
65
66
67
68
69
def phi(self, x, y):
    """Reference phase map at design wavelength."""
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
71
72
73
74
75
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points."""
    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)

Generate optimizer parameters.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Generate optimizer parameters."""
    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.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
def save_ckpt(self, save_path="./fresnel_doe.pth"):
    """Save Fresnel DOE parameters."""
    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.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
102
103
104
105
106
def load_ckpt(self, load_path="./fresnel_doe.pth"):
    """Load Fresnel DOE parameters."""
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.f0 = ckpt["f0"].to(self.device)

surf_dict 

surf_dict()

Return surface parameters.

Source code in deeplens-src/deeplens/phase_surface/fresnel.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.ZernikePhase 

ZernikePhase(r, d, 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

Zernike phase on a plane substrate.

This class implements a diffractive surface using Zernike polynomials to represent the phase profile. It supports up to 37 Zernike terms. Inherits core ray-tracing functionality from Phase class.

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 Zernike phase surface.

Parameters:

Name Type Description Default
r

Radius of the surface

 required
d

Distance to next surface

 required
zernike_order

Number of Zernike terms (default: 37)

 37
norm_radii

Normalization radius (default: r)

 None
mat2

Material on the right side (default: "air")

 'air'
pos_xy

Position in xy plane

 (0.0, 0.0)
vec_local

Local coordinate system vector

 (0.0, 0.0, 1.0)
is_square

Whether the aperture is square

 False
device

Computation device

 'cpu'
Source code in deeplens-src/deeplens/phase_surface/zernike.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def __init__(
    self,
    r,
    d,
    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 Zernike phase surface.

    Args:
        r: Radius of the surface
        d: Distance to next surface
        zernike_order: Number of Zernike terms (default: 37)
        norm_radii: Normalization radius (default: r)
        mat2: Material on the right side (default: "air")
        pos_xy: Position in xy plane
        vec_local: Local coordinate system vector
        is_square: Whether the aperture is square
        device: Computation device
    """
    super().__init__(
        r=r,
        d=d,
        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)

Initialize Zernike phase surface from dictionary.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize Zernike phase surface from dictionary."""
    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"],
        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)

Reference phase map at design wavelength using Zernike polynomials.

Overrides the parent Phase.phi() method to use Zernike polynomial representation.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
101
102
103
104
105
106
107
108
109
110
def phi(self, x, y):
    """Reference phase map at design wavelength using Zernike polynomials.

    Overrides the parent Phase.phi() method to use Zernike polynomial representation.
    """
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Overrides the parent Phase.dphi_dxy() method to use Zernike derivatives.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
112
113
114
115
116
117
118
119
120
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

    Overrides the parent Phase.dphi_dxy() method to use Zernike derivatives.
    """
    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)

Generate optimizer parameters for Zernike coefficients.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
455
456
457
458
459
460
461
462
463
464
465
466
def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Generate optimizer parameters for Zernike coefficients."""
    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 Zernike DOE coefficients.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
471
472
473
474
475
476
477
478
479
480
def save_ckpt(self, save_path="./zernike_doe.pth"):
    """Save Zernike DOE coefficients."""
    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 DOE coefficients.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
482
483
484
485
486
def load_ckpt(self, load_path="./zernike_doe.pth"):
    """Load Zernike DOE coefficients."""
    ckpt = torch.load(load_path)
    self.z_coeff = ckpt["z_coeff"].to(self.device)
    self.zernike_order = ckpt["zernike_order"]

surf_dict 

surf_dict()

Return surface parameters.

Source code in deeplens-src/deeplens/phase_surface/zernike.py 
488
489
490
491
492
493
494
495
496
497
498
499
500
501
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.PolyPhase 

PolyPhase(r, d, 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 on a plane substrate.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def __init__(
    self,
    r,
    d,
    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",
):
    super().__init__(
        r=r,
        d=d,
        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)

Reference phase map at design wavelength.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def phi(self, x, y):
    """Reference phase map at design wavelength."""
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points."""
    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)

Generate optimizer parameters.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def get_optimizer_params(
    self, lrs=[1e-4, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7], optim_mat=False
):
    """Generate optimizer parameters."""
    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 Poly1D DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def save_ckpt(self, save_path="./poly1d_doe.pth"):
    """Save Poly1D DOE parameters."""
    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 Poly1D DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
142
143
144
145
146
147
148
149
150
151
def load_ckpt(self, load_path="./poly1d_doe.pth"):
    """Load Poly1D DOE parameters."""
    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 surface parameters.

Source code in deeplens-src/deeplens/phase_surface/poly.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.CubicPhase 

CubicPhase(r, d, 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 on a plane substrate.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    r,
    d,
    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",
):
    super().__init__(
        r=r,
        d=d,
        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)

Reference phase map at design wavelength.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def phi(self, x, y):
    """Reference phase map at design wavelength."""
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points."""
    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)

Generate optimizer parameters.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def get_optimizer_params(
    self, lrs=[1e-4, 1e-4, 1e-4, 1e-4, 1e-5, 1e-5], optim_mat=False
):
    """Generate optimizer parameters."""
    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 cubic DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def save_ckpt(self, save_path="./cubic_doe.pth"):
    """Save cubic DOE parameters."""
    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 cubic DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
137
138
139
140
141
142
143
144
145
146
def load_ckpt(self, load_path="./cubic_doe.pth"):
    """Load cubic DOE parameters."""
    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 surface parameters.

Source code in deeplens-src/deeplens/phase_surface/cubic.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.GratingPhase 

GratingPhase(r, d, 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

Grating phase on a plane substrate.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(
    self,
    r,
    d,
    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",
):
    super().__init__(
        r=r,
        d=d,
        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 GratingPhase from a dictionary of parameters.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
@classmethod
def init_from_dict(cls, param_dict):
    """
    Initialize GratingPhase from a dictionary of parameters.
    """
    # Extract parameters with defaults matching __init__ signature
    r = param_dict.get("r")
    d = param_dict.get("d")
    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=d,
        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)

Reference phase map at design wavelength.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
71
72
73
74
75
76
77
78
79
80
81
def phi(self, x, y):
    """Reference phase map at design wavelength."""
    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)

Calculate phase derivatives (dphi/dx, dphi/dy) for given points.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
83
84
85
86
87
88
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) for given points."""
    # 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)

Generate optimizer parameters.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def get_optimizer_params(self, lrs=[1e-4, 1e-3], optim_mat=False):
    """Generate optimizer parameters."""
    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 DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
107
108
109
110
111
112
113
114
115
116
def save_ckpt(self, save_path="./grating_doe.pth"):
    """Save grating DOE parameters."""
    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 DOE parameters.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
118
119
120
121
122
123
def load_ckpt(self, load_path="./grating_doe.pth"):
    """Load grating DOE parameters."""
    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 surface parameters.

Source code in deeplens-src/deeplens/phase_surface/grating.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.NURBSPhase 

NURBSPhase(r, d, 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

NURBS phase on a plane substrate.

This class implements a diffractive surface where the phase profile is represented by a NURBS surface. The NURBS surface is defined by control points arranged in a 2D grid, with knot vectors for both u and v directions.

The surface is evaluated using B-spline basis functions and Cox-de Boor recursion algorithm.

Reference

[1] The NURBS Book by Piegl and Tiller [2] https://en.wikipedia.org/wiki/Non-uniform_rational_B-spline

Initialize NURBS phase surface.

Parameters:

Name Type Description Default
r

Radius of the surface

 required
d

Distance to next surface

 required
control_points_u

Number of control points in u direction (default: 8)

 8
control_points_v

Number of control points in v direction (default: 8)

 8
degree_u

Degree of B-spline in u direction (default: 3)

 3
degree_v

Degree of B-spline in v direction (default: 3)

 3
control_points

Optional 3D tensor of shape (control_points_u, control_points_v, 3) containing control point coordinates (x, y, z) where z is phase. If None, initialized with small random values.

 None
weights

Optional 2D tensor of shape (control_points_u, control_points_v) containing weights for rational B-splines. If None, all weights = 1.

 None
norm_radii

Normalization radius (default: r)

 None
mat2

Material on the right side (default: "air")

 'air'
pos_xy

Position in xy plane

 (0.0, 0.0)
vec_local

Local coordinate system vector

 (0.0, 0.0, 1.0)
is_square

Whether the aperture is square

 True
device

Computation device

 'cpu'
Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(
    self,
    r,
    d,
    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 NURBS phase surface.

    Args:
        r: Radius of the surface
        d: Distance to next surface
        control_points_u: Number of control points in u direction (default: 8)
        control_points_v: Number of control points in v direction (default: 8)
        degree_u: Degree of B-spline in u direction (default: 3)
        degree_v: Degree of B-spline in v direction (default: 3)
        control_points: Optional 3D tensor of shape (control_points_u, control_points_v, 3)
                       containing control point coordinates (x, y, z) where z is phase.
                       If None, initialized with small random values.
        weights: Optional 2D tensor of shape (control_points_u, control_points_v)
                containing weights for rational B-splines. If None, all weights = 1.
        norm_radii: Normalization radius (default: r)
        mat2: Material on the right side (default: "air")
        pos_xy: Position in xy plane
        vec_local: Local coordinate system vector
        is_square: Whether the aperture is square
        device: Computation device
    """
    super().__init__(
        r=r,
        d=d,
        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
    if control_points is None:
        # Initialize with small random phase values
        self.control_points = 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')
        self.control_points[..., 0] = u_grid * 2 - 1  # x coordinates
        self.control_points[..., 1] = v_grid * 2 - 1  # y coordinates
    else:
        self.control_points = torch.tensor(control_points, dtype=torch.float32, device=device)
        assert self.control_points.shape == (control_points_u, control_points_v, 3), (
            f"control_points must have shape ({control_points_u}, {control_points_v}, 3)"
        )

    # Initialize weights for rational B-splines
    if weights is None:
        self.weights = torch.ones(control_points_u, control_points_v, device=device)
    else:
        self.weights = torch.tensor(weights, dtype=torch.float32, device=device)
        assert self.weights.shape == (control_points_u, control_points_v), (
            f"weights must have shape ({control_points_u}, {control_points_v})"
        )

    self.param_model = "nurbs"
    self.to(device)

init_from_dict classmethod 

init_from_dict(surf_dict)

Initialize NURBS phase surface from dictionary.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize NURBS phase surface from dictionary."""
    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"],
        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.tensor(control_points, device=obj.device)

    weights = surf_dict.get("weights", None)
    if weights is not None:
        obj.weights = torch.tensor(weights, device=obj.device)

    return obj

phi 

phi(x, y)

Reference phase map at design wavelength using NURBS surface evaluation.

Parameters:

Name Type Description Default
x, y

Coordinate tensors

 required

Returns:

Type Description

Phase values in radians at the specified coordinates
Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
def phi(self, x, y):
    """Reference phase map at design wavelength using NURBS surface evaluation.

    Args:
        x, y: Coordinate tensors

    Returns:
        Phase values in radians at the specified coordinates
    """
    # 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]

    # Flatten for batch processing
    x_flat = x_norm.flatten()
    y_flat = y_norm.flatten()
    batch_size = x_flat.shape[0]

    # Evaluate NURBS surface for each point
    phi_values = []
    for i in range(batch_size):
        point = self._evaluate_nurbs_surface(x_flat[i], y_flat[i])
        phi_values.append(point[2])  # z-coordinate contains phase

    phi = torch.stack(phi_values).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)

Calculate phase derivatives (dphi/dx, dphi/dy) using NURBS surface.

Parameters:

Name Type Description Default
x, y

Coordinate tensors

 required

Returns:

Type Description

dphidx, dphidy: Phase derivatives in x and y directions
Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
def dphi_dxy(self, x, y):
    """Calculate phase derivatives (dphi/dx, dphi/dy) using NURBS surface.

    Args:
        x, y: Coordinate tensors

    Returns:
        dphidx, dphidy: Phase derivatives in x and y directions
    """
    # 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)

Generate optimizer parameters for NURBS control points.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def get_optimizer_params(self, lrs=[1e-4, 1e-2], optim_mat=False):
    """Generate optimizer parameters for NURBS control points."""
    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.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
def save_ckpt(self, save_path="./nurbs_doe.pth"):
    """Save NURBS DOE parameters."""
    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.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
396
397
398
399
400
401
402
403
404
405
406
407
def load_ckpt(self, load_path="./nurbs_doe.pth"):
    """Load NURBS DOE parameters."""
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.control_points = ckpt["control_points"].to(self.device)
    self.weights = ckpt["weights"].to(self.device)
    self.control_points_u = ckpt["control_points_u"]
    self.control_points_v = ckpt["control_points_v"]
    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.

Source code in deeplens-src/deeplens/phase_surface/nurbs.py 
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
def surf_dict(self):
    """Return surface parameters."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

deeplens.phase_surface.VortexPhase 

VortexPhase(r, d, 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:

φ(x, y) = charge · atan2(y, x)  −  π · (x² + y²) / f₀

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 (units: mm²); no design wavelength is assumed — chromatic behaviour emerges naturally from the ray wavelength in diffract(). Setting f0=None disables the Fresnel term (pure vortex).

Reference

Yu et al., "Light Propagation with Phase Discontinuities," Science 2011.

Parameters:

Name Type Description Default
r

Aperture radius [mm].

 required
d

Axial position [mm].

 required
charge

Topological charge (integer). Positive for left-handed, negative for right-handed helix.

 1
f0

Phase curvature parameter [mm²] of the co-centered Fresnel term. Physically equivalent to λ·f for a thin lens at wavelength λ and focal length f. None disables the Fresnel term.

 None
norm_radii

Normalisation radius for phase map display. Defaults to r.

 None
mat2

Material name after the surface.

 'air'
pos_xy

(x, y) position offset in the global frame [mm].

 (0.0, 0.0)
vec_local

Local surface normal direction.

 (0.0, 0.0, 1.0)
is_square

If True, use a square aperture; else circular.

 True
device

Torch device.

 'cpu'
Source code in deeplens-src/deeplens/phase_surface/vortex.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(
    self,
    r,
    d,
    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",
):
    """
    Args:
        r: Aperture radius [mm].
        d: Axial position [mm].
        charge: Topological charge (integer). Positive for left-handed, negative
            for right-handed helix.
        f0: Phase curvature parameter [mm²] of the co-centered Fresnel term.
            Physically equivalent to λ·f for a thin lens at wavelength λ and
            focal length f. ``None`` disables the Fresnel term.
        norm_radii: Normalisation radius for phase map display. Defaults to ``r``.
        mat2: Material name after the surface.
        pos_xy: (x, y) position offset in the global frame [mm].
        vec_local: Local surface normal direction.
        is_square: If True, use a square aperture; else circular.
        device: Torch device.
    """
    super().__init__(
        r=r,
        d=d,
        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 VortexPhase from a parameter dictionary.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize VortexPhase from a parameter dictionary."""
    f0_raw = surf_dict.get("f0", None)
    return cls(
        r=surf_dict["r"],
        d=surf_dict["d"],
        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)

Phase map (wrapped to [0, 2π]).

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
91
92
93
94
95
96
97
98
def phi(self, x, y):
    """Phase map (wrapped to [0, 2π])."""
    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)

Analytical phase gradient (unwrapped) for generalized Snell's law.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
100
101
102
103
104
105
106
107
108
109
110
111
def dphi_dxy(self, x, y):
    """Analytical phase gradient (unwrapped) for generalized Snell's law."""
    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.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
116
117
118
119
120
121
122
123
124
125
126
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.
    """
    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.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
131
132
133
134
135
136
137
138
def save_ckpt(self, save_path="./vortex_doe.pth"):
    """Save VortexPhase parameters."""
    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.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
140
141
142
143
144
145
146
def load_ckpt(self, load_path="./vortex_doe.pth"):
    """Load VortexPhase parameters."""
    ckpt = torch.load(load_path)
    self.param_model = ckpt["param_model"]
    self.charge = int(ckpt["l"])
    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 serialisable dictionary.

Source code in deeplens-src/deeplens/phase_surface/vortex.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def surf_dict(self):
    """Return surface parameters as a serialisable dictionary."""
    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": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }
    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 (in coherent mode) optical path length. All tensor attributes share the same batch shape (*batch_size, num_rays).

Attributes:

Name Type Description
o Tensor

Ray origins, shape (*batch, num_rays, 3) [mm].
d Tensor

Unit ray directions, shape (*batch, num_rays, 3).
wvln Tensor

Wavelength scalar [µm].
is_valid Tensor

Binary validity mask, shape (*batch, num_rays).
en Tensor

Energy weight, shape (*batch, num_rays, 1).
bend_penalty Tensor

Accumulated per-surface bend penalty, shape (*batch, num_rays, 1).
opl Tensor

Optical path length (coherent mode only), shape (*batch, num_rays, 1) [mm].
is_coherent bool

Whether OPL tracking is enabled.

Initialize a ray object.

Parameters:

Name Type Description Default
o Tensor

Ray origin, shape (..., num_rays, 3) [mm].

 required
d Tensor

Ray direction, shape (..., num_rays, 3).

 required
wvln float

Ray wavelength [µm]. Required — must be 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

Compute device. Defaults to "cpu".

 'cpu'
Source code in deeplens-src/deeplens/light/ray.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(self, o, d, wvln, is_coherent=False, device="cpu"):
    """Initialize a ray object.

    Args:
        o (torch.Tensor): Ray origin, shape ``(..., num_rays, 3)`` [mm].
        d (torch.Tensor): Ray direction, shape ``(..., num_rays, 3)``.
        wvln (float): Ray wavelength [µm]. Required — must be passed
            explicitly (the Lens carries ``primary_wvln``/``wvln_rgb``,
            not the Ray).
        is_coherent (bool): Enable optical path length tracking for coherent
            tracing. Defaults to ``False``.
        device (str): Compute device. Defaults to ``"cpu"``.
    """
    # Basic ray parameters - move to device
    self.o = (o if torch.is_tensor(o) else torch.tensor(o)).to(device)
    self.d = (d if torch.is_tensor(d) else torch.tensor(d)).to(device)
    self.shape = self.o.shape[:-1]

    # Wavelength
    assert wvln > 0.1 and wvln < 10.0, "Ray wavelength unit should be [um]"
    self.wvln = torch.tensor(wvln, device=device)

    # Auxiliary ray parameters - create directly on device
    self.is_valid = torch.ones(self.shape, device=device)
    self.en = torch.ones((*self.shape, 1), device=device)
    self.bend_penalty = torch.zeros((*self.shape, 1), device=device)

    # Coherent ray tracing
    self.is_coherent = is_coherent  # bool
    self.opl = torch.zeros((*self.shape, 1), device=device)

    self.device = device
    self.d = F.normalize(self.d, p=2, dim=-1)

prop_to 

prop_to(z, n=1.0)

Ray propagates to a given depth plane.

Parameters:

Name Type Description Default
z float

depth.

 required
n float

refractive index. Defaults to 1.

 1.0
Source code in deeplens-src/deeplens/light/ray.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def prop_to(self, z, n=1.0):
    """Ray propagates to a given depth plane.

    Args:
        z (float): depth.
        n (float, optional): refractive index. Defaults to 1.
    """
    t = (z - self.o[..., 2]) / self.d[..., 2]
    new_o = self.o + self.d * t.unsqueeze(-1)
    valid_mask = (self.is_valid > 0).unsqueeze(-1)
    self.o = torch.where(valid_mask, new_o, self.o)

    if self.is_coherent:
        if t.dtype != torch.float64:
            raise Warning("Should use float64 in coherent ray tracing.")
        else:
            new_opl = self.opl + n * t.unsqueeze(-1)
            self.opl = torch.where(valid_mask, new_opl, self.opl)

    return self

centroid 

centroid()

Calculate the centroid of the ray, shape (..., num_rays, 3)

Returns:

Type Description

torch.Tensor: Centroid of the ray, shape (..., 3)
Source code in deeplens-src/deeplens/light/ray.py 
90
91
92
93
94
95
96
97
98
def centroid(self):
    """Calculate the centroid of the ray, shape (..., num_rays, 3)

    Returns:
        torch.Tensor: Centroid of the ray, shape (..., 3)
    """
    return (self.o * self.is_valid.unsqueeze(-1)).sum(-2) / self.is_valid.sum(
        -1
    ).add(EPSILON).unsqueeze(-1)

rms_error 

rms_error(center_ref=None)

Calculate the RMS error of the ray.

Parameters:

Name Type Description Default
center_ref Tensor

Reference center of the ray, shape (..., 3). If None, use the centroid of the ray as reference.

 None

Returns:

Type Description

torch.Tensor: average RMS error of the ray
Source code in deeplens-src/deeplens/light/ray.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def rms_error(self, center_ref=None):
    """Calculate the RMS error of the ray.

    Args:
        center_ref (torch.Tensor): Reference center of the ray, shape (..., 3). If None, use the centroid of the ray as reference.

    Returns:
        torch.Tensor: average RMS error of the ray
    """
    # 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
    rms_error = ((self.o[..., :2] - center_ref[..., :2]) ** 2).sum(-1)
    rms_error = (rms_error * self.is_valid).sum(-1) / (
        self.is_valid.sum(-1) + EPSILON
    )
    rms_error = rms_error.sqrt()

    # Average RMS error
    return rms_error.mean()

flip_xy 

flip_xy()

Flip the x and y coordinates of the ray.

This function is used when calculating point spread function and wavefront distribution.

Source code in deeplens-src/deeplens/light/ray.py 
126
127
128
129
130
131
132
133
def flip_xy(self):
    """Flip the x and y coordinates of the ray.

    This function is used when calculating point spread function and wavefront distribution.
    """
    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)

Clone the ray.

Can specify which device we want to clone. Sometimes we want to store all rays in CPU, and when using it, we move it to GPU.

Source code in deeplens-src/deeplens/light/ray.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def clone(self, device=None):
    """Clone the ray.

    Can specify which device we want to clone. Sometimes we want to store all rays in CPU, and when using it, we move it to GPU.
    """
    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.is_coherent = self.is_coherent
    ray.device = target_device
    ray.shape = ray.o.shape[:-1]

    return ray

squeeze 

squeeze(dim=None)

Squeeze the ray.

Parameters:

Name Type Description Default
dim int

dimension to squeeze. Defaults to None.

 None
Source code in deeplens-src/deeplens/light/ray.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def squeeze(self, dim=None):
    """Squeeze the ray.

    Args:
        dim (int, optional): dimension to squeeze. Defaults to None.
    """
    self.o = self.o.squeeze(dim)
    self.d = self.d.squeeze(dim)
    # wvln is a single element tensor, no squeeze needed
    self.is_valid = self.is_valid.squeeze(dim)
    self.en = self.en.squeeze(dim)
    self.opl = self.opl.squeeze(dim)
    self.bend_penalty = self.bend_penalty.squeeze(dim)
    return self

unsqueeze 

unsqueeze(dim=None)

Unsqueeze the ray.

Parameters:

Name Type Description Default
dim int

dimension to unsqueeze. Defaults to None.

 None
Source code in deeplens-src/deeplens/light/ray.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def unsqueeze(self, dim=None):
    """Unsqueeze the ray.

    Args:
        dim (int, optional): dimension to unsqueeze. Defaults to None.
    """
    self.o = self.o.unsqueeze(dim)
    self.d = self.d.unsqueeze(dim)
    # wvln is a single element tensor, no unsqueeze needed
    self.is_valid = self.is_valid.unsqueeze(dim)
    self.en = self.en.unsqueeze(dim)
    self.opl = self.opl.unsqueeze(dim)
    self.bend_penalty = self.bend_penalty.unsqueeze(dim)
    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 (ASM, Fresnel, Fraunhofer) 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π / (λ × 10⁻³) [mm⁻¹].
phy_size tuple

Physical aperture size (W, H) [mm].
ps float

Pixel pitch [mm] (must be square).
res tuple

Grid resolution (H, W) in pixels.
z float

Current axial position [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.

 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 
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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*.
        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
    self.phy_size = phy_size  # [mm], physical size
    assert phy_size[0] / self.res[0] == phy_size[1] / self.res[1], (
        "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, -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 x0y plane originating from a point source.

Parameters:

Name Type Description Default
point tuple

Point source position in object space. [mm]. Defaults to (0, 0, -1000.0).

 (0, 0, -1000.0)
wvln float

Wavelength. [um]. Defaults to 0.55.

 0.55
z float

Field z position. [mm]. Defaults to 0.0.

 0.0
phy_size tuple

Valid plane on x0y plane. [mm]. Defaults to (2, 2).

 (4.0, 4.0)
res tuple

Valid plane resoltution. Defaults to (1000, 1000).

 (2000, 2000)
valid_r float

Valid circle radius. [mm]. Defaults to None.

 None

Returns:

Name Type Description
field ComplexWave

Complex field on x0y plane.
Source code in deeplens-src/deeplens/light/wave.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
@classmethod
def point_wave(
    cls,
    point=(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 x0y plane originating from a point source.

    Args:
        point (tuple): Point source position in object space. [mm]. Defaults to (0, 0, -1000.0).
        wvln (float): Wavelength. [um]. Defaults to 0.55.
        z (float): Field z position. [mm]. Defaults to 0.0.
        phy_size (tuple): Valid plane on x0y plane. [mm]. Defaults to (2, 2).
        res (tuple): Valid plane resoltution. Defaults to (1000, 1000).
        valid_r (float): Valid circle radius. [mm]. Defaults to None.

    Returns:
        field (ComplexWave): Complex field on 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
    r = torch.sqrt((x - point[0]) ** 2 + (y - point[1]) ** 2 + (z - point[2]) ** 2)
    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 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θx + y sinθy)) while the amplitude stays uniform.

Parameters:

Name Type Description Default
wvln float

Wavelength. [um].

 0.55
z float

Field z position. [mm].

 0.0
phy_size tuple

Physical size of the field. [mm].

 (4.0, 4.0)
res tuple

Resolution.

 (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

Valid circle radius. [mm].

 None

Returns:

Name Type Description
field ComplexWave

Complex field.
Source code in deeplens-src/deeplens/light/wave.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
@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 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θx + y sinθy))`` while the amplitude stays uniform.

    Args:
        wvln (float): Wavelength. [um].
        z (float): Field z position. [mm].
        phy_size (tuple): Physical size of the field. [mm].
        res (tuple): Resolution.
        theta_x (float): Tilt angle of the wavevector in the x-z plane.
            [rad]. Defaults to 0.0.
        theta_y (float): Tilt angle of the wavevector in the y-z plane.
            [rad]. Defaults to 0.0.
        valid_r (float): Valid circle radius. [mm].

    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.

Parameters:

Name Type Description Default
img Tensor

Input image with shape [H, W] or [B, C, H, W]. Data range is [0, 1].

 required
wvln float

Wavelength. [um].

 0.55
z float

Field z position. [mm].

 0.0
phy_size tuple

Physical size of the field. [mm].

 (4.0, 4.0)

Returns:

Name Type Description
field ComplexWave

Complex field.
Source code in deeplens-src/deeplens/light/wave.py 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
@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.

    Args:
        img (torch.Tensor): Input image with shape [H, W] or [B, C, H, W]. Data range is [0, 1].
        wvln (float): Wavelength. [um].
        z (float): Field z position. [mm].
        phy_size (tuple): Physical size of the field. [mm].

    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 by distance z. Can only propagate planar wave.

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"

Parameters:

Name Type Description Default
prop_dist float

propagation distance, unit [mm].

 required
n float

refractive index.

 1.0

Returns:

Name Type Description
self

propagated complex wave field.
Source code in deeplens-src/deeplens/light/wave.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def prop(self, prop_dist, n=1.0):
    """Propagate the field by distance z. Can only propagate planar wave.

    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"

    Args:
        prop_dist (float): propagation distance, unit [mm].
        n (float): refractive index.

    Returns:
        self: propagated complex wave field.
    """
    # 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 plane z.

Parameters:

Name Type Description Default
z float

destination plane z coordinate.

 required
Source code in deeplens-src/deeplens/light/wave.py 
306
307
308
309
310
311
312
313
314
315
316
def prop_to(self, z, n=1):
    """Propagate the field to plane z.

    Args:
        z (float): destination plane z coordinate.
    """
    # 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 grid.

Source code in deeplens-src/deeplens/light/wave.py 
322
323
324
325
326
327
328
329
def gen_xy_grid(self):
    """Generate the x and y grid."""
    x, y = torch.meshgrid(
        torch.linspace(-0.5 * self.phy_size[1], 0.5 * self.phy_size[1], self.res[0],),
        torch.linspace(0.5 * self.phy_size[0], -0.5 * self.phy_size[0], self.res[1],),
        indexing="xy",
    )
    return x, y

gen_freq_grid 

gen_freq_grid()

Generate the frequency grid.

Source code in deeplens-src/deeplens/light/wave.py 
331
332
333
334
335
336
def gen_freq_grid(self):
    """Generate the frequency grid."""
    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_npz 

load_npz(filepath)

Load data from npz file.

Source code in deeplens-src/deeplens/light/wave.py 
348
349
350
351
352
353
354
355
356
def load_npz(self, filepath):
    """Load data from 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 a npz file.

Source code in deeplens-src/deeplens/light/wave.py 
358
359
360
361
362
363
def save(self, filepath="./wavefield.npz"):
    """Save the complex wave field to a npz file."""
    if filepath.endswith(".npz"):
        self.save_npz(filepath)
    else:
        raise Exception("Unimplemented file format.")

save_npz 

save_npz(filepath='./wavefield.npz')

Save the complex wave field to a npz file.

Source code in deeplens-src/deeplens/light/wave.py 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
def save_npz(self, filepath="./wavefield.npz"):
    """Save the complex wave field to a npz file."""
    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)

show 

show(save_name=None, data='irr')

Save the field as an image.

Source code in deeplens-src/deeplens/light/wave.py 
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
def show(self, save_name=None, data="irr"):
    """Save the field as an image."""
    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)

Pad the input field by (Hpad, Hpad, Wpad, Wpad). This step will also expand physical size of the field.

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

Returns:

Name Type Description
self

Padded complex wave field.
Source code in deeplens-src/deeplens/light/wave.py 
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
def pad(self, Hpad, Wpad):
    """Pad the input field by (Hpad, Hpad, Wpad, Wpad). This step will also expand physical size of the field.

    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.

    Returns:
        self: Padded complex wave field.
    """
    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 horizontally and vertically.

Source code in deeplens-src/deeplens/light/wave.py 
481
482
483
484
485
486
487
def flip(self):
    """Flip the field horizontally and vertically."""
    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)

Render an image batch with one spatially invariant PSF.

Applies a per-channel 2-D convolution using reflect padding so that the output has the same spatial dimensions as the input. The PSF is internally flipped to convert the cross-correlation implemented by F.conv2d into convolution.

Parameters:

Name Type Description Default
img Tensor

Input image batch, shape [B, C, H, W].

 required
psf Tensor

PSF kernel, shape [C, ks, ks]. ks may be odd or even.

 required

Returns:

Type Description

torch.Tensor: Rendered image, shape [B, C, H, W].
Example

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 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def conv_psf(img, psf):
    """Render an image batch with one spatially invariant PSF.

    Applies a per-channel 2-D convolution using reflect padding so that the
    output has the same spatial dimensions as the input. The PSF is internally
    flipped to convert the cross-correlation implemented by ``F.conv2d`` into
    convolution.

    Args:
        img (torch.Tensor): Input image batch, shape ``[B, C, H, W]``.
        psf (torch.Tensor): PSF kernel, shape ``[C, ks, ks]``.  ``ks`` may be
            odd or even.

    Returns:
        torch.Tensor: Rendered image, shape ``[B, C, H, W]``.

    Example:
        >>> psf = lens.psf_rgb(points=torch.tensor([0.0, 0.0, -10000.0]))
        >>> img_blur = conv_psf(img, psf)
    """
    B, C, H, W = img.shape
    C_psf, ks, _ = psf.shape
    assert C_psf == C, f"psf channels ({C_psf}) must match image channels ({C})."

    # Flip the PSF because F.conv2d use cross-correlation
    psf = torch.flip(psf, [1, 2])
    psf = psf.unsqueeze(1)  # shape [C, 1, ks, ks]

    # 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")

    # Convolution
    img_render = F.conv2d(img_pad, psf, groups=C)
    return img_render

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:

Type Description

torch.Tensor: Rendered image, shape [B, C, H, W].
Source code in deeplens-src/deeplens/imgsim/psf.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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:
        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.

 'reflect'

Returns:

Type Description

torch.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 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
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.

    Returns:
        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 / depth.

 'depth'

Returns:

Type Description

torch.Tensor: Rendered image, shape [B, C, H, W].
Source code in deeplens-src/deeplens/imgsim/psf.py 
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
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): ``"depth"`` for linear depth interpolation or
            ``"disparity"`` for linear interpolation in ``1 / depth``.

    Returns:
        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).

 required
psf_kernels Tensor

PSF at each depth layer, shape (num_layers, C, ks, ks).

 required
psf_depths Tensor

Depth values for each layer, shape (num_layers,). Must be negative and sorted ascending (far to near, i.e. -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 
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
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).
        psf_kernels (torch.Tensor): PSF at each depth layer, shape (num_layers, C, ks, ks).
        psf_depths (torch.Tensor): Depth values for each layer, shape (num_layers,).
            Must be negative and sorted ascending (far to near, i.e. -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]

        if mask.sum() == 0:
            continue

        # 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

The image to be blurred (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

Source tile size for memory-efficient rendering. If None, render the whole image at once.

 None

Returns:

Name Type Description
img_render Tensor

Rendered image (B, C, H, W).
Source code in deeplens-src/deeplens/imgsim/psf.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
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 (Tensor): The image to be blurred (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, optional): Source tile size for memory-efficient
            rendering. If ``None``, render the whole image at once.

    Returns:
        img_render (Tensor): Rendered image (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:

Type Description

torch.Tensor: Interpolated packed PSF map, shape [C, grid_new*ks, grid_new*ks].
Source code in deeplens-src/deeplens/imgsim/psf.py 
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
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:
        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:

Type Description

torch.Tensor: Rotated PSFs, shape [N, 3, ks, ks].
Source code in deeplens-src/deeplens/imgsim/psf.py 
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
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:
        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 × ks grid centred on pointc (or the ray centroid when pointc is None). In coherent mode the complex amplitude is accumulated instead of intensity.

All N field points scatter into their own output slices in a single index_put_(accumulate=True) call with 3-D (batch, row, col) indices. This removes the per-field Python loop the earlier implementation used and fuses the scatter across all field points into one kernel launch.

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 for each field point, shape [N, 2]. If None, the valid-ray centroid is used.

 None
interpolate bool

If True (default), 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).

 True

Returns:

Type Description

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.
Source code in deeplens-src/deeplens/imgsim/monte_carlo.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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 × ks`` grid centred on *pointc* (or the
    ray centroid when *pointc* is ``None``). In coherent mode the complex
    amplitude is accumulated instead of intensity.

    All ``N`` field points scatter into their own output slices in a single
    ``index_put_(accumulate=True)`` call with 3-D ``(batch, row, col)``
    indices. This removes the per-field Python loop the earlier
    implementation used and fuses the scatter across all field points into
    one kernel launch.

    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 for each
            field point, shape ``[N, 2]``. If ``None``, the valid-ray
            centroid is used.
        interpolate (bool, optional): If ``True`` (default), 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).

    Returns:
        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:
        amp = torch.sqrt(ray.d[..., 2].abs())           # 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.

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 object. Shape of ray.o is [h, w, spp, 3].

 required
img_obj

[B, C, H, W]. Spatial size H, W is read from this tensor.

 required
ps

pixel size

 required
interpolate

whether to interpolate

 True
energy_correction

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=True the denominator is fixed, so the weight only scales the numerator. None (default) gives uniform per-ray weights.

 None
vignetting

If True, divide by a fixed denominator (torch.numel(ray.is_valid)) instead of the sum of weights; pixels hit by few / attenuated rays therefore appear dimmer (mechanical vignetting). Defaults to False.

 False

Returns:

Name Type Description
output

shape [B, C, h, w]
Source code in deeplens-src/deeplens/imgsim/monte_carlo.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def backward_integral(
    ray,
    img_obj,
    ps,
    interpolate=True,
    energy_correction=None,
    vignetting=False,
):
    """Backward Monte Carlo integration, for ray tracing based rendering.

    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 object. Shape of ``ray.o`` is ``[h, w, spp, 3]``.
        img_obj: [B, C, H, W]. Spatial size ``H, W`` is read from this tensor.
        ps: pixel size
        interpolate: whether to interpolate
        energy_correction: 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=True`` the denominator is
            fixed, so the weight only scales the numerator. ``None``
            (default) gives uniform per-ray weights.
        vignetting: If True, divide by a fixed denominator
            (``torch.numel(ray.is_valid)``) instead of the sum of weights;
            pixels hit by few / attenuated rays therefore appear dimmer
            (mechanical vignetting). Defaults to False.

    Returns:
        output: shape [B, C, h, w]
    """
    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)

Assign points to pixels, supports both incoherent and coherent ray tracing. Use advanced indexing to increment the count for each corresponding pixel.

This function can only compute single point source, constrained by advanced indexing operation.

Parameters:

Name Type Description Default
points

shape [spp, 2]

 required
mask

shape [spp]

 required
ks

kernel size

 required
x_range

x range

 required
y_range

y range

 required
value

shape [spp], values we want to assign to each pixel (intensity or complex amplitude)

 required
interpolate

whether to interpolate

 True

Returns:

Name Type Description
field

intensity or complex amplitude, shape [ks, ks]
Source code in deeplens-src/deeplens/imgsim/monte_carlo.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def assign_points_to_pixels(
    points,
    mask,
    ks,
    x_range,
    y_range,
    value,
    interpolate=True,
):
    """Assign points to pixels, supports both incoherent and coherent ray tracing. Use advanced indexing to increment the count for each corresponding pixel.

    This function can only compute single point source, constrained by advanced indexing operation.

    Args:
        points: shape [spp, 2]
        mask: shape [spp]
        ks: kernel size
        x_range: x range
        y_range: y range
        value: shape [spp], values we want to assign to each pixel (intensity or complex amplitude)
        interpolate: whether to interpolate


    Returns:
        field: intensity or complex amplitude, shape [ks, ks]
    """
    # 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