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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.2.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
@@ -19,6 +19,7 @@ License-File: LICENSE
 Requires-Dist: numpy>=1.22
 Requires-Dist: torch>=2.0
 Requires-Dist: midas-distortion>=0.2.0
+Requires-Dist: midas-stress>=0.8.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"

{midas_diffract-0.2.0 → midas_diffract-0.6.0}/midas_diffract/__init__.py

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

{midas_diffract-0.2.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
 
@@ -34,6 +35,16 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+# Canonical orientation + strain-frame primitives. midas_stress is the single
+# source of truth for this math (Bunge ZXZ orientation algebra, sample<->crystal
+# strain rotation); its torch backend is differentiable end-to-end and
+# device-portable, so the forward model delegates rather than re-porting.
+# NOTE: midas_stress's Voigt is Voigt-MANDEL (sqrt2 on shears); this model uses
+# PLAIN-Voigt / raw 3x3 strain, so we only delegate the *rotation*, never the
+# Voigt packing (see rotate_strain_sample_to_crystal / correct_hkls_latc).
+from midas_stress.orientation import euler_to_orient_mat as _ms_euler_to_orient_mat
+from midas_stress.tensor import strain_lab_to_grain as _ms_strain_lab_to_grain
+
 
 # ---------------------------------------------------------------------------
 #  Configuration data classes
@@ -437,7 +448,13 @@ class HEDMForwardModel(nn.Module):
 
     @staticmethod
     def euler2mat(euler_angles: torch.Tensor) -> torch.Tensor:
-        """Convert ZXZ Euler angles to rotation matrices.
+        """Convert ZXZ (Bunge) Euler angles to crystal->sample rotation matrices.
+
+        Delegates to ``midas_stress.orientation.euler_to_orient_mat`` -- the
+        canonical orientation primitive -- so the convention can never drift
+        from the rest of MIDAS. midas_stress's torch backend is differentiable
+        and vmap-safe; the result is identical to the former in-line ZXZ build
+        (R = Rz(phi1) @ Rx(Phi) @ Rz(phi2)) to ~1e-16.
 
         Parameters
         ----------
@@ -447,35 +464,14 @@ class HEDMForwardModel(nn.Module):
         Returns
         -------
         Tensor (..., 3, 3)
-            Rotation matrices.
+            Rotation matrices (crystal->sample), orthogonalized onto SO(3).
         """
-        c = torch.cos(euler_angles)
-        s = torch.sin(euler_angles)
-
-        c0, c1, c2 = c[..., 0], c[..., 1], c[..., 2]
-        s0, s1, s2 = s[..., 0], s[..., 1], s[..., 2]
-
-        # ZXZ rotation matrix: R = Rz(phi1) @ Rx(Phi) @ Rz(phi2)
-        # Verified element-by-element against nfhedm.py lines 114-120.
-        # Built via torch.stack rather than indexed assignment so the function
-        # composes with torch.func.vmap (in-place writes block vmap).
-        row0 = torch.stack([
-             c0 * c2 - s0 * c1 * s2,
-            -s0 * c1 * c2 - c0 * s2,
-             s0 * s1,
-        ], dim=-1)
-        row1 = torch.stack([
-             s0 * c2 + c0 * c1 * s2,
-             c0 * c1 * c2 - s0 * s2,
-            -c0 * s1,
-        ], dim=-1)
-        row2 = torch.stack([
-            s1 * s2,
-            s1 * c2,
-            c1,
-        ], dim=-1)
-        R = torch.stack([row0, row1, row2], dim=-2)
-
+        if not isinstance(euler_angles, torch.Tensor):
+            euler_angles = torch.as_tensor(euler_angles)
+        R = _ms_euler_to_orient_mat(euler_angles)          # (..., 9), torch
+        R = R.reshape(*R.shape[:-1], 3, 3)
+        # midas_stress already returns a proper rotation; orthogonalize keeps the
+        # historical "exactly on SO(3)" guarantee and is idempotent here.
         return HEDMForwardModel.orthogonalize(R)
 
     # ------------------------------------------------------------------
@@ -525,6 +521,28 @@ class HEDMForwardModel(nn.Module):
         """Numerically stable arccos: clamp to [-1+eps, 1-eps]."""
         return torch.acos(torch.clamp(x, -1.0 + self.epsilon, 1.0 - self.epsilon))
 
+    # ------------------------------------------------------------------
+    #  strain_as_voigt  (accept full 3x3 tensor OR plain-Voigt 6-vector)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def strain_as_voigt(strain: torch.Tensor) -> torch.Tensor:
+        """Normalize a strain input to PLAIN-Voigt [e11,e12,e13,e22,e23,e33].
+
+        Accepts either a plain-Voigt ``(..., 6)`` tensor (returned unchanged) or
+        a full symmetric ``(..., 3, 3)`` strain tensor. The 3x3 path is
+        convention-free -- the natural way to hand a strain field straight from
+        a tensor source (e.g. a midas_stress strain field) into the forward
+        model without picking a Voigt/Mandel packing. The off-diagonals are
+        taken as TRUE tensor components (no factor of 2).
+        """
+        if strain.dim() >= 2 and strain.shape[-1] == 3 and strain.shape[-2] == 3:
+            return torch.stack([
+                strain[..., 0, 0], strain[..., 0, 1], strain[..., 0, 2],
+                strain[..., 1, 1], strain[..., 1, 2], strain[..., 2, 2],
+            ], dim=-1)
+        return strain
+
     # ------------------------------------------------------------------
     #  rotate_strain_sample_to_crystal  (port of C RotateStrainSampleToCrystal)
     # ------------------------------------------------------------------
@@ -536,19 +554,27 @@ class HEDMForwardModel(nn.Module):
     ) -> torch.Tensor:
         """Rotate a symmetric infinitesimal strain from sample to crystal frame.
 
-        Port of ``RotateStrainSampleToCrystal`` from
+        Matches ``RotateStrainSampleToCrystal`` from
         ``FF_HEDM/src/ForwardSimulationCompressed.c:399-419``:
-        eps_crystal = OM^T . eps_sample . OM, in Voigt notation
-        [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33].
+        eps_crystal = OM^T . eps_sample . OM.
+
+        The rotation is delegated to ``midas_stress.tensor.strain_lab_to_grain``
+        (the canonical sample/lab -> crystal/grain strain transform; bit-identical
+        to the former in-line ``OM^T S OM``). PLAIN-Voigt pack/unpack is kept here
+        on purpose -- midas_stress's Voigt is Mandel (sqrt2 on shears), which must
+        not touch the forward model's strain input.
 
         Parameters
         ----------
         orientation_matrices : Tensor (..., 3, 3)
+            Crystal->sample matrices (as returned by :meth:`euler2mat`).
         strain_sample : Tensor (..., 6)
+            PLAIN-Voigt [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33].
 
         Returns
         -------
         strain_crystal : Tensor (..., 6)
+            PLAIN-Voigt, same layout.
         """
         e = strain_sample
         S = torch.stack([
@@ -556,8 +582,7 @@ class HEDMForwardModel(nn.Module):
             torch.stack([e[..., 1], e[..., 3], e[..., 4]], dim=-1),
             torch.stack([e[..., 2], e[..., 4], e[..., 5]], dim=-1),
         ], dim=-2)
-        OM = orientation_matrices
-        C = torch.matmul(torch.matmul(OM.transpose(-1, -2), S), OM)
+        C = _ms_strain_lab_to_grain(S, orientation_matrices)   # OM^T S OM
         return torch.stack([
             C[..., 0, 0], C[..., 0, 1], C[..., 0, 2],
             C[..., 1, 1], C[..., 1, 2], C[..., 2, 2],
@@ -587,10 +612,12 @@ class HEDMForwardModel(nn.Module):
         lattice_params : Tensor (..., 6)
             [a, b, c, alpha, beta, gamma] in Angstroms and degrees.
             The ``...`` dimensions allow per-voxel or per-grain parameters.
-        strain : Tensor (..., 6), optional
-            Crystal-frame symmetric infinitesimal strain in Voigt form
-            [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33]. When supplied,
-            the reciprocal lattice is post-multiplied by (I + eps)^{-1}:
+        strain : Tensor (..., 6) or (..., 3, 3), optional
+            Crystal-frame symmetric infinitesimal strain, either PLAIN-Voigt
+            [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33] or a full symmetric
+            3x3 tensor (normalized via :meth:`strain_as_voigt`; off-diagonals are
+            true tensor components, no factor of 2). When supplied, the
+            reciprocal lattice is post-multiplied by (I + eps)^{-1}:
             B = (I + eps)^{-1} @ B0. Use :meth:`rotate_strain_sample_to_crystal`
             to convert a sample-frame strain into the crystal frame.
 
@@ -671,6 +698,7 @@ class HEDMForwardModel(nn.Module):
         # Voigt layout matches C CorrectHKLsLatCEpsilon:
         #   eps = [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33]
         if strain is not None:
+            strain = self.strain_as_voigt(strain)   # accept full 3x3 too
             e11 = strain[..., 0]
             e12 = strain[..., 1]
             e13 = strain[..., 2]
@@ -711,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.
 
@@ -752,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).
@@ -873,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
@@ -1311,9 +1365,10 @@ class HEDMForwardModel(nn.Module):
         lattice_params : Tensor (..., 6) or (..., N, 6), optional
             Strained lattice parameters [a,b,c,alpha,beta,gamma] in
             Angstroms/degrees. None = use nominal hkls/thetas (no strain).
-        strain : Tensor (..., 6) or (..., N, 6), optional
-            Crystal-frame symmetric infinitesimal strain in Voigt form
-            [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33]. Applied as
+        strain : Tensor (..., 6), (..., N, 6), or (..., 3, 3), optional
+            Crystal-frame symmetric infinitesimal strain, either PLAIN-Voigt
+            [eps_11, eps_12, eps_13, eps_22, eps_23, eps_33] or a full symmetric
+            3x3 tensor (see :meth:`strain_as_voigt`). Applied as
             B = (I + eps)^{-1} @ B0 in addition to any lattice-parameter
             strain expressed through ``lattice_params``. Requires
             ``lattice_params`` to be supplied.
@@ -1326,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)
 
@@ -1354,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.2.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.2.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.2.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
@@ -19,6 +19,7 @@ License-File: LICENSE
 Requires-Dist: numpy>=1.22
 Requires-Dist: torch>=2.0
 Requires-Dist: midas-distortion>=0.2.0
+Requires-Dist: midas-stress>=0.8.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"

{midas_diffract-0.2.0 → midas_diffract-0.6.0}/midas_diffract.egg-info/requires.txt

@@ -1,6 +1,7 @@
 numpy>=1.22
 torch>=2.0
 midas-distortion>=0.2.0
+midas-stress>=0.8.0
 
 [dev]
 pytest>=7.0

{midas_diffract-0.2.0 → midas_diffract-0.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "midas-diffract"
-version = "0.2.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"
@@ -27,6 +27,7 @@ dependencies = [
     "numpy>=1.22",
     "torch>=2.0",
     "midas-distortion>=0.2.0",
+    "midas-stress>=0.8.0",
 ]
 
 [project.optional-dependencies]

{midas_diffract-0.2.0 → midas_diffract-0.6.0}/tests/test_forward.py

@@ -1111,5 +1111,179 @@ class TestCrossValidation:
 import torch.nn.functional as F  # for test_2d_position_backward_compat
 
 
+# ===================================================================
+#  Test: midas_stress is the canonical orientation/strain-frame source
+# ===================================================================
+
+class TestMidasStressDelegation:
+    """Guards that euler2mat / rotate_strain_sample_to_crystal delegate to
+    midas_stress and never silently re-port (convention drift)."""
+
+    def test_euler2mat_matches_midas_stress(self):
+        from midas_stress.orientation import euler_to_orient_mat
+        torch.manual_seed(0)
+        eul = torch.rand(7, 3, dtype=torch.float64) * 2 * math.pi
+        R = HEDMForwardModel.euler2mat(eul)
+        R_ms = euler_to_orient_mat(eul).reshape(7, 3, 3)
+        torch.testing.assert_close(R, R_ms, atol=1e-12, rtol=0)
+
+    def test_euler2mat_vmap_safe(self):
+        from torch.func import vmap
+        eul = torch.rand(5, 3, dtype=torch.float64) * 2 * math.pi
+        R = vmap(HEDMForwardModel.euler2mat)(eul)
+        assert R.shape == (5, 3, 3)
+        torch.testing.assert_close(R, HEDMForwardModel.euler2mat(eul), atol=1e-12, rtol=0)
+
+    def test_euler2mat_differentiable(self):
+        eul = torch.rand(3, 3, dtype=torch.float64, requires_grad=True)
+        HEDMForwardModel.euler2mat(eul).pow(2).sum().backward()
+        assert eul.grad is not None and torch.all(torch.isfinite(eul.grad))
+
+    def test_rotate_strain_matches_midas_stress(self):
+        from midas_stress.tensor import strain_lab_to_grain
+        torch.manual_seed(1)
+        OM = HEDMForwardModel.euler2mat(torch.rand(4, 3, dtype=torch.float64))
+        v = torch.rand(4, 6, dtype=torch.float64) * 1e-3       # plain Voigt
+        out = HEDMForwardModel.rotate_strain_sample_to_crystal(OM, v)
+        # reference: build 3x3, rotate via midas_stress, repack plain Voigt
+        S = torch.zeros(4, 3, 3, dtype=torch.float64)
+        S[:, 0, 0], S[:, 1, 1], S[:, 2, 2] = v[:, 0], v[:, 3], v[:, 5]
+        S[:, 0, 1] = S[:, 1, 0] = v[:, 1]
+        S[:, 0, 2] = S[:, 2, 0] = v[:, 2]
+        S[:, 1, 2] = S[:, 2, 1] = v[:, 4]
+        C = strain_lab_to_grain(S, OM)
+        ref = torch.stack([C[:, 0, 0], C[:, 0, 1], C[:, 0, 2],
+                           C[:, 1, 1], C[:, 1, 2], C[:, 2, 2]], dim=-1)
+        torch.testing.assert_close(out, ref, atol=1e-14, rtol=0)
+
+
+# ===================================================================
+#  Test: full 3x3 strain tensor entry point (convention-free)
+# ===================================================================
+
+class TestStrainTensorInput:
+    """A full symmetric 3x3 strain must give the same result as the
+    equivalent plain-Voigt 6-vector (NOT Voigt-Mandel)."""
+
+    def test_strain_as_voigt_roundtrip(self):
+        v = torch.tensor([1e-3, 2e-4, 3e-4, -5e-4, 1e-4, 2e-4])
+        S = torch.tensor([[1e-3, 2e-4, 3e-4],
+                          [2e-4, -5e-4, 1e-4],
+                          [3e-4, 1e-4, 2e-4]])
+        torch.testing.assert_close(HEDMForwardModel.strain_as_voigt(S), v)
+        # plain Voigt passes through unchanged
+        torch.testing.assert_close(HEDMForwardModel.strain_as_voigt(v), v)
+
+    def test_3x3_strain_equals_voigt_in_correct_hkls_latc(self, nf_geometry, device):
+        model, _, _ = make_model_with_cubic_iron(nf_geometry, device)
+        latc = torch.tensor([2.87, 2.87, 2.87, 90., 90., 90.])
+        v = torch.tensor([1e-3, 2e-4, 3e-4, -5e-4, 1e-4, 2e-4])
+        S = torch.tensor([[1e-3, 2e-4, 3e-4],
+                          [2e-4, -5e-4, 1e-4],
+                          [3e-4, 1e-4, 2e-4]])
+        g_v, t_v = model.correct_hkls_latc(latc, strain=v)
+        g_S, t_S = model.correct_hkls_latc(latc, strain=S)
+        torch.testing.assert_close(g_v, g_S)
+        torch.testing.assert_close(t_v, t_S)
+
+    def test_3x3_strain_equals_voigt_in_forward(self, nf_geometry, device):
+        model, _, _ = make_model_with_cubic_iron(nf_geometry, device)
+        eul = torch.tensor([[0.3, 0.5, 1.1]])
+        pos = torch.zeros(1, 3)
+        latc = torch.tensor([2.87, 2.87, 2.87, 90., 90., 90.]).expand(1, 6)
+        v = torch.tensor([[1e-3, 2e-4, 3e-4, -5e-4, 1e-4, 2e-4]])
+        S = torch.tensor([[[1e-3, 2e-4, 3e-4],
+                           [2e-4, -5e-4, 1e-4],
+                           [3e-4, 1e-4, 2e-4]]])
+        sp_v = model(eul, pos, lattice_params=latc, strain=v)
+        sp_S = model(eul, pos, lattice_params=latc, strain=S)
+        torch.testing.assert_close(sp_v.two_theta, sp_S.two_theta)
+        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"])