Adding trefoil astigmatism to Parallax.guess_common_aberrations() ?

[alouiset]

Hi py4DSTEM team members!
We are struggling with residual trefoil aberrations (A2 and/or D4) in our microscope. I would like to see what would be the effect of a manual trefoil compensation during parallax reconstruction using the guess_common_aberrations() function. Do you think you could add trefoil astigmatism (C23 and phi23) to the list of aberration parameters in this function?
Thank you!

[alouiset]

Actually, it was pretty easy to include. Here is the modified function if you want to test it. Seems to work with the wrapper function as well.

def guess_common_aberrations(
    self,
    rotation_angle_deg=0,
    transpose=False,
    kde_upsample_factor=None,
    kde_sigma_px=0.125,
    kde_lowpass_filter=False,
    lanczos_interpolation_order=None,
    defocus=0,
    astigmatism=0,
    astigmatism_angle_deg=0,
    coma=0,
    coma_angle_deg=0,
    trefoil_astigmatism=0,
    trefoil_astigmatism_angle_deg=0,
    spherical_aberration=0,
    max_batch_size=None,
    plot_shifts_and_aligned_bf=True,
    return_shifts_and_aligned_bf=False,
    plot_arrow_freq=1,
    scale_arrows=1,
    **kwargs,
):
    """
    Generates analytical BF shifts and uses them to align the virtual BF stack,
    based on the experimental geometry (rotation, transpose), and common aberrations.

    Parameters
    ----------
    rotation_angle_deg: float, optional
        Relative rotation between the scan and the diffraction space coordinate systems
    transpose: bool, optional
        Whether the diffraction intensities are transposed
    kde_upsample_factor: int, optional
        Real-space upsampling factor
    kde_sigma_px: float, optional
        KDE gaussian kernel bandwidth in non-upsampled pixels
    kde_lowpass_filter: bool, optional
        If True, the resulting KDE upsampled image is lowpass-filtered using a sinc-function
    lanczos_interpolation_order: int, optional
        If not None, Lanczos interpolation with the specified order is used instead of bilinear
    defocus: float, optional
        Defocus value to use in computing analytical BF shifts
    astigmatism: float, optional
        Astigmatism value to use in computing analytical BF shifts
    astigmatism_angle_deg: float, optional
        Astigmatism angle to use in computing analytical BF shifts
    coma: float, optional
        Coma value to use in computing analytical BF shifts
    coma_angle_deg: float, optional
        Coma angle to use in computing analytical BF shifts
    trefoil_astigmatism: float, optional
        Trefoil astigmatism value to use in computing analytical BF shifts
    trefoil_astigmatism_angle_deg: float, optional
        Trefoil astigmatism angle to use in computing analytical BF shifts
    spherical_aberration: float, optional
        Spherical aberration value to use in computing analytical BF shifts
    max_batch_size: int, optional
        Max number of virtual BF images to use at once in computing cross-correlation
    plot_shifts_and_aligned_bf: bool, optional
        If True, the analytical shifts and the aligned virtual VF image are plotted
    return_shifts_and_aligned_bf: bool, optional
        If True, the analytical shifts and the aligned virtual VF image are returned
    plot_arrow_freq: int, optional
        Frequency of shifts to plot in quiver plot
    scale_arrows: float, optional
        Scale to multiply shifts by

    """
    xp = self._xp
    asnumpy = self._asnumpy

    if not hasattr(self, "_recon_BF"):
        raise ValueError(
            (
                "Aberration guessing is meant to be ran after preprocessing. "
                "Please run the `preprocess()` function first."
            )
        )

    # aberrations_coefs
    aberrations_mn = [
        [1, 0, 0],
        [1, 2, 0],
        [1, 2, 1],
        [2, 1, 0],
        [2, 1, 1],
        [2, 3, 0],
        [2, 3, 1],
        [3, 0, 0],
    ]
    astigmatism_x = astigmatism * np.cos(np.deg2rad(astigmatism_angle_deg) * 2)
    astigmatism_y = astigmatism * np.sin(np.deg2rad(astigmatism_angle_deg) * 2)
    coma_x = coma * np.cos(np.deg2rad(coma_angle_deg) * 1)
    coma_y = coma * np.sin(np.deg2rad(coma_angle_deg) * 1)
    trefoil_astigmatism_x = trefoil_astigmatism * np.cos(np.deg2rad(trefoil_astigmatism_angle_deg) * 3)
    trefoil_astigmatism_y = trefoil_astigmatism * np.sin(np.deg2rad(trefoil_astigmatism_angle_deg) * 3)
    aberrations_coefs = xp.array(
        [
            -defocus,
            astigmatism_x,
            astigmatism_y,
            coma_x,
            coma_y,
            trefoil_astigmatism_x,
            trefoil_astigmatism_y,
            spherical_aberration,
        ]
    )

    # transpose rotation matrix
    if transpose:
        rotation_angle_deg *= -1

    # aberrations_basis
    sampling = 1 / (
        np.array(self._reciprocal_sampling) * self._region_of_interest_shape
    )
    aberrations_basis, aberrations_basis_du, aberrations_basis_dv = (
        calculate_aberration_gradient_basis(
            aberrations_mn,
            sampling,
            self._region_of_interest_shape,
            self._wavelength,
            rotation_angle=np.deg2rad(rotation_angle_deg),
            xp=xp,
        )
    )

    # shifts
    corner_indices = self._xy_inds - xp.array(self._region_of_interest_shape // 2)
    raveled_indices = xp.ravel_multi_index(
        corner_indices.T, self._region_of_interest_shape, mode="wrap"
    )
    gradients = xp.array(
        (
            aberrations_basis_du[raveled_indices, :],
            aberrations_basis_dv[raveled_indices, :],
        )
    )
    shifts_ang = xp.tensordot(gradients, aberrations_coefs, axes=1).T

    # transpose predicted shifts
    if transpose:
        shifts_ang = xp.flip(shifts_ang, axis=1)

    shifts_px = shifts_ang / xp.array(self._scan_sampling)

    # upsampled stack
    if kde_upsample_factor is not None:
        BF_size = np.array(self._stack_BF_unshifted.shape[-2:])
        pixel_output_shape = np.round(BF_size * kde_upsample_factor).astype("int")

        x = xp.arange(BF_size[0], dtype=xp.float32)
        y = xp.arange(BF_size[1], dtype=xp.float32)
        xa_init, ya_init = xp.meshgrid(x, y, indexing="ij")

        # kernel density output the upsampled BF image
        xa = (xa_init + shifts_px[:, 0, None, None]) * kde_upsample_factor
        ya = (ya_init + shifts_px[:, 1, None, None]) * kde_upsample_factor

        pix_output = self._kernel_density_estimate(
            xa,
            ya,
            self._stack_BF_unshifted,
            pixel_output_shape,
            kde_sigma_px * kde_upsample_factor,
            lanczos_alpha=lanczos_interpolation_order,
            lowpass_filter=kde_lowpass_filter,
        )

        # hack since cropping requires "_kde_upsample_factor"
        old_upsample_factor = getattr(self, "_kde_upsample_factor", None)
        self._kde_upsample_factor = kde_upsample_factor
        cropped_image = asnumpy(
            self._crop_padded_object(pix_output, upsampled=True)
        )
        if old_upsample_factor is not None:
            self._kde_upsample_factor = old_upsample_factor
        else:
            del self._kde_upsample_factor

    # shifted stack
    else:
        kde_upsample_factor = 1
        aligned_stack = xp.zeros_like(self._stack_BF_shifted_initial[0])

        if max_batch_size is None:
            max_batch_size = self._num_bf_images

        for start, end in generate_batches(
            self._num_bf_images, max_batch=max_batch_size
        ):
            shifted_BFs = self._stack_BF_shifted_initial[start:end]

            Gs = xp.fft.fft2(shifted_BFs)

            dx = shifts_px[start:end, 0]
            dy = shifts_px[start:end, 1]

            shift_op = xp.exp(
                self._qx_shift[None] * dx[:, None, None]
                + self._qy_shift[None] * dy[:, None, None]
            )
            stack_BF_shifted = xp.real(xp.fft.ifft2(Gs * shift_op))
            aligned_stack += stack_BF_shifted.sum(0)

        cropped_image = asnumpy(
            self._crop_padded_object(aligned_stack, upsampled=False)
        )

    if plot_shifts_and_aligned_bf:
        figsize = kwargs.pop("figsize", (8, 4))
        color = kwargs.pop("color", (1, 0, 0, 1))
        cmap = kwargs.pop("cmap", "magma")

        fig, axs = plt.subplots(1, 2, figsize=figsize)

        self.show_shifts(
            shifts_ang=shifts_ang,
            plot_arrow_freq=plot_arrow_freq,
            scale_arrows=scale_arrows,
            plot_rotated_shifts=False,
            color=color,
            figax=(fig, axs[0]),
        )

        axs[0].set_title("Predicted BF Shifts")

        extent = [
            0,
            self._scan_sampling[1] * cropped_image.shape[1] / kde_upsample_factor,
            self._scan_sampling[0] * cropped_image.shape[0] / kde_upsample_factor,
            0,
        ]

        axs[1].imshow(cropped_image, cmap=cmap, extent=extent, **kwargs)
        axs[1].set_ylabel("x [A]")
        axs[1].set_xlabel("y [A]")
        axs[1].set_title("Predicted Aligned BF Image")

        fig.tight_layout()

    if return_shifts_and_aligned_bf:
        return shifts_ang, cropped_image

Cheers!

[gvarnavi]

Thanks for the suggestion @alouiset!

Actually, I recently re-factored parallax a tiny bit in #695 to allow one to pass arbitrary aberrations during initialization. This means that you can now initialize with a guess of your aberrations (and force the rotation/transpose in preprocess), skip the reconstruct step, and upsample / visualize directly to see the effect.

This is more general but not wrapped up as nicely as guess_common_aberrations. I'll think how best to de-duplicate the functionality.

[alouiset]

Hi @gvarnavi,
Thanks for the feedback! I saw you also added the DirectPtychography class which I would like to test. Is this only available in the phase_contrast branch? Do you plan to release everything in the default branch? Or maybe it is not stable yet.

[gvarnavi]

The direct ptychography functionality (note this has now split into three classes: SSB, OBF, and WDD) is currently only in the phase_contrast branch yes. It is however pretty-much production ready as far as I'm aware -- you can checkout the branch and have a loot at the tutorial notebooks here.

Once #695 gets reviewed and merged it'll also be available in dev -- and I believe we'll plan a release after that.

[gvarnavi]

@alouiset guess_common_aberrations has now been generalized in e6ebf70 to accept any aberration up to 5th order. Note that, similar to the rest of the phase modules this follows the abTEM polar convention and aliases. Meaning you can pass the aberrations using any of the following:

parallax.guess_common_aberrations(
    rotation_angle_deg=-15,
    plot_arrow_freq=4,
    kde_upsample_factor=3,
    defocus=1.5e4, # alias as kwarg
    # C10 = -1.5e4, # polar symbol as kwarg
    # polar_parameters = {'C10': -1.5e4}, # polar_parameters as dictionary
)

Also note that for your specific use-case trefoil is not in polar_aliases, so you have to specify it as C23 and phi23.

#695 should be merged soon, and a new release is planned for next week, so I will close this.