pmace

pmace.pmace

Functions:

PMACE(func, *args, **kwargs)	This class is a decorator that can be used to prepare a function before it is called.
get_data_fit_pt(cur_est, joint_est, y_meas)	Data-fitting point.
pixel_weighted_avg_op(projected_patches, ...)	Calculate the pixel-weighted average of projected patches and reallocates the result.
add_probe_mode(probe_modes, ...[, ...])	Add a new probe mode to the existing list of probe modes and update the probe dictionary.
pmace_recon(y_meas, patch_bounds, init_obj)	Projected Multi-Agent Consensus Equilibrium (PMACE).

class pmace.pmace.PMACE(func, *args, **kwargs)

Bases: object

This class is a decorator that can be used to prepare a function before it is called.

Parameters:

• func (function) – The function to be decorated.

• *args – Positional arguments to be passed to the decorated function.

• **kwargs – Keyword arguments to be passed to the decorated function.

__init__(func, *args, **kwargs)

pmace.pmace.add_probe_mode(probe_modes, projected_patches, y_intsty, probe_dict, energy_ratio=0.05, fresnel_propagation=False, dx=None, wavelength=None, propagation_dist=None)

Add a new probe mode to the existing list of probe modes and update the probe dictionary.

Parameters:

• probe_modes (numpy.ndarray) – Array of existing probe modes.

• projected_patches (numpy.ndarray) – The estimates of projected patches.

• y_intsty (numpy.ndarray) – The intensity measurement.

• probe_dict (dictionary) – Dictionary containing probe modes.

• energy_ratio (float, optional) – Ratio of energy in the new probe mode compared to the existing ones. Default is 0.05 (empirically selected).

• fresnel_propagation (bool, optional) – Flag for performing Fresnel propagation.

• dx (float, optional) – Sampling interval at source plane.

• wavelength (float, optional) – Wavelength of the imaging radiation.

• propagation_dist (float, optional) – Propagation distance.

Returns:

numpy.ndarray – Updated array of probe modes with the newly added probe mode. dictionary: Updated dictionary of probe modes with the newly added probe mode.

pmace.pmace.get_data_fit_pt(cur_est, joint_est, y_meas)

Data-fitting point.

The function calculates the closest data fitting point using given measurement and current estimate.

Parameters:

• cur_est (numpy.ndarray) – Current estimate of projected images or complex probe.

• joint_est (numpy.ndarray) – Current estimate of complex probe or projected images.

• y_meas (numpy.ndarray) – Measurements.

Returns:

numpy.ndarray – Closest data-fitting point.

pmace.pmace.pixel_weighted_avg_op(projected_patches, probe_modes, mode_weights, patch_bounds, image_sz, probe_exp=1, regularization=False, bm3d_psd=0.1, blk_idx=None)

Calculate the pixel-weighted average of projected patches and reallocates the result.

Parameters:

• projected_patches (numpy.ndarray) – Array of projected patches.

• probe_modes (list of numpy.ndarrays) – List of probe modes.

• mode_weights (list) – list of mode weights.

• patch_bounds (numpy.ndarray) – List of patch coordinates.

• image_sz (tuple) – The dimensions of the output image.

• probe_exp (float, optional) – Exponent for probe mode weight calculation. Default is 1.

• regularization (bool, optional) – Apply regularization using BM3D denoising if True. Default is False.

• bm3d_psd (float, optional) – Power spectral density parameter for BM3D denoising. Default is 0.1.

• blk_idx (list, optional) – Block indices for BM3D denoising region. Default is None.

Returns:

numpy.ndarray – Resulting complex image. numpy.ndarray: Updated image patches.

pmace.pmace.pmace_recon(y_meas, patch_bounds, init_obj, init_probe=None, ref_obj=None, ref_probe=None, num_iter=100, joint_recon=False, recon_win=None, save_dir=None, obj_data_fit_prm=0.5, probe_data_fit_prm=0.5, rho=0.5, probe_exp=1.5, obj_exp=0, add_reg=False, sigma=0.02, probe_center_correction=[], reinit=False, scan_loc_refinement_iterations=[], scan_loc_search_step_sz=2, scan_loc_refine_step_sz=1, gt_scan_loc=None, add_mode=[], energy_ratio=0.05, img_px_sz=4.52e-09, wavelength=1.24e-09, propagation_dist=1e-07, orthogonalize_modes=[], gamma=2)

Projected Multi-Agent Consensus Equilibrium (PMACE).

Parameters:

• y_meas (numpy.ndarray) – Pre-processed measurements (diffraction patterns / intensity data).

• patch_bounds (list) – Scan coordinates of projections.

• init_obj (numpy.ndarray) – Formulated initial guess of the complex object.

• init_probe (numpy.ndarray) – Formulated initial guess of the complex probe.

• ref_obj (numpy.ndarray) – Complex reference image for the object.

• ref_probe (numpy.ndarray) – Complex reference image for the probe.

• num_iter (int) – Number of iterations.

• joint_recon (bool) – Option to estimate the complex probe for blind ptychography.

• recon_win (numpy.ndarray) – Pre-defined window for showing and comparing reconstruction results.

• save_dir (str) – Directory to save reconstruction results.

• obj_data_fit_prm (float) – Averaging weight in the object update function. Parameter near 1 gives a closer fit to data.

• probe_data_fit_prm (float) – Averaging weight in the probe update function.

• rho (float) – Mann averaging parameter.

• probe_exp (float) – Exponent of probe weighting in the consensus calculation of the image estimate.

• obj_exp (float) – Exponent of image weighting in the consensus calculation of the probe estimate.

• add_reg (bool) – Option to apply denoiser.

• sigma (float) – Denoising parameter.

• scan_loc_refinement_iterations (list) – List of iterations when scan location refinement is applied.

• scan_loc_search_step_sz (int) – Step size for scan location search.

• scan_loc_refine_step_sz (float) – Step size for refining scan locations.

• gt_scan_loc (numpy.ndarray) – Ground truth scan locations.

• add_mode (list) – The index of reconstruction iterations to add new probe modes.

• gamma (int) – Power parameter for energy weighting.

Returns:

dict –

Reconstructed complex images and NRMSE between reconstructions and reference images.

Keys:

• ’object’ (numpy.ndarray): Reconstructed complex object.

• ’probe’ (numpy.ndarray): Reconstructed complex probe.

• ’err_obj’ (list): NRMSE values for the object reconstructions.

• ’err_probe’ (list): NRMSE values for the probe reconstructions (if joint_recon is True).

• ’err_meas’ (list): NRMSE values for the measured data.

pmace.utils

Functions:

gen_syn_data(cmplx_obj, cmplx_probe, ...[, ...])	Simulate ptychographic intensity measurements.
load_measurement(fpath)	Read measurements from a path and pre-process data.
patch2img(img_patches, patch_coords, img_sz)	Project image patches to a full-sized image with weights.
img2patch(full_img, patch_coords, patch_sz)	Extract image patches from full-sized image.

pmace.utils.gen_syn_data(cmplx_obj, cmplx_probe, patch_bounds, add_noise=True, peak_photon_rate=100000.0, shot_noise_pm=0.5, save_dir=None)

Simulate ptychographic intensity measurements.

Parameters:

• cmplx_obj (numpy.ndarray) – Complex object.

• cmplx_probe (numpy.ndarray) – Complex probe.

• patch_bounds (list of tuple) – Scan coordinates of projections.

• add_noise (bool) – Option to add noise to data.

• peak_photon_rate (float) – Peak rate of photon detection at the detector.

• shot_noise_pm (float) – Expected number of Poisson distributed dark current noise.

• save_dir (str) – Directory for saving generated data.

Returns:

numpy.ndarray – Simulated ptychographic data.

pmace.utils.img2patch(full_img, patch_coords, patch_sz)

Extract image patches from full-sized image.

Parameters:

• img (numpy.ndarray) – Full-sized image.

• coords (numpy.ndarray) – Coordinates of projections.

• patch_sz (tuple) – Size of the output patches (rows, columns).

Returns:

list of numpy.ndarray – Projected image patches.

pmace.utils.load_measurement(fpath)

Read measurements from a path and pre-process data.

Parameters:

fpath (str) – File directory.

Returns:

numpy.ndarray – Pre-processed measurement (square root of non-negative data).

pmace.utils.patch2img(img_patches, patch_coords, img_sz, norm_wgt=None)

Project image patches to a full-sized image with weights.

Parameters:

• img_patches (list of numpy.ndarray) – Projected image patches.

• patch_coords (numpy.ndarray) – Coordinates of projections.

• img_sz (tuple) – Size of the full image (rows, columns).

• norm_wgt (numpy.ndarray) – Normalization weight.

Returns:

numpy.ndarray – The full-sized complex image.

pmace.nrmse

Functions:

compute_nrmse(input_img, ref_img[, cstr])	Compute Normalized Root-Mean-Square-Error (NRMSE) between two images.
pha_err(img, ref_img)	Calculate the phase error between complex images.
phase_norm(img, ref_img[, cstr])	Perform phase normalization on reconstructed image.

pmace.nrmse.compute_nrmse(input_img, ref_img, cstr=None)

Compute Normalized Root-Mean-Square-Error (NRMSE) between two images.

This function calculates the NRMSE between the provided image and reference image.

Parameters:

• input_img (numpy.ndarray) – Complex-valued image for comparison.

• ref_img (numpy.ndarray) – Reference image for comparison.

• cstr (numpy.ndarray, optional) – Area for comparison. If provided, only this region will be considered.

Returns:

float – The computed NRMSE between two images.

pmace.nrmse.pha_err(img, ref_img)

Calculate the phase error between complex images.

The phase error is determined as the minimum of | angle(img) - angle(ref_img) - 2*k*pi| where k belongs to {-1, 0, 1}.

Parameters:

• img (numpy.ndarray) – Complex-valued image for comparison.

• ref_img (numpy.ndarray) – Reference image for comparison.

Returns:

numpy.ndarray – Phase error between two images.

pmace.nrmse.phase_norm(img, ref_img, cstr=None)

Perform phase normalization on reconstructed image.

Since reconstruction is blind to the absolute phase of the ground truth image, this function applies a phase shift to the reconstruction results given the known ground truth image.

Parameters:

• img (numpy.ndarray) – The reconstruction that needs phase normalization.

• ref_img (numpy.ndarray) – The known ground truth image or reference image.

• cstr (numpy.ndarray, optional) – Preconditioning window. If provided, only this region will be considered.

Returns:

numpy.ndarray – The phase-normalized reconstruction as a complex image (dtype: np.complex64).

pmace.display

Functions:

plot_scan_pt(scan_pt, save_dir)	Plot scan points.
plot_nrmse(nrmse_ls, title, label[, ...])	Plot the NRMSE (Normalized Root Mean Squared Error) versus the number of iterations.
plot_cmplx_img(cmplx_img[, img_title, ...])	Plot complex object images and error images compared with a reference image.

pmace.display.plot_scan_pt(scan_pt, save_dir)

Plot scan points.

Parameters:

• scan_pt (numpy.ndarray) – Array of scan points as (x, y) coordinates.

• save_dir (str) – Directory to save the plot.

pmace.display.plot_nrmse(nrmse_ls, title, label, abscissa=None, step_sz=15, fig_sz=[10, 4.8], display=False, save_fname=None)

Plot the NRMSE (Normalized Root Mean Squared Error) versus the number of iterations.

Parameters:

• nrmse_ls (list or array) – List of NRMSE values for each iteration or a dictionary of labels and NRMSE values.

• title (str) – Title for the plot.

• label (list) – List containing the X and Y axis labels and the label for the legend (e.g., [‘X Label’, ‘Y Label’, ‘Legend Label’]).

• abscissa (list or None) – X-axis values corresponding to NRMSE data. If None, it is automatically generated.

• step_sz (int) – Step size for X-axis ticks.

• fig_sz (list) – Size of the figure in inches (width, height).

• display (bool) – Display the plot if True.

• save_fname (str or None) – Save the plot to a file with the specified filename (without extension).

pmace.display.plot_cmplx_img(cmplx_img, img_title='img', ref_img=None, display_win=None, display=False, save_fname=None, fig_sz=[8, 3], mag_vmax=1, mag_vmin=0, phase_vmax=3.141592653589793, phase_vmin=-3.141592653589793, real_vmax=1, real_vmin=0, imag_vmax=0, imag_vmin=-1)

Plot complex object images and error images compared with a reference image.

Parameters:

• cmplx_img (numpy.ndarray) – Complex-valued image.

• img_title (str) – Title for the complex image.

• ref_img (numpy.ndarray or None) – Reference image. If provided, error images will be displayed.

• display_win (numpy.ndarray or None) – Pre-defined window for displaying images.

• display (bool) – Display images if True.

• save_fname (str or None) – Save images to the specified file directory.

• fig_sz (list) – Size of image plots in inches (width, height).

• mag_vmax (float) – Maximum value for showing image magnitude.

• mag_vmin (float) – Minimum value for showing image magnitude.

• phase_vmax (float) – Maximum value for showing image phase.

• phase_vmin (float) – Minimum value for showing image phase.

• real_vmax (float) – Maximum value for showing the real part of the image.

• real_vmin (float) – Minimum value for showing the real part of the image.

• imag_vmax (float) – Maximum value for showing the imaginary part of the image.

• imag_vmin (float) – Minimum value for showing the imaginary part of the image.