"""
Initialize measurements, object, probe, probe positions, tilts, and other variables
This module is completely using NumPy for interoperability,
so users can initialize their data with PtyRAD first and reconstruct with other approaches later.
"""
import os
from copy import deepcopy
from math import floor
import numpy as np
from scipy.io.matlab import matfile_version as get_matfile_version
from scipy.ndimage import gaussian_filter, zoom
from ptyrad.load import load_array_from_file, load_hdf5, load_mat, load_ptyrad
from ptyrad.save import save_array
from ptyrad.utils import (
complex_object_z_resample_torch,
compose_affine_matrix,
create_one_hot_mask,
exponential_decay,
fit_background,
fit_cbed_pattern,
get_default_probe_simu_params,
get_EM_constants,
get_nested,
guess_radius_of_bright_field_disk,
infer_dx_from_params,
make_fzp_probe,
make_mixed_probe,
make_stem_probe,
near_field_evolution,
orthogonalize_modes_vec_np,
power_law,
set_random_seed,
sort_by_mode_int_np,
vprint,
)
[docs]
class Initializer:
[docs]
def __init__(self, init_params, seed=None, verbose=True):
# A deepcopy creates a new object so modifying self.init_params won't affect the original init_params dict that was outside the class
# This is important because self.init_params might get updated if there's cropping, padding, or resampling of the measurements
# The original params file will could be directly saved to the output dir with `save.copy_params_to_dir`,
# while we also keep a digital copy of the original params in self.init_params_original
self.init_params = deepcopy(init_params) # This is the central params dict that will be used for initialization
self.init_params_original = deepcopy(init_params)
self.init_variables = {'random_seed': seed} # This dict stores all the variables that will be used for the later ptychography reconstruction
self.random_seed = seed
self.verbose=verbose
self.print_init_params()
##### Public methods for initializing everything #####
def print_init_params(self):
''' Print the current init_params in the Initialzier object '''
vprint("init_params are displayed below:", verbose=self.verbose)
for key, value in self.init_params.items():
vprint(f" {key}: {value}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_cache(self):
""" Check if the source paths are the same, if so, we may cache that field to reduce file loading time """
# Note:
# For caching, at least 2 out of 3 fields are using the same file path
# Therefore, there's only one possible source for the self.cache_contents
# With 2 file source posibilities, the self.cache_contents is either caching from 'PtyRAD' or 'PtyShv'
# Even we add more file type supports in the future (py4dstem or ptypy), the cache would still be a single file type
vprint("### Initializing cache ###", verbose=self.verbose)
# Initialize flags for cached fields
self.cache_source = None
self.cache_path = None
self.cache_contents = None
self.use_cached_obj = False
self.use_cached_probe = False
self.use_cached_pos = False
# Set cache_source, cache_path, and use_cached_xxx flags iteratively
for source in ['PtyRAD', 'PtyShv', 'py4DSTEM']:
self._set_use_cached_flags(source)
# Set cache_contents
if any([self.use_cached_obj, self.use_cached_probe, self.use_cached_pos]):
if self.cache_source == 'PtyRAD':
vprint(f"Loading 'PtyRAD' file from {self.cache_path} for caching", verbose=self.verbose)
self.cache_contents = load_ptyrad(self.cache_path)
elif self.cache_source == 'PtyShv':
vprint(f"Loading 'PtyShv' file from {self.cache_path} for caching", verbose=self.verbose)
self.cache_contents = load_mat(self.cache_path, key=['object', 'probe', 'outputs.probe_positions'], delimiter='.') # flattend dict with key using delimiter
elif self.cache_source == 'py4DSTEM':
vprint(f"Loading 'py4DSTEM' file from {self.cache_path} for caching", verbose=self.verbose)
self.cache_contents = load_hdf5(self.cache_path, key=None)
else:
raise ValueError(f"File type {source} not implemented for caching yet, please use 'PtyRAD', or 'PtyShv'!")
# Cache is only used when 2 out of 3 fields have the same source and path, so the following flags could only be all false, 1 false 2 true, or 3 true.
vprint(f"use_cached_obj = {self.use_cached_obj}", verbose=self.verbose)
vprint(f"use_cached_probe = {self.use_cached_probe}", verbose=self.verbose)
vprint(f"use_cached_pos = {self.use_cached_pos}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_measurements(self):
vprint("### Initializing measurements ###", verbose=self.verbose)
meas = self._load_meas()
meas = self._process_meas(meas)
meas_avg = meas.mean(0) # This is equivalent to PACBED in electron microscopy. Note that if pad/resample are set to "on_the_fly", this would be different from the final one used for reconstruction.
meas_avg_sum = meas_avg.sum() # This is the total integrated intensity of the averaged diffraction pattern
pad_mode = get_nested(self.init_params, key=['meas_pad', 'mode'], safe=True, default=None)
if pad_mode == 'on_the_fly':
padded = self.init_variables.get('on_the_fly_meas_padded')
padded_int_sum = padded.sum() if padded is not None else 0
vprint(f"Adjusting `meas_avg_sum` by adding {padded_int_sum:.4f} for on_the_fly meas padding", verbose=self.verbose)
meas_avg_sum += padded_int_sum # meas_avg_sum is used to normalize the probe intensity.
# Because the meas could gain intensity during on_the_fly padding,
# we need to consider the extra intensity from the padded region here.
self.init_variables['meas_avg'] = meas_avg
self.init_variables['meas_avg_sum'] = meas_avg_sum
self.init_variables['measurements'] = meas
export_params = self.init_params.get('meas_export') # Ture, False, None, dict (could be {})
if export_params is True or isinstance(export_params, dict):
vprint(f"Exporting measurements with `meas_export` = {export_params}")
self._export_meas(export_params if isinstance(export_params, dict) else {})
# Print out some measurements statistics
vprint(f"meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.4f}, {meas.max():.4f})", verbose=self.verbose)
vprint(f"measurements (N, Ky, Kx) = {meas.dtype}, {meas.shape}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_calibration(self):
vprint("### Setting up calibration ###", verbose=self.verbose)
calib_dict = self.init_params['meas_calibration']
calib_mode = calib_dict['mode'] # One of 'dx', 'dk', 'kMax', 'da', 'angleMax', 'RBF', 'n_alpha', or 'fitRBF'
calib_value = calib_dict.get('value') # fitRBF doesn't need a value here
Npix = self.init_params_original.get('meas_Npix') # Load the original Npix because init_params['meas_Npix'] could have been modified in init_measurements
conv_angle = self.init_params.get('probe_conv_angle')
illum_type = self.init_params.get('probe_illum_type') or 'electron'
vprint(f"meas_calibration mode = '{calib_mode}', value = {calib_value}", verbose=self.verbose) # No need to add :.4f to value because it could be None, also it's user input so won't have too many digits
# Load the meas_raw_avg first to ensure measurement is initialized
try:
meas_raw_avg = self.init_variables['meas_raw_avg'] # This is the averaged measurements with only simple permuting/reshaping/flipping
except KeyError:
vprint("Warning: 'init_variables['meas_raw_avg]' not found. Initializing measurements first for calibration...", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
self.init_measurements()
meas_raw_avg = self.init_variables['meas_raw_avg']
if illum_type == 'electron':
# Get wavelength
energy = self.init_params.get('probe_kv') # kV
wavelength = get_EM_constants(energy, output_type='wavelength') # wavelength in Ang
unit_str = 'Ang'
# Run fitRBF routine for electron ptychography
vprint("Using loaded raw averaged measurement (before crop/pad/resample) to fit RBF as a part of the meas calibration", verbose=self.verbose)
fitRBF = guess_radius_of_bright_field_disk(meas_raw_avg, thresh=calib_dict.get('thresh', 0.5))
vprint(f"Radius of fitted bright field disk (RBF) = {fitRBF:.2f} px with Npix = {meas_raw_avg.shape[-1]}", verbose=self.verbose)
vprint(f"Suggested probe_mask_k radius (RBF*2/Npix) > {(fitRBF * 2 / Npix):.4f}", verbose=self.verbose)
vprint("Fitting raw averaged measurement with center, radius, and Gaussian blur std as a sanity check", verbose=self.verbose)
vprint("Note that the fitted Gaussian blur std (detector blur) would be affected by overlapping Bragg disks", verbose=self.verbose)
_ = fit_cbed_pattern(meas_raw_avg, verbose=self.verbose)
# Actually calculating dx for each calib_mode
if calib_mode == 'fitRBF':
dx = infer_dx_from_params(**{'RBF': fitRBF, 'Npix': Npix, 'wavelength': wavelength, 'conv_angle': conv_angle})
else:
dx = infer_dx_from_params(**{calib_mode: calib_value, 'Npix': Npix, 'wavelength': wavelength, 'conv_angle': conv_angle})
if calib_mode != 'RBF':
inferRBF = conv_angle / 1e3 * Npix * dx / wavelength # We can still infer RBF using the user provided calib value
vprint(f"Using init_params, the inferred RBF (conv_angle / 1e3 * Npix * dx / wavelength) = {inferRBF:.2f} px with Npix = {meas_raw_avg.shape[-1]}", verbose=self.verbose)
if calib_mode in ['fitRBF', 'RBF']:
vprint("WARNING: The 'fitRBF' and 'RBF' calibration methods are highly dependent on the accuracy of user-provided experimental parameters and acquisition conditions,",verbose=self.verbose)
vprint(" including convergence angle, kV, dose, specimen thickness, and collection angle for the estimation of RBF.", verbose=self.verbose)
vprint(" For example, a 5-10% error in convergence angle is fairly common.", verbose=self.verbose)
vprint(" Users are strongly advised to perform proper microscope calibration to ensure accurate results.", verbose=self.verbose)
vprint(" These method should only be used as a rough estimate and not as a substitute for proper experimental calibration.", verbose=self.verbose)
elif illum_type == 'xray':
if calib_mode in ['RBF', 'fitRBF', 'n_alpha']:
raise ValueError(f"Calibration mode '{calib_mode}' is not supported for xray. Use 'dx', 'dk', 'kMax', 'da', 'angleMax'.")
# Get wavelength
energy = self.init_params.get('beam_kev') # keV
wavelength = 1.23984193e-9 / energy # wavelength in m, energy in keV
unit_str = 'm'
# Infer dx calibration from provided values
dx = infer_dx_from_params(**{calib_mode: calib_value,
'Npix': Npix,
'wavelength': wavelength})
else:
raise ValueError(f"'probe_illum_type' = {illum_type} not implemented yet, please use either 'electron' or 'xray'!")
# Print the information
vprint(f"dx (real space pixel size of probe and object) set to {dx:.4f} {unit_str} with Npix = {meas_raw_avg.shape[-1]}", verbose=self.verbose)
Npix_is_modified = False
# Handle additional changes to dx if there's meas_crop
crop_ranges = self.init_params.get('meas_crop')
if crop_ranges is not None and len(crop_ranges) == 4:
if crop_ranges[-1] is not None and len(crop_ranges[-1]) == 2:
kx_i, kx_f = crop_ranges[-1]
Npix_new = kx_f - kx_i
dx = dx * Npix / Npix_new
Npix_is_modified = True
Npix_modified = Npix_new
vprint(f"Update dx to {dx:.4f} {unit_str} due to meas_crop, Npix = {Npix_modified}", verbose=self.verbose)
if illum_type == 'electron':
vprint(f"Suggested probe_mask_k radius (RBF*2/Npix) changes to > {(fitRBF * 2 / Npix_modified):.4f}", verbose=self.verbose)
# Handle additional changes to dx if there's meas_pad
pad_cfg = self.init_params.get('meas_pad')
if pad_cfg is not None and pad_cfg.get('mode') is not None:
mode = pad_cfg['mode'] # 'precompute' or 'on_the_fly'
padding_type = pad_cfg['padding_type']
target_Npix = pad_cfg['target_Npix']
if Npix_is_modified:
Npix = Npix_modified
dx = dx * Npix / target_Npix
vprint(f"Update dx to {dx:.4f} {unit_str} due to meas_pad (mode = {mode}, padding_type = {padding_type}), Npix = {target_Npix}", verbose=self.verbose)
if illum_type == 'electron':
vprint(f"Suggested probe_mask_k radius (RBF*2/Npix) changes to > {(fitRBF * 2 / target_Npix):.4f}", verbose=self.verbose)
# Handle additional change to fitRBF if there's meas_resample
resample_cfg = self.init_params.get('meas_resample')
if resample_cfg is not None and resample_cfg.get('mode') is not None:
mode = resample_cfg['mode'] # 'precompute' or 'on_the_fly'
scale_factors = resample_cfg['scale_factors']
fitRBF_modified = fitRBF * scale_factors[0] # Currently the 2 values need to be the same
final_Npix = self.init_params['meas_Npix']
vprint(f"Update fitRBF to {fitRBF_modified:.4f} due to meas_resample (mode = {mode}, scale_factors = {scale_factors}), Npix = {final_Npix}", verbose=self.verbose)
if illum_type == 'electron':
vprint(f"Suggested probe_mask_k radius (RBF*2/Npix) changes to > {(fitRBF_modified * 2 / final_Npix):.4f}", verbose=self.verbose)
# Set the final dx for internal calibration, this dx would be used for probe, pos, object_extent, H
self.init_params['probe_dx'] = dx
vprint(" ", verbose=self.verbose)
def set_variables_dict(self):
vprint("### Setting init_variables dict ###", verbose=self.verbose)
# Note that the self.init_params can be modified by _meas_crop and other methods
# So this method is called after the entire init_measurements is done
# Keep in mind that crop could modify dx, Npix, scans
# pad could modify dx, Npix
# resample would only modify Npix
probe_illum_type = self.init_params.get('probe_illum_type') or 'electron'
if probe_illum_type == 'electron':
voltage = self.init_params['probe_kv']
wavelength = get_EM_constants(voltage, output_type='wavelength')
unit_str = 'Ang'
conv_angle = self.init_params['probe_conv_angle']
Npix = self.init_params['meas_Npix']
N_scan_slow = self.init_params['pos_N_scan_slow']
N_scan_fast = self.init_params['pos_N_scan_fast']
N_scans = N_scan_slow * N_scan_fast
dx = self.init_params['probe_dx']
dk = 1 / (dx * Npix)
kMax = Npix * dk / 2
da = dk * wavelength * 1e3
angleMax = Npix * da / 2
inferRBF = conv_angle / da
n_alpha = angleMax / conv_angle
# Print some derived values for sanity check
if self.verbose:
vprint("Derived values given input init_params:")
vprint(f' kv = {voltage} kV')
vprint(f' wavelength = {wavelength:.4f} Ang')
vprint(f' conv_angle = {conv_angle} mrad')
vprint(f' Npix = {Npix} px')
vprint(f' dk = {dk:.4f} Ang^-1')
vprint(f' kMax = {kMax:.4f} Ang^-1')
vprint(f' da = {da:.4f} mrad')
vprint(f' angleMax = {angleMax:.4f} mrad')
vprint(f' RBF = {inferRBF:.4f} px (Inferred from the given calibration, NOT necessarily from the loaded measurement data)')
vprint(f' n_alpha = {n_alpha:.4f} (# conv_angle)')
vprint(f' dx = {dx:.4f} Ang, Nyquist-limited dmin = 2*dx = {2*dx:.4f} Ang')
vprint(f' Rayleigh-limited resolution = {(0.61*wavelength/conv_angle*1e3):.4f} Ang (0.61*lambda/alpha for focused probe )')
vprint(f' Real space probe extent = {dx*Npix:.4f} Ang')
elif probe_illum_type == 'xray':
energy = self.init_params['beam_kev']
wavelength = 1.23984193e-9 / energy
unit_str = 'm'
dx = self.init_params['probe_dx']
N_scan_slow = self.init_params['probe_N_scan_slow']
N_scan_fast = self.init_params['probe_N_scan_fast']
N_scans = N_scan_slow * N_scan_fast
Npix = self.init_params['meas_Npix']
dRn = self.init_params['probe_dRn']
Rn = self.init_params['probe_Rn']
D_H = self.init_params['probe_D_H']
D_FZP = self.init_params['probe_D_FZP']
Ls = self.init_params['probe_Ls']
dk = 1/(dx*Npix)
if self.verbose:
vprint("Derived values given input init_params:")
vprint(f' x-ray beam energy = {energy} keV')
vprint(f' wavelength = {wavelength} m')
vprint(f' outmost zone width = {dRn} m')
vprint(f' Rn = {Rn} m')
vprint(f' D_H = {D_H} m')
vprint(f' D_FZP = {D_FZP} m')
vprint(f' Ls = {Ls} m')
vprint(f' Npix = {Npix} px')
vprint(f' dx = {dx} m')
else:
raise ValueError(f"init_params['probe_illum_type'] = {probe_illum_type} not implemented yet, please use either 'electron' or 'xray'!")
# Save general values into init_variables
# While they aren't necessarily "critical" for all initialization scenarios (like some variables aren't needed when we load things),
# But it's better to request these experimental parameters from users, since most of them should come with the measurements.
# We should collect all available experimental parameteres needed for minimal reconstruction from scratch
# And keep them with useful derived values in self.init_variables dict
# TODO: May consider use here to centralize the initalizaiton of all useful/derived variables
self.init_variables['probe_illum_type'] = probe_illum_type
self.init_variables['lambd'] = wavelength # Ang for electron, m for x-ray
self.init_variables['length_unit'] = unit_str
self.init_variables['Npix'] = Npix
self.init_variables['probe_shape'] = np.array([Npix, Npix]).astype(float) # Keep this at float for later init_pos
self.init_variables['N_scan_slow'] = N_scan_slow
self.init_variables['N_scan_fast'] = N_scan_fast
self.init_variables['N_scans'] = N_scans
self.init_variables['scan_step_size'] = self.init_params['pos_scan_step_size']
self.init_variables['dx'] = dx # Ang
self.init_variables['dk'] = dk # 1/Ang
self.init_variables['slice_thickness'] = self.init_params['obj_slice_thickness']
vprint(" ", verbose=self.verbose)
def init_probe(self):
"""
Initialize the probe by loading or simulating and then processing it.
"""
vprint("### Initializing probe ###", verbose=self.verbose)
probe = self._load_probe()
probe = self._process_probe(probe)
self.init_variables['probe'] = probe
# Print summary
vprint(f"probe (pmode, Ny, Nx) = {probe.dtype}, {probe.shape}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_pos(self):
"""
Initialize the probe positions by loading and processing them.
"""
vprint("### Initializing probe positions ###", verbose=self.verbose)
pos = self._load_pos()
pos = self._process_pos(pos)
probe_shape = self.init_variables['probe_shape']
obj_lateral_extent = (1.2 * np.ceil(pos.max(0) - pos.min(0) + probe_shape)).astype(int)
crop_pos = np.round(pos).astype('int16')
probe_pos_shifts = (pos - crop_pos).astype('float32')
# Save the processed positions
self.init_variables['obj_lateral_extent'] = obj_lateral_extent
self.init_variables['crop_pos'] = crop_pos
self.init_variables['probe_pos_shifts'] = probe_pos_shifts
self.init_variables['scan_affine'] = self.init_params['pos_scan_affine']
# Print summary
vprint(f"crop_pos (N,2) = {crop_pos.dtype}, {crop_pos.shape}", verbose=self.verbose)
vprint(f"crop_pos 1st and last px coords (y,x) = {crop_pos[0].tolist(), crop_pos[-1].tolist()}", verbose=self.verbose)
vprint(f"crop_pos extent (Ang) = {(crop_pos.max(0) - crop_pos.min(0))*self.init_variables['dx']}", verbose=self.verbose)
vprint(f"probe_pos_shifts (N,2) = {probe_pos_shifts.dtype}, {probe_pos_shifts.shape}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_obj(self):
"""
Initialize the object by loading and processing it.
"""
vprint("### Initializing object ###", verbose=self.verbose)
obj = self._load_obj()
obj = self._process_obj(obj)
obj = obj.astype('complex64')
self.init_variables['obj'] = obj
# Print summary
dz = self.init_variables['slice_thickness']
dx = self.init_variables['dx']
vprint(f"object (omode, Nz, Ny, Nx) = {obj.dtype}, {obj.shape}", verbose=self.verbose)
vprint(f"object extent (Z, Y, X) (Ang) = {np.round((obj.shape[1]*dz, obj.shape[2]*dx, obj.shape[3]*dx),4)}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_omode_occu(self):
"""
Initialize the mixed-state object mode occupancy so each mode has a fixed weight
"""
# Note: Initially I tried to make it optimizable, but then I noticed the AD algorithm
# tended to entirely shut off the mode by reducing the omode_occu rather than improving the mode
# So I decided to keep it as fixed values for now
omode_occu_params = self.init_params.get('obj_omode_init_occu') or {}
occu_type = omode_occu_params.get('occu_type', 'uniform')
init_occu = omode_occu_params.get('init_occu')
vprint(f"### Initializing omode_occu from '{occu_type}' ###", verbose=self.verbose)
if occu_type == 'custom':
omode_occu = np.array(init_occu)
elif occu_type == 'uniform':
omode = self.init_params['obj_omode_max']
omode_occu = np.ones(omode)/omode
else:
raise ValueError(f"Initialization method {occu_type} not implemented yet, please use 'custom' or 'uniform'!")
omode_occu = omode_occu.astype('float32')
vprint(f"omode_occu (omode) = {omode_occu.dtype}, {omode_occu.shape}", verbose=self.verbose)
self.init_variables['omode_occu'] = omode_occu
vprint(" ", verbose=self.verbose)
def init_H(self):
"""
Initialize the near-field Fresnel propagator for multislice ptychography
"""
vprint("### Initializing H (Fresnel propagator) ###", verbose=self.verbose)
probe_shape = self.init_variables['probe_shape']
dx = self.init_variables['dx']
slice_thickness = self.init_variables['slice_thickness']
lambd = self.init_variables['lambd']
unit_str = self.init_variables['length_unit']
vprint(f"Calculating H with probe_shape = {probe_shape} px, dx = {dx:.4f} {unit_str}, slice_thickness = {slice_thickness:.4f} {unit_str}, lambd = {lambd:.4f} {unit_str}", verbose=self.verbose)
H = near_field_evolution(probe_shape, dx, slice_thickness, lambd)
H = H.astype('complex64')
self.init_variables['H'] = H
vprint(f"H (Ky, Kx) = {H.dtype}, {H.shape}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_obj_tilts(self):
"""
Initialize the object crystal tilts. Tilts can be global tilt (1,2) or pos-dependent tilt (N,2)
"""
try:
tilt_source = self.init_params['tilt_source']
tilt_params = self.init_params['tilt_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
vprint(f"### Initializing obj tilts from = '{tilt_source}' ###", verbose=self.verbose)
if tilt_source == 'custom':
obj_tilts = tilt_params # (1,2) or (N,2) array in unit of mrad
elif tilt_source == 'file':
# Infer file type from extension
file_path = tilt_params.get('path')
key = tilt_params.get('key')
_, ext = os.path.splitext(file_path)
ext = ext.lower()
vprint(f"Detected tilt file type = '{ext}'")
# Warning when there's no key specified
if ext in ('.mat', '.h5', '.hdf5') and key is None:
vprint(f"WARNING: Couldn't find the 'key' in 'tilt_params' with file type = '{ext}'.")
vprint("It is strongly recommended to provide an explicit key to better find the desired dataset, which is often much faster as well.")
vprint("PtyRAD will still try to find the dataset, but you may consider setting 'key': <DATASET_KEY> inside your 'tilt_params' dict.")
if ext == '.raw':
raise ValueError("PtyRAD doesn't support loading object tilt from .raw file yet, please use other tilt_source.")
obj_tilts = np.float32(load_array_from_file(**tilt_params, ndims=[2]))
vprint(f"Initialized obj_tilts with loaded obj_tilts from file, mean obj_tilts = {obj_tilts.mean(0).round(2)} (theta_y, theta_x) mrad", verbose=self.verbose)
elif tilt_source == 'PtyRAD':
pt_path = tilt_params
ckpt = self.cache_contents if pt_path == self.cache_path else load_ptyrad(pt_path)
obj_tilts = np.float32(ckpt['optimizable_tensors']['obj_tilts'])
vprint(f"Initialized obj_tilts with loaded obj_tilts from PtyRAD, mean obj_tilts = {obj_tilts.mean(0).round(2)} (theta_y, theta_x) mrad", verbose=self.verbose)
elif tilt_source == 'simu':
N_scans = self.init_variables['N_scans']
tilt_type = tilt_params.get('tilt_type') or 'all' # Use the specified tilt_type when specified, fall back to 'all' for unspecified or None
init_tilts = tilt_params.get('init_tilts') or [[0,0]] # (1,2) array in unit of mrad
if tilt_type == 'each':
obj_tilts = np.broadcast_to(np.float32(init_tilts), shape=(N_scans,2))
vprint(f"Initialized obj_tilts with init_tilts = {init_tilts} (theta_y, theta_x) mrad", verbose=self.verbose)
elif tilt_type == 'all':
obj_tilts = np.broadcast_to(np.float32(init_tilts), shape=(1,2))
vprint(f"Initialized obj_tilts with init_tilts = {init_tilts} (theta_y, theta_x) mrad", verbose=self.verbose)
else:
raise ValueError(f"Tilt type {tilt_type} not implemented yet, please use either 'each', or 'all' when initializing obj_tilts with 'simu'!")
else:
raise ValueError(f"File type {tilt_source} not implemented yet, please use 'custom', 'PtyRAD', 'file', or 'simu'!")
# Print summary
self.init_variables['obj_tilts'] = obj_tilts
vprint(f"obj_tilts (N, 2) = {obj_tilts.dtype}, {obj_tilts.shape}", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
def init_check(self):
# Although some of the input experimental parameters might not be used directly by the package
# I think it's a good practice to check for overall consistency and remind the user to check carefully
# While these check could be performed within the init methods and achieve early return
# It's more readable to separate the initializaiton logic with the checking logic in this way
vprint("### Checking consistency between input params with the initialized variables ###", verbose=self.verbose)
# Check the consistency of init params with the initialized variables
init_params = self.init_params
Npix = init_params['meas_Npix']
Nlayer = init_params['obj_Nlayer']
N_scans = init_params['pos_N_scans']
N_scan_slow = init_params['pos_N_scan_slow']
N_scan_fast = init_params['pos_N_scan_fast']
# Initialized variables
meas = self.init_variables['measurements']
probe = self.init_variables['probe']
crop_pos = self.init_variables['crop_pos']
probe_pos_shifts = self.init_variables['probe_pos_shifts']
obj = self.init_variables['obj']
omode_occu = self.init_variables['omode_occu']
H = self.init_variables['H']
obj_tilts = self.init_variables['obj_tilts']
if self.init_variables.get('on_the_fly_meas_padded') is not None:
target_Npix = self.init_variables['on_the_fly_meas_padded'].shape[-1]
else:
target_Npix = meas.shape[-1]
if self.init_variables.get('on_the_fly_meas_scale_factors') is not None:
scale_factors = self.init_variables['on_the_fly_meas_scale_factors']
else:
scale_factors = [1,1]
# TODO These checks should probably be refactored a bit with clearer message
# We could duplicate some of them at the specific section to catch them early
# Check DP shape
if Npix == meas.shape[-2] == meas.shape[-1] == probe.shape[-2] == probe.shape[-1] == H.shape[-2] == H.shape[-1]:
vprint(f"Npix, DP measurements, probe, and H shapes are consistent as '{Npix}'", verbose=self.verbose)
elif Npix == target_Npix == probe.shape[-2] == probe.shape[-1] == H.shape[-2] == H.shape[-1]:
vprint(f"Npix, DP measurements, probe, and H shapes will be consistent as '{Npix}' during on-the-fly measurement padding", verbose=self.verbose)
elif Npix == floor(meas.shape[-2]*scale_factors[-2]) == floor(meas.shape[-1]*scale_factors[-1]) == probe.shape[-2] == probe.shape[-1] == H.shape[-2] == H.shape[-1]:
vprint(f"Npix, DP measurements, probe, and H shapes will be consistent as '{Npix}' during on-the-fly measurement resampling", verbose=self.verbose)
elif Npix == floor(target_Npix*scale_factors[-2]) == floor(target_Npix*scale_factors[-1]) == probe.shape[-2] == probe.shape[-1] == H.shape[-2] == H.shape[-1]:
vprint(f"Npix, DP measurements, probe, and H shapes will be consistent as '{Npix}' during on-the-fly measurement padding and then resampling", verbose=self.verbose)
else:
raise ValueError(f"Found inconsistency between Npix({Npix}), DP measurements({meas.shape[-2:]}), probe({probe.shape[-2:]}), and H({H.shape[-2:]}) shape")
# Check scan pattern
if N_scans == len(meas) == N_scan_slow*N_scan_fast == len(crop_pos) == len(probe_pos_shifts):
vprint(f"N_scans, len(meas), N_scan_slow*N_scan_fast, len(crop_pos), and len(probe_pos_shifts) are consistent as '{N_scans}'", verbose=self.verbose)
else:
raise ValueError(f"Found inconstency between N_scans({N_scans}), len(meas)({len(meas)}), N_scan_slow({N_scan_slow})*N_scan_fast({N_scan_fast}), len(crop_pos)({len(crop_pos)}), and len(probe_pos_shifts)({len(probe_pos_shifts)})")
# Check object shape
if obj.shape[0] == len(omode_occu):
vprint(f"obj.shape[0] is consistent with len(omode_occu) as '{obj.shape[0]}'", verbose=self.verbose)
else:
raise ValueError(f"Found inconsistency between obj.shape[0]({obj.shape[0]}) and len(omode_occu)({len(omode_occu)})")
if obj.shape[1] == Nlayer:
vprint(f"obj.shape[1] is consistent with Nlayer as '{Nlayer}'", verbose=self.verbose)
else:
raise ValueError(f"Found inconsistency between obj.shape[1]({obj.shape[1]}) and Nlayer({Nlayer})")
# Check object extent and probe positions
if (crop_pos.min(0) < 0).any():
raise ValueError(f"Found invalid crop position. crop_pos.min(0) {crop_pos.min(0)} must be equal or larger than 0. Please check your position and object initialization.")
if (crop_pos.max(0) + Npix - obj.shape[-2:] > 0).any():
raise ValueError(f"Found invalid crop position. crop_pos.max(0) {crop_pos.max(0)} + Npix ({Npix}) = {crop_pos.max(0) + Npix} must be equal or smaller than object canvas lateral size (Ny, Nx) = {obj.shape[-2:]}. Please check your position and object initialization.")
vprint(f"crop positions (yx_min={crop_pos.min(0)}, yx_max={crop_pos.max(0)+Npix}) are well contained inside object canvas (Ny,Nx) = {obj.shape[-2:]}.", verbose=self.verbose)
# Check obj tilts
if len(obj_tilts) in [1, N_scans]:
vprint("obj_tilts is consistent with either 1 or N_scans", verbose=self.verbose)
else:
raise ValueError(f"Found inconsistency between len(obj_tilts) ({len(obj_tilts)}), 1, and N_scans({N_scans})")
vprint("Pass the consistency check of initialized variables, initialization is done!", verbose=self.verbose)
def init_all(self):
# Run this method to initialize all
self.init_cache()
self.init_measurements()
self.init_calibration()
self.set_variables_dict()
self.init_probe()
self.init_pos()
self.init_obj()
self.init_omode_occu()
self.init_H()
self.init_obj_tilts()
self.init_check()
return self
###### Private methods for setting the cache ######
def _set_use_cached_flags(self, source):
""" Set the flags for each field whether we can cache or not """
# Validate required fields
try:
obj_source = self.init_params['obj_source']
obj_params = self.init_params['obj_params']
probe_source = self.init_params['probe_source']
probe_params = self.init_params['probe_params']
pos_source = self.init_params['pos_source']
pos_params = self.init_params['pos_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
triplets = [
('obj', obj_source, obj_params),
('probe', probe_source, probe_params),
('pos', pos_source, pos_params)]
# Helper for comparison
def same_source_and_params(a, b):
return a[1] == b[1] == source and a[2] == b[2]
# Check if obj, probe, and pos sources are the same
if same_source_and_params(triplets[0], triplets[1]) and same_source_and_params(triplets[1], triplets[2]):
self.use_cached_obj = self.use_cached_probe = self.use_cached_pos = True
self.cache_path = obj_params
self.cache_source = obj_source
return
if same_source_and_params(triplets[0], triplets[1]):
self.use_cached_obj = self.use_cached_probe = True
self.cache_path = obj_params
self.cache_source = obj_source
return
if same_source_and_params(triplets[0], triplets[2]):
self.use_cached_obj = self.use_cached_pos = True
self.cache_path = obj_params
self.cache_source = obj_source
return
if same_source_and_params(triplets[1], triplets[2]):
self.use_cached_probe = self.use_cached_pos = True
self.cache_path = probe_params
self.cache_source = probe_source
return
###### Private methods for initializing measurements ######
def _load_meas(self):
"""Load diffraction data from file or memory according to init_params['meas']."""
# Validate required fields
try:
meas_source = self.init_params['meas_source']
meas_params = self.init_params['meas_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
# Check for 'path' key for all sources
if meas_source != 'custom' and 'path' not in meas_params:
raise KeyError(f"'path' is required in 'meas_params' for source '{meas_source}'. Set 'path': <PATH_TO_YOUR_DATASET> inside your 'meas_params' dict.")
vprint(f"Loading measurements from source = '{meas_source}'", verbose=self.verbose)
if meas_source == 'custom':
if not isinstance(meas_params, np.ndarray): # assume to be a numpy array
raise TypeError(f"'custom' source requires 'meas_params' to be a NumPy array. Got {type(meas_params)}.")
meas = meas_params
elif meas_source in ['file', 'tif', 'tiff', 'mat', 'h5', 'hdf5', 'npy', 'raw']: # Keep the file types for backward compatibility
# Infer file type from extension
file_path = meas_params.get('path')
key = meas_params.get('key')
_, ext = os.path.splitext(file_path)
ext = ext.lower()
vprint(f"Detected measurement file type = '{ext}'")
# Warning when there's no key specified
if ext in ('.mat', '.h5', '.hdf5') and key is None:
vprint(f"WARNING: Couldn't find the 'key' in 'meas_params' with file type = '{ext}'.")
vprint("It is strongly recommended to provide an explicit key to better find the desired dataset, which is often much faster as well.")
vprint("PtyRAD will still try to find the dataset, but you may consider setting 'key': <DATASET_KEY> inside your 'meas_params' dict.")
# Provide default shape for .raw files if it's not specified
if ext == '.raw' and meas_params.get('shape') is None:
vprint(f"WARNING: Couldn't find the 'shape' in 'meas_params' with file type = '{ext}'.")
vprint("It is strongly recommended to provide an explicit shape to better load from .raw files")
vprint("PtyRAD will still try to load the dataset based on the provided 'init_params', but you may consider setting 'shape': (N_scans, Npix, Npix) inside your 'meas_params' dict.")
meas_params['shape'] = (self.init_params['pos_N_scans'],
self.init_params['meas_Npix'],
self.init_params['meas_Npix'])
meas = load_array_from_file(**meas_params)
else:
raise ValueError(f"Unsupported measurement source '{meas_source}'. Use 'custom' or 'file'.")
meas = meas.astype('float32')
vprint("Casting measurements dtype to float32 (single precision) for computational efficiency.")
vprint(f"Imported meausrements shape / dtype = {meas.shape}, dtype = {meas.dtype}", verbose=self.verbose)
vprint(f"Imported meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.4f}, {meas.max():.4f})", verbose=self.verbose)
return meas
def _process_meas(self, meas):
"""
Applies all processing steps to raw loaded measurements.
"""
# If the processing config is None, the methods will skip it internally
# Note that _meas_remove_neg_values and _meas_normalization will always be executed
# If you really want to nullify them, explictly set
# 'meas_remove_neg_values': {'mode': 'subtract_value', 'value': 0}
# 'meas_normalization': {'mode': 'divide_const', 'value': 1}
# Simple geometric operations
meas = self._meas_permute(meas, self.init_params.get('meas_permute'))
meas = self._meas_reshape(meas, self.init_params.get('meas_reshape'))
meas = self._meas_flipT(meas, self.init_params.get('meas_flipT'))
self.init_variables['meas_raw_avg'] = meas.mean(0) # Save this for initial dx calibration. The crop/pad/resample effect would be accounted accordingly in `init_calibration`
# Shape check after flipT (`meas` corresponds to the freshly loaded dataset before anything that could change its shape)
N_scans = self.init_params_original['pos_N_scans']
Npix = self.init_params_original['meas_Npix']
if meas.ndim != 3 or meas.shape[0] != N_scans or meas.shape[1:] != (Npix, Npix):
raise ValueError(
f"Shape mismatch after loading and processing the measurements: expected measurements shape = (N_scans={N_scans}, Npix={Npix}, Npix={Npix}), "
f"but got {meas.shape}. PtyRAD allows you to directly preprocess your loaded measurements with `meas_permute` and `meas_reshape` specified in params files to make it (N_scans, Npix(ky), Npix(kx)). "
f"Please read the comments in demo YAML params files or the documentation for more information about how to set `meas_permute` and `meas_reshape`."
)
# Operations that may change the shape of the measurements
meas = self._meas_crop(meas, self.init_params.get('meas_crop'))
meas = self._meas_remove_neg_values(meas, self.init_params.get('meas_remove_neg_values')) # meas need to be positive before the padding with background fitting mode
meas = self._meas_normalization(meas, self.init_params.get('meas_normalization')) # The normalization is needed because the background is calculated now and it needs to match the level of the final normalized meas
meas = self._meas_pad(meas, self.init_params.get('meas_pad'))
meas = self._meas_resample(meas, self.init_params.get('meas_resample'))
# Operations that add realistic factors to (simulated perfect) measurements
meas = self._meas_add_source_size(meas, self.init_params.get('meas_add_source_size'))
meas = self._meas_add_detector_blur(meas, self.init_params.get('meas_add_detector_blur'))
meas = self._meas_add_poisson_noise(meas, self.init_params.get('meas_add_poisson_noise'))
# Final guard on negative values
meas = self._meas_remove_neg_values(meas, {'mode': 'clip_neg'})
return meas
def _meas_permute(self, meas, order):
if order is not None:
vprint(f"Permuting measurements with order = {order}", verbose=self.verbose)
return meas.transpose(order)
return meas
def _meas_reshape(self, meas, target_shape):
if target_shape is not None:
vprint(f"Reshaping measurements to shape = {target_shape}", verbose=self.verbose)
return meas.reshape(target_shape)
return meas
def _meas_flipT(self, meas, flipT_axes):
"""
Flip and transpose measurement array.
flipT_axes: list of 3 binary/int values [flipud, fliplr, transpose]
"""
if flipT_axes is None:
return meas
# Validate length
if not isinstance(flipT_axes, (list, tuple)) or len(flipT_axes) != 3:
raise ValueError(f"Expected flipT_axes to be a list of 3 values, got: {flipT_axes}")
# Safely cast all entries to int
try:
flipT_axes = [int(v) for v in flipT_axes]
except Exception as e:
raise ValueError(f"flipT_axes must contain values convertible to int (0 or 1). Got: {flipT_axes}") from e
vprint(f"Flipping measurements with [flipud, fliplr, transpose] = {flipT_axes}", verbose=self.verbose)
if flipT_axes[0]:
meas = np.flip(meas, axis=1)
if flipT_axes[1]:
meas = np.flip(meas, axis=2)
if flipT_axes[2]:
meas = np.transpose(meas, (0, 2, 1))
return meas
def _meas_crop(self, meas, crop_ranges):
"""
Crop measurements across 4 dimensions:
[[slow_i, slow_f], [fast_i, fast_f], [ky_i, ky_f], [kx_i, kx_f]]
Allows any entry to be `None` to skip cropping that axis.
Note that this method would also update the `self.init_params`
"""
if crop_ranges is None:
return meas
if len(crop_ranges) != 4:
raise ValueError(f"Expected 4 crop ranges [N_slow, N_fast, ky, kx], got {crop_ranges}")
# Reshape (N, ky, kx) -> (N_slow, N_fast, ky, kx)
Nslow, Nfast = self.init_params['pos_N_scan_slow'], self.init_params['pos_N_scan_fast']
meas = meas.reshape(Nslow, Nfast, *meas.shape[-2:])
vprint(f"Reshaping measurements into {meas.shape} for cropping", verbose=self.verbose)
axes_names = ['N_slow', 'N_fast', 'ky', 'kx']
slices = []
for i, bounds in enumerate(crop_ranges):
if bounds is None:
slices.append(slice(None))
else:
try:
start, stop = bounds
slices.append(slice(start, stop))
vprint(f"Cropping axis {axes_names[i]} from {start} to {stop}", verbose=self.verbose)
except Exception as e:
raise ValueError(f"Invalid crop bounds for axis {axes_names[i]}: {bounds}") from e
meas = meas[slices[0], slices[1], slices[2], slices[3]]
vprint(f"Cropped measurements have shape (N_slow, N_fast, ky, kx) = {meas.shape}", verbose=self.verbose)
# Update self.init_params
vprint("Update (Npix, N_scans, N_scan_slow, N_scan_fast) after the measurements cropping", verbose=self.verbose)
self.init_params['meas_Npix'] = meas.shape[-1]
self.init_params['pos_N_scans'] = meas.shape[0] * meas.shape[1]
self.init_params['pos_N_scan_slow'] = meas.shape[0]
self.init_params['pos_N_scan_fast'] = meas.shape[1]
meas = meas.reshape(-1, meas.shape[-2], meas.shape[-1])
vprint(f"Reshape measurements back to (N, ky, kx) = {meas.shape}", verbose=self.verbose)
return meas
def _meas_remove_neg_values(self, meas, neg_cfg):
"""
Correct negative values in the measurement array based on the specified configuration.
Args:
meas (numpy.ndarray): The measurement array to process.
neg_cfg (dict): Configuration for handling negative values. Expected keys:
- mode (str): Method to handle negative values. Options are 'clip_neg', 'subtract_min',
'clip_value', or 'subtract_value'. Default is 'clip_neg'.
- value (float or None): Value used for 'clip_value' or 'subtract_value' modes. Default is None.
Returns:
numpy.ndarray: The processed measurement array with negative values handled.
"""
# This correction is enforced even the neg_cfg is None (not provided by user)
if neg_cfg is None:
neg_cfg = {}
# Extract configuration with defaults
mode = neg_cfg.get('mode', 'clip_neg') # Default to 'clip_neg'
value = neg_cfg.get('value', None) # Default to None
force = neg_cfg.get('force', False) # Default to False so it'll skip if there's no negative values
# Check if there are negative values
if not (meas < 0).any():
if not force:
vprint("No negative values found in measurements. Skipping non-neg correction.", verbose=self.verbose)
return meas
else:
vprint(f"No negative values found in measurements, but force = '{force}' so continuing measurement negative value correction", verbose=self.verbose)
vprint(f"Removing negative values in measurement with method = {mode} and value = {value}", verbose=self.verbose)
if mode == 'subtract_min':
min_value = meas.min()
meas -= min_value
value = None # Not relevant for this mode
vprint(f"Minimum value of {min_value:.4f} subtracted due to the positive px value constraint of measurements", verbose=self.verbose)
elif mode == 'clip_value':
if value is None:
raise KeyError("Mode 'clip_value' requires a non-None 'value'.")
vprint(f"Minimum value = {meas.min():.4f}, measurements below {value} are clipped to 0 due to the positive px value constraint of measurements", verbose=self.verbose)
meas[meas < value] = 0
elif mode == 'subtract_value':
if value is None:
raise KeyError("Mode 'subtract_value' requires a non-None 'value'.")
vprint(f"Minimum value = {meas.min():.4f}, measurements subtracted by {value} due to the positive px value constraint of measurements", verbose=self.verbose)
meas -= value
elif mode == 'clip_neg': # Default mode
vprint(f"Minimum value = {meas.min():.4f}, negative values are clipped to 0 due to the positive px value constraint of measurements", verbose=self.verbose)
meas[meas < 0] = 0
value = None # Not relevant for clipping
else:
raise ValueError(f"Unsupported mode '{mode}' for handling negative values. Use 'clip_neg', 'subtract_min', 'clip_value', or 'subtract_value'.")
# Final check in case the user specified value is not enough to remove all neg values
if (meas < 0).any():
vprint(f"User specified value = {value} is not enough to remove negative values, applying 0 clipping")
vprint(f"Minimum value of {meas.min():.4f} is clipped to 0 due to the positive px value constraint of measurements", verbose=self.verbose)
meas[meas<0] = 0
return meas
def _meas_normalization(self, meas, norm_cfg):
"""
Normalize the measurement array based on the specified normalization mode.
Args:
meas (numpy.ndarray): The measurement array to normalize, shape (N, ky, kx).
Returns:
numpy.ndarray: The normalized measurement array.
"""
# This correction is enforced even the norm_cfg is None (not provided by user)
if norm_cfg is None:
norm_cfg = {}
norm_mode = norm_cfg.get('mode', 'max_at_one') # Default to 'max_at_one'
norm_const = norm_cfg.get('value', None) # Used for 'divide_const' mode
vprint(f"Normalizing measurements with mode = '{norm_mode}' and value = '{norm_const}'", verbose=self.verbose)
if norm_mode == 'max_at_one':
normalization_const = meas.mean(0).max()
vprint(f"Normalizing by max of the 2D mean pattern intensity: {normalization_const:.8g}", verbose=self.verbose)
elif norm_mode == 'mean_at_one':
normalization_const = meas.mean(0).mean()
vprint(f"Normalizing by mean of the 2D mean pattern intensity: {normalization_const:.8g}", verbose=self.verbose)
elif norm_mode == 'sum_to_one':
normalization_const = meas.mean(0).sum()
vprint(f"Normalizing by sum of the 2D mean pattern intensity: {normalization_const:.8g}", verbose=self.verbose)
elif norm_mode == 'divide_const':
if norm_const is None:
raise KeyError("Mode 'divide_const' requires a non-None 'norm_const'.")
normalization_const = norm_const
vprint(f"Normalizing by user-defined constant: {normalization_const:.8g}", verbose=self.verbose)
else:
raise ValueError(f"Unsupported normalization mode '{norm_mode}'. Use 'max_at_one', 'mean_at_one', 'sum_to_one', or 'divide_const'.")
# Normalize the measurements
meas = meas / normalization_const
meas = meas.astype('float32')
vprint(f"meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.4f}, {meas.max():.4f})", verbose=self.verbose)
return meas
def _meas_pad(self, meas, pad_cfg):
"""
_meas_pad Padd the 3D measurements array to a target size using the specified padding mode and type.
Args:
meas (numpy.ndarray): The measurement array to normalize, shape (N, ky, kx).
pad_cfg (dict): A dictionary containing the padding configuration. Expected keys:
pad_cfg = {'mode': 'on_the_fly', 'padding_type': 'power', 'target_Npix': 256, 'value': 0}
Raises:
KeyError: _description_
KeyError: _description_
Returns:
numpy.ndarray: The padded measurement array.
"""
if pad_cfg is None or pad_cfg.get('mode') is None:
self.init_variables['on_the_fly_meas_padded'] = None
self.init_variables['on_the_fly_meas_padded_idx'] = None
return meas
mode = pad_cfg['mode'] # 'precompute' or 'on_the_fly'. Use `on_the_fly` to save GPU memory
padding_type = pad_cfg['padding_type']
target_Npix = pad_cfg['target_Npix']
value = pad_cfg.get('value', 10) # For constant and linear_ramp padding
threshold = pad_cfg.get('threshold', 70) # For exp and power padding that requires fitting a thresholded mask
vprint(f"Padding measurements with mode='{mode}', padding_type='{padding_type}', target_Npix={target_Npix}", verbose=self.verbose)
# Get amplitude from average DP
meas_avg = meas.mean(axis=0)
meas_int_sum = meas_avg.sum()
amp_avg = np.sqrt(meas_avg)
H, W = amp_avg.shape
# Calculate padding for each dimension
pad_y = max(0, target_Npix - H)
pad_x = max(0, target_Npix - W)
pad_y1, pad_y2 = pad_y // 2, pad_y - pad_y // 2
pad_x1, pad_x2 = pad_x // 2, pad_x - pad_x // 2
pad_h1, pad_h2 = pad_y1, pad_y1 + H
pad_w1, pad_w2 = pad_x1, pad_x1 + W
# Create coordinate grid for radial background fitting
y, x = np.ogrid[:target_Npix, :target_Npix]
center = (H // 2 + pad_y1, W // 2 + pad_x1)
r = np.sqrt((y - center[0])**2 + (x - center[1])**2) + 1e-10 # so r is never 0
# Compute background
if padding_type == 'constant':
amp_padded = np.pad(amp_avg, ((pad_y1, pad_y2), (pad_x1, pad_x2)), mode='constant', constant_values=value)
elif padding_type == 'edge':
amp_padded = np.pad(amp_avg, ((pad_y1, pad_y2), (pad_x1, pad_x2)), mode='edge')
elif padding_type == 'linear_ramp':
amp_padded = np.pad(amp_avg, ((pad_y1, pad_y2), (pad_x1, pad_x2)), mode='linear_ramp', end_values=value)
elif padding_type == 'exp':
mask = create_one_hot_mask(amp_avg, percentile=threshold) # It feels like we probably don't need to normalize meas before padding because the mask is calculated by percentile
popt = fit_background(amp_avg, mask, fit_type='exp')
amp_padded = exponential_decay(r, *popt)
elif padding_type == 'power':
mask = create_one_hot_mask(amp_avg, percentile=threshold)
popt = fit_background(amp_avg, mask, fit_type='power')
amp_padded = power_law(r, *popt)
else:
raise ValueError(f"Unsupported padding_type = '{padding_type}'")
# Square the padded amplitude back to intensity
meas_padded = np.square(amp_padded)[None,] # (1, ky, kx)
meas_padded[..., pad_h1:pad_h2, pad_w1:pad_w2] = 0
padded_int_sum = meas_padded.sum()
vprint(f"Original meas int sum = {meas_int_sum:.4f}, padded region int sum = {padded_int_sum:.4f}, or {padded_int_sum/meas_int_sum:.2%} more intensity after padding.", verbose=self.verbose)
vprint("This percentage should be ideally less than 5%, or you should set a lower threshold to exclude more central region.", verbose=self.verbose)
if mode == 'precompute':
canvas = np.zeros((meas.shape[0], *meas_padded.shape[1:]))
canvas += meas_padded
canvas[..., pad_h1:pad_h2, pad_w1:pad_w2] = meas
meas = canvas
self.init_variables['on_the_fly_meas_padded'] = None
self.init_variables['on_the_fly_meas_padded_idx'] = None
elif mode == 'on_the_fly':
# For on_the_fly padding, we pass the padded 2D pattern (extra background) and padding indices to the model
self.init_variables['on_the_fly_meas_padded'] = meas_padded
self.init_variables['on_the_fly_meas_padded_idx'] = [pad_h1, pad_h2, pad_w1, pad_w2]
else:
raise ValueError(f"meas_pad does not support mode = '{mode}', please choose from 'on_the_fly', 'precompute', or null")
# Update iself.init_params similar to _meas_crop
vprint("Update Npix after the measurements padding", verbose=self.verbose)
self.init_params['meas_Npix'] = meas_padded.shape[-1] # This will update Npix to target_Npix no matter what mode is used
return meas
def _meas_resample(self, meas, resample_cfg):
"""
_meas_resample Resample measurements along the ky, kx dimension
"""
if resample_cfg is None or resample_cfg.get('mode') is None:
self.init_variables['on_the_fly_meas_scale_factors'] = None
return meas
# Validate required fields
try:
mode = resample_cfg['mode']
Npix = self.init_params['meas_Npix']
scale_factors = resample_cfg['scale_factors']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
# Ensure scale_factors is a list or tuple of length 2
if len(scale_factors) != 2:
raise ValueError("scale_factors for resample must be a list or tuple of two elements.")
if scale_factors[0] != scale_factors[1]:
min_scale = min(scale_factors)
vprint(f"Non-uniform scale_factors {scale_factors} detected. Using uniform scale factor: {min_scale}")
scale_factors = [min_scale, min_scale]
# If on-the-fly padding is set, force resample to be on-the-fly as well
if self.init_variables.get('on_the_fly_meas_padded', None) is not None:
mode = 'on_the_fly'
vprint("'meas_resample' is set to 'on_the_fly' mode because 'meas_pad' is also set to 'on_the_fly' mode", verbose=self.verbose)
vprint(f"Resampling measurements with mode = '{mode}', scale_factors = {scale_factors}", verbose=self.verbose)
if mode == 'precompute':
zoom_factors = np.array([1.0, *scale_factors]) # scipy.ndimage.zoom applies to all axes.
meas = zoom(meas, zoom_factors, order=1) # bilinear (order=1) could prevent overshooting. Resampling would change the meas.sum(), but we have normalization at the end of the process.
Npix = meas.shape[-1] # Update Npix
self.init_variables['on_the_fly_meas_scale_factors'] = None
elif mode == 'on_the_fly':
# Don't change `meas`, just update Npix
Npix = floor(Npix * scale_factors[-1]) # To match the rounding logic with torch.nn.functional.interpolate()
self.init_variables['on_the_fly_meas_scale_factors'] = scale_factors
else:
raise ValueError(f"meas_resample does not support mode = '{mode}', please choose from 'on_the_fly', 'precompute', or null")
# Update self.init_params similar to _meas_crop
self.init_params['meas_Npix'] = Npix
vprint(f"Update Npix into '{Npix}' after the measurements resampling", verbose=self.verbose)
vprint(f"Resampled measurements have shape (N_scans, ky, kx) = {meas.shape}", verbose=self.verbose)
return meas
def _meas_add_source_size(self, meas, source_size_std_ang):
if source_size_std_ang is None or source_size_std_ang == 0:
return meas
Nslow, Nfast = self.init_params['pos_N_scan_slow'], self.init_params['pos_N_scan_fast']
meas = meas.reshape(Nslow, Nfast, *meas.shape[-2:])
vprint(f"Reshaping measurements into {meas.shape} for adding partial spatial coherence (source size) induced blurring on measurements", verbose=self.verbose)
# Convert real-space blur in Angstroms to Gaussian std in scan units (px)
source_size_std_px = source_size_std_ang / self.init_params['pos_scan_step_size']
vprint(f"Adding source size (partial spatial coherence) of Gaussian blur std = {source_size_std_px:.4f} scan_step sizes or {source_size_std_ang:.4f} Ang to measurements along the scan directions", verbose=self.verbose)
# Apply blur over scan dimensions (0,1)
meas = gaussian_filter(meas, sigma=source_size_std_px, axes=(0,1)) # Partial spatial coherence is approximated by mixing DPs at nearby probe positions
meas = meas.reshape(-1, meas.shape[-2], meas.shape[-1])
vprint(f"Reshape measurements back to (N, ky, kx) = {meas.shape}", verbose=self.verbose)
return meas
def _meas_add_detector_blur(self, meas, detector_blur_std_px):
"""
Add detector blur (point-spread function of the detector)
"""
if detector_blur_std_px is None or detector_blur_std_px == 0:
return meas
meas = gaussian_filter(meas, sigma=detector_blur_std_px, axes=(-2,-1)) # Detector blur is essentially the Gaussian blur along ky, kx
vprint(f"Adding detector blur (point-spread function of the detector) of Gaussian blur std = {detector_blur_std_px:.4f} px to measurements along the ky, kx directions", verbose=self.verbose)
return meas
def _meas_add_poisson_noise(self, meas, poisson_cfg):
if poisson_cfg is None:
return meas
# Validate required fields
try:
unit = poisson_cfg['unit']
value = poisson_cfg['value']
scan_step_size = self.init_params['pos_scan_step_size']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
# Check negative values before applying Poisson noise
eps = meas.min() / np.abs(meas.mean() + 1e-12)
if meas.min() < 0:
vprint(f"Found negative values in meas, meas.min() = {meas.min():.4g}.", verbose=self.verbose)
if eps > -1e-5:
vprint(f"Negative values ({meas[meas < 0].mean():.4g}) are within relative numerical tolerance (min/mean) 1e-5 , clipping negative values to 0.", verbose=self.verbose)
meas[meas < 0] = 0
else:
raise ValueError(f"meas needs to be positive before applying poisson noise, got meas.min = {meas.min():.4g}. Check your 'meas_remove_neg_values'.")
# Convert units to total electrons per pattern
if unit == 'total_e_per_pattern':
total_electron = value
dose = total_electron / scan_step_size**2
elif unit == 'e_per_Ang2':
dose = value
total_electron = dose * scan_step_size**2
else:
raise ValueError(f"Unsupported unit for Poisson noise: '{unit}'. Use 'total_e_per_pattern' or 'e_per_Ang2'.")
vprint(f"total electron per measurement = dose x scan_step_size^2 = {dose:.3f}(e-/Ang^2) x {scan_step_size:.3f}(Ang)^2 = {total_electron:.3f}", verbose=self.verbose)
# Normalize meas to sum to ~ 1 before applying Poisson noise
vprint(f"Before applying Poisson noise: meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.5f}, {meas.max():.4f})", verbose=self.verbose)
normalization_const = meas.sum() / meas.shape[0]
vprint(f"Normalization constant = {normalization_const:.4f}, this makes each measurement sum to ~ 1.", verbose=self.verbose)
meas = meas / normalization_const # Make each slice of the meas to sum to ~ 1. A global normalization constant keeps the relative intensity.
vprint(f"After applying normalization: meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.5f}, {meas.max():.4f})", verbose=self.verbose)
vprint(f"Mean total electron per pattern = meas.sum((-2,-1)).mean(0) = ({meas.sum((-2,-1)).mean(0):.5f})", verbose=self.verbose)
set_random_seed(seed=self.random_seed)
meas = np.random.poisson(meas * total_electron)
vprint(f"Adding Poisson noise with a total electron per diffraction pattern of {int(total_electron)}", verbose=self.verbose)
vprint(f"After applying Poisson noise: meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.5f}, {meas.max():.4f})", verbose=self.verbose)
meas = meas * normalization_const / total_electron # Un-normalize meas back to the original scale
vprint(f"After un-normalizing back to original scale: meausrements int. statistics (min, mean, max) = ({meas.min():.4f}, {meas.mean():.5f}, {meas.max():.4f})", verbose=self.verbose)
return meas
def _export_meas(self, export_params={}):
meas = self.init_variables['measurements']
file_dir = export_params.get("file_dir")
# Handle the case where file_dir is None
if file_dir in (None, ''):
meas_path = get_nested(self.init_params, key=['meas_params', 'path'], safe=True, default='')
export_params["file_dir"] = os.path.dirname(meas_path)
# Ensure the directory exists if it's not empty
if file_dir and not os.path.exists(file_dir):
vprint(f"User specified 'file_dir' = '{file_dir}' doesn't exist, creating the directory now.", verbose=self.verbose)
os.makedirs(file_dir, exist_ok=True)
save_array(meas, **export_params)
return
###### Private methods for initializing probe ######
def _load_probe(self):
"""
Load the probe from the specified source.
returns:
probe (numpy.ndarray): The loaded probe array would always be casted to (pmode, Ny, Nx).
"""
# Validate required fields
try:
probe_source = self.init_params['probe_source']
probe_params = self.init_params['probe_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
probe_illum_type = self.init_variables['probe_illum_type']
vprint(f"Loading probe from source = '{probe_source}'", verbose=self.verbose)
if probe_source == 'custom':
probe = probe_params
elif probe_source == 'PtyRAD':
probe = self._load_probe_from_ptyrad(probe_params)
elif probe_source == 'PtyShv':
probe = self._load_probe_from_ptyshv(probe_params)
elif probe_source == 'py4DSTEM':
probe = self._load_probe_from_py4dstem(probe_params)
elif probe_source == 'simu':
probe = self._simulate_probe(probe_params, probe_illum_type)
else:
raise ValueError(f"Unsupported probe source '{probe_source}'. Use 'custom', 'PtyRAD', 'PtyShv', 'py4DSTEM', or 'simu'.")
vprint(f"Loaded probe shape = {probe.shape}, dtype = {probe.dtype}", verbose=self.verbose)
return probe
def _load_probe_from_ptyrad(self, params: str):
pt_path = params
ckpt = self.cache_contents if self.use_cached_probe else load_ptyrad(pt_path)
probe = ckpt['optimizable_tensors']['probe']
return probe
def _load_probe_from_ptyshv(self, params: str):
mat_path = params
mat_version = get_matfile_version(mat_path) #https://docs.scipy.org/doc/scipy-1.11.3/reference/generated/scipy.io.matlab.matfile_version.html
use_h5py = (mat_version[0] == 2)
probe = self.cache_contents['probe'] if self.use_cached_probe else load_mat(mat_path, key='probe')
vprint(f"Input PtyShv probe has original shape {probe.shape}, while default PtyShv order is (Ny, Nx, pmode, vp)", verbose=self.verbose)
# First unify the axes order induced by loading with scipy / h5py, now it should be (Ny, Nx, pmode, vp)
if use_h5py:
probe = probe.transpose(range(probe.ndim)[::-1])
vprint(f"Reverse array axes order of probe to {probe.shape} because use_h5py = {use_h5py}, which automatically reverse the order", verbose=self.verbose)
else:
vprint(f"Keep array axes order of probe at {probe.shape} because use_h5py = {use_h5py}", verbose=self.verbose)
# Correct the probe dimension to 3 dimensions, now it should be (Ny, Nx, pmode)
if probe.ndim == 4:
vprint("Import only the 1st variable probe mode to make a final probe with (pmode, Ny, Nx)", verbose=self.verbose) # I don't find variable probe modes are particularly useful for electon ptychography
probe = probe[..., 0]
elif probe.ndim == 2:
vprint("Expanding PtyShv probe dimension to make a final probe with (pmode, Ny, Nx)", verbose=self.verbose)
probe = probe[..., None]
# Final permutation to make it (pmode, Ny, Nx)
probe = probe.transpose(2,0,1)
vprint(f"Permute the array axes order of probe to {probe.shape} make it (pmode, Ny, Nx)", verbose=self.verbose)
return probe
def _load_probe_from_py4dstem(self, params: str):
"""
Load the probe from a py4DSTEM hdf5 file.
Note that the ouput file is expected to be generated by my modified py4DSTEM fork.
https://github.com/chiahao3/py4DSTEM/tree/benchmark
"""
hdf5_path = params
probe = self.cache_contents['probe'] if self.use_cached_probe else load_hdf5(hdf5_path, key='probe')
vprint(f"Input py4DSTEM probe has original shape {probe.shape}", verbose=self.verbose)
if probe.ndim == 2:
vprint("Expanding py4DSTEM probe dimension to make a final probe with (pmode, Ny, Nx)", verbose=self.verbose)
probe = probe[None, ...]
return probe
def _simulate_probe(self, simu_params: dict, probe_illum_type: str):
"""
Simulate the probe based on the specified parameters.
"""
# TODO Can probably improve the params file structure and simulation process for probe
# Currently the probe simu parameters that are not needed when we load existing probes
# are just implictly ignored, like defocus, convergence angle.
# Illumination type is something that can probably live directly under `init_params.probe`
if simu_params is not None:
vprint("Using user-specified parameters in 'init_params['probe_params']' for initial probe simulation.", verbose=self.verbose)
else:
vprint("Using experimental parameters specified by 'init_params' for initial probe simulation.", verbose=self.verbose)
simu_params = get_default_probe_simu_params(self.init_params)
if probe_illum_type == 'electron':
probe = make_stem_probe(simu_params, verbose=self.verbose)[None, ...]
elif probe_illum_type == 'xray':
probe = make_fzp_probe(simu_params, verbose=self.verbose)[None, ...]
else:
raise ValueError(f"Unsupported illumination type '{probe_illum_type}'. Use 'electron' or 'xray'.")
# probe is (1, Ny, Nx) after simulation, expand it to (pmode, Ny, Nx) if needed
if simu_params['pmodes'] > 1:
probe = make_mixed_probe(probe[0], simu_params['pmodes'], simu_params['pmode_init_pows'], verbose=self.verbose)
return probe
def _process_probe(self, probe):
"""
Process the loaded probe, including permutation, setting pmode, and normalization
"""
# If the processing config is None, the methods will skip it internally
pmode_max = self.init_params.get('probe_pmode_max')
pmode_init_pows = self.init_params.get('probe_pmode_init_pows')
probe = self._probe_permute(probe, self.init_params.get('probe_permute'))
probe = self._probe_set_pmode_max(probe, pmode_max, pmode_init_pows, orthogonalize=True, sort=True)
probe = self._probe_z_shift(probe, self.init_params.get('probe_z_shift'))
probe = self._probe_normalize(probe, self.init_params.get('probe_normalize'))
return probe
def _probe_permute(self, probe, order):
"""
Permute the probe dimensions if specified in the parameters.
"""
if order is not None:
vprint(f"Permuting probe with order = {order}", verbose=self.verbose)
probe = probe.transpose(order)
return probe
def _probe_set_pmode_max(self, probe, pmode_max, pmode_init_pows, orthogonalize=True, sort=True):
"""
Either cap or pad the pmode for mixed state probe with optional orthogonalization and sorting
"""
pmode_now = probe.shape[0]
probe_int_sum = np.sum(np.abs(probe)**2)
pmode_init_pow = [min(pmode_init_pows)] # pmode_init_pows is a list of float(s), so we convert it into a list of float
if pmode_now > pmode_max:
vprint(f"pmode_now: {pmode_now} and pmode_max: {pmode_max}, capping the pmode.", verbose=self.verbose)
probe_final = probe[:pmode_max]
elif pmode_now == pmode_max:
vprint(f"pmode_now: {pmode_now} and pmode_max: {pmode_max}, leaving the pmode unchanged.", verbose=self.verbose)
probe_final = probe
else: # pmode_now <= pmode_max: # Need to pad new probe modes
vprint(f"pmode_now: {pmode_now} and pmode_max: {pmode_max}, padding the pmode.", verbose=self.verbose)
num_new_modes = pmode_max - pmode_now
vprint(f"Creating {num_new_modes} new probe modes from the major mode", verbose=self.verbose)
mixed_probe_temp = make_mixed_probe(probe[0], pmode_max, pmode_init_pow, verbose=False) # Take the strongest probe mode and make a temporary new mixed probe (int sum at 1)
new_modes = mixed_probe_temp[-num_new_modes:] * probe_int_sum ** 0.5 # Normalize the new mode intensity with original intensity
probe_final = np.concatenate((probe, new_modes), axis=0) # Total int = 1 + num_new_modes * pmode_init_pow, will normalize it later
# Normalize back to original intensity
normalization_factor = (np.sum(np.abs(probe_final) ** 2) / probe_int_sum) ** 0.5
probe_final = probe_final / normalization_factor
# Optional orthogonalization and sorting
if orthogonalize:
vprint(f"Orthogonalizing {len(probe_final)} pmodes", verbose=self.verbose)
probe_final = orthogonalize_modes_vec_np(probe_final)
if sort:
vprint(f"Sorting {len(probe_final)} pmodes by their intensities", verbose=self.verbose)
probe_final = sort_by_mode_int_np(probe_final)
return probe_final
def _probe_z_shift(self, probe, prop_distance):
"""
Applying user-specified additional axial propagation to the initialized probe. This is used for shifting the reconstructed probe along depth.
Note that prop_distance is defined with propagation direction, so positive value means forward propagation (i.e., increasing depth/z).
"""
if prop_distance is None or prop_distance == 0:
return probe
else:
dx = self.init_variables['dx']
lambd = self.init_variables['lambd']
unit_str = self.init_variables['length_unit']
vprint(f"Applying additional axial propagation (z) = {prop_distance} {unit_str} to the probe. Positive value means forward propagation (i.e., increasing depth/z).", verbose=self.verbose)
H = near_field_evolution(probe.shape[-2:], dx, prop_distance, lambd)
probe_shifted = np.fft.ifft2(H[None,] * np.fft.fft2(probe))
return probe_shifted
def _probe_normalize(self, probe, norm_cfg):
"""
Normalize the probe intensity based on the measurements.
"""
# TODO Extend this method to support other normalization methods
# like the target intensity (vacuum probe), or a scaling factor over meas_avg_sum
# This correction is enforced even the norm_cfg is None (not provided by user)
if norm_cfg is None:
norm_cfg = {}
try:
# Using the pre-calculated meas_avg_sum for probe intensity normalization
# becasue on-the-fly padding could increase the total meas intensity
meas_avg_sum = self.init_variables['meas_avg_sum']
except KeyError:
vprint("WARNING: Measurement average sum ('meas_avg_sum') not found. Initializing measurements first for probe normalization...", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
self.init_measurements()
meas_avg_sum = self.init_variables['meas_avg_sum']
normalization_factor = (np.sum(np.abs(probe) ** 2) / meas_avg_sum) ** 0.5
probe = probe / normalization_factor
vprint(f"sum(|probe_data|**2) = {np.sum(np.abs(probe)**2):.2f}, while meas.mean(0).sum() = {meas_avg_sum:.2f}", verbose=self.verbose)
return probe.astype('complex64')
###### Private methods for initializing positions ######
def _load_pos(self):
"""
Load the probe positions from the specified source.
"""
# Validate required fields
try:
pos_source = self.init_params['pos_source']
pos_params = self.init_params['pos_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
vprint(f"Loading probe positions from source = '{pos_source}'", verbose=self.verbose)
if pos_source == 'custom':
pos = pos_params
elif pos_source == 'PtyRAD':
pos = self._load_pos_from_ptyrad(pos_params)
elif pos_source == 'PtyShv':
pos = self._load_pos_from_ptyshv(pos_params)
elif pos_source == 'py4DSTEM':
pos = self._load_pos_from_py4dstem(pos_params)
elif pos_source == 'simu':
pos = self._simulate_pos(pos_params)
elif pos_source == 'foldslice_hdf5':
pos = self._load_pos_from_foldslice(pos_params)
else:
raise ValueError(f"Unsupported position source '{pos_source}'. Use 'custom', 'PtyRAD', 'PtyShv', 'py4DSTEM', 'simu', or 'foldslice_hdf5'.")
return pos
def _load_pos_from_ptyrad(self, params: str):
pt_path = params
ckpt = self.cache_contents if self.use_cached_pos else load_ptyrad(pt_path)
crop_pos = ckpt['model_attributes']['crop_pos']
probe_pos_shifts = ckpt['optimizable_tensors']['probe_pos_shifts']
pos = crop_pos + probe_pos_shifts
return pos
def _load_pos_from_ptyshv(self, params: str):
mat_path = params
mat_version = get_matfile_version(mat_path) # https://docs.scipy.org/doc/scipy-1.11.3/reference/generated/scipy.io.matlab.matfile_version.html
use_h5py = (mat_version[0] == 2)
mat_contents = self.cache_contents if self.use_cached_pos else load_mat(mat_path, key=['object', 'probe', 'outputs.probe_positions'], delimiter='.')
vprint(f"Input PtyShv probe positions has original shape {mat_contents['outputs.probe_positions'].shape}, while default PtyShv order is (N, 2)", verbose=self.verbose)
# First unify the axes order induced by loading with scipy / h5py, now it should be (N, 2)
if use_h5py:
mat_contents = {key: arr.transpose(range(arr.ndim)[::-1]) for key, arr in mat_contents.items()}
vprint(f"Reverse array axes order because use_h5py = {use_h5py}, which automatically reverse the order", verbose=self.verbose)
probe_positions = mat_contents['outputs.probe_positions']
probe_shape = mat_contents['probe'].shape[:2] # Matlab probe is (Ny,Nx,pmode,vp) or (Ny,Nx,pmode)
obj_shape = mat_contents['object'].shape[:2] # Matlab object is (Ny, Nx, Nz) or (Ny,Nx)
pos_offset = np.ceil((np.array(obj_shape)/2) - (np.array(probe_shape)/2)) - 1 # For Matlab - Python index shift
probe_positions_yx = probe_positions[:, [1,0]] # The first index after shifting is the row index (along vertical axis)
pos = probe_positions_yx + pos_offset
return pos
def _load_pos_from_py4dstem(self, params: str):
hdf5_path = params
hdf5_contents = self.cache_contents if self.use_cached_pos else load_hdf5(hdf5_path)
probe_positions = hdf5_contents['positions_px']
probe_shape = hdf5_contents['probe'].shape[-2:] # py4DSTEM probe is (pmode,Ny,Nx)
pos = probe_positions - np.array(probe_shape)/2
return pos
def _load_pos_from_foldslice(self, params: str):
# This preprocessing routine is equivalent to `p.src_positions='hdf5_pos';` in `fold_slice`
# which was used for many APS instruments
dx = self.init_variables['dx']
probe_shape = self.init_variables['probe_shape']
hdf5_path = params
ppY = load_hdf5(hdf5_path, key='ppY')
ppX = load_hdf5(hdf5_path, key='ppX')
pos = np.stack((-ppY, -ppX), axis=1) / dx
pos = np.flipud(pos) # (N,2) in (pos_y_px, pos_x_px)
obj_shape = 1.2 * np.ceil(pos.max(0) - pos.min(0) + probe_shape)
pos = pos + np.ceil((np.array(obj_shape)/2) - (np.array(probe_shape)/2)) # Shift to obj coordinate
return pos
def _simulate_pos(self, simu_params: dict):
if simu_params is not None:
vprint("Using user-specified parameters in 'init_params['pos_params']' for initial position simulation.", verbose=self.verbose)
else:
simu_params = {}
vprint("Using experimental parameters specified by 'init_params' (dx, scan_step size, N_scan_slow, N_scan_fast) for initial position simulation.", verbose=self.verbose)
# The unspecified parameters will be set to the values specified in self.init_variables
dx = simu_params.get('dx', self.init_variables['dx'])
scan_step_size = simu_params.get('scan_step_size', self.init_variables['scan_step_size'])
N_scan_slow = simu_params.get('N_scan_slow', self.init_variables['N_scan_slow'])
N_scan_fast = simu_params.get('N_scan_fast', self.init_variables['N_scan_fast'])
probe_shape = simu_params.get('probe_shape', self.init_variables['probe_shape'])
vprint(f"Simulating probe positions with dx = {dx:.4f}, scan_step_size = {scan_step_size:.4f}, N_scan_fast = {N_scan_fast}, N_scan_slow = {N_scan_slow}", verbose=self.verbose)
pos = scan_step_size / dx * np.array([(y, x) for y in range(N_scan_slow) for x in range(N_scan_fast)]) # (N,2), each row is (y,x)
pos = pos - pos.mean(0) # Center scan around origin
obj_shape = 1.2 * np.ceil(pos.max(0) - pos.min(0) + probe_shape)
pos = pos + np.ceil((np.array(obj_shape)/2) - (np.array(probe_shape)/2)) # Shift to obj coordinate
return pos
def _process_pos(self, pos):
"""
Process the loaded probe positions, including flipping, affine transformations, and random displacements.
"""
# If the processing config is None, the methods will skip it internally
pos = self._pos_scan_flipT(pos, self.init_params.get('pos_scan_flipT'))
pos = self._pos_scan_affine_transform(pos, self.init_params.get('pos_scan_affine'))
pos = self._pos_scan_add_random_displacement(pos, self.init_params.get('pos_scan_rand_std'))
return pos
def _pos_scan_flipT(self, pos, flipT_axes):
"""
Flip and transpose scan positions.
flipT_axes: list of 3 binary/int values [flipud, fliplr, transpose]
"""
if flipT_axes is None:
return pos
# Validate length
if not isinstance(flipT_axes, (list, tuple)) or len(flipT_axes) != 3:
raise ValueError(f"Expected flipT_axes to be a list of 3 values, got: {flipT_axes}")
# Safely cast all entries to int
try:
flipT_axes = [int(v) for v in flipT_axes]
except Exception as e:
raise ValueError(f"flipT_axes must contain values convertible to int (0 or 1). Got: {flipT_axes}") from e
vprint(f"Flipping scan pattern with [flipup, fliplr, transpose] = {flipT_axes}", verbose=self.verbose)
# Convert the binary code to the indices of non-zero axis. E.g. scan_flipT = [0,1,1] => flip the axes = [1,2]
flipT_axes = np.nonzero(flipT_axes)[0]
if len(flipT_axes) > 0:
pos = pos.reshape(self.init_variables['N_scan_slow'], self.init_variables['N_scan_fast'], 2)
pos = np.flip(pos, flipT_axes)
pos = pos.reshape(-1, 2)
return pos
def _pos_scan_affine_transform(self, pos, scan_affine):
if scan_affine is not None:
(scale, asymmetry, rotation, shear) = scan_affine
vprint(f"Applying affine transformation to scan pattern with (scale, asymmetry, rotation, shear) = {(scale, asymmetry, rotation, shear)}", verbose=self.verbose)
pos = pos - pos.mean(0)
pos = pos @ compose_affine_matrix(scale, asymmetry, rotation, shear)
probe_shape = self.init_variables['probe_shape']
obj_shape = 1.2 * np.ceil(pos.max(0) - pos.min(0) + probe_shape)
pos = pos + np.ceil((np.array(obj_shape) / 2) - (np.array(probe_shape) / 2))
return pos
def _pos_scan_add_random_displacement(self, pos, scan_rand_std):
if scan_rand_std is not None:
vprint(f"Applying Gaussian distributed random displacement with std = {scan_rand_std} px to scan positions", verbose=self.verbose)
set_random_seed(seed=self.random_seed)
pos = pos + scan_rand_std * np.random.randn(*pos.shape)
return pos
###### Private methods for initializing object ######
def _load_obj(self):
"""
Load the object from the specified source.
"""
# Validate required fields
try:
obj_source = self.init_params['obj_source']
obj_params = self.init_params['obj_params']
except KeyError as e:
raise KeyError(f"Missing required configuration field: {e}")
vprint(f"Loading object from source = '{obj_source}'", verbose=self.verbose)
if obj_source == 'custom':
obj = obj_params
elif obj_source == 'PtyRAD':
obj = self._load_obj_from_ptyrad(obj_params)
elif obj_source == 'PtyShv':
obj = self._load_obj_from_ptyshv(obj_params)
elif obj_source == 'py4DSTEM':
obj = self._load_obj_from_py4dstem(obj_params)
elif obj_source == 'simu':
obj = self._simulate_obj(obj_params)
else:
raise ValueError(f"Unsupported object source '{obj_source}'. Use 'custom', 'PtyRAD', 'PtyShv', 'py4DSTEM', or 'simu'.")
return obj
def _load_obj_from_ptyrad(self, params: str):
pt_path = params
ckpt = self.cache_contents if self.use_cached_obj else load_ptyrad(pt_path)
obja = ckpt['optimizable_tensors']['obja']
objp = ckpt['optimizable_tensors']['objp']
obj = obja * np.exp(1j * objp)
return obj
def _load_obj_from_ptyshv(self, params: str):
mat_path = params
mat_version = get_matfile_version(mat_path)
use_h5py = (mat_version[0] == 2)
obj = self.cache_contents['object'] if self.use_cached_obj else load_mat(mat_path, key='object')
vprint(f"Input PtyShv object has original shape {obj.shape}, while default PtyShv order is (Ny, Nx, Nz)", verbose=self.verbose)
# First unify the axes order induced by loading with scipy / h5py, now it should be (Ny, Nx, Nz)
if use_h5py:
obj = obj.transpose(range(obj.ndim)[::-1])
vprint(f"Reverse array axes order because use_h5py = {use_h5py}, which automatically reverse the order", verbose=self.verbose)
vprint("Expanding and permuting PtyShv object dimension to make a final object shape with (omode, Nz, Ny, Nx)", verbose=self.verbose)
if len(obj.shape) == 2: # Single-slice ptycho
obj = obj[None, None, :, :]
elif len(obj.shape) == 3: # Multi-slice ptycho
obj = obj[None,].transpose(0, 3, 1, 2)
return obj
def _load_obj_from_py4dstem(self, params: str):
hdf5_path = params
obj = self.cache_contents['object'] if self.use_cached_obj else load_hdf5(hdf5_path, key='object')
vprint(f"Input py4DSTEM object has original shape {obj.shape}", verbose=self.verbose)
vprint("Expanding py4DSTEM object dimension to (omode, Nz, Ny, Nx)", verbose=self.verbose)
if len(obj.shape) == 2: # Single-slice ptycho
obj = obj[None, None, :, :]
elif len(obj.shape) == 3: # Multi-slice ptycho
obj = obj[None,]
return obj
def _simulate_obj(self, simu_params):
if simu_params is not None:
vprint("Using user-specified parameters in 'init_params['obj_params']' for initial object simulation.", verbose=self.verbose)
obj_shape = simu_params
if len(obj_shape) != 4:
raise ValueError(f"Input `obj_shape` = {obj_shape}, please provide a total dimension of 4 with (omode, Nz, Ny, Nx).")
else:
vprint("Using experimental parameters specified by 'init_params' for initial object simulation.", verbose=self.verbose)
omode = self.init_params['obj_omode_max']
Nz = self.init_params['obj_Nlayer']
try:
Ny, Nx = self.init_variables['obj_lateral_extent']
except KeyError:
vprint("WARNING: 'obj_lateral_extent' not found. Initializing positions first for obj_shape estimation...", verbose=self.verbose)
vprint(" ", verbose=self.verbose)
self.init_pos()
Ny, Nx = self.init_variables['obj_lateral_extent']
obj_shape = (omode, Nz, Ny, Nx)
set_random_seed(seed=self.random_seed)
obj = np.exp(1j * 1e-8 * np.random.rand(*obj_shape))
return obj
def _process_obj(self, obj):
"""
Process the loaded object, including z cropping, padding, resampling, and setting omode
"""
omode_max = self.init_params.get('obj_omode_max')
obj = self._obj_z_crop(obj, self.init_params.get('obj_z_crop'))
obj = self._obj_z_pad(obj, self.init_params.get('obj_z_pad'))
obj = self._obj_z_resample(obj, self.init_params.get('obj_z_resample'))
obj = self._object_set_omode_max(obj, omode_max)
return obj
def _obj_z_crop(self, obj, crop_range):
"""
Crop 4D complex object (omode, Nz, Ny, Nx) across depth (Nz) dimension:
[z_start, z_end]
Note that this method would also update the `self.init_params['obj_Nlayer']`
"""
if crop_range is None:
return obj
if len(crop_range) != 2:
raise ValueError(f"Expected crop range [z_start, z_end], got {crop_range}")
try:
z_start, z_end = crop_range
selected_slices = slice(z_start, z_end)
vprint(f"Cropping object depth from z_start: {z_start} to z_end: {z_end}", verbose=self.verbose)
except Exception as e:
raise ValueError(f"Invalid crop range for object depth: {crop_range}, object shape is {obj.shape}") from e
vprint(f"Current object has shape (omode, Nz, Ny, Nx) = {obj.shape}", verbose=self.verbose)
obj = obj[:,selected_slices,:,:]
vprint(f"Cropped object has shape (omode, Nz, Ny, Nx) = {obj.shape}", verbose=self.verbose)
# Update init_params['obj_Nlayer]
self.init_params['obj_Nlayer'] = obj.shape[1]
return obj
def _obj_z_pad(self, obj, pad_cfg):
"""
Pad 4D complex object (omode, Nz, Ny, Nx) along the depth (Nz) dimension.
Note that this method would also update the `self.init_params['obj_Nlayer']`
"""
if pad_cfg is None:
return obj
pad_layers = pad_cfg['pad_layers']
pad_types = pad_cfg['pad_types']
vprint(f"Current object has shape (omode, Nz, Ny, Nx) = {obj.shape}", verbose=self.verbose)
vprint(f"Padding object along depth with pad_layers = {pad_layers}, pad_types = {pad_types}", verbose=self.verbose)
# Assign variables
pad_layer_top, pad_layer_bottom = pad_layers
pad_type_top, pad_type_bottom = pad_types
# Helper function
def _create_z_pad(obj, num_layers, pad_type, top_or_bottom):
obja = np.abs(obj)
objp = np.angle(obj)
omode, nz, ny, nx = obj.shape
# Return an empty array with the same shape as obj but with nz = 0
if num_layers is None or num_layers == 0:
return np.empty((omode, 0, ny, nx), dtype=obj.dtype)
# Create new layers
else:
new_shape = (omode, num_layers, ny, nx)
if pad_type == 'vacuum':
new_layers_a = np.ones(new_shape)
new_layers_p = np.zeros(new_shape)
elif pad_type == 'mean':
new_layers_a = np.mean(obja, axis=1, keepdims=True) # The variance of amplitude along depth is usually quite small so geometric mean ~= arithmetric mean, although it might feel more natural to use geometric mean.
new_layers_p = np.mean(objp, axis=1, keepdims=True) # Note that this could be a bit biased if phase is not positively constrained. The shape is (omode, 1, Ny, Nx).
elif pad_type == 'edge':
if top_or_bottom == 'top':
edge_idx = [0] # Wrap it with [] to preserve the dimension
elif top_or_bottom == 'bottom':
edge_idx = [-1]
else:
raise ValueError(f"top_or_bottom expects 'top', or 'bottom', got {top_or_bottom}.")
new_layers_a = obja[:, edge_idx] # (omode, 1, Ny, Nx)
new_layers_p = objp[:, edge_idx]
else:
raise ValueError(f"Unsupported pad_type: {pad_type}, please use 'vacuum', 'mean', or 'edge'.")
new_layers = new_layers_a * np.exp(1j * new_layers_p)
return np.broadcast_to(new_layers, new_shape).copy().astype(obj.dtype)
top_layers = _create_z_pad(obj, num_layers=pad_layer_top, pad_type=pad_type_top, top_or_bottom='top')
bottom_layers = _create_z_pad(obj, num_layers=pad_layer_bottom, pad_type=pad_type_bottom, top_or_bottom='bottom')
obj = np.concatenate((top_layers, obj, bottom_layers), axis=1)
vprint(f"Padded object has shape (omode, Nz, Ny, Nx) = {obj.shape}", verbose=self.verbose)
# Update init_params['obj_Nlayer]
self.init_params['obj_Nlayer'] = obj.shape[1]
return obj
def _obj_z_resample(self, obj, resample_cfg):
"""
Resample 4D complex object (omode, Nz, Ny, Nx) along the depth (Nz) dimension.
Note that this method would also update the `self.init_params['obj_Nlayer']`,
`self.init_params['obj_slice_thickness']`, and `self.init_variables['slice_thickness']`
This is currently (v0.1.0b11) the only function in Initializer that uses PyTorch because the scipy.ndimage.zoom is just too slow...
"""
if resample_cfg is None or resample_cfg['mode'] is None:
return obj
# Assign variables
resample_mode = resample_cfg['mode']
resample_value = resample_cfg['value']
Nz_now = obj.shape[1]
dz_now = self.init_variables['slice_thickness'] # This was set by `set_variables_dict` using values in `init_params['obj_slice_thickness]`
length_unit = self.init_variables['length_unit']
# Print current status
vprint(f"Current object has shape (omode, Nz, Ny, Nx) = {(obj.shape)}", verbose=self.verbose)
vprint(f"Current object has slice thickness = {dz_now:.3f} {length_unit}", verbose=self.verbose)
vprint(f"Current object has mean(prod(amp, axis='depth')) = {np.mean(np.prod(np.abs(obj), axis=1)):.3f}, mean(sum(phase, axis='depth')) = {np.mean(np.sum(np.angle(obj), axis=1)):.3g}", verbose=self.verbose)
vprint(f"Resampling object along depth with resampling mode = '{resample_mode}', value = {resample_value}", verbose=self.verbose)
# Get resampled object and infer new slice thickness
obja_resample, objp_resample = complex_object_z_resample_torch(obj, dz_now, resample_mode, resample_value, output_type='amp_phase', return_np=True) # Output amplitude and phase separately so we can check the phase value directly
obj_resample = obja_resample * np.exp(1j * objp_resample) # (omode, Nz, Ny, Nx)
Nz_new = obj_resample.shape[1]
dz_new = dz_now * Nz_now / Nz_new
# Print warning if there's phase wrapping
if objp_resample.max() > 2*np.pi:
vprint(f"Warning: Resampled object phase has a maximum value = {objp_resample.max():.3f} > 2pi, this would cause phase wrapping, try using thinner slices.")
# Update Nlayer and slice thickness
self.init_params['obj_Nlayer'] = obj_resample.shape[1]
self.init_params['obj_slice_thickness'] = dz_new
self.init_variables['slice_thickness'] = dz_new
# Print final status
vprint(f"Resampled object has shape (omode, Nz, Ny, Nx) = {(obj_resample.shape)}", verbose=self.verbose)
vprint(f"Resampled object has slice thickness = {dz_new:.3f} {length_unit}", verbose=self.verbose)
vprint(f"Resampled object has mean(prod(amp, axis='depth')) = {np.mean(np.prod(np.abs(obj_resample), axis=1)):.3f}, mean(sum(phase, axis='depth')) = {np.mean(np.sum(np.angle(obj_resample), axis=1)):.3g}", verbose=self.verbose)
return obj_resample
def _object_set_omode_max(self, obj, omode_max):
"""
Either cap or pad the omode for mixed state object
"""
omode_now = obj.shape[0]
if omode_now > omode_max:
vprint(f"omode_now: {omode_now} and omode_max: {omode_max}, capping the omode.", verbose=self.verbose)
obj_final = obj[:omode_max]
elif omode_now == omode_max:
vprint(f"omode_now: {omode_now} and omode_max: {omode_max}, leaving the omode unchanged.", verbose=self.verbose)
obj_final = obj
else: # omode_now <= omode_max: # Need to pad new probe modes
vprint(f"omode_now: {omode_now} and omode_max: {omode_max}, padding the omode.", verbose=self.verbose)
num_new_modes = omode_max - omode_now
vprint(f"Creating {num_new_modes} new object modes from the mean and std of original object modes", verbose=self.verbose)
# Assign variables
obja = np.abs(obj)
objp = np.angle(obj)
spatial_dims = obj[0].shape # (z,y,x)
obja_mean = np.mean(obja, axis=0, keepdims=True)
objp_mean = np.mean(objp, axis=0, keepdims=True)
if omode_now == 1: # There's no std when omode=1. This is rather rudimentary and we'll introduced some spatially structured noise in future release
obja_std = 5e-4 * obja_mean
objp_std = 0.20 * objp_mean
else:
obja_std = np.std(obja, axis=0, keepdims=True)
objp_std = np.std(objp, axis=0, keepdims=True)
# Create new modes from random variable eps, note that amplitude and phase are perfectly correlated here
set_random_seed(seed=self.random_seed)
eps = np.random.randn(num_new_modes,*spatial_dims) # (num_new_modes, z, y, x)
obja_new = obja_mean + eps * obja_std
objp_new = objp_mean + eps * objp_std
# Check min and max
obja_new = np.clip(obja_new, a_min=np.min(obja), a_max=np.max(obja))
objp_new = np.clip(objp_new, a_min=np.min(objp), a_max=np.max(objp))
# Recombine amplitude and phase back to complex-valued obj
new_modes = obja_new * np.exp(1j * objp_new)
obj_final = np.concatenate((obj, new_modes), axis=0)
return obj_final