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

[wwmills]

define_lattice() --> define_lattice_vectors()

[wwmills]

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

[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

[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.

[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?

[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.

[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.