Imaging refactor part 2 - Atoms: add_atoms(), refine_atoms(), plot(kind='atoms') and tests

[darshan-mali]

What does this PR do?

Part 2 of Imaging refactor

Enables the Lattice class to store atoms and their data. This includes

• add_atoms() function that checks if an atom is present at a given position and stores the atomic positions in both pixel and unit cell coordinates
• refine_atoms() function to refine the atomic positions
• plot() now includes kind='atoms'

Relevant references

Part 1 can be found at : #198

This PR involves the use of the Vector datastructure.
A tutorial for the same can be found at: https://github.com/electronmicroscopy/quantem-tutorials/blob/main/tutorials/core/vector.ipynb

API and logic

After defining the lattice vectors, atoms can be added to Lattice via the add_atoms() method.
The add_atoms() takes in the fractional/unit cell coordinates of all atoms in one unit cell $0 \leq u,v < 1$.
These are tiled across all unit cells and checked against the background to determine if an atom is present.

The refine_atoms() functions optimizes for the position of the atom via 2D Gaussian fitting.

The pixel and unit cell coordinates of the atoms are stored in Lattice.atoms which is a Vector datastructure, alongwith other fitting information.

Both add_atoms() and refine_atoms() set the default_plot='atoms' in Lattice.plot().

Files changed

Updated files

• src/quantem/imaging/lattice.py: The add_atoms() and refine_atoms() functions were added.
• src/quantem/imaging/lattice_visualization.py: Functionality for plot(kind='atoms') was added along with necessary helpers
• tests/imaging/test_lattice.py: Basic pytests for add_atoms() and refine_atoms() based on expected user behaviour

Examples

A testing notebook with examples can be found here:
Lattice_refactor_2_atoms.ipynb

(Note: Some plotting calls have been commented out to reduce file size. Please uncomment them before running)

Example code block:

lat = Lattice.from_data(
    im
).define_lattice_vectors(
    origin = (1626, 1024),
    v = ( 25 ,  0),
    u = ( 0,   25),
    refine_lattice = True,  # Default = True 
    block_size = 5,
).add_atoms(
    positions_frac = [
        [0.0,0.0],
        [0.5,0.5],
    ],
    edge_min_dist_px=100,
    # numbers = [0, 1],
    # intenisty_min = 0.001,
    # intensity_radius = 3.5,
    # contrast_min = 0.001,
).refine_atoms(
    fit_radius = 7,
    max_nfev = 150,
    max_move_px = 2.5,
).plot()

Example output:

PR Checklist

• This PR introduces a public-facing change (e.g., API, CLI input/output).
  • For functional and algorithmic changes, tests are written or updated.
  • Documentation (e.g., tutorials, examples, README) has been updated.

Reviewer checklist

• The notebook provided runs as expected and no bugs were found
• Tests pass and are appropriate for the changes
• Documentation and examples are sufficient and clear
• The implementation matches the stated intent of the contribution

[wwmills]

If None, it appears to use 0, 1, ..., S-1

[darshan-mali]

Fixed this to be consistent

[wwmills]

define_lattice() --> define_lattice_vectors()

[darshan-mali]

Corrected

[wwmills]

Is this what we want? It seems like this ring will include neighboring atoms... This could cause issues for atom position refinement.

[darshan-mali]

Addressing this in main comment.

[wwmills]

not sure why the distance_transform_edt would fail, but if it does, perhaps add an output message for this? otherwise this is a silent failure

[darshan-mali]

This is just a precautionary try except block for all scipy/matrix calculations. Added a warning to let the user know the error occurred and now defaulting to pixel-level mask checking

[wwmills]

we compute mean intensity for all atoms, even though we have a validity mask (we could generate keep right before this) that should rule some out at this point. It probably doesn't cost too much in the way of time, but it is worth making this change.

[darshan-mali]

Rearranged the order of operations to avoid unnecessary calculations

[wwmills]

max_move_px ends up not being the actual maximum movement, it is just the maximum movement in the x and y directions - both moving by max_move_px is not forbidden.

[wwmills]

Additional comment on this - I think that the maximum movement should not be greater than the window that the fit is over. The default is to make it equal, which seems okay, but user can set it greater; do we want to forbid that?

[darshan-mali]

Added a check for max_move_px > r_fit

Implementing a circular max_move_px limit would result in a significant increase in computation for a minimal increase in accuracy. Considering a square patch of side 2*r_fit is simple, fast and is accurate enough.

Also, since the default value of r_fit is 0.5 * nearest_neighbour_distance, it should avoid almost all overlaps. Even for a case when the nearest neighbour is along the diagonal, if the nearest_neighbour_distance is r, then r_fit = 0.5 * r, and the centre to corner distance for the patch would be r/sqrt(2) = 0.7 * r. This should avoid the neighbour influencing the fit especially for lattices without any point defects, as the neighbours on all corners would cancel each others effect and it would just contribute to the background, if any.

[wwmills]

these values seem arbitrary - why 0.25 for sigma min? why amp0*4 for amplitude max? I don't really have a better idea but just want to be thoughtful about these hard-coded values.

[darshan-mali]

Addressing this in main comment

[wwmills]

these tests seem okay, and pass, but there is no test that validates the found positions of the atoms after refinement. That would be a good test.

[darshan-mali]

Added additional pytest to generate synthetic data and check fit accuracy.

[darshan-mali]

@wwmills I have made all the changes that were requested. The changes are ready for review.

Most of the comments were addressed directly in the same comment chain. There are 2 specific ones that I felt needed a more detailed explanation, which is why I am choosing to answer them here.

I also found a bug where an array completely filled with np.nan could be passed through without any errors. Added a check and wrote a pytest to handle it.

Regarding the Gaussian refinement parameters

I tried to come up with multiple different edge cases in which the gaussian fitting could be used, and tried to push the parameters to values beyond which it becomes reasonably unrealistic to attempt to fit a Gaussian.

The x and y initial guesses and bounds are correct and don't need any updates.

The amplitude starts with an initial guess of pmax - bg0, and is constrained by 0 and pmax + (pmax - pmin). The 4.0*amp0 condition was removed as it was arbitrary, while this is a fixed value that no pixel in the patch exceeds.

The background has an initial guess of the median of the part of the patch that is outside the circle of radius r_fit (essentially the region of the square not covered by the inscribed circle), floored by 0.0. It is constrained by the same values as the amplitude.

Regarding sigma, there are a few different scenarios to consider.

1. The lower bound is defined using the minimum size an atom must have in an image to distinguish it from noise / hot pixels. Since, we are assuming it to be a Gaussian, we can approximate the atomic radius as ~3-4 $\sigma$ (higher end of approximation as Norm(3$\sigma$) = 0.44% and Norm (4$\sigma$) = 0.01%). So, the smallest group of pixels, we can classify as an atom would be a 2x2 square, which would have an atomic radius of 1 pixel, and so a minimum sigma value of 0.25

2. Consider the scenario where the user does not specify the fit_radius, so the code defaults to half the interatomic spacing. In this case, we need to calculate the sigma in terms of the the interatomic spacing. Looking only along the direction of the nearest neighbours, an array of Gaussians and their sum was plotted to try and look at when 2 atoms could be considered "just resolved". This does involve some subjective bias, however, when $\sigma$ = 0.55 * interatomic spacing, the individual atoms become extremely difficult to distinguish, even under the assumption that the background is extremely noisy as seen in the figure below. As r_fit is half the interatomic separation, we set the upper bound of sigma to 1.1 * r_fit, when fit_radius is not provided.

3. For the scenario when the fit_radius is provided, if the fit_radius is too large, then the upper bound set by the previous condition holds. However, we need to be concerned about when the fit_radius is too small, as it may result in only part of gaussian being in the patch window, and the fit_radius ends up being smaller than the best Gaussian fit or there is very high background noise. This again involves subjective bias, however, the clipped Gaussian starts becoming undeniably unrecognisable as a Gaussian when sigma exceeds 1.5 * fit_radius, as seen from the figure below.

4. For the initial guess, sigma = 0.5 * fit_radius and sigma = 0.25 * interatomic spacing = 0.5 * auto radius, are both very good estimates as seen in the figures below.

5. The initial guess options of 0.5 px and upper bound of 1.0 px are just safety fallbacks for incorrect user inputs.

Note: For points 2 and 3, this is the upper bound for sigma. If the initial guess is good enough, it should never reach close to this value for most cases. This is just an attempt to get the worst possible scenario in which this would work.

Regarding Atom Detection radii

The intensity radius (atomic disk) and the annulus radii (background annulus) parameters can now be passed as both pixels or in terms of the interatomic separation (auto_radius = 0.5 * interatomic separation), which is the more intuitive unit for the same.

The defaults were changed to 0.8 * auto_radius for the disk and 1.0 - 1.7 * auto_radius for the annulus. These will depend on the sample, imaging conditions and image mode. This attempts to maximise the number of pixels in each while trying to avoid neighbouring atoms. The does a good enough job as seen in the figures below, however, it still does include small parts of the neighbouring atoms, but this is small enough to avoid any errors).

It is very hard to set default radii that satisfy all the criteria, especially since this code is primarily built to deal with polarization measurements which does involve displacement of atoms from expected positions. For the same reason, warnings were chosen instead of raising errors for when the disk radius exceeds 1.0 * auto_radius and when the outer annulus radius exceeds 2.0 * auto_radius, as the values chosen by the user could be more appropriate depending on the data.