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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.1.1
+Version: 0.2.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,7 @@ 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
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
@@ -102,7 +103,7 @@ MIDAS distribution. See the companion paper and the MIDAS repository
 
 ## Scope
 
-`midas-diffract` v0.1.0 is deliberately focused on the forward model and its
+`midas-diffract` v0.1.x is deliberately focused on the forward model and its
 gradient chain. The following capabilities build on this substrate and are
 released separately as they mature:
 

{midas_diffract-0.1.1 → midas_diffract-0.2.0}/README.md

@@ -75,7 +75,7 @@ MIDAS distribution. See the companion paper and the MIDAS repository
 
 ## Scope
 
-`midas-diffract` v0.1.0 is deliberately focused on the forward model and its
+`midas-diffract` v0.1.x is deliberately focused on the forward model and its
 gradient chain. The following capabilities build on this substrate and are
 released separately as they mature:
 

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

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

{midas_diffract-0.1.1 → midas_diffract-0.2.0}/midas_diffract/forward.py

@@ -89,6 +89,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 +352,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}; "
@@ -1163,6 +1192,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

{midas_diffract-0.1.1 → midas_diffract-0.2.0}/midas_diffract.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.1.1
+Version: 0.2.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,7 @@ 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
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
@@ -102,7 +103,7 @@ MIDAS distribution. See the companion paper and the MIDAS repository
 
 ## Scope
 
-`midas-diffract` v0.1.0 is deliberately focused on the forward model and its
+`midas-diffract` v0.1.x is deliberately focused on the forward model and its
 gradient chain. The following capabilities build on this substrate and are
 released separately as they mature:
 

{midas_diffract-0.1.1 → midas_diffract-0.2.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.1 → midas_diffract-0.2.0}/midas_diffract.egg-info/requires.txt

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

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

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

midas_diffract-0.2.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)