HybridLens

Combines a GeoLens with a diffractive optical element (DOE). HybridLens performs coherent ray tracing to the DOE plane, then Angular Spectrum Method (ASM) propagation to the sensor — a hybrid ray–wave model for refractive lenses with DOE or metasurface phase elements.

deeplens.HybridLens

HybridLens(filename=None, device=None, dtype=torch.float64, primary_wvln=DEFAULT_WAVE, wvln_rgb=WAVE_RGB, obj_depth=DEPTH)

Bases: Lens

Hybrid refractive-diffractive lens using a differentiable ray–wave model.

Combines a GeoLens (refractive module) with a diffractive optical element (DOE) placed behind it. The pipeline is:

1. Coherent ray tracing through the embedded GeoLens to obtain a complex wavefront at the DOE plane (including all geometric aberrations).
2. DOE phase modulation applied to the wavefront.
3. Angular Spectrum Method (ASM) propagation from the DOE to the sensor plane to produce the final intensity PSF.

This enables end-to-end gradient flow from image quality metrics back to both refractive surface parameters and the DOE phase profile.

Attributes:

Name	Type	Description
geolens	GeoLens	Embedded refractive module.
doe		Diffractive optical element (one of Binary2, Pixel2D, Fresnel, Zernike, Grating).

Notes

Operates in torch.float64 by default for numerical stability of the wave-propagation step.

References

Xinge Yang et al., "End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model," SIGGRAPH Asia 2024.

Initialize a hybrid refractive-diffractive lens.

Parameters:

Name	Type	Description	Default
filename	str	Path to the lens configuration JSON file. Defaults to None.	None
device	str	Computation device ('cpu' or 'cuda'). Defaults to None.	None
dtype	dtype	Data type for computations. Defaults to torch.float64.	float64
primary_wvln	float	Primary design wavelength [µm]. Used as fallback when a method is called without an explicit wvln. Defaults to DEFAULT_WAVE.	DEFAULT_WAVE
wvln_rgb	sequence of float	Three wavelengths used for RGB computations, ordered [R, G, B] in µm. Defaults to WAVE_RGB.	WAVE_RGB
obj_depth	float	Default object depth [mm], used when a method is called without an explicit depth. Defaults to DEPTH.	DEPTH

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

def __init__(
    self,
    filename=None,
    device=None,
    dtype=torch.float64,
    primary_wvln=DEFAULT_WAVE,
    wvln_rgb=WAVE_RGB,
    obj_depth=DEPTH,
):
    """Initialize a hybrid refractive-diffractive lens.

    Args:
        filename (str, optional): Path to the lens configuration JSON file. Defaults to None.
        device (str, optional): Computation device ('cpu' or 'cuda'). Defaults to None.
        dtype (torch.dtype, optional): Data type for computations. Defaults to torch.float64.
        primary_wvln (float, optional): Primary design wavelength [µm].
            Used as fallback when a method is called without an explicit
            ``wvln``.  Defaults to ``DEFAULT_WAVE``.
        wvln_rgb (sequence of float, optional): Three wavelengths used
            for RGB computations, ordered ``[R, G, B]`` in µm.  Defaults
            to ``WAVE_RGB``.
        obj_depth (float, optional): Default object depth [mm], used
            when a method is called without an explicit ``depth``.
            Defaults to ``DEPTH``.
    """
    super().__init__(
        device=device,
        dtype=dtype,
        primary_wvln=primary_wvln,
        wvln_rgb=wvln_rgb,
        obj_depth=obj_depth,
    )

    # Load lens file
    if filename is not None:
        self.read_lens_json(filename)
    else:
        self.geolens = None
        self.doe = None
        # Set default sensor size and resolution if no file provided
        self.sensor_size = (8.0, 8.0)
        self.sensor_res = (2000, 2000)
        print(
            f"No lens file provided. Using default sensor_size: {self.sensor_size} mm, "
            f"sensor_res: {self.sensor_res} pixels. Use set_sensor() to change."
        )

    self.double()

read_lens_json

read_lens_json(filename)

Read the lens configuration from a JSON file.

Loads a GeoLens and associated DOE from the specified file. A Plane surface is appended to the GeoLens surface list as a placeholder for the DOE plane.

Supported DOE types: binary2, pixel2d, fresnel, zernike, grating.

Parameters:

Name	Type	Description	Default
filename	str	Path to the JSON configuration file. Must contain a "DOE" key with a "type" field.	required

Raises:

Type	Description
ValueError	If the DOE type in the file is not supported.

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

def read_lens_json(self, filename):
    """Read the lens configuration from a JSON file.

    Loads a `GeoLens` and associated DOE from the specified file.
    A ``Plane`` surface is appended to the GeoLens surface list as a
    placeholder for the DOE plane.

    Supported DOE types: ``binary2``, ``pixel2d``, ``fresnel``,
    ``zernike``, ``grating``.

    Args:
        filename (str): Path to the JSON configuration file.  Must
            contain a ``"DOE"`` key with a ``"type"`` field.

    Raises:
        ValueError: If the DOE type in the file is not supported.
    """
    # Load geolens
    geolens = GeoLens(filename=filename, device=self.device)

    # Load DOE (diffractive surface)
    with open(filename, "r") as f:
        data = json.load(f)

        doe_dict = data["DOE"]
        doe_param_model = doe_dict["type"].lower()
        if doe_param_model == "binary2":
            doe = Binary2.init_from_dict(doe_dict)
        elif doe_param_model == "pixel2d":
            doe = Pixel2D.init_from_dict(doe_dict)
        elif doe_param_model == "fresnel":
            doe = Fresnel.init_from_dict(doe_dict)
        elif doe_param_model == "zernike":
            doe = Zernike.init_from_dict(doe_dict)
        elif doe_param_model == "grating":
            doe = Grating.init_from_dict(doe_dict)
        else:
            raise ValueError(f"Unsupported DOE parameter model: {doe_param_model}")
        self.doe = doe

    # Add a Plane/Phase surface to GeoLens (DOE placeholder).
    # Match the DOE's actual aperture (square vs circular) so that rays
    # outside the DOE region are correctly culled at the placeholder.
    geolens.surfaces.append(
        Plane(d=doe.d.item(), r=doe.r, mat2="air", is_square=doe.is_square)
    )
    # r_doe = float(np.sqrt(doe.w**2 + doe.h**2) / 2)
    # geolens.surfaces.append(Phase(r=r_doe, d=doe.d))
    self.geolens = geolens
    self.foclen = geolens.foclen

    # Update hybrid lens sensor resolution and pixel size
    self.set_sensor(sensor_size=geolens.sensor_size, sensor_res=geolens.sensor_res)
    self.to(self.device)

write_lens_json

write_lens_json(lens_path)

Write the lens configuration to a JSON file.

Serialises the GeoLens surfaces (excluding the DOE placeholder) and the DOE configuration into a single JSON file that can be reloaded with read_lens_json.

Parameters:

Name	Type	Description	Default
lens_path	str	Output file path.	required

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

def write_lens_json(self, lens_path):
    """Write the lens configuration to a JSON file.

    Serialises the ``GeoLens`` surfaces (excluding the DOE placeholder)
    and the ``DOE`` configuration into a single JSON file that can be
    reloaded with `read_lens_json`.

    Args:
        lens_path (str): Output file path.
    """
    geolens = self.geolens
    data = {}
    data["info"] = geolens.lens_info if hasattr(geolens, "lens_info") else "None"
    data["foclen"] = round(geolens.foclen, 4)
    data["fnum"] = round(geolens.fnum, 4)
    data["r_sensor"] = round(geolens.r_sensor, 4)
    data["d_sensor"] = round(geolens.d_sensor.item(), 4)
    data["sensor_size"] = [round(i, 4) for i in geolens.sensor_size]
    data["sensor_res"] = geolens.sensor_res

    # Geolens
    data["surfaces"] = []
    for i, s in enumerate(geolens.surfaces[:-1]):
        surf_dict = s.surf_dict()

        # To exclude the last surface (DOE)
        if i < len(geolens.surfaces) - 2:
            surf_dict["d_next"] = round(
                geolens.surfaces[i + 1].d.item() - geolens.surfaces[i].d.item(), 3
            )
        else:
            surf_dict["d_next"] = round(
                geolens.d_sensor.item() - geolens.surfaces[i].d.item(), 3
            )

        data["surfaces"].append(surf_dict)

    # DOE
    data["DOE"] = self.doe.surf_dict()

    with open(lens_path, "w") as f:
        json.dump(data, f, indent=4)

analysis

analysis(save_name='./test.png')

Run a quick visual analysis of the hybrid lens.

Generates two figures: the 2D lens layout (saved to save_name) and the DOE phase map (saved to <save_name>_doe.png).

Parameters:

Name	Type	Description	Default
save_name	str	Base file path for the layout image. The DOE phase-map image is derived by appending _doe before the extension. Defaults to './test.png'.	'./test.png'

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

def analysis(self, save_name="./test.png"):
    """Run a quick visual analysis of the hybrid lens.

    Generates two figures: the 2D lens layout (saved to *save_name*) and
    the DOE phase map (saved to ``<save_name>_doe.png``).

    Args:
        save_name (str, optional): Base file path for the layout image.
            The DOE phase-map image is derived by appending ``_doe``
            before the extension.  Defaults to ``'./test.png'``.
    """
    self.draw_layout(save_name=save_name)
    self.doe.draw_phase_map(save_name=f"{save_name}_doe.png")

double

double()

Convert the GeoLens and DOE to float64 precision.

Double precision is required for numerically stable phase accumulation during coherent ray tracing and ASM propagation. Called automatically by __init__.

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

def double(self):
    """Convert the GeoLens and DOE to ``float64`` precision.

    Double precision is required for numerically stable phase
    accumulation during coherent ray tracing and ASM propagation.
    Called automatically by `__init__`.
    """
    self.geolens.astype(torch.float64)
    self.doe.astype(torch.float64)

refocus

refocus(foc_dist)

Refocus the hybrid lens to a given object distance.

Only the GeoLens sensor-to-last-surface spacing is adjusted; the DOE remains fixed relative to the refractive group (it is physically cemented to the lens barrel).

Parameters:

Name	Type	Description	Default
foc_dist	float	Target focus distance in [mm] (negative, towards the object).	required

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

def refocus(self, foc_dist):
    """Refocus the hybrid lens to a given object distance.

    Only the ``GeoLens`` sensor-to-last-surface spacing is adjusted; the
    DOE remains fixed relative to the refractive group (it is physically
    cemented to the lens barrel).

    Args:
        foc_dist (float): Target focus distance in [mm] (negative,
            towards the object).
    """
    self.geolens.refocus(foc_dist)

calc_scale

calc_scale(depth)

Calculate the object-to-image magnification scale factor.

Delegates to the embedded GeoLens.

Parameters:

Name	Type	Description	Default
depth	float	Object distance in [mm] (negative, towards the object).	required

Returns:

Name	Type	Description
float		Scale factor mapping normalised sensor coordinates [-1, 1] to physical object-space coordinates [mm].

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

def calc_scale(self, depth):
    """Calculate the object-to-image magnification scale factor.

    Delegates to the embedded `GeoLens`.

    Args:
        depth (float): Object distance in [mm] (negative, towards the
            object).

    Returns:
        float: Scale factor mapping normalised sensor coordinates
            ``[-1, 1]`` to physical object-space coordinates [mm].
    """
    return self.geolens.calc_scale(depth)

doe_field

doe_field(point, wvln=None, spp=SPP_COHERENT)

Compute the complex wave field at the DOE plane via coherent ray tracing.

Similar to GeoLens.pupil_field(), but evaluates the field at the last surface (DOE plane) instead of the exit pupil. The returned wavefront encodes amplitude, phase, and all diffraction-order information needed for subsequent DOE modulation and ASM propagation.

Parameters:

Name	Type	Description	Default
point	Tensor	Point source position, shape (3,) or (1, 3) as [x, y, z] in normalised sensor coordinates for x/y and mm for z.	required
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Number of rays to sample. Must be = 1,000,000 for accurate coherent simulation. Defaults to SPP_COHERENT.	SPP_COHERENT

Returns:

Name	Type	Description
tuple		wavefront (torch.Tensor) -- Complex wavefront at the DOE plane, shape [H, W]. psf_center (list[float]) -- Estimated PSF centre on the sensor in normalised coordinates [x, y].

Raises:

Type	Description
AssertionError	If spp < 1,000,000 or the default dtype is not float64.

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

def doe_field(self, point, wvln=None, spp=SPP_COHERENT):
    """Compute the complex wave field at the DOE plane via coherent ray tracing.

    Similar to ``GeoLens.pupil_field()``, but evaluates the field at the
    last surface (DOE plane) instead of the exit pupil.  The returned
    wavefront encodes amplitude, phase, and all diffraction-order
    information needed for subsequent DOE modulation and ASM propagation.

    Args:
        point (torch.Tensor): Point source position, shape ``(3,)`` or
            ``(1, 3)`` as ``[x, y, z]`` in normalised sensor coordinates
            for x/y and mm for z.
        wvln (float, optional): Wavelength in µm.  When ``None`` (default),
            falls back to ``self.primary_wvln``.
        spp (int, optional): Number of rays to sample.  Must be
            >= 1,000,000 for accurate coherent simulation.  Defaults to
            ``SPP_COHERENT``.

    Returns:
        tuple:
            - **wavefront** (*torch.Tensor*) -- Complex wavefront at the
              DOE plane, shape ``[H, W]``.
            - **psf_center** (*list[float]*) -- Estimated PSF centre on
              the sensor in normalised coordinates ``[x, y]``.

    Raises:
        AssertionError: If *spp* < 1,000,000 or the default dtype is not
            ``float64``.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    assert spp >= 1_000_000, (
        "Coherent ray tracing spp is too small, "
        "which may lead to inaccurate simulation."
    )
    assert torch.get_default_dtype() == torch.float64, (
        "Default dtype must be set to float64 for accurate phase tracing."
    )

    geolens, doe = self.geolens, self.doe

    if point.dim() == 1:
        point = point.unsqueeze(0)
    point = point.to(self.device)

    # Calculate ray origin in the object space
    scale = geolens.calc_scale(point[:, 2].item())
    point_obj = point.clone()
    point_obj[:, 0] = point[:, 0] * scale * geolens.sensor_size[1] / 2
    point_obj[:, 1] = point[:, 1] * scale * geolens.sensor_size[0] / 2

    # Determine ray center via chief ray
    pointc_chief_ray = geolens.psf_center(point_obj, method="chief_ray")[
        0
    ]  # shape [2]

    # Ray tracing to the DOE plane
    ray = geolens.sample_from_points(points=point_obj, num_rays=spp, wvln=wvln)
    ray.is_coherent = True
    ray, _ = geolens.trace(ray)
    ray = ray.prop_to(doe.d)

    # Calculate full-resolution complex field for exit-pupil diffraction
    wavefront = forward_integral(
        ray.flip_xy(),
        ps=doe.ps,
        ks=doe.res[0],
        pointc=torch.zeros_like(point[:, :2]),
    ).squeeze(0)  # shape [H, W]

    # Compute PSF center based on chief ray
    psf_center = [
        pointc_chief_ray[0] / geolens.sensor_size[0] * 2,
        pointc_chief_ray[1] / geolens.sensor_size[1] * 2,
    ]

    return wavefront, psf_center

psf

psf(points=[0.0, 0.0, -10000.0], ks=PSF_KS, wvln=None, spp=SPP_COHERENT)

Compute a single-point monochromatic PSF using the ray-wave model.

The returned PSF includes all diffraction orders with physically correct diffraction efficiencies. The pipeline is:

1. Coherent ray tracing through the GeoLens to obtain the complex wavefront at the DOE plane.
2. DOE phase modulation applied to the wavefront.
3. ASM propagation to the sensor, intensity calculation, cropping, and normalisation.

Parameters:

Name	Type	Description	Default
points	list or Tensor	[x, y, z] point source coordinates. x, y are in normalised sensor coordinates [-1, 1]; z is depth in [mm]. Defaults to [0.0, 0.0, -10000.0].	[0.0, 0.0, -10000.0]
ks	int or None	Output PSF patch size. If None, returns the central quarter of the full-sensor intensity. Defaults to PSF_KS.	PSF_KS
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Number of coherent rays to sample. Defaults to SPP_COHERENT.	SPP_COHERENT

Returns:

Type	Description
	torch.Tensor: Normalised PSF patch (sums to 1), shape [ks, ks]. Returned in float32 precision.

Raises:

Type	Description
ValueError	If the default dtype is not float64 (call double first).

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

def psf(
    self,
    points=[0.0, 0.0, -10000.0],
    ks=PSF_KS,
    wvln=None,
    spp=SPP_COHERENT,
):
    """Compute a single-point monochromatic PSF using the ray-wave model.

    The returned PSF includes all diffraction orders with physically
    correct diffraction efficiencies.  The pipeline is:

    1. Coherent ray tracing through the ``GeoLens`` to obtain the complex
       wavefront at the DOE plane.
    2. DOE phase modulation applied to the wavefront.
    3. ASM propagation to the sensor, intensity calculation, cropping, and
       normalisation.

    Args:
        points (list or torch.Tensor, optional): ``[x, y, z]`` point
            source coordinates.  *x, y* are in normalised sensor
            coordinates ``[-1, 1]``; *z* is depth in [mm].  Defaults to
            ``[0.0, 0.0, -10000.0]``.
        ks (int or None, optional): Output PSF patch size.  If ``None``,
            returns the central quarter of the full-sensor intensity.
            Defaults to ``PSF_KS``.
        wvln (float, optional): Wavelength in µm.  When ``None`` (default),
            falls back to ``self.primary_wvln``.
        spp (int, optional): Number of coherent rays to sample.  Defaults
            to ``SPP_COHERENT``.

    Returns:
        torch.Tensor: Normalised PSF patch (sums to 1), shape
            ``[ks, ks]``.  Returned in ``float32`` precision.

    Raises:
        ValueError: If the default dtype is not ``float64`` (call
            `double` first).
    """
    wvln = self.primary_wvln if wvln is None else wvln
    # Check double precision
    if not torch.get_default_dtype() == torch.float64:
        raise ValueError(
            "Please call HybridLens.double() to set the default dtype to float64 for accurate phase tracing."
        )

    # Check lens last surface
    assert isinstance(self.geolens.surfaces[-1], Phase) or isinstance(
        self.geolens.surfaces[-1], Plane
    ), "The last lens surface should be a DOE."
    geolens, doe = self.geolens, self.doe

    # Compute pupil field by coherent ray tracing
    if isinstance(points, list):
        point0 = torch.tensor(points)
    elif isinstance(points, torch.Tensor):
        point0 = points
    else:
        raise ValueError("point should be a list or a torch.Tensor.")

    wavefront, psfc = self.doe_field(point=point0, wvln=wvln, spp=spp)
    wavefront = wavefront.squeeze(0)  # shape of [H, W]

    # DOE phase modulation. We have to flip the phase map because the wavefront has been flipped
    phase_map = torch.flip(doe.get_phase_map(wvln), [-1, -2])
    wavefront = wavefront * torch.exp(1j * phase_map)

    # Propagate wave field to sensor plane
    h, w = wavefront.shape
    wavefront = F.pad(
        wavefront.unsqueeze(0).unsqueeze(0),
        [h // 2, h // 2, w // 2, w // 2],
        mode="constant",
        value=0,
    )
    sensor_field = AngularSpectrumMethod(
        wavefront, z=geolens.d_sensor - doe.d, wvln=wvln, ps=doe.ps, padding=False
    )

    # Compute PSF (intensity distribution)
    psf_inten = sensor_field.abs() ** 2
    psf_inten = (
        F.interpolate(
            psf_inten,
            scale_factor=geolens.sensor_res[0] / h,
            mode="bilinear",
            align_corners=False,
        )
        .squeeze(0)
        .squeeze(0)
    )

    # Calculate PSF center index and crop valid PSF region (Consider both interplation and padding)
    if ks is not None:
        h, w = psf_inten.shape[-2:]
        psfc_idx_i = ((2 - psfc[1]) * h / 4).round().long()
        psfc_idx_j = ((2 + psfc[0]) * w / 4).round().long()

        # Pad to avoid invalid edge region
        psf_inten_pad = F.pad(
            psf_inten,
            [ks // 2, ks // 2, ks // 2, ks // 2],
            mode="constant",
            value=0,
        )
        psf = psf_inten_pad[
            psfc_idx_i : psfc_idx_i + ks, psfc_idx_j : psfc_idx_j + ks
        ]
    else:
        h, w = psf_inten.shape[-2:]
        psf = psf_inten[
            int(h / 2 - h / 4) : int(h / 2 + h / 4),
            int(w / 2 - w / 4) : int(w / 2 + w / 4),
        ]

    # Normalize and convert to float precision
    psf /= psf.sum()  # shape of [ks, ks] or [h, w]
    return diff_float(psf)

draw_layout

draw_layout(save_name='./DOELens.png', depth=-10000.0, ax=None, fig=None, dpi=600)

Draw the hybrid-lens layout with ray paths and wave-propagation arcs.

Renders the refractive elements via GeoLens.draw_lens_2d(), traces rays at three field angles (on-axis, 0.707x, 0.99x full field), and overlays concentric arcs between the DOE and sensor to illustrate the wave-propagation region.

Parameters:

Name	Type	Description	Default
save_name	str	File path to save the figure (used only when ax is None). Defaults to './DOELens.png'.	'./DOELens.png'
depth	float	Object depth [mm] for the traced rays. Defaults to -10000.0.	-10000.0
ax	Axes	Pre-existing axes to draw into. If None, a new figure is created and saved.	None
fig	Figure	Pre-existing figure. Required when ax is provided.	None
dpi	int	Resolution used when saving a new figure. Defaults to 600.	600

Returns:

Type	Description
	tuple or None: (ax, fig) when ax was provided; otherwise the figure is saved to save_name and nothing is returned.

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

@torch.no_grad()
def draw_layout(
    self,
    save_name="./DOELens.png",
    depth=-10000.0,
    ax=None,
    fig=None,
    dpi=600,
):
    """Draw the hybrid-lens layout with ray paths and wave-propagation arcs.

    Renders the refractive elements via ``GeoLens.draw_lens_2d()``, traces
    rays at three field angles (on-axis, 0.707x, 0.99x full field), and
    overlays concentric arcs between the DOE and sensor to illustrate the
    wave-propagation region.

    Args:
        save_name (str, optional): File path to save the figure (used only
            when *ax* is ``None``).  Defaults to ``'./DOELens.png'``.
        depth (float, optional): Object depth [mm] for the traced rays.
            Defaults to ``-10000.0``.
        ax (matplotlib.axes.Axes, optional): Pre-existing axes to draw
            into.  If ``None``, a new figure is created and saved.
        fig (matplotlib.figure.Figure, optional): Pre-existing figure.
            Required when *ax* is provided.
        dpi (int, optional): Resolution used when saving a new figure.
            Defaults to 600.

    Returns:
        tuple or None: ``(ax, fig)`` when *ax* was provided; otherwise
            the figure is saved to *save_name* and nothing is returned.
    """
    geolens = self.geolens

    # Draw lens layout
    if ax is None:
        ax, fig = geolens.draw_lens_2d()
        save_fig = True
    else:
        save_fig = False

    # Draw DOE as orange Fresnel-style widget
    self.doe.draw_widget(ax, color="orange")

    # Draw light path
    color_list = ["#CC0000", "#006600", "#0066CC"]
    views = [
        0.0,
        float(np.rad2deg(geolens.rfov) * 0.707),
        float(np.rad2deg(geolens.rfov) * 0.99),
    ]
    arc_radi_list = [0.1, 0.4, 0.7, 1.0, 1.4, 1.8]
    num_rays = 11
    arc_half_angle = 20
    for i, view in enumerate(views):
        # Draw ray tracing
        ray = geolens.sample_point_source_2D(
            depth=depth,
            fov=view,
            num_rays=num_rays,
            entrance_pupil=True,
            wvln=self.wvln_rgb[2 - i],
        )
        ray.prop_to(-1.0)

        ray, ray_o_record = geolens.trace(ray=ray, record=True)
        ax, fig = geolens.draw_ray_2d(
            ray_o_record, ax=ax, fig=fig, color=color_list[i]
        )

        # Draw wave propagation
        # Calculate ray center for wave propagation visualization
        ray_center_doe = (
            ((ray.o * ray.is_valid.unsqueeze(-1)).sum(dim=0) / ray.is_valid.sum())
            .cpu()
            .numpy()
        )  # shape [3]
        ray.prop_to(geolens.d_sensor)  # shape [num_rays, 3]
        ray_center_sensor = (
            ((ray.o * ray.is_valid.unsqueeze(-1)).sum(dim=0) / ray.is_valid.sum())
            .cpu()
            .numpy()
        )  # shape [3]

        arc_radi = ray_center_sensor[2] - ray_center_doe[2]
        chief_theta = np.rad2deg(
            np.arctan2(
                ray_center_sensor[0] - ray_center_doe[0],
                ray_center_sensor[2] - ray_center_doe[2],
            )
        )
        theta1 = chief_theta - arc_half_angle
        theta2 = chief_theta + arc_half_angle

        for j in arc_radi_list:
            arc_radi_j = arc_radi * j
            arc = patches.Arc(
                (ray_center_sensor[2], ray_center_sensor[0]),
                arc_radi_j,
                arc_radi_j,
                angle=180.0,
                theta1=theta1,
                theta2=theta2,
                color=color_list[i],
            )
            ax.add_patch(arc)

    if save_fig:
        # Save figure
        ax.axis("off")
        ax.set_title("DOE Lens")
        fig.savefig(save_name, bbox_inches="tight", dpi=dpi)
        plt.close()
    else:
        return ax, fig

get_optimizer

get_optimizer(doe_lr=0.0001, lens_lr=[0.0001, 0.0001, 0.01, 1e-05])

Build an Adam optimiser for joint lens + DOE design.

Collects trainable parameters from both the GeoLens (surface thicknesses, curvatures, conic constants, aspheric coefficients) and the DOE phase profile into a single optimiser with per-group learning rates.

Parameters:

Name	Type	Description	Default
doe_lr	float	Learning rate for DOE phase parameters. Defaults to 1e-4.	0.0001
lens_lr	list[float]	Per-parameter-group learning rates for the GeoLens, ordered as [thickness_d, curvature_c, conic_k, aspheric_a]. Defaults to [1e-4, 1e-4, 1e-2, 1e-5].	[0.0001, 0.0001, 0.01, 1e-05]

Returns:

Type	Description
	torch.optim.Adam: Configured optimiser over all trainable parameters.

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

def get_optimizer(
    self, doe_lr=1e-4, lens_lr=[1e-4, 1e-4, 1e-2, 1e-5]
):
    """Build an Adam optimiser for joint lens + DOE design.

    Collects trainable parameters from both the ``GeoLens`` (surface
    thicknesses, curvatures, conic constants, aspheric coefficients) and
    the DOE phase profile into a single optimiser with per-group learning
    rates.

    Args:
        doe_lr (float, optional): Learning rate for DOE phase parameters.
            Defaults to ``1e-4``.
        lens_lr (list[float], optional): Per-parameter-group learning
            rates for the GeoLens, ordered as
            ``[thickness_d, curvature_c, conic_k, aspheric_a]``.
            Defaults to ``[1e-4, 1e-4, 1e-2, 1e-5]``.

    Returns:
        torch.optim.Adam: Configured optimiser over all trainable
            parameters.
    """
    params = []
    params += self.geolens.get_optimizer_params(lrs=lens_lr)
    params += self.doe.get_optimizer_params(lr=doe_lr)

    optimizer = torch.optim.Adam(params)
    return optimizer

DOE Models

The diffractive element is configured by the DOE block in the lens JSON; its type field selects one of the phase parameterizations below. All subclass DiffractiveSurface, which defines the shared phase / propagation interface.

deeplens.diffractive_surface.DiffractiveSurface

DiffractiveSurface(d, res, fab_ps=0.001, fab_step=16, wvln0=0.55, mat='fused_silica', design_ps=None, is_square=True, device='cpu')

Bases: DeepObj

Diffractive (multi-layer diffractive) surface class. Optical properties of diffractive surfaces are simulated with wave optics.

By default the DOE is designed for 0.55um, which means it will have the highest 1st-order diffraction efficiency for 0.55um.

Parameters:

Name	Type	Description	Default
d	float	Distance of the DOE surface. [mm]	required
res	tuple or int	Resolution of the DOE, [w, h]. [pixel]	required
fab_ps	float	Fabrication pixel size. [mm]	0.001
fab_step	int	Fabrication step. Default is 16.	16
wvln0	float	Design wavelength. [um]	0.55
mat	str	Material of the DOE.	'fused_silica'
design_ps	float	Design pixel size. [mm]	None
device	str	Device to run the DOE.	'cpu'

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def __init__(
    self,
    d,
    res,
    fab_ps=0.001,
    fab_step=16,
    wvln0=0.55,
    mat="fused_silica",
    design_ps=None,
    is_square=True,
    device="cpu",
):
    """Diffractive (multi-layer diffractive) surface class. Optical properties of diffractive surfaces are simulated with wave optics.

    By default the DOE is designed for 0.55um, which means it will have the highest 1st-order diffraction efficiency for 0.55um.

    Args:
        d (float): Distance of the DOE surface. [mm]
        res (tuple or int): Resolution of the DOE, [w, h]. [pixel]
        fab_ps (float): Fabrication pixel size. [mm]
        fab_step (int): Fabrication step. Default is 16.
        wvln0 (float): Design wavelength. [um]
        mat (str): Material of the DOE.
        design_ps (float): Design pixel size. [mm]
        device (str): Device to run the DOE.
    """
    # Geometry
    self.d = torch.tensor(d) if not isinstance(d, torch.Tensor) else d
    self.res = (res, res) if isinstance(res, int) else res
    self.ps = fab_ps if design_ps is None else design_ps
    self.w = self.res[0] * self.ps
    self.h = self.res[1] * self.ps
    self.is_square = is_square
    # Surface radius: half-diagonal (circumscribed-circle radius) so it
    # is consistent with Phase / Surface conventions for square apertures.
    self.r = float(np.sqrt(self.w**2 + self.h**2) / 2)

    # Phase map
    self.mat = Material(mat)
    self.wvln0 = wvln0  # [um], design wavelength. Sometimes the maximum working wavelength is preferred.
    self.n0 = self.mat.refractive_index(
        self.wvln0
    )  # refractive index at design wavelength

    # Fabrication for DOE
    self.fab_ps = fab_ps  # [mm], fabrication pixel size
    self.fab_step = fab_step

    # x, y coordinates
    self.x, self.y = torch.meshgrid(
        torch.linspace(-self.w / 2, self.w / 2, self.res[1]),
        torch.linspace(self.h / 2, -self.h / 2, self.res[0]),
        indexing="xy",
    )

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize DOE from a dict.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize DOE from a dict."""
    raise NotImplementedError

phase_func

phase_func()

Calculate raw phase function (no wrapping, no quantization) at design wavelength.

Returns:

Name	Type	Description
phase	tensor	raw phase function at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def phase_func(self):
    """Calculate raw phase function (no wrapping, no quantization) at design wavelength.

    Returns:
        phase (tensor): raw phase function at design wavelength.
    """
    raise NotImplementedError

get_phase_map0

get_phase_map0()

Calculate phase map at design wavelength with phase wrapping and quantization.

In this function, we are actually processing height map. The maximum height is 2pi for design wavelength.

Returns:

Name	Type	Description
phase0	tensor	phase map at design wavelength, range [0, 2pi].

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def get_phase_map0(self):
    """Calculate phase map at design wavelength with phase wrapping and quantization.

    In this function, we are actually processing height map. The maximum height is 2pi for design wavelength.

    Returns:
        phase0 (tensor): phase map at design wavelength, range [0, 2pi].
    """
    # Raw phase map at design wavelength
    phase0 = self.phase_func()

    # Phase wrapping and quantization
    phase0 = torch.remainder(phase0, 2 * torch.pi)
    phase0 = diff_quantize(phase0, levels=self.fab_step)
    return phase0

get_phase_map

get_phase_map(wvln)

Calculate phase map at the given wavelength.

Parameters:

Name	Type	Description	Default
wvln	float	Wavelength. [um].	required

Returns:

Name	Type	Description
phase_map	tensor	Phase map. [1, 1, H, W], range [0, 2pi].

Note

First we should calculate the phase map at 0.55um, then calculate the phase map for the given other wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def get_phase_map(self, wvln):
    """Calculate phase map at the given wavelength.

    Args:
        wvln (float): Wavelength. [um].

    Returns:
        phase_map (tensor): Phase map. [1, 1, H, W], range [0, 2pi].

    Note:
        First we should calculate the phase map at 0.55um, then calculate the phase map for the given other wavelength.
    """
    # Phase map at design wavelength
    phase_map0 = self.get_phase_map0()

    # Phase map at given wavelength (implicitly converted to height map)
    n = self.mat.refractive_index(wvln)
    phase_map = phase_map0 * (self.wvln0 / wvln) * (n - 1) / (self.n0 - 1)

    # Interpolate to the desired resolution (skip if already matching)
    if phase_map.shape[-2:] != (self.res[0], self.res[1]):
        phase_map = (
            F.interpolate(
                phase_map.unsqueeze(0).unsqueeze(0), size=self.res, mode="nearest"
            )
            .squeeze(0)
            .squeeze(0)
        )

    return phase_map

forward

forward(wave)

Propagate wave field to the DOE and apply phase modulation. Input wave field can have different pixel size and physical size with the DOE.

Parameters:

Name	Type	Description	Default
wave	Wave	Input complex wave field. Shape of [B, 1, H, W].	required

Returns:

Name	Type	Description
wave	Wave	Output complex wave field. Shape of [B, 1, H, W].

Reference

[1] https://github.com/vsitzmann/deepoptics function phaseshifts_from_height_map

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def forward(self, wave):
    """Propagate wave field to the DOE and apply phase modulation. Input wave field can have different pixel size and physical size with the DOE.

    Args:
        wave (Wave): Input complex wave field. Shape of [B, 1, H, W].

    Returns:
        wave (Wave): Output complex wave field. Shape of [B, 1, H, W].

    Reference:
        [1] https://github.com/vsitzmann/deepoptics function phaseshifts_from_height_map
    """
    # Propagate to DOE
    wave.prop_to(self.d)

    # Compute phase map at the wave field wavelength, shape of [H, W]
    phase_map = self.get_phase_map(wave.wvln)

    # Consider the different pixel size between the wave field and the DOE
    if self.ps != wave.ps:
        scale = self.ps / wave.ps
        phase_map = (
            F.interpolate(
                phase_map.unsqueeze(0).unsqueeze(0),
                scale_factor=(scale, scale),
                mode="nearest",
            )
            .squeeze(0)
            .squeeze(0)
        )

    # Check if the field and phase map resolution (physical size) are the same
    wave_h, wave_w = wave.u.shape[-2:]
    phase_h, phase_w = phase_map.shape[-2:]
    if phase_h > wave_h or phase_w > wave_w:
        start_h = (phase_h - wave_h) // 2
        start_w = (phase_w - wave_w) // 2
        phase_map = phase_map[
            ..., start_h : start_h + wave_h, start_w : start_w + wave_w
        ]
    elif phase_h < wave_h or phase_w < wave_w:
        pad_top = (wave_h - phase_h) // 2
        pad_bottom = wave_h - phase_h - pad_top
        pad_left = (wave_w - phase_w) // 2
        pad_right = wave_w - phase_w - pad_left
        phase_map = F.pad(
            phase_map,
            (pad_left, pad_right, pad_top, pad_bottom),
            mode="constant",
            value=0,
        )

    wave.u = wave.u * torch.exp(1j * phase_map)
    return wave

__call__

__call__(wave)

Forward function.

Parameters:

Name	Type	Description	Default
wave	Wave	Input complex wave field.	required

Returns:

Name	Type	Description
wave	Wave	Output complex wave field.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def __call__(self, wave):
    """Forward function.

    Args:
        wave (Wave): Input complex wave field.

    Returns:
        wave (Wave): Output complex wave field.
    """
    return self.forward(wave)

pmap_quantize

pmap_quantize(bits=16)

Quantize phase map to bits levels.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def pmap_quantize(self, bits=16):
    """Quantize phase map to bits levels."""
    pmap = self.get_phase_map0()
    pmap_q = torch.round(pmap / (2 * torch.pi / bits)) * (2 * torch.pi / bits)
    return pmap_q

pmap_fab

pmap_fab(bits=16, save_path=None)

Convert to fabricate phase map and save it. This function is used to output DOE_fab file, and it will not change the DOE object itself.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def pmap_fab(self, bits=16, save_path=None):
    """Convert to fabricate phase map and save it. This function is used to output DOE_fab file, and it will not change the DOE object itself."""
    # Fab resolution quantized pmap
    pmap = self.get_phase_map0()
    fab_res = int(self.ps / self.fab_ps * self.res[0])
    pmap = (
        F.interpolate(
            pmap.unsqueeze(0).unsqueeze(0),
            scale_factor=self.ps / self.fab_ps,
            mode="bilinear",
            align_corners=True,
        )
        .squeeze(0)
        .squeeze(0)
    )
    pmap_q = torch.round(pmap / (2 * torch.pi / bits)) * (2 * torch.pi / bits)

    # Save phase map
    if save_path is None:
        save_path = f"./doe_fab_{fab_res}x{fab_res}_{int(self.fab_ps * 1000)}um_{bits}bit.pth"
    self.save_ckpt(save_path=save_path)

    return pmap_q

activate_grad

activate_grad(activate=True)

Activate gradient for phase map parameters.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def activate_grad(self, activate=True):
    """Activate gradient for phase map parameters."""
    raise NotImplementedError

get_optimizer

get_optimizer(lr=None)

Generate optimizer for DOE.

Parameters:

Name	Type	Description	Default
lr	float	Learning rate. Defaults to 1e-3.	None

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def get_optimizer(self, lr=None):
    """Generate optimizer for DOE.

    Args:
        lr (float, optional): Learning rate. Defaults to 1e-3.
    """
    params = self.get_optimizer_params(lr)
    optimizer = torch.optim.Adam(params)

    return optimizer

loss_quantization

loss_quantization(bits=16)

DOE quantization errors.

Reference: Quantization-aware Deep Optics for Diffractive Snapshot Hyperspectral Imaging

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def loss_quantization(self, bits=16):
    """DOE quantization errors.

    Reference: Quantization-aware Deep Optics for Diffractive Snapshot Hyperspectral Imaging
    """
    pmap = self.get_phase_map0()
    step = 2 * torch.pi / bits
    pmap_q = torch.round(pmap / step) * step
    loss = torch.mean(torch.abs(pmap - pmap_q))
    return loss

draw_phase_map

draw_phase_map(bits=None, save_name='./DOE_phase_map.png')

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

Parameters:

Name	Type	Description	Default
bits	int	Number of quantization bits. If provided, quantizes the phase map.	None
save_name	str	Path to save the image.	'./DOE_phase_map.png'

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def draw_phase_map(self, bits=None, save_name="./DOE_phase_map.png"):
    """Draw phase map. Range from [0, 2pi].

    Args:
        bits (int, optional): Number of quantization bits. If provided, quantizes the phase map.
        save_name (str): Path to save the image.
    """
    if bits is not None:
        pmap = self.pmap_quantize(bits)
    else:
        pmap = self.get_phase_map0()
    save_image(pmap, save_name, normalize=True)

draw_phase_map3d

draw_phase_map3d(bits=None, save_name='./DOE_phase_map3d.png')

Draw 3D phase map.

Parameters:

Name	Type	Description	Default
bits	int	Number of quantization bits. If provided, quantizes the phase map.	None
save_name	str	Path to save the image.	'./DOE_phase_map3d.png'

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def draw_phase_map3d(self, bits=None, save_name="./DOE_phase_map3d.png"):
    """Draw 3D phase map.

    Args:
        bits (int, optional): Number of quantization bits. If provided, quantizes the phase map.
        save_name (str): Path to save the image.
    """
    if bits is not None:
        pmap = self.pmap_quantize(bits)
    else:
        pmap = self.get_phase_map0()

    pmap = pmap / 20.0
    x = np.linspace(-self.w / 2, self.w / 2, self.res[0])
    y = np.linspace(-self.h / 2, self.h / 2, self.res[1])
    X, Y = np.meshgrid(x, y)

    fig = plt.figure(figsize=(5, 5))
    ax = fig.add_subplot(111, projection="3d")
    ax.scatter(
        X.flatten(),
        Y.flatten(),
        pmap.cpu().numpy().flatten(),
        marker=".",
        s=0.01,
        c=pmap.cpu().numpy().flatten(),
        cmap="viridis",
    )
    ax.set_aspect("equal")
    ax.axis("off")
    fig.savefig(save_name, dpi=600, bbox_inches="tight")
    plt.close(fig)

draw_phase_map_fab

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

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

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def draw_phase_map_fab(self, save_name="./DOE_phase_map.png"):
    """Draw phase map. Range from [0, 2pi]."""
    pmap = self.get_phase_map0()
    step = 2 * torch.pi / 16
    pmap_q = torch.round(pmap / step) * step

    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
    ax[0].imshow(pmap.cpu().numpy(), vmin=0, vmax=2 * float(np.pi))
    ax[0].set_title(f"Phase map ({self.wvln0}um)", fontsize=10)
    ax[0].grid(False)
    fig.colorbar(ax[0].get_images()[0])

    ax[1].imshow(pmap_q.cpu().numpy(), vmin=0, vmax=2 * float(np.pi))
    ax[1].set_title(f"Quantized phase map ({self.wvln0}um)", fontsize=10)
    ax[1].grid(False)
    fig.colorbar(ax[1].get_images()[0])

    fig.savefig(save_name, dpi=600, bbox_inches="tight")
    plt.close(fig)

draw_cross_section

draw_cross_section(save_name='./DOE_cross_section.png')

Draw cross section of the phase map.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def draw_cross_section(self, save_name="./DOE_cross_section.png"):
    """Draw cross section of the phase map."""
    pmap = self.get_phase_map0()
    pmap = torch.diag(pmap).cpu().numpy()
    r = np.linspace(
        -self.w / 2 * float(np.sqrt(2)), self.w / 2 * float(np.sqrt(2)), self.res[0]
    )

    fig, ax = plt.subplots()
    ax.plot(r, pmap)
    ax.set_title(f"Phase map ({self.wvln0}um) cross section")
    fig.savefig(save_name, dpi=600, bbox_inches="tight")
    plt.close(fig)

draw_widget

draw_widget(ax, color='orange', linestyle='-')

Draw a 2D Fresnel-style widget for the DOE in a layout plot.

Plots the cross-section along the x-axis at y=0. For a square aperture the half-extent is the half-side (w/2); for a circular aperture it is the full radius r (= half-diagonal).

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def draw_widget(self, ax, color="orange", linestyle="-"):
    """Draw a 2D Fresnel-style widget for the DOE in a layout plot.

    Plots the cross-section along the x-axis at y=0. For a square aperture
    the half-extent is the half-side (``w/2``); for a circular aperture it
    is the full radius ``r`` (= half-diagonal).
    """
    d = self.d.item()
    max_offset = d / 100
    roc = self.r * 2
    x_half = self.w / 2 if self.is_square else self.r
    x = np.linspace(-x_half, x_half, 256)
    sag = roc * (1 - np.sqrt(1 - x**2 / roc**2))
    sag = max_offset - np.fmod(sag, max_offset)
    ax.plot(d + sag, x, color=color, linestyle=linestyle, linewidth=0.75)

surf_dict

surf_dict()

Return a dict of surface.

Source code in deeplens-src/deeplens/diffractive_surface/diffractive.py

def surf_dict(self):
    """Return a dict of surface."""
    surf_dict = {
        "type": self.__class__.__name__,
        "(size)": [round(self.w, 4), round(self.h, 4)],
        "d": round(self.d.item(), 4),
        "wvln0": round(self.wvln0, 4),
        "res": self.res,
        "is_square": True,
    }

    return surf_dict

Polynomial (Binary-2) rotationally-symmetric phase profile.

deeplens.diffractive_surface.Binary2

Binary2(d, res=(2000, 2000), mat='fused_silica', wvln0=0.55, fab_ps=0.001, fab_step=16, is_square=True, device='cpu')

Bases: DiffractiveSurface

Initialize Binary DOE.

Source code in deeplens-src/deeplens/diffractive_surface/binary2.py

def __init__(
    self,
    d,
    res=(2000, 2000),
    mat="fused_silica",
    wvln0=0.55,
    fab_ps=0.001,
    fab_step=16,
    is_square=True,
    device="cpu",
):
    """Initialize Binary DOE."""
    super().__init__(
        d=d, res=res, mat=mat, wvln0=wvln0, fab_ps=fab_ps, fab_step=fab_step,
        is_square=is_square, device=device,
    )

    # Initialize with random small values
    self.alpha2 = (torch.rand(1) - 0.5) * 0.02
    self.alpha4 = (torch.rand(1) - 0.5) * 0.002
    self.alpha6 = (torch.rand(1) - 0.5) * 0.0002
    self.alpha8 = (torch.rand(1) - 0.5) * 0.00002
    self.alpha10 = (torch.rand(1) - 0.5) * 0.000002

    self.x, self.y = torch.meshgrid(
        torch.linspace(-self.w / 2, self.w / 2, self.res[1]),
        torch.linspace(self.h / 2, -self.h / 2, self.res[0]),
        indexing="xy",
    )

    # Cache static r² grid (x, y never change after init)
    self.r2 = self.x**2 + self.y**2

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize Binary DOE from a dict.

Source code in deeplens-src/deeplens/diffractive_surface/binary2.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize Binary DOE from a dict."""
    return cls(
        d=doe_dict["d"],
        res=doe_dict["res"],
        mat=doe_dict.get("mat", "fused_silica"),
        wvln0=doe_dict.get("wvln0", 0.55),
        fab_ps=doe_dict.get("fab_ps", 0.001),
        fab_step=doe_dict.get("fab_step", 16),
        is_square=doe_dict.get("is_square", True),
    )

phase_func

phase_func()

Get the phase map at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/binary2.py

def phase_func(self):
    """Get the phase map at design wavelength."""
    # Horner's method: r2*(a2 + r2*(a4 + r2*(a6 + r2*(a8 + r2*a10))))
    r2 = self.r2
    phase = torch.pi * r2 * (
        self.alpha2
        + r2 * (self.alpha4 + r2 * (self.alpha6 + r2 * (self.alpha8 + r2 * self.alpha10)))
    )
    return phase

get_optimizer_params

get_optimizer_params(lr=0.001)

Get parameters for optimization.

Parameters:

Name	Type	Description	Default
lr	float	Base learning rate for alpha2. Learning rates for higher-order parameters will be scaled progressively (10x, 100x, 1000x, 10000x).	0.001

Source code in deeplens-src/deeplens/diffractive_surface/binary2.py

def get_optimizer_params(self, lr=0.001):
    """Get parameters for optimization.

    Args:
        lr (float): Base learning rate for alpha2. Learning rates for higher-order parameters will be scaled progressively (10x, 100x, 1000x, 10000x).
    """
    self.alpha2.requires_grad = True
    self.alpha4.requires_grad = True
    self.alpha6.requires_grad = True
    self.alpha8.requires_grad = True
    self.alpha10.requires_grad = True

    optimizer_params = [
        {"params": [self.alpha2], "lr": lr},
        {"params": [self.alpha4], "lr": lr * 10},
        {"params": [self.alpha6], "lr": lr * 100},
        {"params": [self.alpha8], "lr": lr * 1000},
        {"params": [self.alpha10], "lr": lr * 10000},
    ]

    return optimizer_params

surf_dict

surf_dict()

Return a dict of surface.

Source code in deeplens-src/deeplens/diffractive_surface/binary2.py

def surf_dict(self):
    """Return a dict of surface."""
    surf_dict = super().surf_dict()
    surf_dict["alpha2"] = round(self.alpha2.item(), 6)
    surf_dict["alpha4"] = round(self.alpha4.item(), 6)
    surf_dict["alpha6"] = round(self.alpha6.item(), 6)
    surf_dict["alpha8"] = round(self.alpha8.item(), 6)
    surf_dict["alpha10"] = round(self.alpha10.item(), 6)
    return surf_dict

Free-form, per-pixel phase map.

deeplens.diffractive_surface.Pixel2D

Pixel2D(d, phase_map_path=None, res=(2000, 2000), mat='fused_silica', wvln0=0.55, fab_ps=0.001, fab_step=16, device='cpu')

Bases: DiffractiveSurface

Pixel2D DOE parameterization - direct phase map representation.

Initialize Pixel2D DOE, where each pixel is independent parameter.

Parameters:

Name	Type	Description	Default
d	float	Distance of the DOE surface. [mm]	required
size	tuple or int	Size of the DOE, [w, h]. [mm]	required
res	tuple or int	Resolution of the DOE, [w, h]. [pixel]	(2000, 2000)
mat	str	Material of the DOE.	'fused_silica'
fab_ps	float	Fabrication pixel size. [mm]	0.001
fab_step	int	Fabrication step.	16
device	str	Device to run the DOE.	'cpu'

Source code in deeplens-src/deeplens/diffractive_surface/pixel2d.py

def __init__(
    self,
    d,
    phase_map_path=None,
    res=(2000, 2000),
    mat="fused_silica",
    wvln0=0.55,
    fab_ps=0.001,
    fab_step=16,
    device="cpu",
):
    """Initialize Pixel2D DOE, where each pixel is independent parameter.

    Args:
        d (float): Distance of the DOE surface. [mm]
        size (tuple or int): Size of the DOE, [w, h]. [mm]
        res (tuple or int): Resolution of the DOE, [w, h]. [pixel]
        mat (str): Material of the DOE.
        fab_ps (float): Fabrication pixel size. [mm]
        fab_step (int): Fabrication step.
        device (str): Device to run the DOE.
    """
    super().__init__(d=d, res=res, mat=mat, fab_ps=fab_ps, fab_step=fab_step, wvln0=wvln0, device=device)

    # Initialize phase map with random values
    if phase_map_path is None:
        self.phase_map = torch.randn(self.res, device=self.device) * 1e-3
    elif isinstance(phase_map_path, str):
        self.phase_map = torch.load(phase_map_path, map_location=device, weights_only=True)
    else:
        raise ValueError(f"Invalid phase_map_path: {phase_map_path}")

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize Pixel2D DOE from a dict.

Source code in deeplens-src/deeplens/diffractive_surface/pixel2d.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize Pixel2D DOE from a dict."""
    return cls(
        d=doe_dict["d"],
        res=doe_dict["res"],
        mat=doe_dict.get("mat", "fused_silica"),
        fab_ps=doe_dict.get("fab_ps", 0.001),
        fab_step=doe_dict.get("fab_step", 16),
        phase_map_path=doe_dict.get("phase_map_path", None),
        wvln0=doe_dict.get("wvln0", 0.55),
    )

phase_func

phase_func()

Get the phase map at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/pixel2d.py

def phase_func(self):
    """Get the phase map at design wavelength."""
    return self.phase_map

get_optimizer_params

get_optimizer_params(lr=0.01)

Get parameters for optimization.

Source code in deeplens-src/deeplens/diffractive_surface/pixel2d.py

def get_optimizer_params(self, lr=0.01):
    """Get parameters for optimization."""
    self.phase_map.requires_grad = True
    optimizer_params = [{"params": [self.phase_map], "lr": lr}]
    return optimizer_params

surf_dict

surf_dict(phase_map_path)

Return a dict of surface.

Source code in deeplens-src/deeplens/diffractive_surface/pixel2d.py

def surf_dict(self, phase_map_path):
    """Return a dict of surface."""
    surf_dict = super().surf_dict()
    surf_dict["phase_map_path"] = phase_map_path
    torch.save(self.phase_map.clone().detach().cpu(), phase_map_path)
    return surf_dict

Fresnel-lens (quadratic) phase profile.

deeplens.diffractive_surface.Fresnel

Fresnel(d, f0=None, wvln0=0.55, res=(2000, 2000), mat='fused_silica', fab_ps=0.001, fab_step=16, device='cpu')

Bases: DiffractiveSurface

Initialize Fresnel DOE. A diffractive Fresnel lens shows inverse dispersion property compared to refractive lens.

Parameters:

Name	Type	Description	Default
f0	float	Initial focal length. [mm]	None
d	float	Distance of the DOE surface. [mm]	required
res	tuple or int	Resolution of the DOE, [w, h]. [pixel]	(2000, 2000)
wvln0	float	Design wavelength. [um]	0.55
mat	str	Material of the DOE.	'fused_silica'
fab_ps	float	Fabrication pixel size. [mm]	0.001
fab_step	int	Fabrication step.	16
device	str	Device to run the DOE.	'cpu'

Source code in deeplens-src/deeplens/diffractive_surface/fresnel.py

def __init__(
    self,
    d,
    f0=None,
    wvln0=0.55,
    res=(2000, 2000),
    mat="fused_silica",
    fab_ps=0.001,
    fab_step=16,
    device="cpu",
):
    """Initialize Fresnel DOE. A diffractive Fresnel lens shows inverse dispersion property compared to refractive lens.

    Args:
        f0 (float): Initial focal length. [mm]
        d (float): Distance of the DOE surface. [mm]
        res (tuple or int): Resolution of the DOE, [w, h]. [pixel]
        wvln0 (float): Design wavelength. [um]
        mat (str): Material of the DOE.
        fab_ps (float): Fabrication pixel size. [mm]
        fab_step (int): Fabrication step.
        device (str): Device to run the DOE.
    """
    super().__init__(
        d=d, res=res, wvln0=wvln0, mat=mat, fab_ps=fab_ps, fab_step=fab_step, device=device
    )

    # Initial focal length
    if f0 is None:
        self.f0 = torch.randn(1) * 1e6
    else:
        self.f0 = torch.tensor(f0)

    # Cache static r² grid (x, y never change after init)
    self.r2 = self.x**2 + self.y**2

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize Fresnel DOE from a dict.

Source code in deeplens-src/deeplens/diffractive_surface/fresnel.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize Fresnel DOE from a dict."""
    return cls(
        d=doe_dict["d"],
        res=doe_dict["res"],
        fab_ps=doe_dict.get("fab_ps", 0.001),
        fab_step=doe_dict.get("fab_step", 16),
        f0=doe_dict.get("f0", None),
        wvln0=doe_dict.get("wvln0", 0.55),
        mat=doe_dict.get("mat", "fused_silica"),
    )

phase_func

phase_func()

Get the phase map at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/fresnel.py

def phase_func(self):
    """Get the phase map at design wavelength."""
    wvln0_mm = self.wvln0 * 1e-3
    phase = -2 * torch.pi * self.r2 / (2 * self.f0 * wvln0_mm)
    self._warn_if_undersampled(phase, self.f0, self.wvln0)
    return phase

get_optimizer_params

get_optimizer_params(lr=0.001)

Get parameters for optimization.

Source code in deeplens-src/deeplens/diffractive_surface/fresnel.py

def get_optimizer_params(self, lr=0.001):
    """Get parameters for optimization."""
    self.f0.requires_grad = True
    optimizer_params = [{"params": [self.f0], "lr": lr}]
    return optimizer_params

surf_dict

surf_dict()

Return a dict of surface.

Source code in deeplens-src/deeplens/diffractive_surface/fresnel.py

def surf_dict(self):
    """Return a dict of surface."""
    surf_dict = super().surf_dict()
    surf_dict["f0"] = self.f0.item()
    surf_dict["wvln0"] = self.wvln0
    return surf_dict

Phase parameterized by Zernike polynomials.

deeplens.diffractive_surface.Zernike

Zernike(d, z_coeff=None, zernike_order=37, res=(2000, 2000), mat='fused_silica', fab_ps=0.001, fab_step=16, wvln0=0.55, device='cpu')

Bases: DiffractiveSurface

DOE parameterized by Zernike polynomials.

Initialize Zernike DOE.

Parameters:

Name	Type	Description	Default
d		DOE position	required
res		DOE resolution	(2000, 2000)
z_coeff		Zernike coefficients	None
zernike_order		Number of Zernike coefficients to use	37
fab_ps		Fabrication pixel size	0.001
fab_step		Fabrication step	16
device		Computation device	'cpu'

Source code in deeplens-src/deeplens/diffractive_surface/zernike.py

def __init__(
    self,
    d,
    z_coeff=None,
    zernike_order=37,
    res=(2000, 2000),
    mat="fused_silica",
    fab_ps=0.001,
    fab_step=16,
    wvln0=0.55,
    device="cpu",
):
    """Initialize Zernike DOE.

    Args:
        d: DOE position
        res: DOE resolution
        z_coeff: Zernike coefficients
        zernike_order: Number of Zernike coefficients to use
        fab_ps: Fabrication pixel size
        fab_step: Fabrication step
        device: Computation device
    """
    super().__init__(
        d=d, res=res, mat=mat, fab_ps=fab_ps, fab_step=fab_step, wvln0=wvln0, device=device
    )

    # Initialize Zernike coefficients with random values
    assert zernike_order == 37, "Currently, Zernike DOE only supports 37 orders"
    self.zernike_order = zernike_order
    if z_coeff is None:
        self.z_coeff = torch.randn(zernike_order, device=self.device) * 1e-3
    else:
        self.z_coeff = z_coeff

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize Zernike DOE from a dict.

Source code in deeplens-src/deeplens/diffractive_surface/zernike.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize Zernike DOE from a dict."""
    return cls(
        d=doe_dict["d"],
        res=doe_dict["res"],
        mat=doe_dict.get("mat", "fused_silica"),
        fab_ps=doe_dict.get("fab_ps", 0.001),
        fab_step=doe_dict.get("fab_step", 16),
        z_coeff=doe_dict.get("z_coeff", None),
        zernike_order=doe_dict.get("zernike_order", 37),
        wvln0=doe_dict.get("wvln0", 0.55),
    )

phase_func

phase_func()

Get the phase map at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/zernike.py

def phase_func(self):
    """Get the phase map at design wavelength."""
    return calculate_zernike_phase(self.z_coeff, grid=self.res[0])

get_optimizer_params

get_optimizer_params(lr=0.01)

Get parameters for optimization.

Source code in deeplens-src/deeplens/diffractive_surface/zernike.py

def get_optimizer_params(self, lr=0.01):
    """Get parameters for optimization."""
    self.z_coeff.requires_grad = True
    optimizer_params = [{"params": [self.z_coeff], "lr": lr}]
    return optimizer_params

surf_dict

surf_dict()

Return a dict of surface.

Source code in deeplens-src/deeplens/diffractive_surface/zernike.py

def surf_dict(self):
    """Return a dict of surface."""
    surf_dict = super().surf_dict()
    surf_dict["z_coeff"] = self.z_coeff.clone().detach().cpu()
    surf_dict["zernike_order"] = self.zernike_order
    return surf_dict

Linear / blazed grating phase.

deeplens.diffractive_surface.Grating

Grating(d, res=(2000, 2000), mat='fused_silica', wvln0=0.55, fab_ps=0.001, fab_step=16, theta=0.0, alpha=0.0, device='cpu')

Bases: DiffractiveSurface

Grating diffractive optical element.

A grating introduces a linear phase gradient defined by

phi(x, y) = alpha * (x * sin(theta) + y * cos(theta)) / norm_radii

where

• theta: angle from y-axis to grating vector
• alpha: slope of the grating (phase gradient strength)
• norm_radii: normalization radius

Initialize Grating DOE.

Parameters:

Name	Type	Description	Default
d	float	Distance of the DOE surface. [mm]	required
res	tuple or int	Resolution of the DOE, [w, h]. [pixel]	(2000, 2000)
mat	str	Material of the DOE.	'fused_silica'
wvln0	float	Design wavelength. [um]	0.55
fab_ps	float	Fabrication pixel size. [mm]	0.001
fab_step	int	Fabrication step.	16
theta	float	Angle from y-axis to grating vector. [rad]	0.0
alpha	float	Slope of the grating (phase gradient strength).	0.0
device	str	Device to run the DOE.	'cpu'

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def __init__(
    self,
    d,
    res=(2000, 2000),
    mat="fused_silica",
    wvln0=0.55,
    fab_ps=0.001,
    fab_step=16,
    theta=0.0,
    alpha=0.0,
    device="cpu",
):
    """Initialize Grating DOE.

    Args:
        d (float): Distance of the DOE surface. [mm]
        res (tuple or int): Resolution of the DOE, [w, h]. [pixel]
        mat (str): Material of the DOE.
        wvln0 (float): Design wavelength. [um]
        fab_ps (float): Fabrication pixel size. [mm]
        fab_step (int): Fabrication step.
        theta (float): Angle from y-axis to grating vector. [rad]
        alpha (float): Slope of the grating (phase gradient strength).
        device (str): Device to run the DOE.
    """
    super().__init__(
        d=d, res=res, mat=mat, wvln0=wvln0, fab_ps=fab_ps, fab_step=fab_step, device=device
    )

    # Grating parameters
    self.theta = torch.tensor(theta)  # angle from y-axis to grating vector
    self.alpha = torch.tensor(alpha)  # slope of the grating

    # Normalization radius (use half of the width)
    self.norm_radii = self.w / 2

    self.to(device)

init_from_dict classmethod

init_from_dict(doe_dict)

Initialize Grating DOE from a dict.

Parameters:

Name	Type	Description	Default
doe_dict	dict	Dictionary containing DOE parameters.	required

Returns:

Name	Type	Description
Grating		Initialized Grating DOE object.

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

@classmethod
def init_from_dict(cls, doe_dict):
    """Initialize Grating DOE from a dict.

    Args:
        doe_dict (dict): Dictionary containing DOE parameters.

    Returns:
        Grating: Initialized Grating DOE object.
    """
    return cls(
        d=doe_dict["d"],
        res=doe_dict["res"],
        mat=doe_dict.get("mat", "fused_silica"),
        wvln0=doe_dict.get("wvln0", 0.55),
        fab_ps=doe_dict.get("fab_ps", 0.001),
        fab_step=doe_dict.get("fab_step", 16),
        theta=doe_dict.get("theta", 0.0),
        alpha=doe_dict.get("alpha", 0.0),
    )

phase_func

phase_func()

Get the phase map at design wavelength.

The grating phase is a linear function of position

phi(x, y) = alpha * (x * sin(theta) + y * cos(theta)) / norm_radii

Returns:

Name	Type	Description
phase	tensor	Phase map at design wavelength.

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def phase_func(self):
    """Get the phase map at design wavelength.

    The grating phase is a linear function of position:
        phi(x, y) = alpha * (x * sin(theta) + y * cos(theta)) / norm_radii

    Returns:
        phase (tensor): Phase map at design wavelength.
    """
    # Normalize coordinates
    x_norm = self.x / self.norm_radii
    y_norm = self.y / self.norm_radii

    # Calculate linear phase gradient
    phase = self.alpha * (
        x_norm * torch.sin(self.theta) + y_norm * torch.cos(self.theta)
    )

    return phase

get_optimizer_params

get_optimizer_params(lr=0.001)

Get parameters for optimization.

Parameters:

Name	Type	Description	Default
lr	float	Learning rate for grating parameters.	0.001

Returns:

Name	Type	Description
list		List of parameter groups for optimizer.

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def get_optimizer_params(self, lr=0.001):
    """Get parameters for optimization.

    Args:
        lr (float): Learning rate for grating parameters.

    Returns:
        list: List of parameter groups for optimizer.
    """
    self.theta.requires_grad = True
    self.alpha.requires_grad = True

    optimizer_params = [
        {"params": [self.theta], "lr": lr},
        {"params": [self.alpha], "lr": lr * 10},
    ]

    return optimizer_params

surf_dict

surf_dict()

Return a dict of surface parameters.

Returns:

Name	Type	Description
dict		Dictionary containing surface parameters.

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def surf_dict(self):
    """Return a dict of surface parameters.

    Returns:
        dict: Dictionary containing surface parameters.
    """
    surf_dict = super().surf_dict()
    surf_dict["theta"] = round(self.theta.item(), 6)
    surf_dict["alpha"] = round(self.alpha.item(), 6)
    surf_dict["norm_radii"] = round(self.norm_radii, 6)
    return surf_dict

save_ckpt

save_ckpt(save_path='./grating_doe.pth')

Save grating DOE parameters.

Parameters:

Name	Type	Description	Default
save_path	str	Path to save the checkpoint.	'./grating_doe.pth'

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def save_ckpt(self, save_path="./grating_doe.pth"):
    """Save grating DOE parameters.

    Args:
        save_path (str): Path to save the checkpoint.
    """
    torch.save(
        {
            "param_model": "grating",
            "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.

Parameters:

Name	Type	Description	Default
load_path	str	Path to load the checkpoint from.	'./grating_doe.pth'

Source code in deeplens-src/deeplens/diffractive_surface/grating.py

def load_ckpt(self, load_path="./grating_doe.pth"):
    """Load grating DOE parameters.

    Args:
        load_path (str): Path to load the checkpoint from.
    """
    ckpt = torch.load(load_path)
    self.theta = ckpt["theta"].to(self.device)
    self.alpha = ckpt["alpha"].to(self.device)