Local FFT pseudo-4DSTEM for HRTEM images, with phase and in-plane orientation mapping

src/quantem/diffraction/__init__.py

@@ -6,6 +6,7 @@
 from quantem.diffraction.phase import PhaseMap as PhaseMap
 from quantem.diffraction.reverse_monte_carlo import ReverseMonteCarlo as ReverseMonteCarlo
 from quantem.diffraction.strain import StrainMap as StrainMap
+from quantem.diffraction.systematic_row import SystematicRowMap as SystematicRowMap
 from quantem.diffraction import bloch as bloch
 from quantem.diffraction import calibration as calibration
 from quantem.diffraction import crystal_rotations as crystal_rotations

src/quantem/diffraction/crystal.py

@@ -450,6 +450,9 @@ def from_cif(cls, file_path: str | Path, name: str | None = None, **kwargs) -> "
 
         atoms = read(file_path)
         assert isinstance(atoms, Atoms)
+        # the reader adds an ase Spacegroup object; Crystal finds the space
+        # group itself (Crystal.spacegroup), so we keep only native data
+        atoms.info.pop("spacegroup", None)
         return cls(_expand_partial_occupancy(atoms), name=name, **kwargs)
 
     @property

src/quantem/imaging/__init__.py

@@ -1,3 +1,4 @@
 from quantem.imaging.drift import DriftCorrection as DriftCorrection
 from quantem.imaging.lattice import Lattice as Lattice
 from quantem.imaging.lattice_visualization import PLOT_REGISTRY as PLOT_REGISTRY
+from quantem.imaging.local_fft import LocalFFT as LocalFFT

tests/diffraction/test_crystal.py

@@ -392,6 +392,7 @@ def test_from_cif_keeps_partial_occupancy(tmp_path):
     assert content["Li"] == pytest.approx(8.444, abs=1e-6)
     assert content["O"] == pytest.approx(12.0, abs=1e-6)
     assert xtl.spacegroup.startswith("C2/m")
+    assert "spacegroup" not in xtl.atoms.info  # no ase objects inside Crystal
 
 
 def test_pseudo_symmetry_names_the_breaking_reflection(tmp_path):

tests/diffraction/test_systematic_row.py

@@ -0,0 +1,44 @@
+"""SystematicRowMap: phase and in-plane angle from Friedel-paired rows."""
+
+import numpy as np
+from ase.build import bulk
+
+from quantem.core.datastructures.vector import Vector
+from quantem.diffraction import Crystal, SystematicRowMap
+
+
+def test_row_indexing():
+    lif = Crystal.from_ase(bulk("LiF", "rocksalt", a=4.027, cubic=True), name="LiF", verbose=False)
+    lif.calculate_structure_factors(k_max=1.0)
+    li = Crystal.from_ase(bulk("Li", "bcc", a=3.51, cubic=True), name="Li", verbose=False)
+    li.calculate_structure_factors(k_max=1.0)
+
+    def pair(k, deg, inten):
+        # deg is measured from the row axis toward the column axis
+        g = k * np.array([np.cos(np.deg2rad(deg)), np.sin(np.deg2rad(deg))])
+        return [[g[0], g[1], inten], [-g[0], -g[1], inten]]
+
+    k200 = 2 / 4.027
+    k110 = np.sqrt(2) / 3.51
+    cells = [
+        np.array(pair(k200, 30.0, 5.0)),  # LiF (200) at 30 degrees
+        np.array(pair(k110, 100.0, 4.0) + pair(k200, 10.0, 1.0)),  # Li wins, two directions
+        np.array([[0.3, 0.1, 9.0]]),  # unpaired peak, no row
+        np.array(pair(0.3, 45.0, 2.0)),  # paired but on no ring
+    ]
+    peaks = Vector.from_data(
+        [cells[:2], cells[2:]], fields=["qx", "qy", "intensity"], units=["A^-1"] * 2 + ["counts"]
+    )
+    rm = SystematicRowMap.from_vectors(peaks, [lif, li], k_min=0.2, k_max=0.6).index(k_tol=0.01)
+
+    phase = rm.phase_index()
+    assert phase.tolist() == [[0, 1], [-1, -1]], phase
+    assert rm.ring_label[rm.ring[0, 0, 0]] == "(200)"
+    angle, _ = rm.in_plane_angle("LiF")
+    assert np.isclose(angle[0, 0], 30.0) and np.isclose(angle[0, 1], 10.0), angle
+    assert rm.num_directions()[0, 1] == 2
+    assert rm.intensity[1, 0, 0] == 0, "an unpaired peak must not form a row"
+    assert rm.phase[1, 1, 0] == -1 and rm.intensity[1, 1, 0] > 0, "off-ring row kept unassigned"
+    single = SystematicRowMap.from_vectors(peaks, [lif, li], k_min=0.2, k_max=0.6)
+    single.index(k_tol=0.01, require_pair=False)
+    assert single.intensity[1, 0, 0] == 9.0, "require_pair=False keeps single peaks"

tests/imaging/test_local_fft.py

@@ -0,0 +1,31 @@
+"""LocalFFT: windowed Fourier transforms of an image as a pseudo-4DSTEM dataset."""
+
+import numpy as np
+
+from quantem.core.datastructures.dataset2d import Dataset2d
+from quantem.imaging.local_fft import LocalFFT, window_mask
+
+
+def test_local_fft():
+    # full-radius taper sin^2(pi/2 * (1 - r/R)), on a row through the center
+    r = np.hypot(0.5, np.arange(64) - 31.5)
+    expected = np.sin(np.pi / 2 * np.clip(1 - r / 32, 0, 1)) ** 2
+    assert np.allclose(window_mask(64)[31], expected, atol=1e-6)
+
+    # cosine fringes with g = 0.25 1/A at 60 degrees from the row axis toward the column axis
+    sampling, k, theta = 0.5, 0.25, np.deg2rad(60.0)
+    rr, cc = np.mgrid[:256, :256] * sampling
+    image = 100.0 + 5.0 * np.cos(2 * np.pi * k * (rr * np.cos(theta) + cc * np.sin(theta)))
+    image += np.random.default_rng(0).normal(0, 1.0, image.shape)
+    ds = Dataset2d.from_array(image, sampling=[0.05, 0.05], units=["nm", "nm"])
+    lf = LocalFFT.from_image(ds, window_size=64, step=32, k_max=0.4)
+
+    assert np.isclose(lf.sampling, sampling), "nm must convert to Angstroms"
+    dp = lf.dataset.array[3, 3]
+    h = int(lf.center[0])
+    assert dp[h, h] == 0.0, "the center is removed by the normalization"
+    row, col = np.unravel_index(np.argmax(dp), dp.shape)
+    q = np.array([row - h, col - h]) * lf.pixel_size
+    assert np.isclose(np.hypot(*q), k, atol=lf.pixel_size)
+    assert np.isclose(np.rad2deg(np.arctan2(q[1], q[0])) % 180, 60.0, atol=5.0)
+    assert 0.5 < np.median(dp[dp != 0] + 1.0) < 1.5, "normalized noise must be close to 1"