Add quantem.atoms: 3D atomic model analysis and atom tracing from volumes

[cophus]

What problem this PR addresses

quantem had no tools for analyzing 3D atomic models, the output of atomic electron tomography (AET) or of atomistic simulations, and the atomic tracing on this branch stopped at finding peaks in a volume. This PR adds a quantem.atoms subpackage that takes a set of 3D coordinates through calibration, local structure classification, grain segmentation, strain, and visualization, and completes the tracer in quantem.tomography.atom_trace so that a reconstruction can be converted into such a model. No new required dependencies are added.

Atom tracing (quantem.tomography.Atoms). A volume is modeled as a sum of isotropic 3D Gaussians plus a smooth, non-negative Bezier background, and all site positions, amplitudes, widths and the background are optimized jointly with Adam against the volume. Each Gaussian is splatted only into a window of 3 to 4 sigma around its center, so the cost scales with the number of atoms rather than with the number of voxels. Sites are seeded from a difference-of-Gaussians filter and, during refinement, added at residual peaks, removed when weak or isolated, and merged when closer than a fraction of the nearest-neighbor spacing. On a 280^3 volume with 20,000 sites, 25 iterations take 4 s on an Apple GPU and 9 s on the CPU.

Atomic model analysis (quantem.atoms).

• AtomicModel stores the sites in a Vector with named per-site channels, so every measurement (scores, grain labels, strain components, intensities) is a column that the plots and the viewer can select. Coordinates from any file format are passed in as an (N, 3) array; the class has no format-specific loaders.
• compute_pdf and calibrate measure the radial distribution function with a KD-tree pair count, fit the first peak with an asymmetric generalized Gaussian, and set the physical voxel size from the nearest-neighbor distance of a reference crystal. merge_close_sites removes near-duplicate sites from tracing.
• templates generates polyhedral templates (fcc, hcp, bcc, sc, diamond/zincblende, wurtzite, icosahedral) from crystal definitions for any number of neighbor shells, together with their proper symmetry rotations, which are found numerically.
• matching classifies every site by polyhedral template matching, vectorized in torch. Trial rotations are built from pairs of neighbor vectors, each is scored by the one-to-one match between the rotated template and the measured neighbors, and the best rotation is refined with a batched Kabsch solve; an affine fit gives the deformation gradient. Both fcc and hcp templates take about 1 s for 20,000 sites on a laptop CPU. On synthetic fcc, hcp, bcc, diamond and wurtzite particles the classification accuracy is above 97% for interior sites at 3% bond-length noise and about 99% at 5%.
• segment_grains connects neighboring sites whose disorientation, reduced by the template symmetry, is below a threshold and labels the connected components, then fills the gaps between grains by majority vote; it also produces a boundary channel (interior, grain boundary, twin, surface). Twin boundaries in fcc appear as hcp-classified sites, since a site on a coherent {111} twin plane has the hcp environment.
• fit_icosahedral_centers locates the centers of multiply twinned particles as the least-squares intersection of the twin planes carried by the hcp sites, and layer_positions finds the atomic layers along any direction so that plot("slices") can step through a model one layer at a time.
• compute_strain, surface_distance, sample_volume and classify_species give per-site strain from the affine fit, the distance to the convex hull, the reconstruction intensity around each site, and a species assignment from a one-dimensional Gaussian mixture that leaves low-posterior sites unassigned.
• structures builds ideal Mackay icosahedra, cuboctahedra, Ino decahedra and mirror-twinned double icosahedra with a chosen center separation, with shell and sector channels, for comparison with measured models and for tests.
• visualization provides model.plot(kind=...) with the RDF, channel histograms, shaded-sphere slab projections, layer-indexed slice grids, and a template overlay for one site. explode_grains returns a copy with every grain displaced away from its center and the shared boundary sites duplicated into each adjacent grain.
• model.show() opens an interactive 3D viewer with channel selection, a histogram-driven color range, legend toggles, slab clipping along any normal with a draggable handle, a growth slider, and a site inspector. The viewer itself lives in the companion quantem.widget repository (separate PR); quantem.atoms.show_atoms only re-exports it.

tests/atoms/test_atoms.py adds 28 tests covering the templates (shell counts and symmetry orders), the RDF fit, the matcher on synthetic lattices (classification accuracy and rotation recovery), the measurements, the ideal structures, the center fit and layer indexing on a synthetic twinned particle, grain explosion, species assignment, and the end-to-end pipeline on a synthetic fcc twin. The tracer has no automated tests yet; it was checked on a reconstruction against an independently traced model, recovering 84% of the reference sites with a 0.6-voxel rms offset.

Two small changes outside the new package: Dataset3d.show gains same_scale so that all slices share one intensity range, and quantem/__init__.py imports atoms.

What should the reviewer(s) do

Review the public API of AtomicModel and the module layout, and try the pipeline on a 3D model of your own: from_array, compute_pdf, calibrate, find_neighbors, match_templates(["fcc", "hcp"]), segment_grains, then plot("slab", channel="structure"). Two points I would like feedback on:

1. Whether quantem.atoms should be a top-level subpackage, as here, or live under quantem.tomography next to the tracer. I put it at the top level because the analysis applies to any 3D atomic model, including simulation output.
2. The default classification rule. The soft score is the default; an RMSD-based rule and a neighbor-smoothed rule are available, and none of them is reliable above about 8% bond-length noise.

Open items before merging: the empty tomography/atom_analysis.py should be deleted, atom_trace.py needs tests, and the viewer PR in quantem.widget should land first so that model.show() resolves.

• This PR introduces a public-facing change (e.g., figures, CLI input/output, API).
  • For functional and algorithmic changes, tests are written or updated.
  • Documentation (e.g., tutorials, examples, README) has been updated.
  • A tracking issue or plan to update documentation exists: a tutorial notebook on a synthetic twinned particle will follow once the API settles.

[arthurmccray]

how the heck do i cancel a copilot review? i misclicked and now cant stop it 😢

[Copilot]

🟡 Changes recommended

Critical calibration and viewer integration defects, plus unresolved moderate correctness issues, block approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds quantem.atoms for 3D atomic-model analysis and completes Gaussian-splatting atom tracing.

Changes:

• Adds calibration, matching, grain analysis, strain, structures, measurements, and visualization.
• Adds volume-based atom tracing and public API exports.
• Adds shared-scale dataset visualization and extensive atom-analysis tests.

File summaries

File	Summary and review notes
tests/atoms/test_atoms.py	Adds coverage for atomic-model functionality; tracer refinement remains untested.
src/quantem/tomography/atom_trace.py	Implements tracing and refinement. Critical: viewer class-name mismatch. Moderate: fallback isolation handling can still fail. Nit: missing tracer refinement test.
src/quantem/tomography/atom_analysis.py	Empty legacy placeholder.
src/quantem/tomography/__init__.py	Exports the tracer.
src/quantem/core/datastructures/dataset3d.py	Adds shared slice scaling. Nit: new scaling branches lack test coverage.
src/quantem/atoms/visualization.py	Adds model plotting and grain explosion visualization.
src/quantem/atoms/templates.py	Defines crystal templates and symmetry operations.
src/quantem/atoms/structures.py	Builds ideal nanoparticle structures.
src/quantem/atoms/show_atoms.py	Provides viewer compatibility exports.
src/quantem/atoms/pdf.py	Adds RDF, calibration, and neighbor utilities. Critical: incorrect wurtzite scale conversion. Moderate: one-site neighbor queries fail.
src/quantem/atoms/measurements.py	Adds grain, strain, geometry, and species measurements.
src/quantem/atoms/matching.py	Implements template matching. Moderate: partial-match RMSD can become NaN; CUDA/MPS and gradient-backed tensors fail conversion.
src/quantem/atoms/atomic_model.py	Provides the main pipeline. Moderate: explicit origins, categorical merges, late calibration, tensor-backed sampling, and derived-state invalidation contain correctness issues.
src/quantem/atoms/__init__.py	Exports the atoms public API.
src/quantem/__init__.py	Exposes the new package.

Review details

Suppressed comments (6)

src/quantem/atoms/atomic_model.py:487

• If calibration is called after find_neighbors() or surface_distance(), those channels retain values computed with the old sampling while _units changes to the new unit. Clear or recompute calibrated dependent channels here (or explicitly reject late calibration), otherwise the model exposes stale bond and surface distances under the new calibration.

        self._sampling = np.full(3, scale)
        self._units = units

src/quantem/atoms/atomic_model.py:853

• For a tensor-backed Dataset3d, .array is None, so this assigns arr = None and then passes a 0-D object array into sample_volume; the documented Dataset3d input path fails instead of sampling the tensor. Fall back to .tensor when .array is absent before converting to NumPy.

        arr = getattr(volume, "array", volume)
        if hasattr(arr, "detach"):
            arr = arr.detach().cpu().numpy()
        values = meas.sample_volume(np.asarray(arr), self.positions_native, radius=radius)

src/quantem/atoms/atomic_model.py:348

• Invalidation clears the raw matches but leaves _template_specs, _structure_names, and all derived match/classification channels on the model. After positions_native is edited, templates and structure can therefore describe the old geometry while matches is empty, and subsequent analysis sees inconsistent state. Clear or mark these derived results stale together with _matches.

    def _invalidate(self) -> None:
        self._pdf = None
        self._nn_fit = None
        self._neighbor_distances = None
        self._neighbor_indices = None

src/quantem/atoms/matching.py:274

• When a site has fewer matched neighbors than the template, dmin contains inf for unmatched entries, so dmin**2 * w evaluates inf * 0 to nan. This makes the public rmsd result NaN for partial/surface matches even when valid matches exist.

        sq = (dmin**2 * w).sum(-1) / n_match.clamp_min(1).to(dtype)
        rmsd[sl] = torch.where(n_match > 0, sq.sqrt(), torch.full_like(sq, float("nan"))).cpu()

src/quantem/atoms/pdf.py:168

• For a one-site model, cKDTree.query(..., k=1) returns 1-D arrays, but the following code assumes (N, K) arrays and indexes/sorts along axis 1. find_neighbors(xyz, num_neighbors) therefore raises instead of returning the documented padded (-1, inf) result for small models.

    k = min(num_neighbors + 1, n)
    tree = cKDTree(xyz)
    dist, idx = tree.query(xyz, k=k)

src/quantem/core/datastructures/dataset3d.py:340

• The new default scaling behavior and the explicit vmin/vmax/norm override are not covered by the existing Dataset3d.show tests. Add assertions on the show_2d/axes normalization for both branches, since this changes the displayed intensity interpretation for every caller.

        if same_scale and not any(key in kwargs for key in ("vmin", "vmax", "norm")):
            kwargs["vmin"] = float(np.nanmin(self.array))
            kwargs["vmax"] = float(np.nanmax(self.array))

• Files reviewed: 14/15 changed files
• Comments generated: 7
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.