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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.1.2
+Version: 0.4.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
@@ -18,6 +18,8 @@ Description-Content-Type: text/markdown
 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.1.2 → midas_diffract-0.4.0}/midas_diffract/__init__.py

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

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

@@ -34,6 +34,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
@@ -89,6 +99,16 @@ class HEDMGeometry:
                                    # Default False preserves the existing
                                    # behaviour: NF applies tilts, FF skips.
                                    # Set True for raw multi-panel simulation.
+    # Radial detector distortion (canonical midas_distortion v2 model). Like
+    # tilts, this maps an IDEAL prediction to the RAW detector position and is
+    # OFF by default -- the FF/pf experimental pipeline pre-corrects distortion
+    # at peak-finding time (transforms), so the forward must NOT re-apply it for
+    # the indexer/fit-grain. Raw-pixel-patch consumers (pf_odf, grain_odf) that
+    # never go through transforms set ``apply_distortion=True`` and supply the
+    # calibrated v2 coefficients to predict in the raw frame.
+    apply_distortion: bool = False
+    p_distortion: "list[float] | None" = None   # 15 v2 coeffs (midas_distortion P_COEF_NAMES order); None/zeros => no-op
+    rho_d: "float | None" = None                # distortion radius normalization (um); None => resolve from detector corner
     multi_mode: str = "layered"    # "layered" (default): NF semantics --
                                    # spot must land on the detector at
                                    # EVERY distance (AllDistsFound).
@@ -342,6 +362,25 @@ class HEDMForwardModel(nn.Module):
 
         # Multi-detector / multi-panel configuration
         self.apply_tilts = bool(geometry.apply_tilts)
+
+        # Radial detector distortion (canonical midas_distortion v2 model),
+        # applied ideal->raw in project_to_detector when apply_distortion=True.
+        # OFF by default => identity => indexer/fit-grain output unchanged.
+        self.apply_distortion = bool(getattr(geometry, "apply_distortion", False))
+        p_dist = getattr(geometry, "p_distortion", None)
+        if p_dist is not None:
+            self.p_distortion = nn.Parameter(
+                torch.as_tensor(p_dist, dtype=torch.float64, device=device),
+                requires_grad=False,
+            )
+            self._has_distortion = bool(
+                torch.any(torch.abs(self.p_distortion.detach()) > 0.0).item()
+            )
+        else:
+            self.p_distortion = None
+            self._has_distortion = False
+        self.rho_d = getattr(geometry, "rho_d", None)
+
         if geometry.multi_mode not in ("layered", "panel"):
             raise ValueError(
                 f"Unknown multi_mode {geometry.multi_mode!r}; "
@@ -408,7 +447,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
         ----------
@@ -418,35 +463,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)
 
     # ------------------------------------------------------------------
@@ -496,6 +520,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)
     # ------------------------------------------------------------------
@@ -507,19 +553,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([
@@ -527,8 +581,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],
@@ -558,10 +611,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.
 
@@ -642,6 +697,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]
@@ -1163,6 +1219,37 @@ class HEDMForwardModel(nn.Module):
             ydet_d = torch.stack(out_y, dim=0)
             zdet_d = torch.stack(out_z, dim=0)
 
+        # Ideal->raw radial distortion (canonical midas_distortion v2 model),
+        # applied on the BC-relative detector-plane coords (um) before pixel
+        # conversion -- mirrors midas_calibrate_v2.forward.geometry. Gated OFF
+        # by default so the indexer/fit-grain (ideal frame) are byte-unchanged;
+        # raw-patch consumers (pf_odf, grain_odf) opt in. R is BC-relative, so
+        # the distortion is frame-flip invariant; the eta convention (and any
+        # phase offset between the calibration frame and this frame) is the one
+        # thing to validate empirically (see implementation_plan_distortion_layer).
+        if self.apply_distortion and self._has_distortion:
+            from midas_distortion import apply_distortion as _apply_dist, \
+                v2_term_layout as _v2_terms, resolve_rho_d_um as _resolve_rho_d
+            eps = torch.tensor(1e-9, dtype=ydet_d.dtype, device=ydet_d.device)
+            R = torch.sqrt(ydet_d * ydet_d + zdet_d * zdet_d).clamp(min=eps)
+            # eta convention matches calibrate_v2 forward: atan2(-y, z), degrees.
+            eta_deg_d = self.RAD2DEG * torch.atan2(-ydet_d, zdet_d)
+            # resolve_rho_d_um passes a supplied rho_d through, or computes the
+            # max BC-relative corner distance (um) when None.
+            rho_d_val, _rho_how = _resolve_rho_d(
+                self.rho_d,
+                NrPixelsY=self.n_pixels_y, NrPixelsZ=self.n_pixels_z,
+                BC_y=float(self._y_BC.reshape(-1)[0]),
+                BC_z=float(self._z_BC.reshape(-1)[0]),
+                pxY=self.px,
+            )
+            rho_d_t = torch.as_tensor(float(rho_d_val), dtype=R.dtype, device=R.device)
+            p_v2 = self.p_distortion.to(R.dtype)
+            R_corr = _apply_dist(R, eta_deg_d, p_v2, rho_d_t, terms=_v2_terms())
+            scale = R_corr / R
+            ydet_d = ydet_d * scale
+            zdet_d = zdet_d * scale
+
         # FF/PF: y-axis on detector flipped (yBC - ydet/px), validated against C
         # NF:    not flipped (yBC + ydet/px), validated against C
         y_sign = -1.0 if self.flip_y else 1.0
@@ -1251,9 +1338,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.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.1.2
+Version: 0.4.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
@@ -18,6 +18,8 @@ Description-Content-Type: text/markdown
 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.1.2 → midas_diffract-0.4.0}/midas_diffract.egg-info/SOURCES.txt

@@ -13,6 +13,7 @@ midas_diffract.egg-info/dependency_links.txt
 midas_diffract.egg-info/requires.txt
 midas_diffract.egg-info/top_level.txt
 tests/test_c_comparison.py
+tests/test_distortion_layer.py
 tests/test_forward.py
 tests/test_hkls.py
 tests/test_losses.py

{midas_diffract-0.1.2 → midas_diffract-0.4.0}/midas_diffract.egg-info/requires.txt

@@ -1,5 +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.1.2 → midas_diffract-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "midas-diffract"
-version = "0.1.2"
+version = "0.4.0"
 description = "End-to-end differentiable forward model for High-Energy Diffraction Microscopy (FF, NF, pf-HEDM)"
 readme = "README.md"
 license = "BSD-3-Clause"
@@ -26,6 +26,8 @@ classifiers = [
 dependencies = [
     "numpy>=1.22",
     "torch>=2.0",
+    "midas-distortion>=0.2.0",
+    "midas-stress>=0.8.0",
 ]
 
 [project.optional-dependencies]

midas_diffract-0.4.0/tests/test_distortion_layer.py

@@ -0,0 +1,78 @@
+"""Tests for the gated ideal->raw radial-distortion layer (midas_distortion v2).
+
+Guarantees:
+  1. default (apply_distortion=False) is BYTE-IDENTICAL to the pre-layer forward
+     -- proves the indexer/fit-grain (ideal-frame) path is untouched.
+  2. apply_distortion=True with zero coeffs is also identity.
+  3. nonzero coeffs shift predicted spots (ideal->raw) by a sane, radius-growing
+     amount.
+  4. the layer is differentiable w.r.t. the distortion coefficients.
+"""
+import numpy as np
+import pytest
+import torch
+
+from midas_diffract.forward import HEDMForwardModel, HEDMGeometry
+
+_GD = dict(
+    Lsd=752000.0, y_BC=695.0, z_BC=874.0, px=172.0,
+    omega_start=180.0, omega_step=-0.25, n_frames=1440,
+    n_pixels_y=1679, n_pixels_z=1679, min_eta=6.0, wavelength=0.172979,
+)
+# A representative calibrated v2 coefficient set (Bucsek Pilatus 2M CeO2).
+_COEFFS = [0.00707, -0.01, 0.00624, 0.01, -34.76, 0.00234, 81.47,
+           -0.00369, -12.29, -0.00727, -5.29, -0.00863, -1.51, -0.00446, -7.79]
+
+
+def _model(**extra):
+    rng = np.random.default_rng(0)
+    hk = torch.tensor(rng.standard_normal((40, 3)), dtype=torch.float64)
+    th = torch.tensor(np.abs(rng.standard_normal(40)) * 0.04 + 0.04, dtype=torch.float64)
+    hi = torch.tensor(rng.integers(-3, 4, (40, 3)), dtype=torch.float64)
+    g = HEDMGeometry(**_GD, **extra)
+    return HEDMForwardModel(hk, th, g, hkls_int=hi).double()
+
+
+_EUL = torch.tensor([[0.3, 0.5, 0.2]], dtype=torch.float64)
+_POS = torch.zeros(1, 3, dtype=torch.float64)
+
+
+def test_default_off_is_unchanged():
+    """No distortion fields == apply_distortion=False == zero coeffs (identity)."""
+    o0 = _model().forward(_EUL, _POS)
+    o_off = _model(apply_distortion=False, p_distortion=[0.0] * 15).forward(_EUL, _POS)
+    o_zero = _model(apply_distortion=True, p_distortion=[0.0] * 15, rho_d=2.0e5).forward(_EUL, _POS)
+    assert torch.equal(o0.y_pixel, o_off.y_pixel)
+    assert torch.equal(o0.z_pixel, o_off.z_pixel)
+    # zero coeffs with the layer ACTIVE must also be a no-op (D == 1).
+    assert torch.allclose(o0.y_pixel, o_zero.y_pixel, atol=1e-9)
+    assert torch.allclose(o0.z_pixel, o_zero.z_pixel, atol=1e-9)
+
+
+def test_distortion_shifts_spots():
+    o0 = _model().forward(_EUL, _POS)
+    o2 = _model(apply_distortion=True, p_distortion=_COEFFS, rho_d=2.0e5).forward(_EUL, _POS)
+    vm = (o0.valid > 0.5) & (o2.valid > 0.5)
+    assert int(vm.sum()) > 5
+    dy = (o2.y_pixel - o0.y_pixel)[vm].abs()
+    assert float(dy.median()) > 0.05   # measurable
+    assert float(dy.max()) < 50.0      # but physical (sub-module)
+
+
+def test_distortion_is_differentiable():
+    m = _model(apply_distortion=True, p_distortion=_COEFFS, rho_d=2.0e5)
+    m.p_distortion.requires_grad_(True)
+    out = m.forward(_EUL, _POS)
+    loss = (out.y_pixel * out.valid).sum() + (out.z_pixel * out.valid).sum()
+    loss.backward()
+    assert m.p_distortion.grad is not None
+    assert float(m.p_distortion.grad.norm()) > 0.0
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_distortion_cpu_cuda_parity():
+    o_cpu = _model(apply_distortion=True, p_distortion=_COEFFS, rho_d=2.0e5).forward(_EUL, _POS)
+    m = _model(apply_distortion=True, p_distortion=_COEFFS, rho_d=2.0e5).to("cuda")
+    o_gpu = m.forward(_EUL.cuda(), _POS.cuda())
+    assert torch.allclose(o_cpu.y_pixel, o_gpu.y_pixel.cpu(), atol=1e-7)
+    assert torch.allclose(o_cpu.z_pixel, o_gpu.z_pixel.cpu(), atol=1e-7)

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

@@ -1111,5 +1111,95 @@ 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)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])