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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.4.0
+Version: 0.6.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.6.0}/midas_diffract/__init__.py

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

{midas_diffract-0.4.0 → midas_diffract-0.6.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
 
@@ -738,6 +739,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 +781,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 +920,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
@@ -1354,6 +1381,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 +1414,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.6.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.6.0}/midas_diffract.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.4.0
+Version: 0.6.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.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "midas-diffract"
-version = "0.4.0"
+version = "0.6.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.6.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"])