Skip to content

Commit

Permalink
splitting up constraints
Browse files Browse the repository at this point in the history
  • Loading branch information
@gvarnavi
gvarnavi committed Jan 5, 2024
1 parent 6223139 commit 85d1b76
Show file tree
Hide file tree
Showing 7 changed files with 336 additions and 967 deletions.
9 changes: 9 additions & 0 deletions py4DSTEM/process/phase/iterative_base_class.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1980,6 +1980,15 @@ def _report_reconstruction_summary(
)
)

def _constraints(self, current_object, current_probe, current_positions, **kwargs):
"""Wrapper function for all classes to inherit"""

current_object = self._object_constraints(current_object, **kwargs)
current_probe = self._probe_constraints(current_probe, **kwargs)
current_positions = self._positions_constraints(current_positions, **kwargs)

return current_object, current_probe, current_positions

@property
def angular_sampling(self):
"""Angular sampling [mrad]"""
Expand Down
209 changes: 1 addition & 208 deletions py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -647,213 +647,6 @@ def preprocess(

return self

def _constraints(
self,
current_object,
current_probe,
current_positions,
fix_com,
fit_probe_aberrations,
fit_probe_aberrations_max_angular_order,
fit_probe_aberrations_max_radial_order,
constrain_probe_amplitude,
constrain_probe_amplitude_relative_radius,
constrain_probe_amplitude_relative_width,
constrain_probe_fourier_amplitude,
constrain_probe_fourier_amplitude_max_width_pixels,
constrain_probe_fourier_amplitude_constant_intensity,
fix_probe_aperture,
initial_probe_aperture,
fix_positions,
global_affine_transformation,
gaussian_filter,
gaussian_filter_sigma,
butterworth_filter,
q_lowpass,
q_highpass,
butterworth_order,
kz_regularization_filter,
kz_regularization_gamma,
identical_slices,
object_positivity,
shrinkage_rad,
object_mask,
pure_phase_object,
tv_denoise_chambolle,
tv_denoise_weight_chambolle,
tv_denoise_pad_chambolle,
tv_denoise,
tv_denoise_weights,
tv_denoise_inner_iter,
orthogonalize_probe,
):
"""
Ptychographic constraints operator.
Parameters
--------
current_object: np.ndarray
Current object estimate
current_probe: np.ndarray
Current probe estimate
current_positions: np.ndarray
Current positions estimate
fix_com: bool
If True, probe CoM is fixed to the center
fit_probe_aberrations: bool
If True, fits the probe aberrations to a low-order expansion
fit_probe_aberrations_max_angular_order: bool
Max angular order of probe aberrations basis functions
fit_probe_aberrations_max_radial_order: bool
Max radial order of probe aberrations basis functions
constrain_probe_amplitude: bool
If True, probe amplitude is constrained by top hat function
constrain_probe_amplitude_relative_radius: float
Relative location of top-hat inflection point, between 0 and 0.5
constrain_probe_amplitude_relative_width: float
Relative width of top-hat sigmoid, between 0 and 0.5
constrain_probe_fourier_amplitude: bool
If True, probe aperture is constrained by fitting a sigmoid for each angular frequency.
constrain_probe_fourier_amplitude_max_width_pixels: float
Maximum pixel width of fitted sigmoid functions.
constrain_probe_fourier_amplitude_constant_intensity: bool
If True, the probe aperture is additionally constrained to a constant intensity.
fix_probe_aperture: bool
If True, probe Fourier amplitude is replaced by initial_probe_aperture
initial_probe_aperture: np.ndarray
Initial probe aperture to use in replacing probe Fourier amplitude
fix_positions: bool
If True, positions are not updated
gaussian_filter: bool
If True, applies real-space gaussian filter in A
gaussian_filter_sigma: float
Standard deviation of gaussian kernel
butterworth_filter: bool
If True, applies fourier-space butterworth filter
q_lowpass: float
Cut-off frequency in A^-1 for low-pass butterworth filter
q_highpass: float
Cut-off frequency in A^-1 for high-pass butterworth filter
butterworth_order: float
Butterworth filter order. Smaller gives a smoother filter
kz_regularization_filter: bool
If True, applies fourier-space arctan regularization filter
kz_regularization_gamma: float
Slice regularization strength
identical_slices: bool
If True, forces all object slices to be identical
object_positivity: bool
If True, forces object to be positive
shrinkage_rad: float
Phase shift in radians to be subtracted from the potential at each iteration
object_mask: np.ndarray (boolean)
If not None, used to calculate additional shrinkage using masked-mean of object
pure_phase_object: bool
If True, object amplitude is set to unity
tv_denoise_chambolle: bool
If True, performs TV denoising along z
tv_denoise_weight_chambolle: float
weight of tv denoising constraint
tv_denoise_pad_chambolle: int
If not None, pads object at top and bottom with this many zeros before applying denoising
tv_denoise: bool
If True, applies TV denoising on object
tv_denoise_weights: [float,float]
Denoising weights[z weight, r weight]. The greater `weight`,
the more denoising.
tv_denoise_inner_iter: float
Number of iterations to run in inner loop of TV denoising
orthogonalize_probe: bool
If True, probe will be orthogonalized
Returns
--------
constrained_object: np.ndarray
Constrained object estimate
constrained_probe: np.ndarray
Constrained probe estimate
constrained_positions: np.ndarray
Constrained positions estimate
"""

if gaussian_filter:
current_object = self._object_gaussian_constraint(
current_object, gaussian_filter_sigma, pure_phase_object
)

if butterworth_filter:
current_object = self._object_butterworth_constraint(
current_object,
q_lowpass,
q_highpass,
butterworth_order,
)

if identical_slices:
current_object = self._object_identical_slices_constraint(current_object)
elif kz_regularization_filter:
current_object = self._object_kz_regularization_constraint(
current_object,
kz_regularization_gamma,
z_padding=1,
)
elif tv_denoise:
current_object = self._object_denoise_tv_pylops(
current_object,
tv_denoise_weights,
tv_denoise_inner_iter,
z_padding=1,
)
elif tv_denoise_chambolle:
current_object = self._object_denoise_tv_chambolle(
current_object,
tv_denoise_weight_chambolle,
axis=0,
padding=tv_denoise_pad_chambolle,
)

if shrinkage_rad > 0.0 or object_mask is not None:
current_object = self._object_shrinkage_constraint(
current_object,
shrinkage_rad,
object_mask,
)

if self._object_type == "complex":
current_object = self._object_threshold_constraint(
current_object, pure_phase_object
)
elif object_positivity:
current_object = self._object_positivity_constraint(current_object)

if fix_com:
current_probe = self._probe_center_of_mass_constraint(current_probe)

# These constraints don't _really_ make sense for mixed-state
if fix_probe_aperture:
raise NotImplementedError()
elif constrain_probe_fourier_amplitude:
raise NotImplementedError()
if fit_probe_aberrations:
raise NotImplementedError()
if constrain_probe_amplitude:
raise NotImplementedError()

if orthogonalize_probe:
current_probe = self._probe_orthogonalization_constraint(current_probe)

if not fix_positions:
current_positions = self._positions_center_of_mass_constraint(
current_positions
)

if global_affine_transformation:
current_positions = self._positions_affine_transformation_constraint(
self._positions_px_initial, current_positions
)

return current_object, current_probe, current_positions

def reconstruct(
self,
max_iter: int = 8,
Expand Down Expand Up @@ -1058,8 +851,8 @@ def reconstruct(
projection_b,
projection_c,
normalization_min,
step_size,
max_batch_size,
step_size,
)

# batching
Expand Down
Loading

0 comments on commit 85d1b76

Please sign in to comment.