Lens

Abstract base class for every lens model in DeepLens. Lens defines the shared interface — psf(), psf_rgb(), render(), sensor configuration, and file I/O — that GeoLens, HybridLens, DiffractiveLens, PSFNetLens, and DefocusLens all inherit.

deeplens.Lens

Lens(
    dtype=torch.float32,
    device=None,
    primary_wvln=DEFAULT_WAVE,
    wvln_rgb=WAVE_RGB,
    obj_depth=DEPTH,
)

Bases: DeepObj

Abstract base class for all lens models in DeepLens.

Lens defines the shared interface — PSF computation (psf, psf_rgb), image rendering (render), sensor configuration, and JSON file I/O — that GeoLens, HybridLens, DiffractiveLens, PSFNetLens, and DefocusLens all inherit. Subclasses override the core optical methods (e.g. psf) with their own differentiable implementations.

Initialize a lens class.

Parameters:

Name	Type	Description	Default
dtype	dtype	Data type. Defaults to torch.float32.	float32
device	str	Device to run the lens. Defaults to None.	None
primary_wvln	float	Primary design wavelength [µm]. Used as fallback when a method is called without an explicit wvln. Defaults to DEFAULT_WAVE (0.587, d-line).	DEFAULT_WAVE
wvln_rgb	sequence of float	Three wavelengths used for RGB (polychromatic) computations, ordered [R, G, B] in µm. Defaults to WAVE_RGB.	WAVE_RGB
obj_depth	float	Default object depth [mm] used as fallback when a method is called without an explicit depth. Should be negative (object in front of the lens). Defaults to DEPTH (−20 000 mm, practical infinity).	DEPTH

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

def __init__(
    self,
    dtype=torch.float32,
    device=None,
    primary_wvln=DEFAULT_WAVE,
    wvln_rgb=WAVE_RGB,
    obj_depth=DEPTH,
):
    """Initialize a lens class.

    Args:
        dtype (torch.dtype, optional): Data type. Defaults to torch.float32.
        device (str, optional): Device to run the lens. Defaults to None.
        primary_wvln (float, optional): Primary design wavelength [µm].
            Used as fallback when a method is called without an explicit
            ``wvln``.  Defaults to ``DEFAULT_WAVE`` (0.587, d-line).
        wvln_rgb (sequence of float, optional): Three wavelengths used for
            RGB (polychromatic) computations, ordered ``[R, G, B]`` in
            µm.  Defaults to ``WAVE_RGB``.
        obj_depth (float, optional): Default object depth [mm] used as
            fallback when a method is called without an explicit
            ``depth``.  Should be negative (object in front of the lens).
            Defaults to ``DEPTH`` (−20 000 mm, practical infinity).
    """
    # Lens device
    if device is None:
        self.device = init_device()
    else:
        self.device = torch.device(device)

    # Lens default dtype
    self.dtype = dtype

    primary_wvln = torch.as_tensor(primary_wvln, dtype=torch.float64)
    wvln_rgb = torch.as_tensor(wvln_rgb, dtype=torch.float64)
    obj_depth = torch.as_tensor(obj_depth, dtype=torch.float64)

    if primary_wvln.numel() != 1:
        raise ValueError("primary_wvln must be a scalar wavelength in [µm].")
    if wvln_rgb.numel() != 3:
        raise ValueError("wvln_rgb must contain exactly three wavelengths in [µm].")
    if obj_depth.numel() != 1:
        raise ValueError("obj_depth must be a scalar depth in [mm].")

    if not (primary_wvln.item() > 0.1 and primary_wvln.item() < 10.0):
        raise ValueError(
            "primary_wvln must be in [µm] and satisfy 0.1 < primary_wvln < 10."
        )
    if not torch.all((wvln_rgb > 0.1) & (wvln_rgb < 10.0)):
        raise ValueError(
            "wvln_rgb must be in [µm] and every value must satisfy 0.1 < wvln < 10."
        )
    if not obj_depth.item() < 0.0:
        raise ValueError(
            "obj_depth must be negative [mm], with the object in front of the lens."
        )

    # Design wavelengths [µm].  IO may override.
    self.primary_wvln = float(primary_wvln.item())
    self.wvln_rgb = [float(w) for w in wvln_rgb.tolist()]

    # Default object depth [mm].
    self.obj_depth = float(obj_depth.item())

read_lens_json

read_lens_json(filename)

Read the lens from a JSON file. Must be overridden by subclasses.

Parameters:

Name	Type	Description	Default
filename	str	Path to the JSON lens file.	required

Raises:

Type	Description
NotImplementedError	This base implementation must be overridden.

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

def read_lens_json(self, filename):
    """Read the lens from a JSON file. Must be overridden by subclasses.

    Args:
        filename (str): Path to the JSON lens file.

    Raises:
        NotImplementedError: This base implementation must be overridden.
    """
    raise NotImplementedError

write_lens_json

write_lens_json(filename)

Write the lens to a JSON file. Must be overridden by subclasses.

Parameters:

Name	Type	Description	Default
filename	str	Destination path for the JSON lens file.	required

Raises:

Type	Description
NotImplementedError	This base implementation must be overridden.

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

def write_lens_json(self, filename):
    """Write the lens to a JSON file. Must be overridden by subclasses.

    Args:
        filename (str): Destination path for the JSON lens file.

    Raises:
        NotImplementedError: This base implementation must be overridden.
    """
    raise NotImplementedError

set_sensor

set_sensor(sensor_size, sensor_res)

Set sensor size and resolution.

Parameters:

Name	Type	Description	Default
sensor_size	tuple	Sensor size (w, h) in [mm].	required
sensor_res	tuple	Sensor resolution (W, H) in [pixels].	required

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

def set_sensor(self, sensor_size, sensor_res):
    """Set sensor size and resolution.

    Args:
        sensor_size (tuple): Sensor size (w, h) in [mm].
        sensor_res (tuple): Sensor resolution (W, H) in [pixels].
    """
    assert sensor_size[0] * sensor_res[1] == sensor_size[1] * sensor_res[0], (
        "Sensor resolution aspect ratio does not match sensor size aspect ratio."
    )
    self.sensor_size = sensor_size
    self.sensor_res = sensor_res
    self.pixel_size = self.sensor_size[0] / self.sensor_res[0]
    self.r_sensor = float(np.sqrt(sensor_size[0] ** 2 + sensor_size[1] ** 2)) / 2
    self.calc_fov()

set_sensor_res

set_sensor_res(sensor_res)

Set sensor resolution (and aspect ratio) while keeping sensor radius unchanged.

Parameters:

Name	Type	Description	Default
sensor_res	tuple	Sensor resolution (W, H) in [pixels].	required

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

def set_sensor_res(self, sensor_res):
    """Set sensor resolution (and aspect ratio) while keeping sensor radius unchanged.

    Args:
        sensor_res (tuple): Sensor resolution (W, H) in [pixels].
    """
    # Change sensor resolution
    self.sensor_res = sensor_res

    # Change sensor size (r_sensor is fixed)
    diam_res = float(np.sqrt(self.sensor_res[0] ** 2 + self.sensor_res[1] ** 2))
    self.sensor_size = (
        2 * self.r_sensor * self.sensor_res[0] / diam_res,
        2 * self.r_sensor * self.sensor_res[1] / diam_res,
    )
    self.pixel_size = self.sensor_size[0] / self.sensor_res[0]
    self.calc_fov()

create_dummy_sensor

create_dummy_sensor()

Fill in any sensor geometry the lens file did not provide.

A lens file may omit part of the sensor description. The sensor radius is derived from sensor_size when absent, and defaults are supplied for sensor_size and sensor_res so the lens is usable without an explicit set_sensor() call.

Raises:

Type	Description
ValueError	If neither r_sensor nor a two-element sensor_size is available, or if r_sensor is not finite and positive.

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

def create_dummy_sensor(self):
    """Fill in any sensor geometry the lens file did not provide.

    A lens file may omit part of the sensor description. The sensor radius
    is derived from `sensor_size` when absent, and defaults are supplied for
    `sensor_size` and `sensor_res` so the lens is usable without an explicit
    `set_sensor()` call.

    Raises:
        ValueError: If neither `r_sensor` nor a two-element `sensor_size`
            is available, or if `r_sensor` is not finite and positive.
    """
    if not hasattr(self, "r_sensor"):
        if hasattr(self, "sensor_size") and len(self.sensor_size) == 2:
            self.r_sensor = math.hypot(*self.sensor_size) / 2.0
        else:
            raise ValueError(
                "Lens file is missing sensor geometry; provide an image surface "
                "or r_sensor/sensor_size."
            )
    if not math.isfinite(float(self.r_sensor)) or float(self.r_sensor) <= 0:
        raise ValueError("Lens sensor radius must be a finite positive number.")

    # Complete sensor size and resolution if not set from lens file
    if not hasattr(self, "sensor_size"):
        self.sensor_size = (8.0, 8.0)
        print(
            f"Sensor_size not found in lens file. Using default: {self.sensor_size} mm. "
            "Consider specifying sensor_size in the lens file or using set_sensor()."
        )

    if not hasattr(self, "sensor_res"):
        self.sensor_res = (2000, 2000)
        print(
            f"Sensor_res not found in lens file. Using default: {self.sensor_res} pixels. "
            "Consider specifying sensor_res in the lens file or using set_sensor()."
        )
        self.set_sensor_res(self.sensor_res)

calc_fov

calc_fov()

Compute FoV (radian) of the lens.

Reference

[1] https://en.wikipedia.org/wiki/Angle_of_view_(photography)

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

@torch.no_grad()
def calc_fov(self):
    """Compute FoV (radian) of the lens.

    Reference:
        [1] https://en.wikipedia.org/wiki/Angle_of_view_(photography)
    """
    if not hasattr(self, "foclen"):
        return

    self.hfov = 2 * float(np.arctan(self.sensor_size[0] / 2 / self.foclen))
    self.vfov = 2 * float(np.arctan(self.sensor_size[1] / 2 / self.foclen))
    self.dfov = 2 * float(np.arctan(self.r_sensor / self.foclen))
    self.rfov_eff = self.dfov / 2  # effective (paraxial) half-diagonal FoV
    self.rfov = (
        self.rfov_eff
    )  # default to effective; GeoLens overrides with ray-traced value

psf

psf(points, wvln=None, ks=PSF_KS, **kwargs)

Compute the monochromatic PSF for one or more point sources.

Subclasses must override this method with a differentiable implementation. Three computation models are common in practice: geometric ray binning, coherent ray-wave, and Huygens spherical-wave integration.

Parameters:

Name	Type	Description	Default
points	Tensor	Point source coordinates, shape [N, 3] or [3]. x, y are normalised to [-1, 1] (relative to the sensor half-diagonal); z is depth in mm (must be negative, i.e. in front of the lens).	required
wvln	float	Wavelength in micrometers. When None (default), falls back to self.primary_wvln.	None
ks	int	Output PSF kernel size in pixels. Defaults to PSF_KS (64).	PSF_KS
**kwargs		Additional keyword arguments forwarded to the underlying PSF computation (e.g. spp, model, recenter).	{}

Returns:

Name	Type	Description
psf	Tensor	PSF intensity map, shape [ks, ks] for a single point or [N, ks, ks] for a batch.

Raises:

Type	Description
NotImplementedError	This base implementation must be overridden.

Note

The method is differentiable with respect to all optimisable lens parameters so it can be used directly inside a training loop.

Examples:

point = torch.tensor([0.0, 0.0, -10000.0])
psf = lens.psf(points=point, ks=64, model="geometric")
print(psf.shape)  # torch.Size([64, 64])

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

def psf(self, points, wvln=None, ks=PSF_KS, **kwargs):
    """Compute the monochromatic PSF for one or more point sources.

    Subclasses must override this method with a differentiable
    implementation.  Three computation models are common in practice:
    geometric ray binning, coherent ray-wave, and Huygens spherical-wave
    integration.

    Args:
        points (torch.Tensor): Point source coordinates, shape ``[N, 3]``
            or ``[3]``.  ``x, y`` are normalised to ``[-1, 1]``
            (relative to the sensor half-diagonal); ``z`` is depth in mm
            (must be negative, i.e. in front of the lens).
        wvln (float, optional): Wavelength in micrometers.  When ``None``
            (default), falls back to ``self.primary_wvln``.
        ks (int, optional): Output PSF kernel size in pixels.  Defaults
            to ``PSF_KS`` (64).
        **kwargs: Additional keyword arguments forwarded to the underlying
            PSF computation (e.g. ``spp``, ``model``, ``recenter``).

    Returns:
        psf (torch.Tensor): PSF intensity map, shape ``[ks, ks]`` for a single
            point or ``[N, ks, ks]`` for a batch.

    Raises:
        NotImplementedError: This base implementation must be overridden.

    Note:
        The method is differentiable with respect to all optimisable lens
        parameters so it can be used directly inside a training loop.

    Example:
        ```python
        point = torch.tensor([0.0, 0.0, -10000.0])
        psf = lens.psf(points=point, ks=64, model="geometric")
        print(psf.shape)  # torch.Size([64, 64])
        ```
    """
    raise NotImplementedError

psf_rgb

psf_rgb(points, ks=PSF_KS, **kwargs)

Compute the RGB (tri-chromatic) PSF by stacking three wavelength calls.

Calls psf three times for the RGB primary wavelengths stored in self.wvln_rgb and stacks the results along the channel axis.

Parameters:

Name	Type	Description	Default
points	Tensor	Point source coordinates, shape [N, 3] or [3]. Same convention as psf.	required
ks	int	PSF kernel size. Defaults to PSF_KS.	PSF_KS
**kwargs		Forwarded to psf (e.g. spp, model).	{}

Returns:

Name	Type	Description
psf_rgb	Tensor	RGB PSF, shape [3, ks, ks] for a single point or [N, 3, ks, ks] for a batch.

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

def psf_rgb(self, points, ks=PSF_KS, **kwargs):
    """Compute the RGB (tri-chromatic) PSF by stacking three wavelength calls.

    Calls `psf` three times for the RGB primary wavelengths stored
    in ``self.wvln_rgb`` and stacks the results along the channel axis.

    Args:
        points (torch.Tensor): Point source coordinates, shape ``[N, 3]``
            or ``[3]``.  Same convention as `psf`.
        ks (int, optional): PSF kernel size. Defaults to ``PSF_KS``.
        **kwargs: Forwarded to `psf` (e.g. ``spp``, ``model``).

    Returns:
        psf_rgb (torch.Tensor): RGB PSF, shape ``[3, ks, ks]`` for a single point
            or ``[N, 3, ks, ks]`` for a batch.
    """
    psfs = []
    for wvln in self.wvln_rgb:
        psfs.append(self.psf(points=points, ks=ks, wvln=wvln, **kwargs))
    psf_rgb = torch.stack(psfs, dim=-3)  # shape [3, ks, ks] or [N, 3, ks, ks]
    return psf_rgb

point_source_grid

point_source_grid(depth, grid=(9, 9), normalized=True, quater=False, center=True)

Generate a grid of point sources for PSF calculation.

Parameters:

Name	Type	Description	Default
depth	float	Depth (z-coordinate) of the point sources [mm] (negative, in front of the lens).	required
grid	tuple	Grid size (grid_w, grid_h). Defaults to (9, 9), meaning a 9x9 grid.	(9, 9)
normalized	bool	If True, return normalized object-space xy coordinates in [-1, 1]; if False, scale to physical positions [mm]. Defaults to True.	True
quater	bool	If True, return only one quarter of the grid to save memory. Defaults to False.	False
center	bool	If True, place points at the center of each patch; otherwise sample to the field corners. Defaults to True.	True

Returns:

Name	Type	Description
point_source	Tensor	Object source coordinates, shape [grid_h, grid_w, 3] with the last dim ordered (x, y, z). When quater is True, the first two dimensions are reduced to the returned quarter.

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

def point_source_grid(
    self, depth, grid=(9, 9), normalized=True, quater=False, center=True
):
    """Generate a grid of point sources for PSF calculation.

    Args:
        depth (float): Depth (z-coordinate) of the point sources [mm]
            (negative, in front of the lens).
        grid (tuple, optional): Grid size (grid_w, grid_h). Defaults to
            (9, 9), meaning a 9x9 grid.
        normalized (bool, optional): If True, return normalized object-space
            xy coordinates in [-1, 1]; if False, scale to physical positions
            [mm]. Defaults to True.
        quater (bool, optional): If True, return only one quarter of the grid
            to save memory. Defaults to False.
        center (bool, optional): If True, place points at the center of each
            patch; otherwise sample to the field corners. Defaults to True.

    Returns:
        point_source (torch.Tensor): Object source coordinates, shape
            [grid_h, grid_w, 3] with the last dim ordered (x, y, z). When
            `quater` is True, the first two dimensions are reduced to the
            returned quarter.
    """
    # Compute point source grid
    if grid[0] == 1:
        x = torch.tensor([[0.0]], device=self.device, dtype=self.dtype)
        y = torch.tensor([[0.0]], device=self.device, dtype=self.dtype)
        assert not quater, "Quater should be False when grid is 1."
    else:
        if center:
            # Use center of each patch
            half_bin_size = 1 / 2 / (grid[0] - 1)
            x, y = torch.meshgrid(
                torch.linspace(
                    -1 + half_bin_size,
                    1 - half_bin_size,
                    grid[0],
                    device=self.device,
                    dtype=self.dtype,
                ),
                torch.linspace(
                    1 - half_bin_size,
                    -1 + half_bin_size,
                    grid[1],
                    device=self.device,
                    dtype=self.dtype,
                ),
                indexing="xy",
            )
        else:
            # Use corner of image sensor
            x, y = torch.meshgrid(
                torch.linspace(
                    -0.98, 0.98, grid[0], device=self.device, dtype=self.dtype
                ),
                torch.linspace(
                    0.98, -0.98, grid[1], device=self.device, dtype=self.dtype
                ),
                indexing="xy",
            )

    z = torch.full_like(x, depth)
    point_source = torch.stack([x, y, z], dim=-1)

    # Use quater of the sensor plane to save memory
    if quater:
        bound_i = grid[0] // 2 if grid[0] % 2 == 0 else grid[0] // 2 + 1
        bound_j = grid[1] // 2
        point_source = point_source[0:bound_i, bound_j:, :]

    # De-normalize object source coordinates to physical coordinates
    if not normalized:
        scale = self.calc_scale(depth)
        point_source[..., 0] *= scale * self.sensor_size[0] / 2
        point_source[..., 1] *= scale * self.sensor_size[1] / 2

    return point_source

psf_map

psf_map(grid=(5, 5), wvln=None, depth=None, ks=PSF_KS, **kwargs)

Compute monochrome PSF map.

Parameters:

Name	Type	Description	Default
grid	tuple	Grid size (grid_w, grid_h). Defaults to (5, 5), meaning 5x5 grid.	(5, 5)
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
depth	float	Depth of the object. When None (default), falls back to self.obj_depth.	None
ks	int	Kernel size. Defaults to PSF_KS.	PSF_KS

Returns:

Name	Type	Description
psf_map	Tensor	Monochrome PSF map, shape [grid_h, grid_w, 1, ks, ks].

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

def psf_map(self, grid=(5, 5), wvln=None, depth=None, ks=PSF_KS, **kwargs):
    """Compute monochrome PSF map.

    Args:
        grid (tuple): Grid size (grid_w, grid_h). Defaults to (5, 5), meaning 5x5 grid.
        wvln (float): Wavelength in µm. When ``None`` (default), falls back
            to ``self.primary_wvln``.
        depth (float): Depth of the object. When ``None`` (default), falls
            back to ``self.obj_depth``.
        ks (int): Kernel size. Defaults to PSF_KS.

    Returns:
        psf_map (torch.Tensor): Monochrome PSF map, shape
            [grid_h, grid_w, 1, ks, ks].
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth

    # PSF map grid
    points = self.point_source_grid(depth=depth, grid=grid, center=True)
    points = points.reshape(-1, 3)

    # Compute PSF map
    psfs = []
    for i in range(points.shape[0]):
        point = points[i, ...]
        psf = self.psf(points=point, wvln=wvln, ks=ks)
        psfs.append(psf)
    psf_map = torch.stack(psfs).unsqueeze(1)  # shape [grid_h * grid_w, 1, ks, ks]

    # Reshape PSF map from [grid_h * grid_w, 1, ks, ks] -> [grid_h, grid_w, 1, ks, ks]
    psf_map = psf_map.reshape(grid[1], grid[0], 1, ks, ks)
    return psf_map

psf_map_rgb

psf_map_rgb(grid=(5, 5), ks=PSF_KS, depth=None, **kwargs)

Compute RGB PSF map.

Parameters:

Name	Type	Description	Default
grid	tuple	Grid size (grid_w, grid_h). Defaults to (5, 5), meaning 5x5 grid.	(5, 5)
ks	int	Kernel size. Defaults to PSF_KS, meaning PSF_KS x PSF_KS kernel size.	PSF_KS
depth	float	Depth of the object. When None (default), falls back to self.obj_depth.	None
**kwargs		Additional arguments for psf_map().	{}

Returns:

Name	Type	Description
psf_map	Tensor	Shape of [grid_h, grid_w, 3, ks, ks].

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

def psf_map_rgb(self, grid=(5, 5), ks=PSF_KS, depth=None, **kwargs):
    """Compute RGB PSF map.

    Args:
        grid (tuple): Grid size (grid_w, grid_h). Defaults to (5, 5), meaning 5x5 grid.
        ks (int): Kernel size. Defaults to PSF_KS, meaning PSF_KS x PSF_KS kernel size.
        depth (float): Depth of the object. When ``None`` (default), falls
            back to ``self.obj_depth``.
        **kwargs: Additional arguments for psf_map().

    Returns:
        psf_map (torch.Tensor): Shape of [grid_h, grid_w, 3, ks, ks].
    """
    depth = self.obj_depth if depth is None else depth
    psfs = []
    for wvln in self.wvln_rgb:
        psf_map = self.psf_map(grid=grid, ks=ks, depth=depth, wvln=wvln, **kwargs)
        psfs.append(psf_map)
    psf_map = torch.cat(psfs, dim=2)  # shape [grid_h, grid_w, 3, ks, ks]
    return psf_map

draw_psf_map

draw_psf_map(
    grid=(7, 7),
    ks=PSF_KS,
    depth=None,
    log_scale=False,
    save_name="./psf_map.png",
    show=False,
)

Draw the RGB PSF map of the lens and save it (or return the figure).

Parameters:

Name	Type	Description	Default
grid	tuple	Grid size (grid_w, grid_h). Defaults to (7, 7).	(7, 7)
ks	int	PSF kernel size in pixels. Defaults to PSF_KS.	PSF_KS
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
log_scale	bool	If True, normalize each PSF in log scale for better visualization. Defaults to False.	False
save_name	str	Output image path. Defaults to "./psf_map.png".	'./psf_map.png'
show	bool	If True, return (fig, ax) instead of saving. Defaults to False.	False

Returns:

Name	Type	Description
result	tuple or None	(fig, ax) if show is True, otherwise None (the figure is saved to save_name).

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

@torch.no_grad()
def draw_psf_map(
    self,
    grid=(7, 7),
    ks=PSF_KS,
    depth=None,
    log_scale=False,
    save_name="./psf_map.png",
    show=False,
):
    """Draw the RGB PSF map of the lens and save it (or return the figure).

    Args:
        grid (tuple, optional): Grid size (grid_w, grid_h). Defaults to (7, 7).
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        log_scale (bool, optional): If True, normalize each PSF in log scale
            for better visualization. Defaults to False.
        save_name (str, optional): Output image path. Defaults to
            "./psf_map.png".
        show (bool, optional): If True, return (fig, ax) instead of saving.
            Defaults to False.

    Returns:
        result (tuple or None): (fig, ax) if `show` is True, otherwise None
            (the figure is saved to `save_name`).
    """
    depth = self.obj_depth if depth is None else depth
    # Calculate RGB PSF map, shape [grid_h, grid_w, 3, ks, ks]
    psf_map = self.psf_map_rgb(depth=depth, grid=grid, ks=ks)

    # Create a grid visualization (vis_map: shape [3, grid_h * ks, grid_w * ks])
    grid_w, grid_h = grid if isinstance(grid, tuple) else (grid, grid)
    h, w = grid_h * ks, grid_w * ks
    vis_map = torch.zeros((3, h, w), device=psf_map.device, dtype=psf_map.dtype)

    # Put each PSF into the vis_map
    for i in range(grid_h):
        for j in range(grid_w):
            # Extract the PSF at this grid position
            psf = psf_map[i, j]  # shape [3, ks, ks]

            # Normalize the PSF
            if log_scale:
                # Log scale normalization for better visualization
                psf = torch.log(psf + 1e-4)  # 1e-4 is an empirical value
                psf = (psf - psf.min()) / (psf.max() - psf.min() + 1e-8)
            else:
                # Linear normalization
                local_max = psf.max()
                if local_max > 0:
                    psf = psf / local_max

            # Place the normalized PSF in the visualization map
            y_start, y_end = i * ks, (i + 1) * ks
            x_start, x_end = j * ks, (j + 1) * ks
            vis_map[:, y_start:y_end, x_start:x_end] = psf

    # Create the figure and display
    fig, ax = plt.subplots(figsize=(10, 10))

    # Convert to numpy for plotting
    vis_map = vis_map.permute(1, 2, 0).cpu().numpy()
    ax.imshow(vis_map)

    # Add scale bar near bottom-left
    H, W, _ = vis_map.shape
    scale_bar_length = 100
    arrow_length = scale_bar_length / (self.pixel_size * 1e3)
    y_position = H - 20  # a little above the lower edge
    x_start = 20
    x_end = x_start + arrow_length

    ax.annotate(
        "",
        xy=(x_start, y_position),
        xytext=(x_end, y_position),
        arrowprops=dict(arrowstyle="-", color="white"),
        annotation_clip=False,
    )
    ax.text(
        x_end + 5,
        y_position,
        f"{scale_bar_length} μm",
        color="white",
        fontsize=12,
        ha="left",
        va="center",
        clip_on=False,
    )

    # Clean up axes and save
    ax.axis("off")
    plt.tight_layout(pad=0)

    if show:
        return fig, ax
    else:
        plt.savefig(save_name, dpi=300, bbox_inches="tight", pad_inches=0)
        plt.close(fig)

point_source_radial

point_source_radial(depth, grid=9, center=False, direction='diagonal', normalized=True)

Generate radial point sources from center to edge of the field.

Produces grid evenly-spaced points along a chosen radial direction (diagonal, meridional, or sagittal) in normalized or physical object-space coordinates.

Parameters:

Name	Type	Description	Default
depth	float	Object depth (z-coordinate) in mm.	required
grid	int	Number of sample points. Defaults to 9.	9
center	bool	If True, offset positions to bin centers. Defaults to False.	False
direction	str	Sampling direction — "diagonal" (x = y, 45°, default), "y" (meridional, x = 0), "x" (sagittal, y = 0).	'diagonal'
normalized	bool	If True, return coordinates in [0, 1]. If False, scale to physical object-space positions (mm). Defaults to True.	True

Returns:

Name	Type	Description
point_source	Tensor	Point source positions, shape [grid, 3].

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

def point_source_radial(
    self, depth, grid=9, center=False, direction="diagonal", normalized=True
):
    """Generate radial point sources from center to edge of the field.

    Produces ``grid`` evenly-spaced points along a chosen radial direction
    (diagonal, meridional, or sagittal) in normalized or physical object-space
    coordinates.

    Args:
        depth (float): Object depth (z-coordinate) in mm.
        grid (int): Number of sample points. Defaults to 9.
        center (bool): If ``True``, offset positions to bin centers.
            Defaults to ``False``.
        direction (str): Sampling direction —
            ``"diagonal"`` (x = y, 45°, default),
            ``"y"`` (meridional, x = 0),
            ``"x"`` (sagittal, y = 0).
        normalized (bool): If ``True``, return coordinates in [0, 1].
            If ``False``, scale to physical object-space positions (mm).
            Defaults to ``True``.

    Returns:
        point_source (torch.Tensor): Point source positions, shape ``[grid, 3]``.
    """
    if grid == 1:
        r = torch.tensor([0.0], device=self.device, dtype=self.dtype)
    else:
        # Select center of bin to calculate PSF
        if center:
            half_bin_size = 1 / 2 / (grid - 1)
            r = torch.linspace(
                0, 1 - half_bin_size, grid, device=self.device, dtype=self.dtype
            )
        else:
            r = torch.linspace(0, 0.98, grid, device=self.device, dtype=self.dtype)

    # Map radial coordinate to (x, y) based on direction
    if direction == "diagonal":
        px, py = r, r
    elif direction == "y":
        px, py = torch.zeros_like(r), r
    elif direction == "x":
        px, py = r, torch.zeros_like(r)
    else:
        raise ValueError(
            f"Invalid direction: {direction!r}. Use 'diagonal', 'x', or 'y'."
        )

    z = torch.full_like(px, depth)
    point_source = torch.stack([px, py, z], dim=-1)

    if not normalized:
        scale = self.calc_scale(depth)
        point_source[..., 0] = (
            point_source[..., 0] * scale * self.sensor_size[0] / 2
        )
        point_source[..., 1] = (
            point_source[..., 1] * scale * self.sensor_size[1] / 2
        )

    return point_source

draw_psf_radial

draw_psf_radial(
    M=3, depth=None, ks=PSF_KS, log_scale=False, save_name="./psf_radial.png"
)

Draw a radial (45 deg, diagonal) sequence of RGB PSFs and save it.

Draws M PSFs evenly spaced from the field center to the corner, each of size ks x ks, arranged in a single row.

Parameters:

Name	Type	Description	Default
M	int	Number of PSFs to draw. Defaults to 3.	3
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
ks	int	PSF kernel size in pixels. Defaults to PSF_KS.	PSF_KS
log_scale	bool	If True, normalize each PSF in log scale for better visualization. Defaults to False.	False
save_name	str	Output image path. Defaults to "./psf_radial.png".	'./psf_radial.png'

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

@torch.no_grad()
def draw_psf_radial(
    self, M=3, depth=None, ks=PSF_KS, log_scale=False, save_name="./psf_radial.png"
):
    """Draw a radial (45 deg, diagonal) sequence of RGB PSFs and save it.

    Draws M PSFs evenly spaced from the field center to the corner, each of
    size ks x ks, arranged in a single row.

    Args:
        M (int, optional): Number of PSFs to draw. Defaults to 3.
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        log_scale (bool, optional): If True, normalize each PSF in log scale
            for better visualization. Defaults to False.
        save_name (str, optional): Output image path. Defaults to
            "./psf_radial.png".
    """
    from torchvision.utils import make_grid, save_image

    depth = self.obj_depth if depth is None else depth
    x = torch.linspace(0, 1, M, device=self.device, dtype=self.dtype)
    y = torch.linspace(0, 1, M, device=self.device, dtype=self.dtype)
    z = torch.full_like(x, depth)
    points = torch.stack((x, y, z), dim=-1)

    psfs = []
    for i in range(M):
        # Scale PSF for a better visualization
        psf = self.psf_rgb(points=points[i], ks=ks, recenter=True, spp=SPP_PSF)
        psf /= psf.max()

        if log_scale:
            psf = torch.log(psf + EPSILON)
            psf = (psf - psf.min()) / (psf.max() - psf.min())

        psfs.append(psf)

    psf_grid = make_grid(psfs, nrow=M, padding=1, pad_value=0.0)
    save_image(psf_grid, save_name, normalize=True)

render

render(img_obj, depth=None, method=None, **kwargs)

Differentiable image simulation for a 2D (flat) scene.

Performs only the optical component of image simulation and is fully differentiable.

For incoherent imaging the intensity PSF is convolved with the object-space image. For coherent imaging the complex PSF is convolved with the complex object image before squaring for intensity.

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Input image in linear (raw) space, shape [B, C, H, W].	required
depth	float	Object depth in mm (negative value). When None (default), falls back to self.obj_depth.	None
method	str	Rendering method. When None (default), falls back to self._default_render_method ("psf_patch" for the base Lens). One of: "psf_patch" – convolve a single PSF evaluated at patch_center. "psf_map" – spatially-varying PSF block convolution.	None
**kwargs		Method-specific keyword arguments: For "psf_map": psf_grid (tuple, default (10, 10)), psf_ks (int, default PSF_KS), psf_spp (int, default SPP_PSF). For "psf_patch": patch_center (tuple or Tensor, default (0.0, 0.0)), psf_ks (int).	{}

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Raises:

Type	Description
AssertionError	If method is "psf_map" and the image resolution does not match the sensor resolution.
Exception	If method is not recognised.

Reference

[1] "Optical Aberration Correction in Postprocessing using Imaging Simulation", TOG 2021. [2] "Efficient depth- and spatially-varying image simulation for defocus deblur", ICCVW 2025.

Examples:

img_rendered = lens.render(
    img, depth=-10000.0, method="psf_patch",
    patch_center=(0.3, 0.0), psf_ks=64,
)

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

def render(self, img_obj, depth=None, method=None, **kwargs):
    """Differentiable image simulation for a 2D (flat) scene.

    Performs only the optical component of image simulation and is fully
    differentiable.

    For incoherent imaging the intensity PSF is convolved with the
    object-space image.  For coherent imaging the complex PSF is convolved
    with the complex object image before squaring for intensity.

    Args:
        img_obj (torch.Tensor): Input image in linear (raw) space,
            shape ``[B, C, H, W]``.
        depth (float, optional): Object depth in mm (negative value).
            When ``None`` (default), falls back to ``self.obj_depth``.
        method (str, optional): Rendering method.  When ``None`` (default),
            falls back to ``self._default_render_method``
            (``"psf_patch"`` for the base `Lens`).  One of:

            * ``"psf_patch"`` – convolve a single PSF evaluated at
              *patch_center*.
            * ``"psf_map"`` – spatially-varying PSF block convolution.

        **kwargs: Method-specific keyword arguments:

            * For ``"psf_map"``: ``psf_grid`` (tuple, default ``(10, 10)``),
              ``psf_ks`` (int, default ``PSF_KS``), ``psf_spp`` (int,
              default ``SPP_PSF``).
            * For ``"psf_patch"``: ``patch_center`` (tuple or Tensor,
              default ``(0.0, 0.0)``), ``psf_ks`` (int).

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

    Raises:
        AssertionError: If *method* is ``"psf_map"`` and the image
            resolution does not match the sensor resolution.
        Exception: If *method* is not recognised.

    Reference:
        [1] "Optical Aberration Correction in Postprocessing using Imaging Simulation", TOG 2021.
        [2] "Efficient depth- and spatially-varying image simulation for defocus deblur", ICCVW 2025.

    Example:
        ```python
        img_rendered = lens.render(
            img, depth=-10000.0, method="psf_patch",
            patch_center=(0.3, 0.0), psf_ks=64,
        )
        ```
    """
    method = self._default_render_method if method is None else method
    depth = self.obj_depth if depth is None else depth
    # Check sensor resolution
    B, C, Himg, Wimg = img_obj.shape
    Wsensor, Hsensor = self.sensor_res

    # Image simulation (in RAW space)
    if method == "psf_map":
        # Render full resolution image with PSF map convolution
        assert Wimg == Wsensor and Himg == Hsensor, (
            f"Sensor resolution {Wsensor}x{Hsensor} must match input image {Wimg}x{Himg}."
        )
        psf_grid = kwargs.get("psf_grid", (10, 10))
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        psf_spp = kwargs.get("psf_spp", SPP_PSF)
        img_render = self.render_psf_map(
            img_obj,
            depth=depth,
            psf_grid=psf_grid,
            psf_ks=psf_ks,
            psf_spp=psf_spp,
        )

    elif method == "psf_patch":
        # Render an image patch with its corresponding PSF
        patch_center = kwargs.get("patch_center", (0.0, 0.0))
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        img_render = self.render_psf_patch(
            img_obj,
            depth=depth,
            patch_center=patch_center,
            psf_ks=psf_ks,
            method=kwargs.get("conv_method", "conv"),
        )

    elif method == "psf_pixel":
        raise NotImplementedError(
            "Per-pixel PSF convolution has not been implemented."
        )

    else:
        raise Exception(f"Image simulation method {method} is not supported.")

    return img_render

render_psf

render_psf(img_obj, depth=None, patch_center=(0, 0), psf_ks=PSF_KS, method='conv')

Render an image patch using PSF convolution (deprecated alias).

Thin wrapper around render_psf_patch. Prefer calling render_psf_patch directly to avoid confusion.

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Input image in raw space, shape [B, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
patch_center	tuple	Patch center (x, y) in normalized object coordinates. Defaults to (0, 0).	(0, 0)
psf_ks	int	PSF kernel size in pixels. Defaults to PSF_KS.	PSF_KS

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

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

def render_psf(
    self,
    img_obj,
    depth=None,
    patch_center=(0, 0),
    psf_ks=PSF_KS,
    method="conv",
):
    """Render an image patch using PSF convolution (deprecated alias).

    Thin wrapper around `render_psf_patch`. Prefer calling `render_psf_patch`
    directly to avoid confusion.

    Args:
        img_obj (torch.Tensor): Input image in raw space, shape [B, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        patch_center (tuple, optional): Patch center (x, y) in normalized
            object coordinates. Defaults to (0, 0).
        psf_ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.

    Returns:
        img_render (torch.Tensor): Rendered image, shape [B, C, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    return self.render_psf_patch(
        img_obj,
        depth=depth,
        patch_center=patch_center,
        psf_ks=psf_ks,
        method=method,
    )

render_psf_patch

render_psf_patch(
    img_obj, depth=None, patch_center=(0, 0), psf_ks=PSF_KS, method="conv"
)

Render an image patch using a single PSF evaluated at the patch center.

Computes the RGB PSF at patch_center and convolves it with the input image. All pixels in the patch share the same PSF (valid for a small, roughly isoplanatic patch).

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Input image in raw space, shape [B, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
patch_center	tuple or Tensor	Patch center (x, y) in normalized object coordinates, shape [2] or [B, 2].	(0, 0)
psf_ks	int	PSF kernel size in pixels. Defaults to PSF_KS.	PSF_KS
method	str	Convolution backend, "conv" or "fft". Defaults to "conv".	'conv'

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

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

def render_psf_patch(
    self,
    img_obj,
    depth=None,
    patch_center=(0, 0),
    psf_ks=PSF_KS,
    method="conv",
):
    """Render an image patch using a single PSF evaluated at the patch center.

    Computes the RGB PSF at `patch_center` and convolves it with the input
    image. All pixels in the patch share the same PSF (valid for a small,
    roughly isoplanatic patch).

    Args:
        img_obj (torch.Tensor): Input image in raw space, shape [B, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        patch_center (tuple or torch.Tensor): Patch center (x, y) in
            normalized object coordinates, shape [2] or [B, 2].
        psf_ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        method (str, optional): Convolution backend, ``"conv"`` or
            ``"fft"``. Defaults to ``"conv"``.

    Returns:
        img_render (torch.Tensor): Rendered image, shape [B, C, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    # Convert patch_center to tensor
    if isinstance(patch_center, (list, tuple)):
        points = (patch_center[0], patch_center[1], depth)
        points = torch.as_tensor(
            points, device=self.device, dtype=self.dtype
        ).unsqueeze(0)
    elif isinstance(patch_center, torch.Tensor):
        patch_center = patch_center.to(device=self.device, dtype=self.dtype)
        if patch_center.shape[-1] != 2 or patch_center.ndim not in (1, 2):
            raise ValueError("patch_center tensor must have shape [2] or [B, 2].")
        depth = torch.full_like(patch_center[..., 0], depth)
        points = torch.stack(
            [patch_center[..., 0], patch_center[..., 1], depth], dim=-1
        )
    else:
        raise Exception(
            f"Patch center must be a list or tuple or tensor, but got {type(patch_center)}."
        )

    # Compute PSF and perform PSF convolution
    psf = self.psf_rgb(points=points, ks=psf_ks)
    if psf.ndim == 4 and psf.shape[0] == 1:
        psf = psf.squeeze(0)
    elif psf.ndim == 4 and psf.shape[0] != img_obj.shape[0]:
        raise ValueError(
            f"Per-batch patch centers ({psf.shape[0]}) must match image "
            f"batch size ({img_obj.shape[0]})."
        )
    img_render = conv_psf(img_obj, psf=psf, method=method)
    return img_render

render_psf_map

render_psf_map(img_obj, depth=None, psf_grid=7, psf_ks=PSF_KS, psf_spp=SPP_PSF)

Render a full-resolution image using spatially-varying PSF block convolution.

Note

Larger psf_grid and psf_ks give more accurate rendering but are slower.

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Input image in raw space, shape [B, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
psf_grid	int or tuple	PSF grid size. Defaults to 7.	7
psf_ks	int	PSF kernel size in pixels. Defaults to PSF_KS.	PSF_KS
psf_spp	int	Samples per point for PSF computation. Defaults to SPP_PSF.	SPP_PSF

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

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

def render_psf_map(
    self,
    img_obj,
    depth=None,
    psf_grid=7,
    psf_ks=PSF_KS,
    psf_spp=SPP_PSF,
):
    """Render a full-resolution image using spatially-varying PSF block convolution.

    Note:
        Larger `psf_grid` and `psf_ks` give more accurate rendering but are slower.

    Args:
        img_obj (torch.Tensor): Input image in raw space, shape [B, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        psf_grid (int or tuple, optional): PSF grid size. Defaults to 7.
        psf_ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        psf_spp (int, optional): Samples per point for PSF computation.
            Defaults to SPP_PSF.

    Returns:
        img_render (torch.Tensor): Rendered image, shape [B, C, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    psf_map = self.psf_map_rgb(grid=psf_grid, ks=psf_ks, depth=depth, spp=psf_spp)
    img_render = conv_psf_map(img_obj, psf_map)
    return img_render

render_rgbd

render_rgbd(img_obj, depth_map, method='psf_patch', **kwargs)

Render RGBD image.

TODO: add obstruction-aware image simulation.

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Object image, shape [B, C, H, W].	required
depth_map	Tensor	Depth map [mm], shape [B, 1, H, W] (also accepts [B, H, W]). Values should be positive.	required
method	str	Image simulation method, one of "psf_patch", "psf_map", or "psf_pixel". Defaults to "psf_patch".	'psf_patch'
**kwargs		Method-specific keyword arguments, e.g. interp_mode (str): "depth" or "disparity", defaults to "disparity"; num_layers (int): number of depth layers, defaults to 16.	{}

Returns:

Name	Type	Description
img_render	Tensor	Rendered image, shape [B, C, H, W].

Raises:

Type	Description
ValueError	If depth_map contains negative values.
Exception	If method is not recognised.

Reference

[1] "Aberration-Aware Depth-from-Focus", TPAMI 2023. [2] "Efficient Depth- and Spatially-Varying Image Simulation for Defocus Deblur", ICCVW 2025.

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

def render_rgbd(self, img_obj, depth_map, method="psf_patch", **kwargs):
    """Render RGBD image.

    TODO: add obstruction-aware image simulation.

    Args:
        img_obj (torch.Tensor): Object image, shape [B, C, H, W].
        depth_map (torch.Tensor): Depth map [mm], shape [B, 1, H, W] (also
            accepts [B, H, W]). Values should be positive.
        method (str, optional): Image simulation method, one of "psf_patch",
            "psf_map", or "psf_pixel". Defaults to "psf_patch".
        **kwargs: Method-specific keyword arguments, e.g. interp_mode (str):
            "depth" or "disparity", defaults to "disparity"; num_layers (int):
            number of depth layers, defaults to 16.

    Returns:
        img_render (torch.Tensor): Rendered image, shape [B, C, H, W].

    Raises:
        ValueError: If depth_map contains negative values.
        Exception: If method is not recognised.

    Reference:
        [1] "Aberration-Aware Depth-from-Focus", TPAMI 2023.
        [2] "Efficient Depth- and Spatially-Varying Image Simulation for Defocus Deblur", ICCVW 2025.
    """
    if len(depth_map.shape) == 3:
        # [B, H, W] -> [B, 1, H, W]
        depth_map = depth_map.unsqueeze(1)
    elif len(depth_map.shape) != 4:
        raise ValueError("Depth map must have shape [B, H, W] or [B, 1, H, W].")
    if depth_map.shape[1] != 1:
        raise ValueError("Depth map must have exactly one depth channel.")
    if not bool(torch.isfinite(depth_map).all().item()):
        raise ValueError("Depth map must contain only finite values.")
    if bool((depth_map <= 0).any().item()):
        raise ValueError("Depth map values must be strictly positive [mm].")

    if method == "psf_patch":
        # Render an image patch (same FoV, different depth)
        patch_center = kwargs.get("patch_center", (0.0, 0.0))
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        depth_min = kwargs.get("depth_min", depth_map.min())
        depth_max = kwargs.get("depth_max", depth_map.max())
        num_layers = kwargs.get("num_layers", 16)
        interp_mode = kwargs.get("interp_mode", "disparity")

        # Calculate PSF at different depths, (num_layers, 3, ks, ks)
        disp_ref, depths_ref = self._sample_depth_layers(
            depth_min, depth_max, num_layers
        )

        points = torch.stack(
            [
                torch.full_like(depths_ref, patch_center[0]),
                torch.full_like(depths_ref, patch_center[1]),
                depths_ref,
            ],
            dim=-1,
        )
        psfs = self.psf_rgb(points=points, ks=psf_ks)  # (num_layers, 3, ks, ks)

        # Image simulation
        img_render = conv_psf_depth_interp(
            img_obj, -depth_map, psfs, depths_ref, interp_mode=interp_mode
        )
        return img_render

    elif method == "psf_map":
        # Render full resolution image with PSF map convolution (different FoV, different depth)
        psf_grid = kwargs.get("psf_grid", (8, 8))  # (grid_w, grid_h)
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        depth_min = kwargs.get("depth_min", depth_map.min())
        depth_max = kwargs.get("depth_max", depth_map.max())
        num_layers = kwargs.get("num_layers", 16)
        interp_mode = kwargs.get("interp_mode", "disparity")

        # Calculate PSF map at different depths (convert to negative for PSF calculation)
        disp_ref, depths_ref = self._sample_depth_layers(
            depth_min, depth_max, num_layers
        )

        psf_maps = []
        from tqdm import tqdm

        for depth in tqdm(depths_ref):
            psf_map = self.psf_map_rgb(grid=psf_grid, ks=psf_ks, depth=depth)
            psf_maps.append(psf_map)
        psf_map = torch.stack(
            psf_maps, dim=2
        )  # shape [grid_h, grid_w, num_layers, 3, ks, ks]

        # Image simulation
        img_render = conv_psf_map_depth_interp(
            img_obj, -depth_map, psf_map, depths_ref, interp_mode=interp_mode
        )
        return img_render

    elif method == "psf_pixel":
        # Render full resolution image with per-pixel PSF splatting. This method is computationally expensive.
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        assert img_obj.shape[0] == 1, "Now only support batch size 1"

        # Calculate points in the object space
        points_xy = torch.meshgrid(
            torch.linspace(
                -1, 1, img_obj.shape[-1], device=self.device, dtype=self.dtype
            ),
            torch.linspace(
                1, -1, img_obj.shape[-2], device=self.device, dtype=self.dtype
            ),
            indexing="xy",
        )
        points_xy = torch.stack(points_xy, dim=0).unsqueeze(0)
        points = torch.cat([points_xy, -depth_map], dim=1)  # shape [B, 3, H, W]

        # Calculate PSF at different pixels. This step is the most time-consuming.
        points = points.permute(0, 2, 3, 1).reshape(-1, 3)  # shape [H*W, 3]
        psfs = self.psf_rgb(points=points, ks=psf_ks)  # shape [H*W, 3, ks, ks]
        psfs = psfs.reshape(
            img_obj.shape[-2], img_obj.shape[-1], 3, psf_ks, psf_ks
        )  # shape [H, W, 3, ks, ks]

        # Image simulation
        img_render = splat_psf_per_pixel(img_obj, psfs)  # shape [1, C, H, W]
        return img_render

    else:
        raise Exception(f"Image simulation method {method} is not supported.")

activate_grad

activate_grad(activate=True)

Activate (or deactivate) gradients for each surface.

Must be overridden by subclasses.

Parameters:

Name	Type	Description	Default
activate	bool	Whether to enable gradients. Defaults to True.	True

Raises:

Type	Description
NotImplementedError	This base implementation must be overridden.

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

def activate_grad(self, activate=True):
    """Activate (or deactivate) gradients for each surface.

    Must be overridden by subclasses.

    Args:
        activate (bool, optional): Whether to enable gradients. Defaults to True.

    Raises:
        NotImplementedError: This base implementation must be overridden.
    """
    raise NotImplementedError

get_optimizer_params

get_optimizer_params(lr=[0.0001, 0.0001, 0.1, 0.001])

Build per-parameter-group optimizer params. Must be overridden.

Parameters:

Name	Type	Description	Default
lr	list	Per-group learning rates for the different lens parameters. Defaults to [1e-4, 1e-4, 1e-1, 1e-3].	[0.0001, 0.0001, 0.1, 0.001]

Returns:

Name	Type	Description
params	list	List of parameter-group dicts for a torch optimizer.

Raises:

Type	Description
NotImplementedError	This base implementation must be overridden.

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

def get_optimizer_params(self, lr=[1e-4, 1e-4, 1e-1, 1e-3]):
    """Build per-parameter-group optimizer params. Must be overridden.

    Args:
        lr (list, optional): Per-group learning rates for the different lens
            parameters. Defaults to [1e-4, 1e-4, 1e-1, 1e-3].

    Returns:
        params (list): List of parameter-group dicts for a torch optimizer.

    Raises:
        NotImplementedError: This base implementation must be overridden.
    """
    raise NotImplementedError

get_optimizer

get_optimizer(lr=[0.0001, 0.0001, 0, 0.001])

Build an Adam optimizer over the lens parameter groups.

Parameters:

Name	Type	Description	Default
lr	list	Per-group learning rates passed to get_optimizer_params. Defaults to [1e-4, 1e-4, 0, 1e-3].	[0.0001, 0.0001, 0, 0.001]

Returns:

Name	Type	Description
optimizer	Adam	Configured Adam optimizer.

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

def get_optimizer(self, lr=[1e-4, 1e-4, 0, 1e-3]):
    """Build an Adam optimizer over the lens parameter groups.

    Args:
        lr (list, optional): Per-group learning rates passed to
            `get_optimizer_params`. Defaults to [1e-4, 1e-4, 0, 1e-3].

    Returns:
        optimizer (torch.optim.Adam): Configured Adam optimizer.
    """
    params = self.get_optimizer_params(lr)
    optimizer = torch.optim.Adam(params)
    return optimizer