CombinedConstraint#

class ptyrad.constraints.CombinedConstraint(constraint_params, device='cuda', verbose=True)[source]#

Bases: Module

Applies iteration-wise in-place constraints on optimizable tensors.

This class is designed to apply various constraints to a model’s parameters during the optimization process. The constraints are applied at specific iteration frequencies, as determined by the constraint_params dictionary. These constraints include orthogonality, probe amplitude constraints in Fourier space, intensity constraints, Gaussian blurring, Fourier filtering, and more.

Parameters:
  • constraint_params (dict) – A dictionary containing the configuration for each constraint. Each constraint should have a frequency and other parameters necessary for its application.

  • device (str, optional) – The device on which the tensors are located (e.g., ‘cuda’ or ‘cpu’). Defaults to ‘cuda’.

  • verbose (bool, optional) – If True, prints messages during the application of constraints. Defaults to True.

__init__(constraint_params, device='cuda', verbose=True)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Methods

__init__(constraint_params[, device, verbose])

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.

apply_complex_ratio(model, niter)

Apply complex constraint on object

apply_fix_probe_int(model, niter)

Apply probe intensity constraint

apply_kr_filter(model, niter)

Apply kr Fourier filter constraint on object

apply_kz_filter(model, niter)

Apply kz Fourier filter constraint on object

apply_mirrored_amp(model, niter)

Apply mirrored amplitude constraint on obja at voxel level

apply_obj_rblur(model, niter)

Apply Gaussian blur to object, this only applies to the last 2 dimension (...,H,W)

apply_obj_z_recenter(model, niter)

Apply object z-recentering along depth dimension

apply_obj_zblur(model, niter)

Apply Gaussian blur to object, this only applies to the last dimension (...,L)

apply_obja_thresh(model, niter)

Apply thresholding on obja at voxel level

apply_objp_postiv(model, niter)

Apply positivity constraint on objp at voxel level

apply_ortho_pmode(model, niter)

Apply orthogonality constraint to probe modes

apply_pos_recenter(model, niter)

Apply position recentering constraint on probe positions

apply_probe_mask_k(model, niter)

Apply probe amplitude constraint in Fourier space

apply_tilt_smooth(model, niter)

Apply Gaussian blur to object tilts

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.

compile(*args, **kwargs)

Compile this Module's forward using torch.compile().

cpu()

Move all model parameters and buffers to the CPU.

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(model, niter)

Define the computation performed at every call.

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

get_extra_state()

Return any extra state to include in the module's state_dict.

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

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.

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.

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