ptyrad.core.constraints

Physical constraints that directly modify optimizable tensors with specified intervals of iterations

Functions

complex_ratio_constraint(model, alpha1, alpha2)
dct_threshold_filter(x[, threshold_ratio])	Applies hard-threshold filtering in the DCT domain.
get_obj_z_shift(obj_phase[, threshold, ...])	Compute z-shift from the center-of-mass (CoM) of the object phase.
kr_filter(obj, radius, width)	Apply kr_filter using the 2D sigmoid filter
kz_filter(obj[, beta_regularize_layers, ...])	Apply kz_filter using the arctan filter
orthogonalize_modes_vec(modes[, sort])	Orthogonalize probe modes via Gram matrix eigendecomposition.
shift_obj_along_z(objc, z_shift)	Apply a subpixel shift along z using Fourier shift theorem.
sort_by_mode_int(modes)

Classes

CombinedConstraint(constraint_params[, device])

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

class ptyrad.core.constraints.CombinedConstraint(constraint_params, device='cuda')

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’.

apply_probe_mask_k(model, niter)

Apply probe amplitude constraint in Fourier space

apply_probe_mask_r(model, niter)

Apply probe amplitude constraint in Real space

apply_ortho_pmode(model, niter)

Apply orthogonality constraint to probe modes

apply_fix_probe_int(model, niter)

Apply probe intensity constraint

apply_obj_rblur(model, niter)

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

apply_obj_zblur(model, niter)

Apply Gaussian blur to object along the z-axis (slice dimension)

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

Apply kr threshold constraint on object

apply_complex_ratio(model, niter)

Apply complex constraint on object

apply_mirrored_amp(model, niter)

Apply mirrored amplitude constraint on obja at voxel level

apply_obj_z_recenter(model, niter)

Apply object z-recentering along depth dimension

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

Apply position recentering constraint on probe positions

apply_tilt_smooth(model, niter)

Apply Gaussian blur to object tilts

forward(model, niter)

Applies constraints to the optimizable model parameters if niter satisfies the pre-determined conditions (start_iter, step, end_iter)

ptyrad.core.constraints.sort_by_mode_int(modes)

ptyrad.core.constraints.orthogonalize_modes_vec(modes, sort=False)

Orthogonalize probe modes via Gram matrix eigendecomposition.

Uses eigh (Hermitian eigensolver) for cross-platform numerical stability. eigh dispatches to cheev/dsyev on all LAPACK backends (Accelerate on macOS, MKL on Windows/Linux), unlike eig which used the numerically weaker cgeev. A is upcasted to complex128 for the small (Nmode x Nmode) decomposition to further guard against float32 precision loss on any backend. Falls back silently to the original modes if the result is invalid.

Note:

• MPS does not support complex128 or eigh; the entire double-precision

computation is moved to CPU, which routes through the stable Hermitian LAPACK path on macOS. Results are cast back to the original device/dtype only after validation.

• The use of eigh and A = 0.5 * (A + A.conj().T) are suggested in PR #34 by

@SoverHHH, @EdwardPooh, and @dong-zehao

• This is a highly vectorized PyTorch implementation of ptycho\+core\probe_modes_ortho.m

from PtychoShelves. The expected shape of modes input is (pmode, Ny, Nx) to be consistent with ptyrad.

• Matlab’s dot(p2,p1) for complex input would implictly apply with the complex conjugate,

so Matlab’s dot() != torch.dot because torch.dot doesn’t automatically apply the complex conjugate. This is pointed out by @dong-zehao in issue #11.

ptyrad.core.constraints.dct_threshold_filter(x, threshold_ratio=0.05)

Applies hard-threshold filtering in the DCT domain.

Keeps only the largest-magnitude DCT coefficients according to threshold_ratio and zeros out the rest.

Works for any input shape (…, H, W).

Parameters:

• x (torch.Tensor) – Input real-valued tensor of shape (…, H, W).

• threshold_ratio (float) – Fraction of coefficients to keep. Must be in [0, 1]. For example, 0.05 retains the top 5% coefficients.

Returns:

Filtered tensor of same shape as x.

Return type:

torch.Tensor

Raises:

ValueError – If threshold_ratio is not in [0, 1].

ptyrad.core.constraints.kr_filter(obj, radius, width)

Apply kr_filter using the 2D sigmoid filter

ptyrad.core.constraints.kz_filter(obj, beta_regularize_layers=1, alpha_gaussian=1, obj_type='phase')

Apply kz_filter using the arctan filter

ptyrad.core.constraints.get_obj_z_shift(obj_phase, threshold=95, scale=1, max_shift=10)

Compute z-shift from the center-of-mass (CoM) of the object phase.

Parameters:

• obj_phase – tensor (omode, z, y, x), phase values in radians

• threshold – threshold factor used to remove weak intensities in the image

• scale – scaling factor applied to the measured shift

• max_shift – maximum allowed shift in pixels

Returns:

float, signed shift (positive = shift down in z)

ptyrad.core.constraints.shift_obj_along_z(objc, z_shift)

Apply a subpixel shift along z using Fourier shift theorem.

Parameters:

• objc – tensor (omode, z, y, x), complex

• z_shift – float, shift in pixels along +z direction

Returns:

shifted tensor, same shape

ptyrad.core.constraints.complex_ratio_constraint(model, alpha1, alpha2)