ptyrad.solver.reconstruction

Reconstruction workflow related functions

Functions

compute_loss(batch, model, model_instance, ...)	Compute the model output and loss.
create_optimizer(optimizer_params, ...)
create_scheduler(scheduler_params, optimizer)	Creates a PyTorch LR scheduler from the given params and binds it to the optimizer.
loss_logger(batch_losses, niter, iter_t)	Logs and summarizes the loss values for an iteration during the ptychographic reconstruction.
make_batches(indices, pos, batch_size[, ...])	Make batches from input indices
parse_torch_compile_configs(configs)	Convert user-facing CompilerConfigs to dict suitable for torch.compile
prepare_recon(model, init, params)	Prepares the indices, batches, and output path for ptychographic reconstruction.
recon_loop(model, init, params, optimizer, ...)	Executes the iterative optimization loop for ptychographic reconstruction.
recon_step(batches, grad_accumulation, ...)	Performs one iteration (or step) of the ptychographic reconstruction in the optimization loop.
select_scan_indices(N_scan_slow, N_scan_fast)
time_sync([device])
toggle_grad_requires(model, niter)	Toggle requires_grad based on start and end iteration for each optimizable tensor.

ptyrad.solver.reconstruction.create_optimizer(optimizer_params, optimizable_params)

ptyrad.solver.reconstruction.create_scheduler(scheduler_params, optimizer)

Creates a PyTorch LR scheduler from the given params and binds it to the optimizer.

Converts param-group LRs (and betas[0] for OneCycleLR/CyclicLR) to CPU tensors so that torch.compile guards on tensor identity rather than scalar value, preventing recompilation when LR or momentum changes each step. Patches scheduler.step() to restore tensor identity in-place after each call. Optionally loads a previously saved scheduler state.

Parameters:

• scheduler_params (dict or None) – Scheduler configuration dict with keys ‘name’, ‘configs’, and ‘load_state’. If None, returns None.

• optimizer (torch.optim.Optimizer) – The optimizer whose param-group LRs will be managed.

Returns:

The configured scheduler, or None.

Return type:

torch.optim.lr_scheduler.LRScheduler or None

ptyrad.solver.reconstruction.prepare_recon(model, init, params)

Prepares the indices, batches, and output path for ptychographic reconstruction.

This function parses the necessary parameters and generates the indices for scanning, creates batches based on the probe positions, and sets up the output directory for saving results. It also plots and saves a figure illustrating the grouping of probe positions.

Parameters:

• model (PtychoModel) – The ptychographic model containing the object, probe, probe positions, and other relevant parameters.

• init (Initializer) – The initializer object containing the initialized variables needed for reconstruction.

• params (dict) – A dictionary containing various parameters needed for the reconstruction process, including experimental parameters, loss parameters, constraint parameters, and reconstruction settings.

Returns:

A tuple containing the following:

• indices (numpy.ndarray): Array of indices for scanning positions.

• batches (list of numpy.ndarray): List of batches where each batch contains indices grouped according to the selected grouping mode.

• output_path (str): The path to the directory where reconstruction results and figures will be saved.

Return type:

tuple

ptyrad.solver.reconstruction.select_scan_indices(N_scan_slow, N_scan_fast, subscan_slow=None, subscan_fast=None, mode='full')

ptyrad.solver.reconstruction.make_batches(indices, pos, batch_size, mode='random', seed=None)

Make batches from input indices

ptyrad.solver.reconstruction.recon_loop(model, init, params, optimizer, scheduler, loss_fn, constraint_fn, indices, batches, output_path, acc=None)

Executes the iterative optimization loop for ptychographic reconstruction.

This function performs the iterative reconstruction process by optimizing the model parameters over a specified number of iterations. During each iteration, it applies the loss and constraint functions, updates the model, and logs the loss values. Intermediate results are saved at specified intervals, and a summary is plotted.

Parameters:

• model (PtychoModel) – The ptychographic model containing the parameters and variables to be optimized.

• init (Initializer) – The initializer object containing the initialized variables needed for reconstruction.

• params (dict) – A dictionary containing various parameters for the reconstruction process, including experimental parameters, source parameters, loss parameters, constraint parameters, and reconstruction settings.

• optimizer (torch.optim.Optimizer) – The optimizer used to update the model parameters.

• scheduler (torch.optim.lr_scheduler.LRScheduler or None) – The LR scheduler that adjusts the optimizer’s learning rate each iteration. Pass None to disable.

• loss_fn (CombinedLoss) – The loss function object used to compute the loss during each iteration.

• constraint_fn (CombinedConstraint) – The constraint function object applied during each iteration to enforce specific constraints on the model.

• indices (numpy.ndarray) – Array of indices for scanning positions.

• batches (list of numpy.ndarray) – List of batches where each batch contains indices grouped according to the selected grouping mode.

• output_path (str) – The path to the directory where reconstruction results and figures will be saved.

Returns:

A list of tuples, where each tuple contains the iteration number, the loss

value for that iteration, and the time taken for that iteration.

Return type:

list

ptyrad.solver.reconstruction.recon_step(batches, grad_accumulation, model, optimizer, scheduler, loss_fn, constraint_fn, niter, acc=None, compute_loss_fn=None, scheduler_step_unit='iter')

Performs one iteration (or step) of the ptychographic reconstruction in the optimization loop.

This function executes a single iteration of the reconstruction process, including: - Computing the forward model to generate diffraction patterns. - Calculating the loss by comparing the modeled and measured diffraction patterns. - Performing a backward pass to compute gradients and update the model parameters using the optimizer. - Applying iteration-wise constraints after all batches are processed.

Parameters:

• batches (list of numpy.ndarray) – List of batches where each batch contains indices grouped according to the selected grouping mode.

• grad_accumulation (int) – Number of mini-batches to accumulate gradients over before stepping.

• model (PtychoModel) – The ptychographic model containing the parameters and variables to be optimized.

• optimizer (torch.optim.Optimizer) – The optimizer used to update the model parameters.

• scheduler (torch.optim.lr_scheduler.LRScheduler or None) – LR scheduler stepped according to scheduler_step_unit after the optimizer step. Pass None to disable.

• loss_fn (CombinedLoss) – The loss function object used to compute the loss for each batch.

• constraint_fn (CombinedConstraint) – The constraint function object applied after each iteration to enforce specific constraints on the model.

• niter (int) – The current iteration number in the optimization loop.

• acc (Accelerator or None) – HuggingFace Accelerator for multi-GPU support. None for single-GPU.

• compute_loss_fn (callable or None) – Compiled or uncompiled loss function; defaults to compute_loss.

• scheduler_step_unit (str) – ‘iter’ to step the scheduler once per outer iteration (default), or ‘batch’ to step after every optimizer.step() call (grad-accumulation boundary). ReduceLROnPlateau always steps per iteration regardless of this setting.

Returns:

A dictionary where each key corresponds to a loss component name and the value is

a list of loss values computed for each batch in the iteration.

Return type:

dict

ptyrad.solver.reconstruction.time_sync(device=None)

ptyrad.solver.reconstruction.parse_torch_compile_configs(configs)

Convert user-facing CompilerConfigs to dict suitable for torch.compile

Note

The params.yaml defines as ‘enable’: bool = False, while torch.compile takes only ‘disable’: bool, so a conversion is needed.

ptyrad.solver.reconstruction.toggle_grad_requires(model, niter)

Toggle requires_grad based on start and end iteration for each optimizable tensor.

ptyrad.solver.reconstruction.compute_loss(batch, model, model_instance, measured_DP, loss_fn)

Compute the model output and loss.

ptyrad.solver.reconstruction.loss_logger(batch_losses, niter, iter_t)

Logs and summarizes the loss values for an iteration during the ptychographic reconstruction.

This function computes the average loss for each loss component across all batches in the current iteration. It then logs the total loss, the individual loss components, and the time taken for the iteration. The function also returns the total loss for the iteration.

Parameters:

• batch_losses (dict) – A dictionary where each key corresponds to a loss component name, and the value is a list of loss values computed for each batch in the iteration.

• niter (int) – The current iteration number in the optimization loop.

• iter_t (float) – The total time taken to complete the iteration, in seconds.

Returns:

The total loss for the current iteration, computed as the sum of the average loss values for each component.

Return type:

float