diff --git a/src/quantem/diffractive_imaging/direct_ptycho_utils.py b/src/quantem/diffractive_imaging/direct_ptycho_utils.py index fcc76601..ce08bfeb 100644 --- a/src/quantem/diffractive_imaging/direct_ptycho_utils.py +++ b/src/quantem/diffractive_imaging/direct_ptycho_utils.py @@ -37,6 +37,12 @@ # fmt: on +def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None: + if rotation_angle is None: + return None + return math.radians(float(rotation_angle)) + + def create_edge_window(shape, edge_blend_pixels, device="cpu"): """ Create a smooth edge window that transitions from 0 at edges to 1 in center. @@ -495,7 +501,10 @@ def fit_aberrations_from_shifts( gpts: tuple[int, int], sampling: tuple[float, float], ) -> dict[str, float]: - """ """ + """Fit low-order aberrations from lateral shifts. + + Returns ``rotation_angle`` in degrees. + """ device = shifts_ang.device # Get spatial frequencies at BF positions @@ -534,7 +543,7 @@ def fit_aberrations_from_shifts( "C10": C10.item(), "C12": C12.item(), "phi12": phi12.item(), - "rotation_angle": rotation_rad.item(), + "rotation_angle": torch.rad2deg(rotation_rad).item(), } diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py index ecbeece0..912b4bd2 100644 --- a/src/quantem/diffractive_imaging/direct_ptychography.py +++ b/src/quantem/diffractive_imaging/direct_ptychography.py @@ -19,6 +19,7 @@ validate_int, validate_tensor, ) +from quantem.core.visualization import show_2d from quantem.diffractive_imaging.complex_probe import ( aberration_surface, aberration_surface_cartesian_basis, @@ -44,6 +45,7 @@ from quantem.diffractive_imaging.direct_ptycho_utils import ( ABERRATION_PRESETS, _crop_corner_centered_mask, + _rotation_degrees_to_radians, align_vbf_stack_multiscale, create_edge_window, fit_aberrations_from_shifts, @@ -156,13 +158,13 @@ def add(name: str, value): if self.initial_aberrations: add("initial_aberrations", self.initial_aberrations) if self.initial_rotation_angle is not None: - add("initial_rotation_angle", self.initial_rotation_angle) + add("initial_rotation_angle_deg", self.initial_rotation_angle) elif which == "optimized": if self.optimized_aberrations: add("optimized_aberrations", self.optimized_aberrations) if self.optimized_rotation_angle is not None: - add("optimized_rotation_angle", self.optimized_rotation_angle) + add("optimized_rotation_angle_deg", self.optimized_rotation_angle) elif which == "current": current_abers = self.current_aberrations(override_aberration_coefs) @@ -171,17 +173,17 @@ def add(name: str, value): if current_abers: add("current_aberrations", current_abers) if current_rot is not None: - add("current_rotation_angle", current_rot) + add("current_rotation_angle_deg", current_rot) elif which == "all": if self.initial_aberrations: add("initial_aberrations", self.initial_aberrations) if self.initial_rotation_angle is not None: - add("initial_rotation_angle", self.initial_rotation_angle) + add("initial_rotation_angle_deg", self.initial_rotation_angle) if self.optimized_aberrations: add("optimized_aberrations", self.optimized_aberrations) if self.optimized_rotation_angle is not None: - add("optimized_rotation_angle", self.optimized_rotation_angle) + add("optimized_rotation_angle_deg", self.optimized_rotation_angle) else: raise ValueError( @@ -335,7 +337,7 @@ def from_dataset4d( if rotation_angle is None: origin.estimate_detector_rotation() - rotation_angle = origin.detector_rotation_deg / 180 * math.pi + rotation_angle = origin.detector_rotation_deg # shift to origin origin.shift_origin_to( @@ -494,6 +496,7 @@ def bf_mask(self, value: torch.Tensor): @property def rotation_angle(self) -> float: + """Current detector rotation angle in degrees.""" return self.hyperparameter_state.current_rotation_angle() @property @@ -679,7 +682,7 @@ def reconstruct( upsampling_factor : int, optional Factor by which to upsample the reconstruction override_rotation_angle : float, optional - Rotation angle for coordinate system + Rotation angle for coordinate system, in degrees max_batch_size : int, optional Maximum batch size for processing deconvolution_kernel : str, one of ['ssb', 'obf', 'mf','prlx','icom'] @@ -742,7 +745,10 @@ def reconstruct( # Get k-space grid kxa, kya = spatial_frequencies( - self.gpts, self.sampling, rotation_angle=rotation_angle, device=self.device + self.gpts, + self.sampling, + rotation_angle=_rotation_degrees_to_radians(rotation_angle), + device=self.device, ) k, phi = polar_coordinates(kxa, kya) @@ -900,6 +906,62 @@ def obj(self) -> np.ndarray: obj = to_numpy(self.corrected_bf) return obj + def visualize( + self, + return_fig: bool = False, + show_obj_fft: bool = True, + apply_hanning_window: bool = False, + **kwargs, + ): + """ + Show the reconstructed object and its Hann-windowed Fourier transform. + + Parameters + ---------- + cbar : bool, optional + Whether to show colorbars, by default True. + return_fig : bool, optional + If True, return ``(fig, axs)``. + fft_norm : str | dict, optional + Normalization passed to ``show_2d`` for the object FFT. + **kwargs + Additional arguments passed to ``show_2d``. + """ + if self.corrected_bf is None: + raise RuntimeError("Run reconstruct() before visualize().") + + obj = self.obj + obj_scalebar = {"sampling": self.scan_sampling[1], "units": "Å"} + + if show_obj_fft: + if apply_hanning_window: + window = np.hanning(obj.shape[-2])[:, None] * np.hanning(obj.shape[-1])[None, :] + obj_fft = np.fft.fftshift(np.abs(np.fft.fft2(obj * window))) + else: + obj_fft = np.fft.fftshift(np.abs(np.fft.fft2(obj))) + + fft_sampling = 1 / (self.scan_sampling[1] * obj.shape[-1]) + fft_scalebar = {"sampling": fft_sampling, "units": r"$\mathrm{A^{-1}}$"} + + fig, axs = show_2d( + [obj, obj_fft], + title=["Object phase", "Object phase FFT"], + scalebar=[obj_scalebar, fft_scalebar], + **kwargs, + ) + axs[1].set_aspect(obj.shape[-1] / obj.shape[-2]) + else: + fig, axs = show_2d( + obj, + title="Object phase", + scalebar=obj_scalebar, + **kwargs, + ) + + if return_fig: + return fig, axs + return None + def optimize_hyperparameters( self, aberration_coefs: dict[str, float | OptimizationParameter] | None = None, @@ -917,7 +979,7 @@ def optimize_hyperparameters( aberration_coefs : dict[str, float|OptimizationParameter] Dict of aberration names to either fixed values or optimization ranges. rotation_angle : float|OptimizationParameter - Fixed rotation or optimization range. + Fixed rotation or optimization range, in degrees. n_trials : int Number of Optuna trials. sampler : optuna.samplers.BaseSampler, optional @@ -1107,7 +1169,10 @@ def _return_lateral_shifts( ): # Get initial shifts kxa, kya = spatial_frequencies( - self.gpts, self.sampling, rotation_angle=rotation_angle, device=self.device + self.gpts, + self.sampling, + rotation_angle=_rotation_degrees_to_radians(rotation_angle), + device=self.device, ) k, phi = polar_coordinates(kxa, kya) @@ -1335,7 +1400,7 @@ def _fit_hyperparameters_least_squares_inner( kxa, kya = spatial_frequencies( self.gpts, self.sampling, - rotation_angle=rotation_angle, + rotation_angle=_rotation_degrees_to_radians(rotation_angle), device=device, ) @@ -1511,7 +1576,7 @@ def fit_hyperparameters_least_squares( Args: aberration_coefs: Initial aberration coefficients to deconvolve - rotation_angle: Rotation angle for basis functions + rotation_angle: Rotation angle for basis functions, in degrees cartesian_basis: Aberration basis to fit. Can be: - str: preset name like "low_order" - list[str]: explicit list like ["C10", "C12_a", "C12_b", ...]