feat: azimuthal average of a centred diffraction pattern

Sources/SimulatoDetector/AzimuthalAverage.swift

@@ -0,0 +1,111 @@
+//
+//  AzimuthalAverage.swift
+//  Role: abTEM 1.0.10's `DiffractionPatterns.azimuthal_average` (measurements.py, `_azimuthal_average`) on a centred
+//        diffraction pattern: rings at radii 0, Δ, 2Δ, … below a maximum angle, each the mean of the pixels a step or a
+//        Gaussian weight selects around that radius. The angles follow abTEM bit for bit where it matters: the axis
+//        coordinates are a float32 `linspace` between the pattern limits [mrad], the radius is float32, the ring test is
+//        float64, so a pixel on a ring edge falls on the side abTEM puts it. The sums are float64; abTEM sums a float32
+//        pattern in float32, so on float32 data the two differ at about 1e-7 relative.
+//
+
+import Foundation
+import SimulatoCore
+
+/// How `azimuthalAverage` weights the pixels around each ring (abTEM's `weighting_function`).
+public enum AzimuthalWeighting: Sendable, Hashable {
+    /// The pixels with |α − ring| < width, each weight 1 (abTEM "step").
+    case step
+    /// exp(−(α − ring)² / (width² / 2)) over every pixel (abTEM "gaussian").
+    case gaussian
+}
+
+extension RealField {
+    /// The azimuthal average of a centred diffraction pattern (abTEM `DiffractionPatterns.azimuthal_average`): for each
+    /// ring radius 0, Δ, 2Δ, … below `maxAngle`, Σ P · w / Σ w with w the weighting around that radius, and 0 where
+    /// Σ w is 0. The pattern is centred as `Waves.diffractionPattern` and `PixelatedDetector` return it: sampling in 1/Å,
+    /// the zero frequency at index n / 2 of each axis (abTEM `fftshift=True`).
+    /// - Parameters:
+    ///   - beam: the beam that sets the angles (α = k λ 10³ mrad).
+    ///   - maxAngle: the last ring lies below it [mrad]; `nil` is abTEM's default, the largest |lowest angle| of the two
+    ///     axes (the edge of the axis that reaches further, not the corner).
+    ///   - radialSampling: the ring spacing Δ, in units of the finer angular sampling of the two axes (default 1).
+    ///   - weighting: `.step` (default) or `.gaussian`.
+    ///   - width: the step half-width or the Gaussian width, in units of the finer angular sampling (default 1).
+    /// - Returns: one reciprocal-space profile, sampling Δ / (λ 10³) [1/Å].
+    /// - Throws: `invalidDetector` unless the radial sampling, the width and a given maximum angle are finite and
+    ///   positive (abTEM returns zeros, NaN or an empty profile), or when the rings would number a million or more.
+    public func azimuthalAverage(beam: ElectronBeam, maxAngle: Double? = nil, radialSampling: Double = 1,
+                                 weighting: AzimuthalWeighting = .step, width: Double = 1) throws
+        -> ReciprocalSpaceLineProfiles {
+        guard radialSampling.isFinite, radialSampling > 0 else {
+            throw SimulatoError.invalidDetector("the radial sampling must be finite and positive, got \(radialSampling)")
+        }
+        guard width.isFinite, width > 0 else {
+            throw SimulatoError.invalidDetector("the width must be finite and positive, got \(width)")
+        }
+        if let m = maxAngle, !(m.isFinite && m > 0) {
+            throw SimulatoError.invalidDetector("the maximum angle must be finite and positive, got \(m)")
+        }
+        let wavelength = beam.wavelength
+        // abTEM `limits`, then `angular_limits` (k × λ × 10³, in that order), then `angular_sampling`.
+        func limits(_ n: Int, _ s: Double) -> (lower: Double, upper: Double) {
+            let l = n % 2 == 1 ? (Double(-((n - 1) / 2)) * s, Double((n - 1) / 2) * s)
+                               : (Double(-(n / 2)) * s, Double(n / 2 - 1) * s)
+            return (l.0 * wavelength * 1e3, l.1 * wavelength * 1e3)
+        }
+        let lx = limits(nx, samplingX), ly = limits(ny, samplingY)
+        let minAngular = Swift.min(samplingX * wavelength * 1e3, samplingY * wavelength * 1e3)
+        // abTEM: -min(min(angular_limits)), the lower limit of the lexicographically smaller (lower, upper) pair.
+        let maximum = maxAngle ?? -Swift.min(lx.lower, ly.lower)
+        let step = radialSampling * minAngular
+        let w = width * minAngular
+
+        // NumPy `arange(0, maximum, step)`: ceil(maximum / step) rings at i × step.
+        let ringsDouble = (maximum / step).rounded(.up)
+        guard ringsDouble < 1e6 else {
+            throw SimulatoError.invalidDetector("a maximum angle of \(maximum) mrad at \(step) mrad gives \(ringsDouble) rings")
+        }
+        let rings = Swift.max(Int(ringsDouble), 0)
+
+        // NumPy `linspace(lower, upper, n, dtype=float32)`: i × step + lower in float64, the last point set to upper,
+        // then rounded to float32.
+        func linspace32(_ l: (lower: Double, upper: Double), _ n: Int) -> [Float] {
+            guard n > 1 else { return [Float(l.lower)] }
+            let d = (l.upper - l.lower) / Double(n - 1)
+            var out = (0..<n).map { Float(Double($0) * d + l.lower) }
+            out[n - 1] = Float(l.upper)
+            return out
+        }
+        let ax = linspace32(lx, nx), ay = linspace32(ly, ny)
+        var radius = [Double](repeating: 0, count: nx * ny)
+        for ix in 0..<nx {
+            for iy in 0..<ny {
+                let r: Float = (ax[ix] * ax[ix] + ay[iy] * ay[iy]).squareRoot()
+                radius[ix * ny + iy] = Double(r)
+            }
+        }
+
+        var out = [Double](repeating: 0, count: rings)
+        let gaussianScale = w * w / 2
+        for i in 0..<rings {
+            let center = Double(i) * step
+            var weight = 0.0, sum = 0.0
+            switch weighting {
+            case .step:
+                for k in radius.indices where abs(radius[k] - center) < w {
+                    weight += 1
+                    sum += values[k]
+                }
+            case .gaussian:
+                for k in radius.indices {
+                    let d = radius[k] - center
+                    let m = exp(-(d * d) / gaussianScale)
+                    weight += m
+                    sum += values[k] * m
+                }
+            }
+            out[i] = weight > 0 ? sum / weight : 0
+        }
+        return ReciprocalSpaceLineProfiles(count: rings, values: out, sampling: step / (wavelength * 1e3))
+    }
+}

Sources/SimulatoDetector/SimulatoDetector.docc/SimulatoDetector.md

@@ -34,6 +34,7 @@ Detectors for scanned-probe simulations and the STEM driver that applies them ov
 - ``LineProfiles``
 - ``RealSpaceLineProfiles``
 - ``ReciprocalSpaceLineProfiles``
+- ``AzimuthalWeighting``
 
 ### Detector input
 

Tests/SimulatoDetectorTests/AzimuthalAverageTests.swift

@@ -0,0 +1,76 @@
+//
+//  AzimuthalAverageTests.swift
+//  Role: the azimuthal average of a centred diffraction pattern (abTEM's DiffractionPatterns.azimuthal_average), against
+//        numbers abTEM 1.0.10 printed (tools/pins/azimuthal/pins.py) for planted 9 x 8 and 8 x 6 patterns, values
+//        (k + 1)^1.3 with k the flat index, at 200 keV; and its refusals.
+//
+
+import Foundation
+import Testing
+import SimulatoCore
+import SimulatoDetector
+
+@Suite("Azimuthal average")
+struct AzimuthalAverageTests {
+    static func pattern(_ nx: Int, _ ny: Int, _ sx: Double, _ sy: Double) -> RealField {
+        RealField(nx: nx, ny: ny, values: (0..<(nx * ny)).map { pow(Double($0) + 1, 1.3) }, samplingX: sx, samplingY: sy)
+    }
+
+    static let beam = try! ElectronBeam(energy: 200e3)
+
+    /// assert_allclose defaults (rtol 1e-7, atol 0): abTEM states no tolerance for the azimuthal average.
+    static func close(_ a: [Double], _ b: [Double], rtol: Double = 1e-7) -> Bool {
+        a.count == b.count && zip(a, b).allSatisfy { abs($0 - $1) <= rtol * abs($1) }
+    }
+
+    @Test func stepRingsMatchAbTEMOnAnOddEvenPattern() throws {
+        let p = Self.pattern(9, 8, 0.05, 0.06)
+        let a = try p.azimuthalAverage(beam: Self.beam)
+        #expect(a.count == 5 && a.ensembleShape == [])
+        #expect(abs(a.sampling - 0.05000000000000001) <= 1e-7 * 0.05)
+        #expect(Self.close(a.values, [109.31059072723754, 110.0733829078507, 111.05800473544934, 113.83482973524923,
+                                      116.75029096542004]))
+        let b = try p.azimuthalAverage(beam: Self.beam, maxAngle: 9, radialSampling: 0.7, width: 1.5)
+        #expect(b.count == 11 && abs(b.sampling - 0.035) <= 1e-7 * 0.035)
+        #expect(Self.close(b.values, [109.71729296135845, 110.60442943790507, 110.97479454458214, 112.91164715276787,
+                                      115.2149506722666, 115.6100936043095, 116.37171872576145, 116.20544576168864,
+                                      116.3659299808196, 111.20251111350319, 109.37593155122985]))
+    }
+
+    @Test func gaussianRingsMatchAbTEM() throws {
+        let p = Self.pattern(9, 8, 0.05, 0.06)
+        let a = try p.azimuthalAverage(beam: Self.beam, weighting: .gaussian)
+        #expect(Self.close(a.values, [109.52747242355026, 110.19514603212777, 111.79352436215716, 114.5200943034101,
+                                      117.77733984988313]))
+        let b = try p.azimuthalAverage(beam: Self.beam, maxAngle: 6, radialSampling: 1.3, weighting: .gaussian, width: 0.8)
+        #expect(b.count == 4 && abs(b.sampling - 0.06500000000000002) <= 1e-7 * 0.065)
+        #expect(Self.close(b.values, [109.39182869735909, 110.29310407876528, 112.87926982867883, 117.66284648473913]))
+    }
+
+    @Test func evenPatternMatchesAbTEM() throws {
+        let p = Self.pattern(8, 6, 0.04, 0.04)
+        let a = try p.azimuthalAverage(beam: Self.beam)
+        #expect(a.count == 4 && abs(a.sampling - 0.039999999999999994) <= 1e-7 * 0.04)
+        #expect(Self.close(a.values, [76.36746409305131, 76.84996447773611, 77.6798227805934, 75.22480102040664]))
+        #expect(Self.close(try p.azimuthalAverage(beam: Self.beam, width: 0.5).values,
+                           [76.08612050106568, 76.6139088690763, 77.74902662196723, 77.65019303476393]))
+    }
+
+    @Test func aRingWithNoPixelIsZero() throws {
+        // A 1 x 1 pattern: one pixel at the centre; rings beyond a step width of it hold nothing.
+        let p = RealField(nx: 1, ny: 1, values: [3], samplingX: 0.05, samplingY: 0.05)
+        let a = try p.azimuthalAverage(beam: Self.beam, maxAngle: 5, width: 0.5)
+        #expect(a.count == 4 && a.values == [3, 0, 0, 0])
+    }
+
+    @Test func invalidArgumentsAreRefused() {
+        let p = Self.pattern(9, 8, 0.05, 0.06)
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, radialSampling: 0) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, radialSampling: .nan) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, width: 0) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, width: -1) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, maxAngle: 0) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, maxAngle: .infinity) }
+        #expect(throws: SimulatoError.self) { try p.azimuthalAverage(beam: Self.beam, maxAngle: 1e9) }
+    }
+}

docs/open-items.md

@@ -15,9 +15,10 @@ coverage: `docs/parity.md`.
 ## Work items
 
 - **Detector and noise follow-ups from 0.4.0**: line profiles for segmented and
-  pixelated line scans, the azimuthal average, `LineScan` defaults, the "spline" method of `Images.interpolate`, tilt
-  distributions (#9). `integrate_disc(interpolate:)` and `Images.interpolate` (method "fft") are in; the DPC measurement of a
-  `SegmentedDetector` is under Polar measurements.
+  pixelated line scans, `LineScan` defaults, the "spline" method of `Images.interpolate`, tilt distributions (#9).
+  `integrate_disc(interpolate:)` and `Images.interpolate` (method "fft") are in; the DPC measurement of a `SegmentedDetector`
+  is under Polar measurements. The azimuthal average of a single pattern is in (`RealField.azimuthalAverage`), over a
+  pixelated scan it is not.
 - **Potential API**: `CrystalPotential` (a unit potential repeated, frozen phonons per cell), its API undecided (#7).
 - **Cells**: non-orthogonal and rotated cells, `periodic=False`, `orthogonalize_cell` in `AtomicStructure` (the kit keeps
   its own structure type: `docs/adr/0001-standalone-structure-types.md`). Finite projection, the Kirkland, Peng and Ewald

tools/pins/azimuthal/pins.py

@@ -0,0 +1,36 @@
+"""SimulatoKit azimuthal average: DiffractionPatterns.azimuthal_average (abTEM 1.0.10, numpy).
+
+Prints the numbers that Tests/SimulatoDetectorTests/AzimuthalAverageTests.swift holds inline (no fixture file): planted
+centred patterns (fftshift=True, as PixelatedDetector and Waves.diffractionPattern return them) of 9 x 8 and 8 x 6
+pixels, values (k + 1)^1.3 with k the flat index, float64, at 200 keV. abTEM builds the angular coordinates in float32
+and sums in float64; the test holds the profiles at rtol 1e-7 (assert_allclose defaults; abTEM states no tolerance).
+
+Run from the repo root:  .venv/bin/python -I -B tools/pins/azimuthal/pins.py
+"""
+
+import numpy as np
+from abtem.measurements import DiffractionPatterns
+
+
+def pattern(nx, ny, sampling):
+    arr = (np.arange(nx * ny, dtype=np.float64).reshape(nx, ny) + 1.0) ** 1.3
+    return DiffractionPatterns(arr, sampling=sampling, fftshift=True, metadata={"energy": 200e3})
+
+
+def show(name, p):
+    print(name, "count:", p.array.shape[-1], "sampling:", repr(p.sampling))
+    print(name, "values:", p.array.ravel().tolist())
+
+
+a = pattern(9, 8, (0.05, 0.06))
+b = pattern(8, 6, (0.04, 0.04))
+show("odd default", a.azimuthal_average())
+show("odd max 9 rs 0.7 w 1.5", a.azimuthal_average(max_angle=9.0, radial_sampling=0.7, width=1.5))
+show("odd gaussian", a.azimuthal_average(weighting_function="gaussian"))
+show("odd gaussian max 6 rs 1.3 w 0.8", a.azimuthal_average(max_angle=6.0, radial_sampling=1.3,
+                                                           weighting_function="gaussian", width=0.8))
+show("even default", b.azimuthal_average())
+show("even width 0.5", b.azimuthal_average(width=0.5))
+one = DiffractionPatterns(np.array([[3.0]]), sampling=0.05, fftshift=True, metadata={"energy": 200e3})
+show("one pixel max 5 w 0.5", one.azimuthal_average(max_angle=5.0, width=0.5))
+print("max angles default:", -min(min(a.angular_limits)), -min(min(b.angular_limits)))