midas-diffract 0.4.0__tar.gz → 0.7.0__tar.gz

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.4.0
+Version: 0.7.0
 Summary: End-to-end differentiable forward model for High-Energy Diffraction Microscopy (FF, NF, pf-HEDM)
 Author-email: Hemant Sharma <hsharma@anl.gov>
 License-Expression: BSD-3-Clause

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/midas_diffract/__init__.py

@@ -28,7 +28,7 @@ Quick start
     loss.backward()
 """
 
-__version__ = "0.4.0"
+__version__ = "0.7.0"
 
 from .forward import (
     HEDMForwardModel,

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/midas_diffract/forward.py

@@ -27,6 +27,7 @@ Reference C code:
 """
 
 import math
+import warnings
 from dataclasses import dataclass, field
 from typing import Optional, Tuple
 
@@ -116,6 +117,26 @@ class HEDMGeometry:
                                    # spot is valid if it lands on at least
                                    # one detector; output gains a det_id
                                    # field naming which one.
+    # Paired OmegaRange / BoxSize acceptance filter (the ``OmegaRange`` and
+    # ``BoxSize`` paramfile keys). Ports the KeepSpot gate in
+    # NF_HEDM/src/CalcDiffractionSpots.c:225-236 -- a spot is kept only if
+    # SOME index i satisfies BOTH omega in (omega_ranges[i]) AND the nominal
+    # ring position (yl, zl) inside box_sizes[i]; a spot that fails every
+    # pair is dropped from the theoretical spot list entirely, i.e. it is
+    # excluded from BOTH numerator and denominator of the overlap fraction.
+    #
+    # Entry i of ``omega_ranges`` pairs with entry i of ``box_sizes``, so the
+    # two lists must have equal length. Both default to None => the filter is
+    # OFF and every existing FF/pf/NF caller is bit-unchanged.
+    #
+    # Units: omega_ranges in degrees; box_sizes in MICROMETRES, as
+    # (y_min, y_max, z_min, z_max) on the detector plane relative to the beam
+    # centre, evaluated at the FIRST distance (Lsd[0]) with no grain
+    # displacement, no tilt and no distortion -- exactly the C's
+    # CalcSpotPosition(RingRadius = Lsd[0]*tan(2*theta), eta). Bounds are
+    # STRICT (a spot exactly on an edge is rejected), matching the C.
+    omega_ranges: "list[tuple[float, float]] | None" = None
+    box_sizes: "list[tuple[float, float, float, float]] | None" = None
 
     @property
     def n_distances(self) -> int:
@@ -398,6 +419,38 @@ class HEDMForwardModel(nn.Module):
         )
         self._has_wedge = abs(float(geometry.wedge)) > 0.0
 
+        # Paired OmegaRange / BoxSize acceptance filter. OFF unless BOTH
+        # lists are supplied and non-empty (see HEDMGeometry.box_sizes).
+        om_ranges = getattr(geometry, "omega_ranges", None) or []
+        bx_sizes = getattr(geometry, "box_sizes", None) or []
+        if bool(om_ranges) != bool(bx_sizes):
+            raise ValueError(
+                "omega_ranges and box_sizes are a paired filter: supply both "
+                f"or neither (got {len(om_ranges)} omega_ranges and "
+                f"{len(bx_sizes)} box_sizes)."
+            )
+        if len(om_ranges) != len(bx_sizes):
+            raise ValueError(
+                f"omega_ranges (len {len(om_ranges)}) must match box_sizes "
+                f"(len {len(bx_sizes)}); entry i of one pairs with entry i "
+                "of the other."
+            )
+        self._has_omega_box = bool(bx_sizes)
+        self.register_buffer(
+            "_omega_ranges",
+            torch.tensor(
+                [[float(a), float(b)] for a, b in om_ranges],
+                dtype=torch.float64, device=device,
+            ).reshape(-1, 2),
+        )
+        self.register_buffer(
+            "_box_sizes",
+            torch.tensor(
+                [[float(v) for v in bs] for bs in bx_sizes],
+                dtype=torch.float64, device=device,
+            ).reshape(-1, 4),
+        )
+
         # Scan config
         self.scan_config = scan_config
         if scan_config is not None:
@@ -738,6 +791,7 @@ class HEDMForwardModel(nn.Module):
         orientation_matrices: torch.Tensor,
         hkls_cart: Optional[torch.Tensor] = None,
         thetas: Optional[torch.Tensor] = None,
+        per_grain: bool = False,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """Core Bragg geometry: orientations + G-vectors -> angles.
 
@@ -779,14 +833,32 @@ class HEDMForwardModel(nn.Module):
         # batch, (b) per-voxel hkls_cart shape (..., M, 3) for strained
         # rendering. Both flow through the same einsum via leading-dim
         # broadcasting on the second arg.
-        G_C = torch.einsum("...nij,...mj->...nmi", orientation_matrices, hkls_cart)
+        #
+        # per_grain=True: ELEMENT-WISE pairing of grain i's orientation with
+        # grain i's strained hkls -- NO orientation x strain cross-product.
+        # Requires orientation_matrices (N,3,3); hkls_cart (N,M,3) per-grain or
+        # (M,3) shared. Used by forward_per_grain() for the O(N*M) fast path.
+        if per_grain:
+            if hkls_cart.dim() == 2:                 # (M,3) shared lattice
+                G_C = torch.einsum("nij,mj->nmi", orientation_matrices, hkls_cart)
+            else:                                    # (N,M,3) per-grain strain
+                G_C = torch.einsum("nij,nmj->nmi", orientation_matrices, hkls_cart)
+        else:
+            G_C = torch.einsum("...nij,...mj->...nmi", orientation_matrices, hkls_cart)
 
         # v = sin(theta)*|G| -- C precomputes Gs from the UNROTATED G-vector norm
         # (rotation preserves norm in exact arithmetic but not in float64).
         # Match C: use |hkls_cart| (pre-rotation), not |R @ hkls_cart|.
         len_hkl = torch.norm(hkls_cart, dim=-1)  # (M,) or (..., M)
         v_no_wedge = torch.sin(thetas) * len_hkl  # (M,) or (..., M)
-        v_no_wedge = v_no_wedge.unsqueeze(-2).expand_as(G_C[..., 0])  # (..., N, M)
+        # Broadcast v to G_C's (..., N, M) grid. Layout-agnostic: in the
+        # cross-product layout v gains an orientation axis at -2; in the
+        # per_grain / shared layouts it already matches. Numerically identical
+        # to the former ``unsqueeze(-2).expand_as`` for those existing paths.
+        gc0 = G_C[..., 0]
+        while v_no_wedge.dim() < gc0.dim():
+            v_no_wedge = v_no_wedge.unsqueeze(-2)
+        v_no_wedge = v_no_wedge.expand_as(gc0)
 
         # ---- Wedge: rigorous geometric formulation -------------------
         # The rotation axis tilts from z to n_hat = (sin W, 0, cos W).
@@ -900,9 +972,16 @@ class HEDMForwardModel(nn.Module):
         eta = self.safe_arccos(Gz_lab / r_yz)
         eta = -torch.sign(Gy_lab) * eta
 
-        # 2*theta  (broadcast thetas to match 2N dimension)
-        two_theta_single = 2.0 * thetas.unsqueeze(-2)  # (..., 1, M) or (1, M)
-        two_theta = two_theta_single.expand_as(all_omega)
+        # 2*theta -- broadcast to the single-branch (pre-cat) grid, then double
+        # along the grain axis to match all_omega's 2N. Layout-agnostic (handles
+        # the per_grain layout where the 2N axis is grain-doubled) and
+        # numerically identical to the former expand for the cross-product path
+        # (thetas does not depend on the orientation axis).
+        tt = 2.0 * thetas
+        while tt.dim() < omega_p.dim():
+            tt = tt.unsqueeze(-2)
+        tt = tt.expand_as(omega_p)
+        two_theta = torch.cat([tt, tt], dim=-2)
 
         # Validity mask
         valid_p = disc_valid & coswp_valid
@@ -917,8 +996,93 @@ class HEDMForwardModel(nn.Module):
                  ((math.pi - torch.abs(eta)) >= self.min_eta)
         valid = valid * eta_ok.float()
 
+        # Paired OmegaRange / BoxSize gate (no-op unless both were supplied).
+        if self._has_omega_box:
+            valid = valid * self.omega_box_mask(
+                all_omega, eta, two_theta
+            ).to(valid.dtype)
+
         return all_omega, eta, two_theta, valid
 
+    # ------------------------------------------------------------------
+    #  omega_box_mask  (port of the KeepSpot gate in CalcDiffrSpots_Furnace)
+    # ------------------------------------------------------------------
+
+    def omega_box_mask(
+        self,
+        omega: torch.Tensor,
+        eta: torch.Tensor,
+        two_theta: torch.Tensor,
+    ) -> torch.Tensor:
+        """Paired ``OmegaRange`` / ``BoxSize`` acceptance mask.
+
+        Direct port of the ``KeepSpot`` gate in
+        ``NF_HEDM/src/CalcDiffrSpots_Furnace`` (CalcDiffractionSpots.c:225-236,
+        identical code at MakeDiffrSpots.c:248-259)::
+
+            CalcSpotPosition(RingRadius, eta, &yl, &zl);   // C:224
+            for (OmegaRangeNo = 0; OmegaRangeNo < NOmegaRanges; OmegaRangeNo++) {
+              KeepSpot = 0;
+              if ((Omega > OmegaRange[i][0]) && (Omega < OmegaRange[i][1]) &&
+                  (yl > BoxSizes[i][0]) && (yl < BoxSizes[i][1]) &&
+                  (zl > BoxSizes[i][2]) && (zl < BoxSizes[i][3])) {
+                KeepSpot = 1; break;
+              }
+            }
+
+        with ``RingRadius = Lsd[0] * tan(2*theta)`` (C:209, and the caller
+        passes ``Lsd[0]`` -- SharedFuncsFit.c:830), and
+        ``yl = -sin(eta)*RingRadius``, ``zl = cos(eta)*RingRadius``
+        (CalcDiffractionSpots.c:80-85).
+
+        Semantics that matter:
+
+        - ``(yl, zl)`` are the **nominal ring** coordinates in micrometres at
+          the FIRST distance, with **no** grain displacement, **no** tilt and
+          **no** distortion applied. They are not pixel coordinates and carry
+          no beam-centre offset.
+        - Bounds are **strict** on all six comparisons (edge-exact spots are
+          rejected).
+        - Range ``i`` pairs with box ``i``; a spot is kept if **any** pair
+          accepts (``break`` on first hit).
+        - A rejected spot never enters ``TheorSpots``, so ``nTspots`` shrinks
+          and the spot is dropped from **both** the numerator and the
+          denominator of ``CalcFracOverlap`` (SharedFuncsFit.c:645, 648-650) --
+          it is *not* counted as a miss. Folding the mask into ``valid`` here
+          reproduces that, because ``valid`` is the denominator weight in
+          ``ObsVolume.hard_fraction`` / ``soft_fraction``.
+
+        Parameters
+        ----------
+        omega, eta, two_theta : Tensor
+            Radians, all broadcastable to a common shape.
+
+        Returns
+        -------
+        Tensor (bool), same broadcast shape.
+        """
+        dtype = omega.dtype
+        # C uses Lsd[0] for the ring radius regardless of nDistances.
+        Lsd0 = self._Lsd_eff.to(dtype).reshape(-1)[0]
+        ring_radius = Lsd0 * torch.tan(two_theta)
+        yl = -torch.sin(eta) * ring_radius
+        zl = torch.cos(eta) * ring_radius
+        omega_deg = omega * self.RAD2DEG
+
+        om = self._omega_ranges.to(dtype)
+        bx = self._box_sizes.to(dtype)
+        accept = torch.zeros(
+            torch.broadcast_shapes(omega_deg.shape, yl.shape, zl.shape),
+            dtype=torch.bool, device=omega.device,
+        )
+        for i in range(om.shape[0]):
+            accept = accept | (
+                (omega_deg > om[i, 0]) & (omega_deg < om[i, 1])
+                & (yl > bx[i, 0]) & (yl < bx[i, 1])
+                & (zl > bx[i, 2]) & (zl < bx[i, 3])
+            )
+        return accept
+
     # ------------------------------------------------------------------
     #  Tilt rotation matrix (RotationTilts in SharedFuncsFit.c:230-266)
     # ------------------------------------------------------------------
@@ -1354,6 +1518,11 @@ class HEDMForwardModel(nn.Module):
         if positions.shape[-1] == 2:
             positions = F.pad(positions, (0, 1), value=0.0)
 
+        # Footgun guard: per-grain lattice/strain + N>1 orientations forms an
+        # N x N orientation x strain cross-product (output (N, 2N, M)); callers
+        # simulating a fixed polycrystal almost always want only the diagonal.
+        self._warn_if_cross_product(euler_angles, lattice_params, strain)
+
         # 1. Orientation matrices
         orientation_matrices = self.euler2mat(euler_angles)
 
@@ -1382,6 +1551,102 @@ class HEDMForwardModel(nn.Module):
 
         return spots
 
+    @staticmethod
+    def _warn_if_cross_product(euler_angles, lattice_params, strain):
+        """Warn when forward() would form an orientation x strain cross-product.
+
+        Fires only when there are N>1 orientations AND lattice_params/strain
+        carry a matching per-grain axis -- the case where the (N, 2N, M) output
+        is an N x N cross-product and the caller likely wanted the diagonal.
+        Shared lattice/strain (no grain axis) is the correct (2N, M) path and
+        does not warn.
+        """
+        n_orient = euler_angles.shape[-2] if euler_angles.dim() >= 2 else 1
+        if n_orient <= 1:
+            return
+
+        def _has_grain_axis(t):
+            if t is None:
+                return False
+            if t.shape[-1] == 6 and t.dim() >= 2:          # Voigt lattice or strain
+                return t.shape[-2] == n_orient
+            if tuple(t.shape[-2:]) == (3, 3) and t.dim() >= 3:  # full-tensor strain
+                return t.shape[-3] == n_orient
+            return False
+
+        if _has_grain_axis(lattice_params) or _has_grain_axis(strain):
+            warnings.warn(
+                f"forward() called with {n_orient} orientations and per-grain "
+                "lattice_params/strain forms an orientation x strain "
+                f"cross-product (output shape ({n_orient}, {2 * n_orient}, M)); "
+                "only the diagonal [i, i] and [i, i+N] is physical. For a fixed "
+                "polycrystal use forward_per_grain() (element-wise, O(N*M)), "
+                "or index the diagonal of this output.",
+                stacklevel=3,
+            )
+
+    def forward_per_grain(
+        self,
+        euler_angles: torch.Tensor,
+        positions: torch.Tensor,
+        lattice_params: Optional[torch.Tensor] = None,
+        strain: Optional[torch.Tensor] = None,
+    ) -> SpotDescriptors:
+        """Element-wise per-grain forward simulation -- the fast path.
+
+        Grain ``i`` is simulated with orientation ``i``, lattice/strain ``i``
+        and position ``i``, WITHOUT the orientation x strain cross-product that
+        :meth:`forward` forms when both are per-grain. The output has leading
+        shape ``(2N, M)`` (the two omega branches doubled along the grain axis),
+        which is exactly the diagonal of :meth:`forward`'s ``(N, 2N, M)`` output
+        -- so gradient and gradient-free callers agree bit-for-bit.
+
+        Cost is O(N*M) rather than O(N^2 * M), matching the algorithm of the C
+        reference ``ForwardSimulationCompressed.c``. Fully differentiable and
+        device-portable; for pure forward simulation wrap the call in
+        ``torch.inference_mode()``.
+
+        Parameters
+        ----------
+        euler_angles : Tensor (N, 3)
+            Bunge ZXZ Euler angles (radians), one per grain. No leading batch.
+        positions : Tensor (N, 3) or (N, 2)
+            Real-space grain positions (micrometers).
+        lattice_params : Tensor (6,) or (N, 6), optional
+            Shared or per-grain lattice [a,b,c,alpha,beta,gamma] (Ang/deg).
+        strain : Tensor (6,), (N, 6), (3, 3), or (N, 3, 3), optional
+            Shared or per-grain crystal-frame strain (plain-Voigt or full 3x3).
+
+        Returns
+        -------
+        SpotDescriptors with leading shape ``(2N, M)``: grain ``i``'s two omega
+        solutions live at axis-(-2) indices ``i`` and ``i + N``.
+        """
+        if positions.shape[-1] == 2:
+            positions = F.pad(positions, (0, 1), value=0.0)
+
+        orientation_matrices = self.euler2mat(euler_angles)        # (N, 3, 3)
+
+        hkls_cart = None
+        thetas = None
+        if lattice_params is not None:
+            hkls_cart, thetas = self.correct_hkls_latc(lattice_params, strain=strain)
+        elif strain is not None:
+            raise ValueError(
+                "strain was supplied but lattice_params is None; strain "
+                "requires a reference lattice to apply (I + eps)^{-1} @ B0."
+            )
+
+        omega, eta, two_theta, valid = self.calc_bragg_geometry(
+            orientation_matrices, hkls_cart, thetas, per_grain=True
+        )
+        spots = self.project_to_detector(omega, eta, two_theta, positions, valid)
+
+        if self.scan_config is not None:
+            spots = self.filter_by_scan(spots, positions)
+
+        return spots
+
     # ------------------------------------------------------------------
     #  filter_by_scan  (beam proximity for pf-HEDM)
     # ------------------------------------------------------------------

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/midas_diffract/simulate_panel_zarrs.py

@@ -613,18 +613,13 @@ def simulate_panel_zarrs(
             f"BC=({panel.y_bc:.1f},{panel.z_bc:.1f}), "
             f"tx={panel.tx:.2f}° ty={panel.ty:.2f}° tz={panel.tz:.2f}°")
         with torch.no_grad():
-            spots = model(eulers_t, positions_t,
-                          lattice_params=latc, strain=strain_t)
-
-        valid_np = (spots.valid > 0.5).cpu().numpy()
-        # Strain × orientation diagonal mask
-        G_strain, Kdim, Mdim = valid_np.shape
-        if G_strain == n_grains and Kdim == 2 * n_grains:
-            diag_mask = np.zeros((G_strain, Kdim, Mdim), dtype=bool)
-            for gi in range(G_strain):
-                diag_mask[gi, gi, :] = True
-                diag_mask[gi, gi + n_grains, :] = True
-            valid_np = valid_np & diag_mask
+            # Per-grain fast path: O(N*M), no orientation x strain cross-product.
+            # Output is (2N, M) -- already the per-grain diagonal, no mask needed.
+            spots = model.forward_per_grain(eulers_t, positions_t,
+                                            lattice_params=latc, strain=strain_t)
+
+        valid_np = (spots.valid > 0.5).cpu().numpy()   # (2N, M)
+        Kdim, Mdim = valid_np.shape                     # Kdim == 2 * n_grains
 
         # First-wins: drop spots that an earlier panel already took
         # (avoids double-counting in the rare overlap).
@@ -652,10 +647,11 @@ def simulate_panel_zarrs(
         y_pix    = spots.y_pixel.cpu().numpy()
         z_pix    = spots.z_pixel.cpu().numpy()
         frame_nr = spots.frame_nr.cpu().numpy()
-        grain_ids = np.broadcast_to(np.arange(G_strain)[:, None, None],
-                                    (G_strain, Kdim, Mdim))
-        hkl_ids = np.broadcast_to(np.arange(Mdim)[None, None, :],
-                                  (G_strain, Kdim, Mdim))
+        # (2N, M): row k holds grain (k % n_grains); rows [0:N]/[N:2N] are the
+        # two omega branches.
+        grain_ids = np.broadcast_to((np.arange(Kdim) % n_grains)[:, None],
+                                    (Kdim, Mdim))
+        hkl_ids = np.broadcast_to(np.arange(Mdim)[None, :], (Kdim, Mdim))
 
         rec = {
             "grain_id":     grain_ids.reshape(-1)[flat_idx].astype(np.int32),

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/midas_diffract.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.4.0
+Version: 0.7.0
 Summary: End-to-end differentiable forward model for High-Energy Diffraction Microscopy (FF, NF, pf-HEDM)
 Author-email: Hemant Sharma <hsharma@anl.gov>
 License-Expression: BSD-3-Clause

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/midas_diffract.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ tests/test_forward.py
 tests/test_hkls.py
 tests/test_losses.py
 tests/test_multi_detector.py
+tests/test_omega_box_filter.py
 tests/test_strain_tensor.py
 tests/test_tilts.py
 tests/test_wedge.py

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "midas-diffract"
-version = "0.4.0"
+version = "0.7.0"
 description = "End-to-end differentiable forward model for High-Energy Diffraction Microscopy (FF, NF, pf-HEDM)"
 readme = "README.md"
 license = "BSD-3-Clause"

{midas_diffract-0.4.0 → midas_diffract-0.7.0}/tests/test_forward.py

@@ -1201,5 +1201,89 @@ class TestStrainTensorInput:
         torch.testing.assert_close(sp_v.omega, sp_S.omega)
 
 
+# ===================================================================
+#  Test: forward_per_grain (fast path) == diagonal of forward()
+# ===================================================================
+
+class TestForwardPerGrain:
+    """The O(N*M) per-grain path must equal the diagonal of forward()'s
+    O(N^2*M) orientation x strain cross-product, bit-for-bit, so the
+    gradient and gradient-free callers agree."""
+
+    def _setup(self, nf_geometry, device, N=6, strained=True):
+        model, _, _ = make_model_with_cubic_iron(nf_geometry, device)
+        torch.manual_seed(7)
+        euler = torch.rand(N, 3, dtype=torch.float64) * 2 * math.pi
+        pos = torch.arange(N * 3, dtype=torch.float64).reshape(N, 3) * 5.0
+        latc = torch.tensor([2.87, 2.87, 2.87, 90., 90., 90.],
+                            dtype=torch.float64).expand(N, 6)
+        if strained:
+            s = torch.randn(N, 3, 3, dtype=torch.float64) * 1e-3
+            strain = (s + s.transpose(-1, -2)) / 2
+        else:
+            strain = None
+        return model, euler, pos, latc, strain, N
+
+    def test_per_grain_equals_forward_diagonal_strained(self, nf_geometry, device):
+        # Compare per-grain (no global sort, which would reorder near-degenerate
+        # spots): forward() diagonal rows [gi,gi] & [gi,gi+N] vs
+        # forward_per_grain() rows [gi] & [gi+N]. Element-wise einsum differs
+        # from the cross-product einsum only by fp64 reduction order (~1e-12).
+        model, euler, pos, latc, strain, N = self._setup(nf_geometry, device)
+        spf = model(euler, pos, lattice_params=latc, strain=strain)   # (N,2N,M)
+        spg = model.forward_per_grain(euler, pos, lattice_params=latc, strain=strain)  # (2N,M)
+
+        vf = spf.valid > 0.5
+        vg = spg.valid > 0.5
+        max_diff = 0.0
+        for fld in ("two_theta", "eta", "omega", "y_pixel", "z_pixel", "frame_nr"):
+            F = getattr(spf, fld)
+            G = getattr(spg, fld)
+            for gi in range(N):
+                for kf, kg in ((gi, gi), (gi + N, gi + N)):
+                    # validity must agree exactly on the diagonal
+                    assert torch.equal(vf[gi, kf], vg[kg])
+                    max_diff = max(max_diff, float((F[gi, kf] - G[kg]).abs().max()))
+        assert max_diff < 1e-9, f"per-grain vs forward-diagonal max diff {max_diff:.2e}"
+
+    def test_per_grain_shape_is_2N(self, nf_geometry, device):
+        model, euler, pos, latc, strain, N = self._setup(nf_geometry, device)
+        sp = model.forward_per_grain(euler, pos, lattice_params=latc, strain=strain)
+        assert sp.valid.shape[0] == 2 * N
+
+    def test_per_grain_shared_lattice_no_strain(self, nf_geometry, device):
+        # nominal (no lattice) path must also work and avoid cross-product
+        model, euler, pos, latc, _, N = self._setup(nf_geometry, device, strained=False)
+        sp = model.forward_per_grain(euler, pos)
+        assert sp.valid.shape[0] == 2 * N
+
+    def test_per_grain_differentiable(self, nf_geometry, device):
+        model, euler, pos, latc, strain, N = self._setup(nf_geometry, device)
+        euler = euler.clone().requires_grad_(True)
+        sp = model.forward_per_grain(euler, pos, lattice_params=latc, strain=strain)
+        # omega/eta/pixels depend on orientation (two_theta does not)
+        (sp.omega * sp.valid).sum().backward()
+        assert euler.grad is not None and torch.all(torch.isfinite(euler.grad))
+
+    def test_per_grain_fp32(self, nf_geometry, device):
+        model, euler, pos, latc, strain, N = self._setup(nf_geometry, device)
+        sp = model.forward_per_grain(euler.float(), pos.float(),
+                                     lattice_params=latc.float(), strain=strain.float())
+        assert sp.valid.dtype == torch.float32 or sp.valid.dtype == torch.float64
+
+    def test_forward_warns_on_cross_product(self, nf_geometry, device):
+        model, euler, pos, latc, strain, N = self._setup(nf_geometry, device)
+        with pytest.warns(UserWarning, match="cross-product"):
+            model(euler, pos, lattice_params=latc, strain=strain)
+
+    def test_forward_no_warn_shared_lattice(self, nf_geometry, device):
+        import warnings as _w
+        model, euler, pos, _, _, N = self._setup(nf_geometry, device, strained=False)
+        shared = torch.tensor([2.87, 2.87, 2.87, 90., 90., 90.], dtype=torch.float64)
+        with _w.catch_warnings():
+            _w.simplefilter("error")  # any cross-product warning becomes an error
+            model(euler, pos, lattice_params=shared)  # shared -> (2N,M), no warn
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

midas_diffract-0.7.0/tests/test_omega_box_filter.py

@@ -0,0 +1,352 @@
+"""Tests for the paired ``OmegaRange`` / ``BoxSize`` acceptance gate.
+
+Reference C: ``CalcDiffrSpots_Furnace`` in
+``NF_HEDM/src/CalcDiffractionSpots.c:185-248`` (identical code in
+``NF_HEDM/src/MakeDiffrSpots.c:216-259``)::
+
+    RealType RingRadius = distance * tan(2 * deg2rad * Thetas[indexhkl]);  # :209
+    CalcSpotPosition(RingRadius, etas[i], &yl, &zl);                       # :224
+      # yl = -(sin(eta) * RingRadius);  zl = cos(eta) * RingRadius;        # :83-84
+    for (OmegaRangeNo = 0; OmegaRangeNo < NOmegaRanges; OmegaRangeNo++) {  # :225
+      KeepSpot = 0;
+      if ((Omega > OmegaRange[i][0]) && (Omega < OmegaRange[i][1]) &&      # :227-232
+          (yl > BoxSizes[i][0]) && (yl < BoxSizes[i][1]) &&
+          (zl > BoxSizes[i][2]) && (zl < BoxSizes[i][3])) {
+        KeepSpot = 1; break;
+      }
+    }
+    if (KeepSpot == 1) { spots[spotnr*3 + ...] = ...; spotnr++; }          # :237-243
+
+``distance`` is ``Lsd[0]`` -- the caller passes it explicitly
+(``NF_HEDM/src/SharedFuncsFit.c:830``). A rejected spot is never written into
+``TheorSpots``, so it is dropped from both the numerator and the denominator of
+``CalcFracOverlap`` (``SharedFuncsFit.c:645`` increments ``TotalPixels`` only
+for spots present in the list).
+
+Run with:
+    cd packages/midas_diffract
+    python -m pytest tests/test_omega_box_filter.py -v
+"""
+
+import math
+
+import pytest
+import torch
+
+from midas_diffract.forward import HEDMForwardModel, HEDMGeometry
+
+
+LSD = 1_000_000.0
+
+
+def _geom(**kw):
+    base = dict(
+        Lsd=LSD,
+        y_BC=1024.0,
+        z_BC=1024.0,
+        px=200.0,
+        omega_start=-180.0,
+        omega_step=1.0,
+        n_frames=360,
+        n_pixels_y=2048,
+        n_pixels_z=2048,
+        min_eta=6.0,
+        wavelength=0.172979,
+        flip_y=False,          # NF convention
+        multi_mode="layered",
+    )
+    base.update(kw)
+    return HEDMGeometry(**base)
+
+
+def _model(geom):
+    """Gold (a = 4.08 A) reflections, matching the NF fit-orientation path."""
+    hkls_int = torch.tensor(
+        [[1, 0, 0], [1, 1, 0], [1, 1, 1], [2, 0, 0], [2, 1, 1]],
+        dtype=torch.float64,
+    )
+    B = torch.eye(3, dtype=torch.float64) / 4.08
+    hkls_cart = hkls_int @ B.T
+    g = torch.linalg.norm(hkls_cart, dim=-1)
+    thetas = torch.asin(g * geom.wavelength / 2.0)
+    return HEDMForwardModel(
+        hkls=hkls_cart, thetas=thetas, geometry=geom,
+        hkls_int=hkls_int, device=torch.device("cpu"),
+    )
+
+
+EUL = torch.tensor([[0.3, 0.5, 0.7]], dtype=torch.float64)
+POS = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64)
+
+
+def _nominal_ring_coords(model, eul=EUL):
+    """(yl, zl) in um exactly as the C CalcSpotPosition computes them."""
+    om, eta, tt, valid = model.calc_bragg_geometry(model.euler2mat(eul))
+    R = LSD * torch.tan(tt)
+    return -torch.sin(eta) * R, torch.cos(eta) * R, om, valid
+
+
+# ---------------------------------------------------------------------------
+#  Default: no BoxSize => filter off, behaviour unchanged
+# ---------------------------------------------------------------------------
+
+class TestDefaultOff:
+    def test_flag_off_and_buffers_empty(self):
+        m = _model(_geom())
+        assert m._has_omega_box is False
+        assert m._omega_ranges.shape == (0, 2)
+        assert m._box_sizes.shape == (0, 4)
+
+    def test_valid_mask_identical_to_wide_open_box(self):
+        """A box wider than any possible ring must not change ``valid``."""
+        m_off = _model(_geom())
+        m_open = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, -1e9, 1e9)],
+        ))
+        s_off = m_off(EUL, POS)
+        s_open = m_open(EUL, POS)
+        assert torch.equal(s_off.valid, s_open.valid)
+        assert float(s_off.valid.sum()) > 0     # the test is not vacuous
+
+    def test_ff_geometry_unaffected(self):
+        """FF/pf callers never pass the keys -> byte-identical forward."""
+        m = _model(_geom(flip_y=True, multi_mode="panel"))
+        assert m._has_omega_box is False
+        s = m(EUL, POS)
+        assert float(s.valid.sum()) > 0
+
+
+# ---------------------------------------------------------------------------
+#  Inside is kept, outside is dropped
+# ---------------------------------------------------------------------------
+
+class TestInsideOutside:
+    def test_spot_inside_box_is_kept_outside_is_dropped(self):
+        m_off = _model(_geom())
+        yl, zl, _om, _v = _nominal_ring_coords(m_off)
+        s_off = m_off(EUL, POS)
+        base = s_off.valid.bool()
+        n_base = int(base.sum())
+        assert n_base > 2
+
+        # Cut just above the smallest valid zl -> drop exactly those spots.
+        z_cut = float(zl[base].min()) + 1.0
+        n_below = int((base & (zl <= z_cut)).sum())
+        assert n_below > 0
+
+        m_box = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, z_cut, 1e9)],
+        ))
+        kept = m_box(EUL, POS).valid.bool()
+
+        assert int(kept.sum()) == n_base - n_below
+        # Every kept spot was valid before and is inside the box.
+        assert bool((kept <= base).all())
+        assert bool((zl[kept] > z_cut).all())
+
+    def test_box_is_micrometres_at_lsd0_not_pixels(self):
+        """A 0..2048 box read as *pixels* would keep everything; read as um
+        (the C convention) it keeps nothing for this geometry."""
+        m_off = _model(_geom())
+        yl, zl, _om, _v = _nominal_ring_coords(m_off)
+        base = m_off(EUL, POS).valid.bool()
+        # Sanity: every valid ring here is far outside a +/-2048 um window.
+        assert float(zl[base].abs().min()) > 2048.0
+
+        m_px = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-2048.0, 2048.0, -2048.0, 2048.0)],
+        ))
+        assert float(m_px(EUL, POS).valid.sum()) == 0.0
+
+    def test_y_sign_convention(self):
+        """yl = -sin(eta) * R: an eta = +90 deg spot lands at negative y."""
+        m = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, -1e9, 1e9)],
+        ))
+        two_theta = torch.tensor([[0.1]], dtype=torch.float64)
+        eta = torch.tensor([[math.pi / 2]], dtype=torch.float64)
+        omega = torch.zeros_like(eta)
+        R = float(LSD * math.tan(0.1))
+        # Positive-y half plane rejects it, negative-y half plane accepts it.
+        m._box_sizes = torch.tensor([[0.0, 1e9, -1e9, 1e9]], dtype=torch.float64)
+        assert not bool(m.omega_box_mask(omega, eta, two_theta).item())
+        m._box_sizes = torch.tensor([[-1e9, 0.0, -1e9, 1e9]], dtype=torch.float64)
+        assert bool(m.omega_box_mask(omega, eta, two_theta).item())
+        assert R > 0     # ring radius is positive, so the sign came from -sin
+
+
+# ---------------------------------------------------------------------------
+#  Strict (exclusive) bounds -- C uses > and <, never >= / <=
+# ---------------------------------------------------------------------------
+
+class TestStrictBounds:
+    @staticmethod
+    def _one_spot(box):
+        """Evaluate the mask for a single spot at eta = 0, omega = 0."""
+        m = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)], box_sizes=[box],
+        ))
+        two_theta = torch.tensor([[0.1]], dtype=torch.float64)
+        eta = torch.zeros_like(two_theta)
+        omega = torch.zeros_like(two_theta)
+        # zl = cos(0) * Lsd*tan(2theta); reproduce bit-for-bit.
+        zl = float(
+            torch.cos(eta) * (torch.tensor(LSD, dtype=torch.float64)
+                              * torch.tan(two_theta))
+        )
+        return m, omega, eta, two_theta, zl
+
+    def test_zmin_edge_exact_is_rejected(self):
+        m, om, eta, tt, zl = self._one_spot((-1e6, 1e6, 0.0, 1e9))
+        m._box_sizes = torch.tensor([[-1e6, 1e6, zl, 1e9]], dtype=torch.float64)
+        assert not bool(m.omega_box_mask(om, eta, tt).item())
+
+    def test_zmax_edge_exact_is_rejected(self):
+        m, om, eta, tt, zl = self._one_spot((-1e6, 1e6, 0.0, 1e9))
+        m._box_sizes = torch.tensor([[-1e6, 1e6, -1e9, zl]], dtype=torch.float64)
+        assert not bool(m.omega_box_mask(om, eta, tt).item())
+
+    def test_strictly_inside_is_accepted(self):
+        m, om, eta, tt, zl = self._one_spot((-1e6, 1e6, 0.0, 1e9))
+        m._box_sizes = torch.tensor(
+            [[-1e6, 1e6, zl - 1e-6, zl + 1e-6]], dtype=torch.float64,
+        )
+        assert bool(m.omega_box_mask(om, eta, tt).item())
+
+    def test_omega_edge_exact_is_rejected(self):
+        m, om, eta, tt, zl = self._one_spot((-1e6, 1e6, -1e9, 1e9))
+        # omega == 0 exactly; a range starting at 0 must reject it.
+        m._omega_ranges = torch.tensor([[0.0, 180.0]], dtype=torch.float64)
+        assert not bool(m.omega_box_mask(om, eta, tt).item())
+        m._omega_ranges = torch.tensor([[-1e-9, 180.0]], dtype=torch.float64)
+        assert bool(m.omega_box_mask(om, eta, tt).item())
+
+
+# ---------------------------------------------------------------------------
+#  Paired any-pair-accepts logic
+# ---------------------------------------------------------------------------
+
+class TestPairing:
+    def test_any_pair_accepts(self):
+        """Pair 0 accepts on omega only, pair 1 on box only; each spot needs
+        BOTH halves of the SAME pair, and either pair is enough."""
+        m_off = _model(_geom())
+        yl, zl, om, _v = _nominal_ring_coords(m_off)
+        base = m_off(EUL, POS).valid.bool()
+        om_deg = om * (180.0 / math.pi)
+
+        # Split the valid spots by omega sign.
+        neg = base & (om_deg < 0)
+        pos = base & (om_deg > 0)
+        assert int(neg.sum()) > 0 and int(pos.sum()) > 0
+
+        z_hi = float(zl[base].max()) + 1.0
+        z_lo = float(zl[base].min()) - 1.0
+
+        # Pair 0: negative omega, box open.  Pair 1: positive omega, box shut.
+        m2 = _model(_geom(
+            omega_ranges=[(-180.0, 0.0), (0.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, z_lo, z_hi), (-1e9, 1e9, 1e8, 1e9)],
+        ))
+        kept = m2(EUL, POS).valid.bool()
+        assert torch.equal(kept, neg)
+
+        # Swapping the boxes swaps which half survives -- the pairing is by
+        # index, not a global OR of all omegas with all boxes.
+        m3 = _model(_geom(
+            omega_ranges=[(-180.0, 0.0), (0.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, 1e8, 1e9), (-1e9, 1e9, z_lo, z_hi)],
+        ))
+        assert torch.equal(m3(EUL, POS).valid.bool(), pos)
+
+    def test_both_pairs_open_is_union(self):
+        m_off = _model(_geom())
+        yl, zl, om, _v = _nominal_ring_coords(m_off)
+        base = m_off(EUL, POS).valid.bool()
+        z_hi = float(zl[base].max()) + 1.0
+        z_lo = float(zl[base].min()) - 1.0
+        m2 = _model(_geom(
+            omega_ranges=[(-180.0, 0.0), (0.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, z_lo, z_hi), (-1e9, 1e9, z_lo, z_hi)],
+        ))
+        assert torch.equal(m2(EUL, POS).valid.bool(), base)
+
+
+# ---------------------------------------------------------------------------
+#  Arity validation
+# ---------------------------------------------------------------------------
+
+class TestValidation:
+    def test_box_without_omega_range_raises(self):
+        with pytest.raises(ValueError, match="paired filter"):
+            _model(_geom(box_sizes=[(-1e6, 1e6, -1e6, 1e6)]))
+
+    def test_omega_range_without_box_raises(self):
+        with pytest.raises(ValueError, match="paired filter"):
+            _model(_geom(omega_ranges=[(-180.0, 180.0)]))
+
+    def test_length_mismatch_raises(self):
+        with pytest.raises(ValueError, match="must match box_sizes"):
+            _model(_geom(
+                omega_ranges=[(-180.0, 180.0)],
+                box_sizes=[(-1e6, 1e6, -1e6, 1e6), (-1e6, 1e6, -1e6, 1e6)],
+            ))
+
+
+# ---------------------------------------------------------------------------
+#  The mask reaches every downstream consumer of ``valid``
+# ---------------------------------------------------------------------------
+
+class TestPropagation:
+    def _pair(self):
+        m_off = _model(_geom())
+        yl, zl, _om, _v = _nominal_ring_coords(m_off)
+        base = m_off(EUL, POS).valid.bool()
+        z_cut = float(zl[base].min()) + 1.0
+        m_box = _model(_geom(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, z_cut, 1e9)],
+        ))
+        return m_off, m_box
+
+    def test_layer_valid_inherits_in_multi_distance(self):
+        hkls_int = torch.tensor([[1, 1, 1], [2, 0, 0]], dtype=torch.float64)
+        B = torch.eye(3, dtype=torch.float64) / 4.08
+        hkls_cart = hkls_int @ B.T
+        g = torch.linalg.norm(hkls_cart, dim=-1)
+        thetas = torch.asin(g * 0.172979 / 2.0)
+
+        def build(**kw):
+            geom = _geom(
+                Lsd=[LSD, 1.05 * LSD],
+                y_BC=[1024.0, 1024.0],
+                z_BC=[1024.0, 1024.0],
+                **kw,
+            )
+            return HEDMForwardModel(
+                hkls=hkls_cart, thetas=thetas, geometry=geom,
+                hkls_int=hkls_int, device=torch.device("cpu"),
+            )
+
+        m_off = build()
+        s_off = m_off(EUL, POS)
+        assert s_off.layer_valid is not None
+        assert float(s_off.valid.sum()) > 0
+
+        m_box = build(
+            omega_ranges=[(-180.0, 180.0)],
+            box_sizes=[(-1e9, 1e9, 1e8, 1e9)],     # shut
+        )
+        s_box = m_box(EUL, POS)
+        assert float(s_box.valid.sum()) == 0.0
+        assert float(s_box.layer_valid.sum()) == 0.0
+
+    def test_forward_per_grain_agrees_with_forward(self):
+        _m_off, m_box = self._pair()
+        s_fwd = m_box(EUL, POS)
+        s_pg = m_box.forward_per_grain(EUL, POS)
+        assert torch.equal(s_fwd.valid.reshape(-1), s_pg.valid.reshape(-1))