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