pypty.utils

fourier_clean_3d

Apply a 3D Fourier filter to the input array.

PARAMETER	DESCRIPTION
array	Input 3D array to be filtered. TYPE: array_like
cutoff	Cutoff frequency (default is 0.66). TYPE: float DEFAULT: 0.66
mask	Predefined mask to apply. If None, a mask is generated. TYPE: array_like or None DEFAULT: None
rolloff	Rolloff parameter for smoothing the mask (default is 0). TYPE: float DEFAULT: 0
default_float	Data type for computations (default is cp.float32). TYPE: data - type DEFAULT: float32
xp	Array module (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
array_like	The filtered array after applying the Fourier filter.

fourier_clean

Apply a Fourier filter to the input array. Supports 2D or 3D arrays.

PARAMETER	DESCRIPTION
array	Input array (2D or 3D) to be filtered. TYPE: array_like
cutoff	Cutoff frequency (default is 0.66). TYPE: float DEFAULT: 0.66
mask	Predefined mask to apply. If None, a mask is generated. TYPE: array_like or None DEFAULT: None
rolloff	Rolloff parameter for smoothing the mask (default is 0). TYPE: float DEFAULT: 0
default_float	Data type for computations (default is cp.float32). TYPE: data - type DEFAULT: float32
xp	Array module (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
array_like	The filtered array after applying the Fourier filter.

create_spatial_frequencies

Generate spatial frequency grids and corresponding masks for multislice simulations.

PARAMETER	DESCRIPTION
px	Pixel size in the x-direction. TYPE: float
py	Pixel size in the y-direction. TYPE: float
shape	Size of the grid. TYPE: int
damping_cutoff_multislice	Damping cutoff factor for multislice simulations. TYPE: float
smooth_rolloff	Smoothing rolloff parameter. TYPE: float
default_float	Data type for computations. TYPE: data - type
xp	Array module, e.g., numpy or cupy (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
tuple	Tuple containing: - q2: 2D array of squared spatial frequencies. - qx: 2D array of spatial frequencies in x. - qy: 2D array of spatial frequencies in y. - exclude_mask: Mask in Fourier space. - exclude_mask_ishift: Unshifted mask.

shift_probe_fourier

Shift a probe in Fourier space by applying a phase ramp.

PARAMETER	DESCRIPTION
probe	The input probe array. TYPE: array_like
shift_px	Shift in pixels (y, x). TYPE: tuple of float

RETURNS	DESCRIPTION
tuple	Tuple containing the shifted probe, the phase mask, the Fourier transform of the probe, and the spatial frequency grids (maskx, masky).

generate_mask_for_grad_from_pos

Construct a binary mask from given positions and footprint dimensions.

PARAMETER	DESCRIPTION
shapex	Width of the mask. TYPE: int
shapey	Height of the mask. TYPE: int
positions_list	List of (y, x) positions where the mask should be activated. TYPE: list of tuple
shape_footprint_x	Footprint width. TYPE: int
shape_footprint_y	Footprint height. TYPE: int
shrink	Shrink factor to adjust the footprint (default is 0). TYPE: int DEFAULT: 0

RETURNS	DESCRIPTION
array_like	The constructed binary mask.

complex_grad_to_phase_grad

Convert a Wirtinger derivative to the gradient with respect to the phase.

PARAMETER	DESCRIPTION
grad	The Wirtinger derivative (dL/dz*). TYPE: array_like
array	The complex array (z = |z| exp(i*phase)). TYPE: array_like

RETURNS	DESCRIPTION
array_like	The phase gradient (dL/dp).

complex_grad_to_phase_abs_grad

Compute the phase gradient and negative amplitude gradient from a complex gradient.

PARAMETER	DESCRIPTION
grad	The Wirtinger derivative (dL/dz*). TYPE: array_like
array	The complex array (z = exp(-a + i*phase)). TYPE: array_like
xp	Array module, e.g., numpy or cupy (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
tuple of array_like	A tuple containing: - Phase gradient (dL/dp). - Negative amplitude gradient (dL/da).

complex_grad_to_mag_grad

Calculate a magnitude gradient from a complex gradient and separate magnitude and phase arrays.

PARAMETER	DESCRIPTION
grad	The complex gradient. TYPE: array_like
abs	The magnitude array. TYPE: array_like
phase	The phase array. TYPE: array_like
xp	Array module, e.g., numpy or cupy (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
array_like	The magnitude gradient.

wolfe_1

Check the Armijo condition (Wolfe condition 1) for line search.

PARAMETER	DESCRIPTION
value	The current function value. TYPE: float
new_value	The function value after the step. TYPE: float
d_value	The directional derivative at the current point. TYPE: float
step	Step size. TYPE: float
wolfe_c1	Armijo condition constant (default is 0.5). TYPE: float DEFAULT: 0.5

RETURNS	DESCRIPTION
bool	True if the condition is satisfied, False otherwise.

wolfe_2

Check the curvature condition (Wolfe condition 2) for line search.

PARAMETER	DESCRIPTION
d_value	The directional derivative at the current point. TYPE: float
new_d_value	The directional derivative after the step. TYPE: float
wolfe_c2	Curvature condition constant (default is 0.9). TYPE: float DEFAULT: 0.9

RETURNS	DESCRIPTION
bool	True if the condition is satisfied, False otherwise.

upsample_something_3d

Upsample a 3D array along the last two axes.

PARAMETER	DESCRIPTION
something	The 3D array to be upsampled. TYPE: ndarray
upsample	Upsampling factor. TYPE: int
scale	If True, scales the upsampled result to conserve total sum (default is True). TYPE: bool DEFAULT: True
xp	Array module, e.g., numpy or cupy (default is numpy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
ndarray	The upsampled 3D array.

downsample_something_3d

Downsample a 3D array along the last two axes.

PARAMETER	DESCRIPTION
something	The 3D array to be downsampled. TYPE: ndarray
upsample	Downsampling factor. TYPE: int
xp	Array module, e.g., numpy or cupy. TYPE: module

RETURNS	DESCRIPTION
ndarray	The downsampled 3D array.

upsample_something

Upsample a 2D array.

PARAMETER	DESCRIPTION
something	The 2D array to be upsampled. TYPE: ndarray
upsample	Upsampling factor. TYPE: int
scale	If True, scales the result to conserve total sum (default is True). TYPE: bool DEFAULT: True
xp	Array module (default is numpy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
ndarray	The upsampled array.

downsample_something

Downsample a 2D array.

PARAMETER	DESCRIPTION
something	The 2D array to be downsampled. TYPE: ndarray
upsample	Downsampling factor. TYPE: int
xp	Array module, e.g., numpy or cupy. TYPE: module

RETURNS	DESCRIPTION
ndarray	The downsampled array.

get_cupy_memory_usage

Print current CuPy GPU memory usage statistics.

RETURNS	DESCRIPTION
None

get_compute_batch

Estimate the optimal compute batch size based on GPU memory usage.

PARAMETER	DESCRIPTION
compute_batch	Initial guess or default. TYPE: int
load_one_by_one	Whether data is streamed instead of fully loaded. TYPE: bool
hist_size	History size for optimizers. TYPE: int
measured_data_shape	Shape of the input dataset. TYPE: tuple
memory_saturation	Proportion of GPU memory to use. TYPE: float
smart_memory	User-provided memory strategy. TYPE: callable or bool
data_pad	Padding applied to data. TYPE: int
obj_shape	Shape of the object array. TYPE: tuple
probe_shape	Shape of the probe array. TYPE: tuple
dtype	Data type string ('single' or 'double'). TYPE: str
propmethod	Propagation method name. TYPE: str
print_flag	Verbosity. TYPE: int

RETURNS	DESCRIPTION
tuple	Suggested batch size, load_one_by_one flag, and memory strategy.

preprocess_dataset

Apply preprocessing steps to the dataset including shifting, binning, padding, and scaling.

PARAMETER	DESCRIPTION
dataset	The input dataset. TYPE: ndarray
load_one_by_one	Whether data is loaded incrementally. TYPE: bool
algorithm_type	Type of reconstruction algorithm. TYPE: str
recon_type	Type of reconstruction (e.g., near_field, far_field). TYPE: str
data_shift_vector	Vector indicating pixel shift in y and x. TYPE: list of int
data_bin	Binning factor. TYPE: int
data_pad	Padding size. TYPE: int
upsample_pattern	Upsampling factor for the pattern. TYPE: int
data_multiplier	Factor to scale data intensity. TYPE: float
xp	Array module, e.g., numpy or cupy. TYPE: module
force_pad	If True, apply forced padding. TYPE: bool

RETURNS	DESCRIPTION
tuple	Tuple containing: - preprocessed dataset - data_shift_vector - data_bin - data_pad - data_multiplier

get_window

Create a circular cosine-tapered window mask.

PARAMETER	DESCRIPTION
shape	Size of the square window. TYPE: int
r0	Inner radius where tapering begins (normalized). TYPE: float
r_max	Outer radius where mask falls to zero (normalized). TYPE: float
inverted	If True, returns 1 - mask (default is True). TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
ndarray	A 2D mask array of the specified shape.

create_static_background_from_nothing

Generate an initial static background if none is provided.

PARAMETER	DESCRIPTION
static_background	Initial static background value or None. TYPE: float or ndarray
probe	Probe wavefunction. TYPE: ndarray
damping_cutoff_multislice	Maximum spatial frequency used. TYPE: float
data_pad	Padding to be applied. TYPE: int
upsample_pattern	Upsampling factor used. TYPE: int
default_float	Data type for output. TYPE: dtype
recon_type	Type of reconstruction ('near_field' or 'far_field'). TYPE: str

RETURNS	DESCRIPTION
ndarray	Initialized static background.

create_probe_from_nothing

Generate an initial probe guess when no valid probe is provided.

Depending on the input, this function either uses an aperture mask, computes a mean pattern from the dataset, or adjusts an existing mean pattern to generate a probe. It applies shifting, binning, padding, and scaling to produce a probe suitable for the specified reconstruction type.

PARAMETER	DESCRIPTION
probe	Input probe. If set to "aperture", the aperture mask is used. If None, the probe is generated based on the mean pattern. TYPE: ndarray, str, or None
data_pad	Padding size applied to the data. TYPE: int
mean_pattern	Mean pattern used to generate the probe if no probe is provided. TYPE: ndarray or None
aperture_mask	Aperture mask used when probe is set to "aperture". TYPE: ndarray
tilt_mode	Flag indicating if tilt mode is active. TYPE: bool
tilts	Tilt values. TYPE: ndarray
dataset	Measured dataset. TYPE: ndarray
estimate_aperture_based_on_binary	Factor used to threshold the dataset for aperture estimation. TYPE: bool or float
pixel_size_x_A	Pixel size in the x-direction in angstroms. TYPE: float
acc_voltage	Acceleration voltage in keV. TYPE: float
data_multiplier	Factor to scale the data intensity. TYPE: float
masks	Optional masks to apply to the mean pattern. TYPE: ndarray or None
data_shift_vector	Vector indicating the shift to be applied to the data. TYPE: list or tuple of int
data_bin	Binning factor. TYPE: int
upsample_pattern	Upsampling factor applied to the pattern. TYPE: int
default_complex_cpu	Complex data type for CPU computations. TYPE: dtype
print_flag	Flag controlling verbosity. TYPE: int
algorithm	Identifier for the reconstruction algorithm. TYPE: str
measured_data_shape	Shape of the measured data. TYPE: tuple
n_obj_modes	Number of object modes. TYPE: int
probe_marker	Marker array for probe scenarios. TYPE: ndarray or None
recon_type	Type of reconstruction ("near_field" or "far_field"). TYPE: str
defocus_array	Array of defocus values. TYPE: ndarray
Cs	Spherical aberration coefficient. TYPE: float
skip_preprocessing	Default False. If True, preprocessing steps (rescaling) will be skipped

RETURNS	DESCRIPTION
ndarray	The generated probe as a complex array.

generate_hermite_modes

Generate Hermite polynomial-based probe modes from a main mode.

PARAMETER	DESCRIPTION
main_mode	The main probe mode. TYPE: ndarray
n_herm_x	Max Degree of Hermite polynomials in x. TYPE: int
n_herm_y	Max Degree of Hermite polynomials in y. TYPE: int
default_complex	Complex data type to use. TYPE: dtype
xp	Array module, e.g., numpy or cupy TYPE: module

RETURNS	DESCRIPTION
ndarray	Stack of Hermite-based probe modes.

apply_probe_modulation

Apply defocus, aberrations, Hermite mode generation, and other modulations to the probe.

PARAMETER	DESCRIPTION
probe	Initial probe. TYPE: ndarray
extra_probe_defocus	Defocus distance to apply. TYPE: float
acc_voltage	Accelerating voltage in keV. TYPE: float
pixel_size_x_A	Pixel size in x (Å). TYPE: float
pixel_size_y_A	Pixel size in y (Å). TYPE: float
aberrations	List of aberration coefficients. TYPE: list or ndarray
print_flag	Whether to print info. TYPE: bool
beam_ctf	Optional beam CTF to apply. TYPE: ndarray or None
n_hermite_probe_modes	Number of Hermite modes in (y, x). TYPE: tuple or None
defocus_spread_modes	Defocus values to generate additional modes. TYPE: list or None
probe_marker	Probe assignment array for multi-scenario. TYPE: ndarray or None
default_complex	Complex type. TYPE: dtype
default_float	Float type. TYPE: dtype
xp	Array module, e.g., numpy or cupy TYPE: module

RETURNS	DESCRIPTION
ndarray	Modulated probe array.

resample_slices

Resample the object along the slice (z) axis and optionally pad slices.

This utility optionally prepends/appends empty slices (N_before, N_after). If refactor_thickness != 1, the slice sampling is refined by linear interpolation along z (requires homogeneous slice_distances). After resampling, the object is adjusted by a power-law to keep the effective thickness scaling.

PARAMETER	DESCRIPTION
o	Object array of shape (Ny, Nx, Nz, Nm). TYPE: ndarray
params	Parameter dictionary containing at least "slice_distances". TYPE: dict
refactor_thickness	Factor by which to refine slice sampling (default is 1). Values > 1 increase the number of slices. TYPE: float DEFAULT: 1
N_before	Number of empty slices to prepend (default is 0). TYPE: int DEFAULT: 0
N_after	Number of empty slices to append (default is 0). TYPE: int DEFAULT: 0

RETURNS	DESCRIPTION
o	Resampled/padded object array. TYPE: ndarray
params	Updated parameter dictionary. Keys updated: slice_distances total_thickness obj TYPE: dict

RAISES	DESCRIPTION
ValueError	If refactor_thickness != 1 but slice_distances are not homogeneous.

probe_modes_from_spreads

Generate a coherent-mode expansion of a partially coherent probe via MC sampling + SVD.

The probe is perturbed by random transverse shifts (spatial coherence) and defocus offsets (focal spread) in reciprocal space. The resulting ensemble is reshaped into a matrix and decomposed with SVD; the leading left singular vectors are returned as coherent modes in real space.

PARAMETER	DESCRIPTION
probe	Probe array. Expected shape is (Ny, Nx, 1) (real-space). Only the first mode is used as the starting point. TYPE: ndarray or ndarray
wavelength	Electron wavelength in Å. TYPE: float
rez_px_size_rad	Reciprocal-space sampling in radians per pixel. TYPE: float
focal_spread	Standard deviation of defocus spread (Å). TYPE: float
spatial_coh_x_A	RMS spatial coherence in x (Å). Default is 0. TYPE: float DEFAULT: 0
spatial_coh_y_A	RMS spatial coherence in y (Å). Default is 0. TYPE: float DEFAULT: 0
spatial_coh_xy	Correlation coefficient between x/y coherence in [-1, 1]. Default is 0. TYPE: float DEFAULT: 0
n_modes	Number of coherent modes to return (default is 4). TYPE: int DEFAULT: 4
n_samples	Number of Monte-Carlo samples used to build the ensemble (default is 256). TYPE: int DEFAULT: 256
seed	Random seed for sampling (default is 0). TYPE: int DEFAULT: 0
data_pad	Padding (pixels) to crop before SVD and then restore afterwards for far-field reconstructions (default is 0). TYPE: int DEFAULT: 0
recon_type	If not "near_field", the function optionally crops/pads the probe in reciprocal space by data_pad (default is "far_field"). TYPE: (far_field, near_field) DEFAULT: "far_field"
xp	Array module, typically numpy or cupy. TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
modes_r	Real-space coherent probe modes of shape (Ny, Nx, n_modes). TYPE: ndarray or ndarray

Notes

The mode weights are absorbed into the returned amplitudes; i.e., the returned modes are already scaled by sqrt(lambda_i).

convert_num_to_nmab

Convert a number of aberration terms to (n, m, ab) indices based on Krivanek notation.

PARAMETER	DESCRIPTION
num_abs	Number of aberration coefficients. TYPE: int

RETURNS	DESCRIPTION
tuple of lists	Lists of n, m, and ab strings ('', 'a', or 'b') for each aberration mode.

nmab_to_strings

Convert aberration indices into string identifiers in Krivanek notation.

PARAMETER	DESCRIPTION
possible_n	List of radial indices. TYPE: list of int
possible_m	List of azimuthal indices. TYPE: list of int
possible_ab	List of aberration mode types ('', 'a', 'b'). TYPE: list of str

RETURNS	DESCRIPTION
list of str	List of formatted aberration identifiers like 'C10', 'C12a', etc.

get_ctf_matrix

Generate a matrix of phase contributions for all aberration modes.

PARAMETER	DESCRIPTION
kx	Spatial frequency in x-direction. TYPE: ndarray
ky	Spatial frequency in y-direction. TYPE: ndarray
num_abs	Number of aberration coefficients. TYPE: int
wavelength	Electron wavelength. TYPE: float
xp	backend (default is cupy). TYPE: module DEFAULT: numpy

RETURNS	DESCRIPTION
ndarray	list of Zernike polynomials (num_abs, height, width) with phase contributions.

get_ctf

Compute the scalar contrast transfer function (CTF) from aberrations.

PARAMETER	DESCRIPTION
aberrations	List of aberration coefficients. TYPE: list or ndarray
kx	Spatial frequency in x-direction. TYPE: ndarray
ky	Spatial frequency in y-direction. TYPE: ndarray
wavelength	Electron wavelength. TYPE: float
angle_offset	Additional rotation angle in radians (default is 0). TYPE: float DEFAULT: 0

RETURNS	DESCRIPTION
ndarray	The computed CTF.

get_ctf_derivatives

Compute spatial derivatives of the CTF with respect to kx and ky.

PARAMETER	DESCRIPTION
aberrations	List of aberration coefficients. TYPE: list or ndarray
kx	Spatial frequency in x-direction. TYPE: ndarray
ky	Spatial frequency in y-direction. TYPE: ndarray
wavelength	Electron wavelength. TYPE: float
angle_offset	Additional rotation angle (default is 0). TYPE: float DEFAULT: 0

RETURNS	DESCRIPTION
tuple of ndarray	Derivatives of CTF with respect to kx and ky.

get_ctf_gradient_rotation_angle

Compute the gradient of the phase with respect to rotation angle.

PARAMETER	DESCRIPTION
aberrations	List of aberration coefficients. TYPE: list or ndarray
kx	Spatial frequency in x-direction. TYPE: ndarray
ky	Spatial frequency in y-direction. TYPE: ndarray
wavelength	Electron wavelength. TYPE: float
angle_offset	Additional angular offset (default is 0). TYPE: float DEFAULT: 0

RETURNS	DESCRIPTION
tuple of ndarray	Derivatives of the CTF gradient in x and y directions with respect to angular change.

apply_defocus_probe

Apply a defocus phase shift to a probe in Fourier space.

PARAMETER	DESCRIPTION
probe	The input probe wavefunction. TYPE: ndarray
distance	Defocus distance in meters. TYPE: float
acc_voltage	Acceleration voltage in kiloelectronvolts (keV). TYPE: float
pixel_size_x_A	Pixel size along x in angstroms. TYPE: float
pixel_size_y_A	Pixel size along y in angstroms. TYPE: float
default_complex	Complex data type for computation. TYPE: dtype
default_float	Float data type for computation. xp : module, optional Array module, e.g., numpy or cupy TYPE: dtype

RETURNS	DESCRIPTION
ndarray	The defocused probe.

padfft

Pad the input array in Fourier space by padding its FFT.

PARAMETER	DESCRIPTION
array	Input array to be padded. TYPE: ndarray
pad	Number of pixels to pad on each side. TYPE: int

RETURNS	DESCRIPTION
ndarray	The padded array in spatial domain.

padprobetodatafarfield

Pad or crop a probe in Fourier space to match far-field data dimensions.

PARAMETER	DESCRIPTION
probe	The probe wavefunction. TYPE: ndarray
measured_data_shape	Shape of the measured data. TYPE: tuple
data_pad	Padding applied to the data. TYPE: int
upsample_pattern	Upsampling factor used in the reconstruction. TYPE: int

RETURNS	DESCRIPTION
ndarray	Adjusted probe wavefunction.

padprobetodatanearfield

Pad or crop a probe for near-field reconstruction.

This function adjusts the probe wavefunction by padding or cropping it to match the near-field measured data dimensions after upsampling and padding.

PARAMETER	DESCRIPTION
probe	The input probe wavefunction. TYPE: ndarray
measured_data_shape	Shape of the measured data. TYPE: tuple
data_pad	Padding size applied to the data. TYPE: int
upsample_pattern	Upsampling factor applied to the measured data. TYPE: int

RETURNS	DESCRIPTION
ndarray	The adjusted probe wavefunction.

prepare_main_loop_params

Prepare main loop parameters for reconstruction.

This function adjusts scan positions, pads the object if necessary, handles subpixel corrections, and computes the electron wavelength based on the accelerating voltage.

PARAMETER	DESCRIPTION
algorithm	Identifier for the reconstruction algorithm. TYPE: any
probe	The probe array. TYPE: ndarray
obj	The object array. TYPE: ndarray
positions	Array of scan positions. TYPE: ndarray
tilts	Array of tilt angles. TYPE: ndarray
measured_data_shape	Shape of the measured data. TYPE: tuple
acc_voltage	Accelerating voltage in keV. TYPE: float
allow_subPixel_shift	If True, compute subpixel corrections (default is True). TYPE: bool DEFAULT: True
sequence	Sequence of indices for positions (default is None, which uses full range). TYPE: list or callable DEFAULT: None
use_full_FOV	If True, use full field-of-view adjustments (default is False). TYPE: bool DEFAULT: False
print_flag	Verbosity flag (default is 0). TYPE: int DEFAULT: 0
default_float_cpu	Float data type for CPU computations (default is np.float64). TYPE: data - type DEFAULT: float64
default_complex_cpu	Complex data type for CPU computations (default is np.complex128). TYPE: data - type DEFAULT: complex128
default_int_cpu	Integer data type for CPU computations (default is np.int64). TYPE: data - type DEFAULT: int64
probe_constraint_mask	Optional mask for probe constraints. TYPE: ndarray or None DEFAULT: None
aperture_mask	Optional aperture mask. TYPE: ndarray or None DEFAULT: None
extra_space_on_side_px	Padding to add around the object. Default is 0. TYPE: int DEFAULT: 0
ignore_positions	If True, skip position adjustments. Default is False. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
obj	Padded object array. TYPE: ndarray
positions	Adjusted integer scan positions. TYPE: ndarray
zero_placeholder	Placeholder zero (reserved for future use). TYPE: int
sequence	Sequence of indices used for processing. TYPE: list
wavelength	Computed electron wavelength in angstroms. TYPE: float
full_pos_correction	Sub-pixel corrections for scan positions. TYPE: ndarray
tilts_correction	Zero array matching tilts shape. TYPE: ndarray
aperture_mask	Selected probe constraint or aperture mask. TYPE: ndarray or None

try_to_initialize_beam_current

Initialize beam current array or pad if it's too short.

PARAMETER	DESCRIPTION
beam_current	Existing beam current values. TYPE: ndarray or None
measured_data_shape	Shape of measured dataset. TYPE: tuple
default_float	Float type for the array. TYPE: dtype
xp	Array module, e.g., numpy or cupy TYPE: module

RETURNS	DESCRIPTION
ndarray	Initialized or padded beam current.

save_updated_arrays

Save current reconstruction state and log loss metrics during training.

This function saves checkpoints for object, probe, tilts, scan positions, static background, aberration coefficients, and beam current if specified. It also logs loss and constraint contributions in a CSV file if logging is enabled.

PARAMETER	DESCRIPTION
output_folder	Directory where files will be saved. TYPE: str
epoch	Current epoch number. TYPE: int
current_probe_step	Whether to save the current probe. TYPE: bool
current_probe_pos_step	Whether to save current scan positions. TYPE: bool
current_tilts_step	Whether to save current tilts. TYPE: bool
current_obj_step	Whether to save the current object. TYPE: bool
obj	Object array to save. TYPE: ndarray or ndarray
probe	Probe array to save. TYPE: ndarray or ndarray
tilts_correction	Tilt correction values. TYPE: ndarray
full_pos_correction	Sub-pixel scan position correction. TYPE: ndarray
positions	Integer scan positions. TYPE: ndarray
tilts	Original tilt values. TYPE: ndarray
static_background	Static background array. TYPE: ndarray or ndarray
current_aberrations_array_step	Whether to save aberration array. TYPE: bool
current_static_background_step	Whether to save static background. TYPE: bool
count	Not used inside the function. TYPE: int
current_loss	Current loss value. TYPE: float
current_sse	Current sum of squared errors. TYPE: float
aberrations	Array of aberration coefficients. TYPE: ndarray or ndarray
beam_current	Array of beam current values. TYPE: ndarray or ndarray
current_beam_current_step	Whether to save beam current. TYPE: bool
save_flag	Whether to trigger checkpoint saving. TYPE: bool
save_loss_log	Whether to log loss. If set to 2, log full breakdown of constraints. TYPE: bool or int
constraint_contributions	List of constraint term contributions to the loss. TYPE: list
actual_step	Step size applied in the optimizer. TYPE: float
count_linesearch	Number of line search iterations. TYPE: int
d_value	Initial directional derivative. TYPE: float
new_d_value	New directional derivative after the step. TYPE: float
current_update_step_bfgs	Step size suggested by BFGS or optimizer. TYPE: float
t0	Start time of the epoch (used for timing). TYPE: float
xp	NumPy or CuPy module used for computation. TYPE: module
warnings	Warning string to be logged. TYPE: str

RETURNS	DESCRIPTION
None

prepare_saving_stuff

Prepare folder and loss CSV for saving training logs.

PARAMETER	DESCRIPTION
output_folder	Directory for results. TYPE: str
save_loss_log	Whether to save loss values. TYPE: bool
epoch_prev	Previous epoch index. TYPE: int

RETURNS	DESCRIPTION
None

print_pypty_header

Print formatted header announcing start of reconstruction.

PARAMETER	DESCRIPTION
data_path	Path to the dataset. TYPE: str
output_folder	Directory where results are saved. TYPE: str
save_loss_log	Whether loss logging is enabled. TYPE: bool

RETURNS	DESCRIPTION
None

save_current_checkpoint_obj_probe

Save intermediate reconstruction data as checkpoints.

This function saves the current state of the object, probe, tilt corrections, scan positions, static background, and aberrations to disk. It is intended to allow resuming reconstruction from the last checkpoint.

PARAMETER	DESCRIPTION
output_folder	Directory where checkpoint files will be saved. TYPE: str
obj	The current object array. TYPE: ndarray or GPU array
probe	The current probe array. TYPE: ndarray or GPU array
tilts_correction	Correction values for tilt angles. TYPE: ndarray
full_pos_correction	Sub-pixel correction values for scan positions. TYPE: ndarray
positions	Scan positions array. TYPE: ndarray
tilts	Tilt angles array. TYPE: ndarray
static_background	Static background array. TYPE: ndarray
current_probe_step	Flag indicating whether to save the probe. TYPE: bool
current_obj_step	Flag indicating whether to save the object. TYPE: bool
current_probe_pos_step	Flag indicating whether to save the scan positions. TYPE: bool
current_tilts_step	Flag indicating whether to save the tilt angles. TYPE: bool
current_static_background_step	Flag indicating whether to save the static background. TYPE: bool
current_aberrations_array_step	Flag indicating whether to save the aberrations array. TYPE: bool
aberrations_array	The current aberrations array. TYPE: ndarray or GPU array
beam_current	The current beam current array. TYPE: ndarray or GPU array or None
bcstep	Flag indicating whether to save the beam current. TYPE: bool
xp	Array module (e.g., numpy or cupy). TYPE: module

RETURNS	DESCRIPTION
None

print_recon_state

Display current reconstruction progress including loss, optimization state, and updates.

PARAMETER	DESCRIPTION
t0	Start time of the epoch (Unix timestamp). TYPE: float
algorithm	Name of the loss or optimization algorithm used. TYPE: str
epoch	Current training epoch. TYPE: int
current_loss	Loss value at current epoch. TYPE: float
current_sse	Sum of squared errors. TYPE: float
current_obj_step	Whether the object is being updated. TYPE: bool
current_probe_step	Whether the probe is being updated. TYPE: bool
current_probe_pos_step	Whether the scan grid is being updated. TYPE: bool
current_tilts_step	Whether tilt corrections are being updated. TYPE: bool
current_static_background_step	Whether static background is being updated. TYPE: bool
current_aberrations_array_step	Whether aberration coefficients are being updated. TYPE: bool
current_beam_current_step	Whether beam current is being updated. TYPE: bool
current_hist_length	Optimizer memory length (0=GD, 1=CG, >1=BFGS). TYPE: int
print_flag	Verbosity flag: 0 = silent, 1 = single-line print, 2 = verbose. TYPE: int

RETURNS	DESCRIPTION
None

try_to_gpu

Convert all key reconstruction variables to GPU arrays if using CuPy.

PARAMETER	DESCRIPTION
obj	Object array. TYPE: ndarray
probe	Probe array. TYPE: ndarray
positions	Integer scan positions. TYPE: ndarray
full_pos_correction	Sub-pixel scan grid correction. TYPE: ndarray
tilts	Tilt values. TYPE: ndarray
tilts_correction	Tilt corrections. TYPE: ndarray
masks	Optional segmentation or region masks. TYPE: ndarray or None
defocus_array	Array of defocus values per position. TYPE: ndarray
slice_distances	Slice spacing in multislice simulations. TYPE: ndarray
aperture_mask	Probe aperture mask. TYPE: ndarray or None
dataset	Measured dataset. TYPE: ndarray
load_one_by_one	Whether dataset is streamed from disk. TYPE: bool
static_background	Static background array. TYPE: ndarray or None
aberrations_array	Array of aberration coefficients. TYPE: ndarray or None
beam_current	Beam current scaling factor. TYPE: ndarray or None
default_float	Float precision dtype for casting. TYPE: dtype
default_complex	Complex precision dtype for casting. TYPE: dtype
default_int	Integer dtype for casting. TYPE: dtype
xp	Array module, e.g., numpy or cupy TYPE: module

RETURNS	DESCRIPTION
tuple	The same variables in GPU format (if using CuPy), with proper types.

get_value_for_epoch

Evaluate a list of values or functions at the current epoch.

PARAMETER	DESCRIPTION
func_or_value	List of fixed values or callables. TYPE: list
epoch	Current epoch number. TYPE: int
default_float	Float precision type. TYPE: dtype

RETURNS	DESCRIPTION
list	Evaluated values.

get_steps_epoch

Evaluate step values for the current epoch.

PARAMETER	DESCRIPTION
steps	List of (multiplier, callable) or fixed values. TYPE: list
epoch	Current training epoch. TYPE: int
default_float	Float precision type. TYPE: dtype

RETURNS	DESCRIPTION
list	List of step values.

lambda_to_string

Extract lambda function source as a string.

PARAMETER	DESCRIPTION
f	Lambda function. TYPE: function

RETURNS	DESCRIPTION
str	Extracted string source of the lambda.

convert_to_string

Convert parameter dictionary to string format, including lambda serialization.

PARAMETER	DESCRIPTION
dicti2	Original parameter dictionary. TYPE: dict
strip_dataset_from_params	Whether to exclude 'dataset' key (default is True). TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
dict	Dictionary with string values.

string_to_lambda

Convert stringified lambda expression to a Python function.

PARAMETER	DESCRIPTION
lambda_string	Lambda string to evaluate. TYPE: str

RETURNS	DESCRIPTION
callable or str	The resulting function or original string if evaluation fails.

load_params

Load parameter dictionary from a .pkl file.

PARAMETER	DESCRIPTION
path	Path to the .pkl parameter file. TYPE: str

RETURNS	DESCRIPTION
dict	Loaded parameters.

string_params_to_usefull_params

Convert string-encoded lambdas in parameter dictionary back to callables.

PARAMETER	DESCRIPTION
params	Parameter dictionary possibly containing lambda strings. TYPE: dict

RETURNS	DESCRIPTION
dict	Updated dictionary with callables.

save_params

Save parameters to a .pkl file, optionally removing the dataset.

PARAMETER	DESCRIPTION
params_path	Output path for the parameter file. TYPE: str
params	Parameter dictionary to save. TYPE: dict
strip_dataset_from_params	If True, remove the dataset entry. TYPE: bool

RETURNS	DESCRIPTION
None

change_params

Update one or more entries in a saved parameter file.

PARAMETER	DESCRIPTION
path_params	Path to the .pkl parameter file. TYPE: str
keys	Key(s) to change. TYPE: str or list of str
new_values	New value(s) corresponding to keys. TYPE: any or list

RETURNS	DESCRIPTION
None

detect_changed_params

Check for manual hyperparameter edits and update in-memory list.

PARAMETER	DESCRIPTION
list_hyperparams	Current hyperparameter values. TYPE: list
keys_hyperparams	Names of hyperparameters to monitor. TYPE: list of str
path_params	Path to the .pkl parameter file. TYPE: str
warnings	Existing warning messages. TYPE: str
print_flag	If True, print a warning when changes are detected. TYPE: bool

RETURNS	DESCRIPTION
list_hyperparams	list with changed hyperparameters TYPE: list
warnings	Updated warning messages. TYPE: str

load_nexus_params

Load reconstruction parameters from a NeXus (.nxs) HDF5 file.

PARAMETER	DESCRIPTION
path_nexus	Path to the .nxs file. TYPE: str

RETURNS	DESCRIPTION
dict	Dictionary of extracted parameters.

delete_dataset_from_params

Delete the 'dataset' key from saved parameter file.

PARAMETER	DESCRIPTION
params_path	Path to the pickled parameters file. TYPE: str

RETURNS	DESCRIPTION
None

phase_cross_corr_align

Align two FFT-transformed images using phase cross-correlation.

PARAMETER	DESCRIPTION
im_ref_fft	Reference image FFT. TYPE: ndarray
im_2_fft	FFT of the image to be aligned. TYPE: ndarray
refine_box_dim	Size of the interpolation box for sub-pixel alignment. TYPE: int
upsample	Upsampling factor for interpolation. TYPE: int
x_real	Real space x grid. TYPE: ndarray
y_real	Real space y grid. TYPE: ndarray
shift_y	Predefined shift in y (optional). TYPE: float or None DEFAULT: None
shift_x	Predefined shift in x (optional). TYPE: float or None DEFAULT: None

RETURNS	DESCRIPTION
ndarray	Shifted FFT of the second image.

segment_regions_from_image

Segment regions in an image based on edge detection and boundary blocking.

This function applies Gaussian smoothing to the input image, thresholds it to detect edges, dilates the edges to thicken boundaries, and inverts the result to identify regions. Finally, it labels connected regions in the image.

PARAMETER	DESCRIPTION
image	Input 2D grayscale image. TYPE: ndarray
threshold	Intensity threshold for edge detection. TYPE: float
sigma	Standard deviation for Gaussian blur. TYPE: float DEFAULT: 1.0
dilation_size	Size of the dilation kernel for boundaries. TYPE: int

RETURNS	DESCRIPTION
labeled_regions	2D array where each region has a unique label. TYPE: ndarray
num_regions	Total number of regions. TYPE: int
edge_mask	Binary edge map. TYPE: ndarray

interpolate_label0_points

Interpolate positions for points with label == 0 using neighboring labeled samples.

The input points array stores measured coordinates (x, y) along with integer grid indices (i, j) and a region label. For entries with label == 0, x/y are interpolated from labeled entries in (i, j)-space using scipy.interpolate.griddata.

PARAMETER	DESCRIPTION
points	Array of shape (N, 5) with columns [x, y, i, j, label]. TYPE: ndarray
method	Primary interpolation method for griddata (default is "linear"). TYPE: (linear, cubic, nearest) DEFAULT: "linear"
fallback	Backup method used where the primary interpolation returns NaNs (default is "nearest"). TYPE: (nearest, None) DEFAULT: "nearest"

RETURNS	DESCRIPTION
ndarray	Copy of points with x/y updated for rows where label == 0.

optimize_2d_block_shifts

Optimize per-region (block) shifts to enforce smooth scan-grid geometry.

The function assigns an (dx, dy) shift to each non-zero label such that:

• The discrete Laplacian of the resulting 2D displacement field is minimized (smoothness / low curvature), and
• Points in each scan column/row align well to a line and preserve ordering.

Points with label == 0 are treated as unknowns and are filled by interpolate_label0_points during optimization.

PARAMETER	DESCRIPTION
points	Array of shape (N, 5) with columns [x, y, i, j, label]. TYPE: ndarray

RETURNS	DESCRIPTION
shifted_points	Array of shape (N, 5) with updated x/y after optimal block shifts. TYPE: ndarray
block_shifts	Mapping label -> (dx, dy) for each optimized block. TYPE: dict

position_puzzling

Detect and correct block-wise scan-position discontinuities by segmentation + puzzling.

The function computes a Laplacian-based score on the scan-position field. If the score exceeds score_threshold, it segments the displacement map into regions and then solves for per-region shifts using optimize_2d_block_shifts.

PARAMETER	DESCRIPTION
points	Scan positions of shape (N, 2) in (y, x) order. TYPE: ndarray or ndarray
scan_size	Scan grid size as (Ny, Nx). TYPE: tuple of int
sigma	Smoothing parameter used during segmentation (default is 0.4). TYPE: float DEFAULT: 0.4
score_threshold	Threshold on the Laplacian score above which puzzling is applied (default is 1.2). TYPE: float DEFAULT: 1.2

RETURNS	DESCRIPTION
puzzling_worked	True if segmentation found multiple regions and puzzling was applied. TYPE: bool
points_np	Corrected scan positions of shape (N, 2) (CPU array). TYPE: ndarray

fit_vector_field

Fit a smooth displacement field from calibration points.

PARAMETER	DESCRIPTION
x	Nominal and observed coordinates. TYPE: array_like
y	Nominal and observed coordinates. TYPE: array_like
u	Observed coordinates. TYPE: array_like
v	Observed coordinates. TYPE: array_like
Xi	Grid coordinates for mapping. TYPE: array_like
Yi	Grid coordinates for mapping. TYPE: array_like

RETURNS	DESCRIPTION
tuple of numpy.ndarray	Displacement fields (dV, dU).

warp_image

Warp image(s) based on a displacement field.

PARAMETER	DESCRIPTION
image	Input image(s). TYPE: array_like
vector_field	Displacement fields (U, V). TYPE: tuple of numpy.ndarray
order	Interpolation order. Default is 1. TYPE: int DEFAULT: 1
extra_angle_rotation_compensation	Rotation angle in degrees. Default is 0. TYPE: float

RETURNS	DESCRIPTION
ndarray	Warped image(s).

intorplate_puzzled_obj

Interpolate a puzzled object array using a displacement field.

PARAMETER	DESCRIPTION
array	Input object array of shape (Y, X, slices, modes). TYPE: ndarray or ndarray
pos_old	Original scan positions, shape (N, 2). TYPE: ndarray
pos_new	New scan positions, shape (N, 2). TYPE: ndarray
order	Interpolation order. Default is 1. TYPE: int DEFAULT: 1
method	Interpolation method ('linear', 'nearest', etc.). Default is 'linear'. TYPE: str DEFAULT: 'linear'
xp	Array module, e.g. numpy or cupy DEFAULT: numpy

RETURNS	DESCRIPTION
ndarray	Interpolated object array with same shape as array.