PtychoAD

PtychoAD#

class ptyrad.models.PtychoAD(init_variables, model_params, device='cuda', verbose=True)[source]#

Bases: 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.

device#

Device to run computations on (‘cuda:0’ by default).

Type:: str

verbose#

If True, prints model summary (True by default).

Type:: bool

detector_blur_std#

Standard deviation for detector blur, or None if no blur.

Type:: float

obj_preblur_std#

Standard deviation for object pre-blur, or None if no pre-blur.

Type:: float

lr_params#

Learning rate parameters for optimizable tensors.

Type:: dict

opt_obja#

Amplitude of the object.

Type:: torch.Tensor

opt_objp#

Phase of the object.

Type:: torch.Tensor

opt_obj_tilts#

Tilts of the object.

Type:: torch.Tensor

opt_probe#

Probe function.

Type:: torch.Tensor

opt_probe_pos_shifts#

Shifts for the probe positions.

Type:: torch.Tensor

omode_occu#

Occupation mode.

Type:: torch.Tensor

H#

Propagator matrix.

Type:: torch.Tensor

measurements#

Measurements for the ptychographic reconstruction.

Type:: torch.Tensor

N_scan_slow#

Number of scans in the slow direction.

Type:: torch.Tensor

N_scan_fast#

Number of scans in the fast direction.

Type:: torch.Tensor

crop_pos#

Cropping positions.

Type:: torch.Tensor

slice_thickness#

slice thickness (dz) parameter.

Type:: torch.Tensor

dx#

Pixel size in the x direction.

Type:: torch.Tensor

dk#

K-space sampling interval.

Type:: torch.Tensor

scan_affine#

Affine transformation for scan.

Type:: affine.Affine

tilt_obj#

Whether object tilts are being optimized.

Type:: bool

shift_probes#

Whether probe shifts are being optimized.

Type:: bool

probe_int_sum#

Sum of squared probe intensities.

Type:: torch.Tensor

optimizable_tensors#

Dictionary of tensors that can be optimized.

Type:: dict

Parameters:

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.

__init__(init_variables, model_params, device='cuda', verbose=True)[source]#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

Methods

`__init__`(init_variables, model_params[, ...])	Initialize internal Module state, shared by both nn.Module and ScriptModule.
`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`clear_cache`()	Clear temporary attributes like cached object patches.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`create_grids`()	Create the grids for shifting probes, selecting obj ROI, and Fresnel propagator in a vectorized approach
`create_optimizable_params_dict`(lr_params[, ...])	Sets the optimizer with lr_params
`cuda`([device])	Move all model parameters and buffers to the GPU.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(indices)	Doing the forward pass and get an output diffraction pattern for each input index
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_complex_probe_view`()	Retrieve complex view of the probe
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_forward_meas`(object_patches, probes, ...)
`get_measurements`([indices])	Get measurements for each position
`get_obj_ROI`(indices)	Get object ROI with integer coordinates
`get_obj_patches`(indices)	Get object patches from specified indices
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_probes`(indices)	Get probes for each position
`get_propagated_probe`(index)
`get_propagators`(indices)	Get propagators for each position
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`init_compilation_iters`()	Initialize iteration numbers that require torch.compile
`init_propagator_vars`()	Initialize propagator related variables
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`modules`()	Return an iterator over all modules in the network.
`mtia`([device])	Move all model parameters and buffers to the MTIA.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`print_model_summary`()	Prints a summary of the model's optimizable variables and statistics.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post-hook to be run after module's `load_state_dict()` is called.
`register_load_state_dict_pre_hook`(hook)	Register a pre-hook to be run before module's `load_state_dict()` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_post_hook`(hook)	Register a post-hook for the `state_dict()` method.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`set_submodule`(target, module)	Set the submodule given by `target` if it exists, otherwise throw an error.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

Attributes

`T_destination`
`call_super_init`
`dump_patches`
`training`

PtychoAD

Contents

PtychoAD#