"""
Optimizable model of the ptychographic reconstruction using automatic differentiation (AD)
This is the PyTorch model that holds optimizable tensors and interacts with loss and constraints.
"""
from math import prod
import torch
from torch.fft import fft2, ifft2
import torch.nn as nn
from torchvision.transforms.functional import gaussian_blur
from ptyrad.forward import multislice_forward_model_vec_all
from ptyrad.utils import imshift_batch, torch_phasor, vprint
# The obj_ROI_grid is modified from precalculation to on-the-fly generation for memory consumption
# It has very little performance impact but saves lots of memory for large 4D-STEM data
# Added create_grids and print_model_summary for readability and decoupling
# set_optimizer function is called at the end of the initializaiton, while this can also be called if you want to update the optimizer params without initializing the object
# obj optimization is now split into objp and obja
# mixed object modes are normalized by the init_omode_occu. By design this is a fixed value because optimizing omode_occu with obj simultaneously could be unstable
# obj_ROI cropping is done with vectorization and the obj_ROI_grid is only generated once
# probe with sub-px shifts are calculated only when probe_pos_shifts are enables for optimization
# All the sub-px shifted probes in a batch are processed together with vectorizaiton
# Likewise, the multislice forward model is also fully vectorized across samples (in batch), pmode, and omode
# Note that it's possible to reduce the peak-memory consumption by reducing the level of vectorizaiton and roll back to for loops
# Lastly, the forward pass of this model would output the dp_fwd (N, Ky, Kx) and objp_patches (N, omode, Nz, Ny, Nx) in float32 for later loss calculation
[docs]
class PtychoAD(torch.nn.Module):
"""
Main optimization class for ptychographic reconstruction using automatic differentiation (AD).
This class is responsible for initializing the model parameters, setting up the optimizer, and
performing forward passes to compute diffraction patterns based on the given input indices.
Attributes:
device (str): Device to run computations on ('cuda:0' by default).
verbose (bool): If True, prints model summary (True by default).
detector_blur_std (float): Standard deviation for detector blur, or None if no blur.
obj_preblur_std (float): Standard deviation for object pre-blur, or None if no pre-blur.
lr_params (dict): Learning rate parameters for optimizable tensors.
opt_obja (torch.Tensor): Amplitude of the object.
opt_objp (torch.Tensor): Phase of the object.
opt_obj_tilts (torch.Tensor): Tilts of the object.
opt_probe (torch.Tensor): Probe function.
opt_probe_pos_shifts (torch.Tensor): Shifts for the probe positions.
omode_occu (torch.Tensor): Occupation mode.
H (torch.Tensor): Propagator matrix.
measurements (torch.Tensor): Measurements for the ptychographic reconstruction.
N_scan_slow (torch.Tensor): Number of scans in the slow direction.
N_scan_fast (torch.Tensor): Number of scans in the fast direction.
crop_pos (torch.Tensor): Cropping positions.
slice_thickness (torch.Tensor): slice thickness (dz) parameter.
dx (torch.Tensor): Pixel size in the x direction.
dk (torch.Tensor): K-space sampling interval.
scan_affine (affine.Affine): Affine transformation for scan.
tilt_obj (bool): Whether object tilts are being optimized.
shift_probes (bool): Whether probe shifts are being optimized.
probe_int_sum (torch.Tensor): Sum of squared probe intensities.
optimizable_tensors (dict): Dictionary of tensors that can be optimized.
Args:
init_variables (dict): Dictionary of initial variables required for the model.
model_params (dict): Dictionary of model parameters including learning rates and blur stds.
device (str): Device to run computations on. Default is 'cuda:0'.
verbose (bool): If True, prints model summary. Default is True.
"""
[docs]
def __init__(self, init_variables, model_params, device='cuda', verbose=True):
super(PtychoAD, self).__init__()
with torch.no_grad():
vprint('### Initializing PtychoAD model ###', verbose=verbose)
# Setup model behaviors
self.device = device
self.verbose = verbose
self.detector_blur_std = model_params['detector_blur_std']
self.obj_preblur_std = model_params['obj_preblur_std']
if init_variables.get('on_the_fly_meas_padded', None) is not None:
self.meas_padded = torch.tensor(init_variables['on_the_fly_meas_padded'], dtype=torch.float32, device=device)
self.meas_padded_idx = torch.tensor(init_variables['on_the_fly_meas_padded_idx'], dtype=torch.int32, device=device)
else:
self.meas_padded = None
self.meas_scale_factors = init_variables.get('on_the_fly_meas_scale_factors', None)
# Parse the learning rate and start iter for optimizable tensors
start_iter_dict = {}
end_iter_dict = {}
lr_dict = {}
for key, params in model_params['update_params'].items():
start_iter_dict[key] = params.get('start_iter')
end_iter_dict[key] = params.get('end_iter')
lr_dict[key] = params['lr']
self.optimizer_params = model_params['optimizer_params']
self.start_iter = start_iter_dict
self.end_iter = end_iter_dict
self.lr_params = lr_dict
# Optimizable parameters
self.opt_obja = nn.Parameter(torch.abs(torch.tensor(init_variables['obj'], device=device)).to(torch.float32))
self.opt_objp = nn.Parameter(torch.angle(torch.tensor(init_variables['obj'], device=device)).to(torch.float32))
self.opt_obj_tilts = nn.Parameter(torch.tensor(init_variables['obj_tilts'], dtype=torch.float32, device=device))
self.opt_slice_thickness = nn.Parameter(torch.tensor(init_variables['slice_thickness'], dtype=torch.float32, device=device))
self.opt_probe = nn.Parameter(torch.view_as_real(torch.tensor(init_variables['probe'], dtype=torch.complex64, device=device))) # The `torch.view_as_real` allows correct handling of DDP via NCCL even in PyTorch 2.4
self.opt_probe_pos_shifts = nn.Parameter(torch.tensor(init_variables['probe_pos_shifts'], dtype=torch.float32, device=device))
# Buffers are used during forward pass
self.register_buffer ('omode_occu', torch.tensor(init_variables['omode_occu'], dtype=torch.float32, device=device))
self.register_buffer ('H', torch.tensor(init_variables['H'], dtype=torch.complex64, device=device))
self.register_buffer ('measurements', torch.tensor(init_variables['measurements'], dtype=torch.float32, device=device))
self.register_buffer ('N_scan_slow', torch.tensor(init_variables['N_scan_slow'], dtype=torch.int32, device=device))# Saving this for reference, the cropping is based on self.obj_ROI_grid.
self.register_buffer ('N_scan_fast', torch.tensor(init_variables['N_scan_fast'], dtype=torch.int32, device=device))# Saving this for reference, the cropping is based on self.obj_ROI_grid.
self.register_buffer ('crop_pos', torch.tensor(init_variables['crop_pos'], dtype=torch.int32, device=device))# Saving this for reference, the cropping is based on self.obj_ROI_grid.
self.register_buffer ('slice_thickness', torch.tensor(init_variables['slice_thickness'], dtype=torch.float32, device=device))# Saving this for reference
self.register_buffer ('dx', torch.tensor(init_variables['dx'], dtype=torch.float32, device=device))# Saving this for reference
self.register_buffer ('dk', torch.tensor(init_variables['dk'], dtype=torch.float32, device=device))# Saving this for reference
self.register_buffer ('lambd', torch.tensor(init_variables['lambd'], dtype=torch.float32, device=device))
self.random_seed = init_variables['random_seed'] # Saving this for reference
self.length_unit = init_variables['length_unit'] # Saving this for reference
self.scan_affine = init_variables['scan_affine'] # Saving this for reference
self.tilt_obj = bool(self.lr_params['obj_tilts'] != 0 or torch.any(self.opt_obj_tilts)) # Set tilt_obj to True if lr_params['obj_tilts'] is not 0 or we have any none-zero tilt values
self.shift_probes = bool(self.lr_params['probe_pos_shifts'] != 0) # Set shift_probes to True if lr_params['probe_pos_shifts'] is not 0
self.change_thickness = bool(self.lr_params['slice_thickness'] != 0)
self.probe_int_sum = self.get_complex_probe_view().abs().pow(2).sum() # This is only used for the `fix_probe_int`
self.loss_iters = []
self.iter_times = []
self.dz_iters = []
self.avg_tilt_iters = []
# Create grids for shifting
self.create_grids()
# Create a dictionary to store the optimizable tensors
self.optimizable_tensors = {
'obja' : self.opt_obja,
'objp' : self.opt_objp,
'obj_tilts' : self.opt_obj_tilts,
'slice_thickness' : self.opt_slice_thickness,
'probe' : self.opt_probe,
'probe_pos_shifts': self.opt_probe_pos_shifts}
self.create_optimizable_params_dict(self.lr_params, self.verbose)
# Initialize propagator-related variables
self.init_propagator_vars()
# Initialize iteration numbers that require torch.compile
self.init_compilation_iters()
vprint('### Done initializing PtychoAD model ###', verbose=verbose)
vprint(' ', verbose=verbose)
def get_complex_probe_view(self):
""" Retrieve complex view of the probe """
# This is a post-processing to ensure minimal code changes in PtyRAD for the DDP (multiGPU) via NCCL due to limited support for Complex value
return torch.view_as_complex(self.opt_probe)
def create_grids(self):
""" Create the grids for shifting probes, selecting obj ROI, and Fresnel propagator in a vectorized approach """
# Note that the shift_object_grid is pre-generated for potential future usage of sub-px object shifts
device = self.device
probe = self.get_complex_probe_view()
Npy, Npx = probe.shape[-2:] # Number of probe pixels in y and x directions
Noy, Nox = self.opt_objp.shape[-2:] # Number of object pixels in y and x directions
# Grids for Fresnel propagator
# Note that this grid has a half-bin shift that avoids exact Nyquist frequency (k = -0.5) which would generate NaNs inside sqrt.
# Due to the half-bin shift, the 0th element is also not exactly 0.
# If we decided to go with the small angle approximated propagator, then we may unify this grid with the reciprocal probe gird.
ygrid = (torch.arange(-Npy // 2, Npy // 2, device=device) + 0.5) / Npy
xgrid = (torch.arange(-Npx // 2, Npx // 2, device=device) + 0.5) / Npx
ky = torch.fft.ifftshift(2 * torch.pi * ygrid / self.dx) # Use ifftshift to shift the 0 frequency to the corner
kx = torch.fft.ifftshift(2 * torch.pi * xgrid / self.dx)
Ky, Kx = torch.meshgrid(ky, kx, indexing="ij")
self.propagator_grid = torch.stack([Ky,Kx], dim=0) # (2,Ky,Kx), k-space grid for Fresnel propagator with 2pi absorbed
# Grids for obj_ROI selection
rpy, rpx = torch.meshgrid(torch.arange(Npy, dtype=torch.int32, device=device),
torch.arange(Npx, dtype=torch.int32, device=device), indexing='ij') # real space grid for probe in y and x directions
self.rpy_grid = rpy # real space grid with y-indices spans across probe extent
self.rpx_grid = rpx
# Grids for shifting probes and objects
# This is the typical fftfreq grid that ranges from [-0.5, 0.5) and the 0th element is exactly 0
kpy, kpx = torch.meshgrid(torch.fft.fftfreq(Npy, dtype=torch.float32, device=device),
torch.fft.fftfreq(Npx, dtype=torch.float32, device=device), indexing='ij')
koy, kox = torch.meshgrid(torch.fft.fftfreq(Noy, dtype=torch.float32, device=device),
torch.fft.fftfreq(Nox, dtype=torch.float32, device=device), indexing='ij')
self.shift_probes_grid = torch.stack([kpy, kpx], dim=0) # (2,Npy,Npx), normalized k-space grid stack for sub-px probe shifting
self.shift_object_grid = torch.stack([koy, kox], dim=0) # (2,Noy,Nox), normalized k-space grid stack for sub-px object shifting (Implemented for completeness, not used in PtyRAD)
def create_optimizable_params_dict(self, lr_params, verbose=True):
""" Sets the optimizer with lr_params """
# # Use this to edit learning rate if needed some refinement
# model.set_optimizer(lr_params={'obja' : 5e-4,
# 'objp' : 5e-4,
# 'obj_tilts' : 1e-4,
# 'probe' : 1e-4,
# 'probe_pos_shifts': 1e-4})
# optimizer=torch.optim.Adam(model.optimizer_params)
self.lr_params = lr_params
self.optimizable_params = []
for param_name, lr in lr_params.items():
if param_name not in self.optimizable_tensors:
raise ValueError(f"WARNING: '{param_name}' is not a valid parameter name, check your `update_params` and choose from 'obja', 'objp', 'obj_tilts', 'slice_thickness', 'probe', and 'probe_pos_shifts'")
else:
self.optimizable_tensors[param_name].requires_grad = (lr != 0) and (self.start_iter[param_name] ==1) # Set requires_grad based on learning rate and start_iter
if lr != 0:
self.optimizable_params.append({'params': [self.optimizable_tensors[param_name]], 'lr': lr})
if verbose:
self.print_model_summary()
def init_propagator_vars(self):
""" Initialize propagator related variables """
# Initialize propagator for fixed non-zero tilts and fixed thickness that could be position dependent
# It's better to calculate the full one during initialization and slice it later given indices so we can use torch.compile later
dz = self.opt_slice_thickness.detach()
Ky, Kx = self.propagator_grid
tilts_y_full = self.opt_obj_tilts[:,0,None,None] / 1e3 # mrad, tilts_y = (N,Y,X)
tilts_x_full = self.opt_obj_tilts[:,1,None,None] / 1e3
self.H_fixed_tilts_full = self.H * torch_phasor(dz * (Ky * torch.tan(tilts_y_full) + Kx * torch.tan(tilts_x_full))) # (1,Y,X) or (N,Y,X)
# Initialize other relevant variables
self.k = 2 * torch.pi / self.lambd
self.Kz = torch.sqrt(self.k ** 2 - Kx ** 2 - Ky ** 2) # Upper case K indicates it's a 2D tensor (Y,X)
def init_compilation_iters(self):
""" Initialize iteration numbers that require torch.compile """
compilation_iters = {1} # Always compile at first iteration
for param_name in self.optimizable_tensors.keys():
start_iter = self.start_iter.get(param_name)
end_iter = self.end_iter.get(param_name)
# Add start_iter compilation points
if start_iter is not None and start_iter >= 1:
compilation_iters.add(start_iter)
# Add end_iter compilation points
if end_iter is not None and end_iter >= 1:
# Compile at end_iter to handle the transition. end_iter is exclusive for grad calculation.
compilation_iters.add(end_iter)
# Store as sorted list
self.compilation_iters = sorted(compilation_iters)
def print_model_summary(self):
""" Prints a summary of the model's optimizable variables and statistics. """
# Set all the print as vprint so that it'll only print once in DDP, the actual `if verbose` is set outside of the function
vprint('### PtychoAD optimizable variables ###')
for name, tensor in self.optimizable_tensors.items():
vprint(f"{name.ljust(16)}: {str(tensor.shape).ljust(32)}, {str(tensor.dtype).ljust(16)}, device:{tensor.device}, grad:{str(tensor.requires_grad).ljust(5)}, lr:{self.lr_params[name]:.0e}")
total_var = sum(tensor.numel() for _, tensor in self.optimizable_tensors.items() if tensor.requires_grad)
# When you create a new model, make sure to pass the optimizer_params to optimizer using "optimizer = torch.optim.Adam(model.optimizer_params)"
vprint(" ")
vprint('### Optimizable variables statitsics ###')
vprint(f"Total measurement values : {self.measurements.numel():,d}")
vprint(f"Total optimizing variables: {total_var:,d}")
vprint(f"Overdetermined ratio : {self.measurements.numel()/total_var:.2f}")
vprint(" ")
vprint('### Model behavior ###')
vprint(f"Obj preblur : {True if self.obj_preblur_std is not None else False}")
vprint(f"Tilt propagator : {self.tilt_obj}")
vprint(f"Change slice thickness : {self.change_thickness}")
vprint(f"Sub-px probe shift : {self.shift_probes}")
vprint(f"Detector blur : {True if self.detector_blur_std is not None else False}")
vprint(f"On-the-fly meas padding : {True if self.meas_padded is not None else False}")
vprint(f"On-the-fly meas resample : {True if self.meas_scale_factors is not None else False}")
vprint(" ")
def get_obj_ROI(self, indices):
""" Get object ROI with integer coordinates """
# It's strongly recommended to do integer version of get_obj_ROI
# opt_obj.shape = (B,D,H,W,C) = (omode,D,H,W,2)
# object_roi = (N,B,D,H,W,2), N is the additional sample index within the input batch, B is now used for omode.
# rpy_grid is the y-grid (Ny,Nx), by adding the y coordinates from init_crop_pos (N,1) in a broadcast way, it becomes (N,Ny,Nx)
# obj_ROI_grid_y = (N,Ny,Nx)
opt_obj = torch.stack([self.opt_obja, self.opt_objp], dim=-1)
obj_ROI_grid_y = self.rpy_grid[None,:,:] + self.crop_pos[indices, None, None, 0]
obj_ROI_grid_x = self.rpx_grid[None,:,:] + self.crop_pos[indices, None, None, 1]
object_roi = opt_obj[:,:,obj_ROI_grid_y,obj_ROI_grid_x,:].permute(2,0,1,3,4,5)
return object_roi
def get_obj_patches(self, indices):
"""
Get object patches from specified indices
"""
object_patches = self.get_obj_ROI(indices)
if self.obj_preblur_std is None or self.obj_preblur_std == 0:
return object_patches
else:
# Permute and reshape approach, this is much faster than the stack/list comprehension version
obj = object_patches.permute(5,0,1,2,3,4) # Move the r/i to the front so it's (2,N,B,D,H,W)
obj_shape = obj.shape
obj = obj.reshape(-1, obj_shape[-2], obj_shape[-1])
object_patches = gaussian_blur(obj, kernel_size=5, sigma=self.obj_preblur_std).reshape(obj_shape).permute(1,2,3,4,5,0) # The torchvision Gaussian blur only acts on last 2 dimensions
return object_patches
def get_probes(self, indices):
""" Get probes for each position """
# This function will return a probe tensor with (N, pmode, Ny, Nx)
# If you're not trying to optimize probe positions, there's not much point using sub-px shifted stationary probes
# So the function would broadcast the same probe across the batch dimension,
# and would only be returning multiple sub-px shifted probes if you're optimizing self.opt_probe_pos_shifts
probe = self.get_complex_probe_view()
if self.shift_probes:
probes = imshift_batch(probe, shifts = self.opt_probe_pos_shifts[indices], grid = self.shift_probes_grid)
else:
probes = torch.broadcast_to(probe, (indices.shape[0], *probe.shape)) # Broadcast a batch dimension, essentially using same probe for all samples
return probes
def get_propagators(self, indices):
""" Get propagators for each position """
# self.tilt_obj is True as long as we're optimizing the opt_obj_tilts or we have non-zero initial tilt values
# This function will return a single propagator (H) if self.opt_obj_tilts has shape = (1,2) (single tilt_y, tilt_x)
# If self.opt_obj_tilts has shape = (N,2), it'll return multiple propagtors stacked at axis 0 (N,Y,X)
# Note that 0 tilts is numerically equivalent to the H and can be verified by "torch.allclose(model.H, model.get_propagators([0]))"
# The exp(2pi * i * sqrt(k^2 - kx^2 - ky^2)) approach is equivalent to the common exp(-i * pi * lambda * dz * k^2) for small angles,
# see J. Goodman, Introduction to Fourier Optics (McGraw-Hill, 1968) (PDF page 88, eqn 4-20, 4-21 as attached).
# https://www.hlevkin.com/hlevkin/90MathPhysBioBooks/Physics/Physics/Mix/Introduction%20to%20Fourier%20Optics.pdf
# Note that torch.exp(1j*phase_shift) is not compatible with torch.compile because the 1j is a Python built-in and not a tensor,
# so I've replaced them with torch.polar(torch.ones_like(phase), phase), which is wrapped as a util function `torch_phasor(phase)`
# Setup boolean flags
tilt_obj = self.tilt_obj # Whether we need to apply tilt to the Fresnel propagator
global_tilt = (self.opt_obj_tilts.shape[0] == 1) # 'tilt_type' = 'all' or 'each'
change_tilt = (self.lr_params['obj_tilts'] != 0) # Whether tilts are optimized or not
change_thickness = self.change_thickness # Whether thickness is optimized or not
# Setup tilts and other variables
dz = self.opt_slice_thickness
Kz = self.Kz # kz = torch.sqrt(k ** 2 - Kx ** 2 - Ky ** 2), k = 2pi/lambda.
Ky, Kx = self.propagator_grid
# tilts can be either (1,2) or (N,2) depends on global_tilt flag
if global_tilt:
tilts = self.opt_obj_tilts
else:
tilts = self.opt_obj_tilts[indices]
tilts_y = tilts[:,0,None,None] / 1e3 # mrad, tilts_y = (N,Y,X)
tilts_x = tilts[:,1,None,None] / 1e3
if tilt_obj and change_thickness:
# Case 1: Tilts are either non-zero or optimizing, while thickness is optimizing
H_opt_dz = torch_phasor(dz * Kz) # H has zero frequency at the corner in k-space
return H_opt_dz * torch_phasor(dz * (Ky * torch.tan(tilts_y) + Kx * torch.tan(tilts_x)))
elif tilt_obj and not change_thickness:
if change_tilt:
# Case 2A: Tilts are optimizing, while thickness is fixed
return self.H * torch_phasor(dz * (Ky * torch.tan(tilts_y) + Kx * torch.tan(tilts_x)))
else:
# Case 2B: Tilts are fixed non-zero values (1,2) or (N,2), while thickness is fixed
# Propagator is pre-calculated during init_propagator_vars
return self.H_fixed_tilts_full if global_tilt else self.H_fixed_tilts_full[indices]
elif not tilt_obj and change_thickness:
# Case 3: Tilt is fixed at 0 and thickness is optimizing
H_opt_dz = torch_phasor(dz * Kz)
return H_opt_dz[None,]
else:
# Case 4: Tilt is fixed at 0 and thickness is fixed
return self.H[None,]
def get_propagated_probe(self, index):
probe = self.get_probes(index)[0].detach() # (pmode, Ny, Nx), just grab the probe at 1st index
H = self.get_propagators(index)[[0]].detach() # (1, Ny, Nx) or (N, Ny, Nx) depends on tilt_type ('all' or 'each'), so we need to grab the 1st index without reducing dimension
n_slices = self.opt_objp.shape[1]
probe_prop = torch.zeros((n_slices, *probe.shape), dtype=probe.dtype, device=probe.device)
psi = probe # (z, pmode, Ny, Nx)
for n in range(n_slices):
probe_prop[n] = psi
psi = ifft2(H[None,] * fft2(psi))
return probe_prop
def get_forward_meas(self, object_patches, probes, propagators):
dp_fwd = multislice_forward_model_vec_all(object_patches, probes, propagators, omode_occu=self.omode_occu)
if self.detector_blur_std is not None and self.detector_blur_std != 0:
dp_fwd = gaussian_blur(dp_fwd, kernel_size=5, sigma=self.detector_blur_std)
return dp_fwd
def get_measurements(self, indices=None):
""" Get measurements for each position """
# Return the selected measurements based on input indices
# If no indices are passed, return the entire measurements ignoring any "on-the-fly" padding/resampling
measurements = self.measurements
device = self.device
dtype = measurements.dtype
if self.meas_padded is not None:
meas_padded = self.meas_padded
meas_padded_idx = self.meas_padded_idx
pad_h1, pad_h2, pad_w1, pad_w2 = meas_padded_idx
scale_factor = tuple(self.meas_scale_factors) if self.meas_scale_factors is not None else None
if indices is not None:
measurements = self.measurements[indices]
if self.meas_padded is not None:
canvas = torch.zeros((measurements.shape[0], *meas_padded.shape[-2:]), dtype=dtype, device=device)
canvas += meas_padded
canvas[..., pad_h1:pad_h2, pad_w1:pad_w2] = measurements # Replace the center part with the original meas
measurements = canvas
if self.meas_scale_factors is not None and any(factor != 1 for factor in scale_factor):
measurements = torch.nn.functional.interpolate(measurements[None,], scale_factor=scale_factor, mode='bilinear')[0] # 2D interpolate requires 4D input (N, C, H, W)
measurements = measurements / prod(scale_factor) # This ensures the intensity scale and the integrated intensity are unchanged
else: # Skip the "on-the-fly" operations so it won't throw any CUDA out-of-memory error. All typical PtyRAD usuage would pass get_measurements(batch) so this should be ok.
if self.meas_padded is not None or self.meas_scale_factors is not None:
vprint(f"WARNING: 'on-the-fly' measurements padding/resampling detected, but they are ignored because it may cause 'CUDA out-of-memory' when 'get_measurements()' is called without any indices. The original measurement with shape = {self.measurements.shape} is returned instead.")
measurements = self.measurements
return measurements
def clear_cache(self):
"""Clear temporary attributes like cached object patches."""
self._current_object_patches = None
def forward(self, indices):
""" Doing the forward pass and get an output diffraction pattern for each input index """
# The indices are passed in as an array and representing the whole batch
# Note that detector blur is a physical process and should be included in the forward method
# It's a design choice to put it here, instead of putting it under optimization.py
object_patches = self.get_obj_patches(indices)
probes = self.get_probes(indices)
propagators = self.get_propagators(indices)
dp_fwd = self.get_forward_meas(object_patches, probes, propagators)
# Keep the object_patches for later object-specific loss
self._current_object_patches = object_patches
return dp_fwd