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

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

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

@@ -28,7 +28,7 @@ Quick start
     loss.backward()
 """
 
-__version__ = "0.6.0"
+__version__ = "0.8.0"
 
 from .forward import (
     HEDMForwardModel,
@@ -39,7 +39,7 @@ from .forward import (
 )
 from .hkls import hkls_for_forward_model
 from .losses import SpotMatchingLoss
-from .optimize import optimize_single_grain, evaluate_recovery
+from .optimize import optimize_single_grain, evaluate_recovery, NoMatchError
 
 __all__ = [
     "HEDMForwardModel",
@@ -50,6 +50,7 @@ __all__ = [
     "SpotMatchingLoss",
     "hkls_for_forward_model",
     "optimize_single_grain",
+    "NoMatchError",
     "evaluate_recovery",
     "__version__",
 ]

{midas_diffract-0.6.0 → midas_diffract-0.8.0}/midas_diffract/forward.py

@@ -117,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:
@@ -204,6 +224,90 @@ class SpotDescriptors:
     det_id: Optional[torch.Tensor] = None       # (..., K, M) int64 panel index (panel mode only)
 
 
+# ---------------------------------------------------------------------------
+#  Omega solver
+# ---------------------------------------------------------------------------
+
+def solve_omega(Gx: torch.Tensor, Gy: torch.Tensor, v: torch.Tensor):
+    """Solve ``-Gx cos(w) + Gy sin(w) = v`` for both omega branches.
+
+    Returns ``(omega_p, omega_n, valid)`` with omega in radians on ``(-pi, pi]``
+    and ``valid`` True where a solution exists.
+
+    Closed form, no quadratic::
+
+        -Gx cos w + Gy sin w  ==  rho cos(w - phi),
+            rho = |(Gx, Gy)|,   phi = atan2(Gy, -Gx)
+
+    so the constraint is ``rho cos(w - phi) = v`` and the solutions are
+    ``w = phi +- acos(v / rho)``.
+
+    This replaces a Gy^2-divided quadratic together with everything that route
+    dragged in: ``+ 1e-7`` in each denominator, a discriminant, a clamped
+    ``safe_arccos``, a +-w branch chosen by comparing residuals, and a separate
+    ``|Gy| < 1e-12`` special case. All were consequences of solving for cos w
+    first rather than for w.
+
+    Why it is not merely tidier. MEASURED over 200k random coefficient sets in
+    float64, as the residual of the defining equation above:
+
+    ===========================  ====================  =================
+    regime                       old (eps = 1e-7)      this form
+    ===========================  ====================  =================
+    ordinary ``|Gy| ~ 0.35``     median 1.2e-07        median 2.8e-17
+    small ``|Gy| <= 3e-4``       median 1.7e-04        median 1.4e-17
+    tiny ``|Gy| <= 1e-8``        median 3.0e-04        median 1.4e-17
+    ``Gy`` exactly 0             median 3.0e-04        median 1.4e-17
+    ===========================  ====================  =================
+
+    and, recovering a KNOWN omega within 0.02 deg of zero -- the band the old
+    ``safe_arccos`` clamp froze -- median error 1.56e-02 deg before, 7.2e-15
+    deg now. Note the epsilon was never confined to the small-|Gy| spots the
+    audit identified: ``y2 + 1e-7`` perturbs every coefficient, so essentially
+    every spot carried a small bias.
+
+    The one remaining singularity is ``acos`` at ``|v/rho| = 1``: exact
+    tangency of the Ewald sphere, a genuinely grazing reflection and genuinely
+    non-differentiable. That is where a singularity belongs. The old form
+    additionally put one at ``omega ~ 0``, which is an ordinary spot.
+
+    Verified against the C reference (``ForwardSimulationCompressed`` and
+    ``simulateNF``) via ``tests/test_c_comparison.py``.
+    """
+    rho = torch.sqrt(Gx * Gx + Gy * Gy + 1e-300)
+    phi = torch.atan2(Gy, -Gx)
+
+    # A solution exists iff |v| <= rho -- exactly the old condition: the old
+    # discriminant was 4 Gy^2 (rho^2 - v^2), so ``disc >= 0`` meant
+    # ``|v| <= rho`` for Gy != 0, and the Gy ~ 0 branch's ``|-v/Gx| <= 1`` is
+    # the same statement when rho = |Gx|.
+    cos_off = v / rho
+    valid = torch.abs(cos_off) <= 1.0
+
+    # acos'(+-1) is infinite. Guard the INPUT: torch sums per-spot gradients,
+    # so one tangential spot would otherwise NaN the whole batch. The boundary
+    # value (0 or pi) is restored exactly by the outer where.
+    at_edge = torch.abs(cos_off) >= 1.0
+    cos_safe = torch.where(at_edge, torch.zeros_like(cos_off),
+                           cos_off.clamp(-1.0, 1.0))
+    d_off = torch.acos(cos_safe)
+    d_off = torch.where(
+        at_edge,
+        torch.where(cos_off > 0, torch.zeros_like(d_off),
+                    torch.full_like(d_off, math.pi)),
+        d_off,
+    )
+
+    zero_w = torch.zeros_like(phi)
+    omega_p = torch.where(valid, phi + d_off, zero_w)
+    omega_n = torch.where(valid, phi - d_off, zero_w)
+    # Wrap into (-pi, pi], the range the old acos-based form produced.
+    two_pi = 2.0 * math.pi
+    omega_p = omega_p - two_pi * torch.round(omega_p / two_pi)
+    omega_n = omega_n - two_pi * torch.round(omega_n / two_pi)
+    return omega_p, omega_n, valid
+
+
 # ---------------------------------------------------------------------------
 #  Forward model
 # ---------------------------------------------------------------------------
@@ -399,6 +503,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:
@@ -518,8 +654,21 @@ class HEDMForwardModel(nn.Module):
     # ------------------------------------------------------------------
 
     def safe_arccos(self, x: torch.Tensor) -> torch.Tensor:
-        """Numerically stable arccos: clamp to [-1+eps, 1-eps]."""
-        return torch.acos(torch.clamp(x, -1.0 + self.epsilon, 1.0 - self.epsilon))
+        """Numerically stable arccos: clamp to [-1+eps, 1-eps].
+
+        The clamp sits at the float64 boundary, NOT at ``self.epsilon``.
+        ``self.epsilon`` is 1e-7, and ``acos(1 - 1e-7) = 4.47e-4 rad``, so
+        clamping there froze both the value and the gradient in a +-0.0256 deg
+        band -- 0.029% of valid spots pinned to the band edge. At 1e-15 the
+        residual band is 2.6e-6 deg, four orders smaller and well below any
+        angular resolution this model is used at.
+
+        Where a sine is available, prefer ``atan2``: it has no band at all.
+        The omega solver does exactly that; this is kept for the ``Gy ~ 0``
+        branch and for eta, where the constraint supplies no sine.
+        """
+        eps = 1e-15 if x.dtype == torch.float64 else 1e-7
+        return torch.acos(torch.clamp(x, -1.0 + eps, 1.0 - eps))
 
     # ------------------------------------------------------------------
     #  strain_as_voigt  (accept full 3x3 tensor OR plain-Voigt 6-vector)
@@ -833,67 +982,9 @@ class HEDMForwardModel(nn.Module):
         v = v_no_wedge + sin_W * Gz_p
         # ---------------------------------------------------------------
 
-        # Quadratic solver for omega
-        # -Gx*cos(w) + Gy*sin(w) = v   (with the effective Gx, Gy, v above)
-        # Rearranged: a*cos^2(w) + b*cos(w) + c = 0
-        # C uses almostzero=1e-12 for the Gy≈0 branch (see
-        # NF_HEDM/src/CalcDiffractionSpots.c:96 and
-        # FF_HEDM/src/ForwardSimulationCompressed.c:168). Match exactly.
-        almostzero = 1e-12
-        x2 = Gx * Gx
-        y2 = Gy * Gy
-        a = 1.0 + x2 / (y2 + self.epsilon)
-        b_coeff = 2.0 * v * Gx / (y2 + self.epsilon)
-        c_coeff = v * v / (y2 + self.epsilon) - 1.0
-        discriminant = b_coeff * b_coeff - 4.0 * a * c_coeff
-
-        sqrt_disc = torch.sqrt(torch.abs(discriminant))
-
-        coswp = (-b_coeff + sqrt_disc) / (2.0 * a)
-        coswn = (-b_coeff - sqrt_disc) / (2.0 * a)
-
-        wap = self.safe_arccos(coswp)
-        wan = self.safe_arccos(coswn)
-        wbp = -wap
-        wbn = -wan
-
-        # Select correct branch: the one satisfying -Gx*cos(w)+Gy*sin(w)=v
-        eqap = -Gx * torch.cos(wap) + Gy * torch.sin(wap)
-        eqbp = -Gx * torch.cos(wbp) + Gy * torch.sin(wbp)
-        eqan = -Gx * torch.cos(wan) + Gy * torch.sin(wan)
-        eqbn = -Gx * torch.cos(wbn) + Gy * torch.sin(wbn)
-
-        Dap = torch.abs(eqap - v)
-        Dbp = torch.abs(eqbp - v)
-        Dan = torch.abs(eqan - v)
-        Dbn = torch.abs(eqbn - v)
-
-        all_wp = torch.where(Dap < Dbp, wap, wbp)
-        all_wn = torch.where(Dan < Dbn, wan, wbn)
-
-        # Special case: Gy ~ 0 (C uses almostzero=1e-12)
-        # C code (CalcDiffractionSpots.c:97-106):
-        #   cosome1 = -v / x;
-        #   if (|cosome1| <= 1) { ome = acos(cosome1); solutions: +ome, -ome }
-        gy_zero = torch.abs(Gy) < almostzero
-        cosome_special = -v / (Gx + self.epsilon)
-        cosome_special_valid = (torch.abs(cosome_special) <= 1.0) & gy_zero & (torch.abs(Gx) > self.epsilon)
-        special_w = self.safe_arccos(cosome_special)  # positive omega solution
-        # Two solutions: +ome and -ome
-        special_wp = special_w   # positive
-        special_wn = -special_w  # negative
-
-        # When |Gy| < almostzero, use the special case; otherwise use the quadratic
-        disc_valid = (discriminant >= 0) & (~gy_zero)
-        coswp_valid = (coswp >= -1.0) & (coswp <= 1.0)
-        coswn_valid = (coswn >= -1.0) & (coswn <= 1.0)
-
-        omega_p = torch.where(cosome_special_valid, special_wp,
-                              torch.where(disc_valid & coswp_valid, all_wp,
-                                          torch.zeros_like(all_wp)))
-        omega_n = torch.where(cosome_special_valid, special_wn,
-                              torch.where(disc_valid & coswn_valid, all_wn,
-                                          torch.zeros_like(all_wn)))
+        # Omega solver -- see module-level ``solve_omega``.
+        omega_p, omega_n, sol_valid = solve_omega(Gx, Gy, v)
+
 
         # Concatenate two solutions: (..., 2N, M)
         all_omega = torch.cat([omega_p, omega_n], dim=-2)
@@ -931,21 +1022,102 @@ class HEDMForwardModel(nn.Module):
         tt = tt.expand_as(omega_p)
         two_theta = torch.cat([tt, tt], dim=-2)
 
-        # Validity mask
-        valid_p = disc_valid & coswp_valid
-        valid_n = disc_valid & coswn_valid
-        # For gy_zero special case, valid only if cosome is in [-1, 1]
-        valid_p = valid_p | cosome_special_valid
-        valid_n = valid_n | cosome_special_valid
-        valid = torch.cat([valid_p, valid_n], dim=-2).float()
+        # Validity mask. Both omega solutions exist under the same condition,
+        # |v| <= rho -- the closed form has no separate Gy ~ 0 case to union in.
+        valid = torch.cat([sol_valid, sol_valid], dim=-2).float()
 
         # Eta bounds
         eta_ok = (torch.abs(eta) >= self.min_eta) & \
                  ((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)
     # ------------------------------------------------------------------

{midas_diffract-0.6.0 → midas_diffract-0.8.0}/midas_diffract/optimize.py

@@ -36,6 +36,19 @@ DEG2RAD = math.pi / 180.0
 RAD2DEG = 180.0 / math.pi
 
 
+class NoMatchError(RuntimeError):
+    """Too few observed spots matched for the optimisation to be defined.
+
+    Raised instead of returning a sentinel loss: with no matched spots there is
+    no residual, hence no gradient and no information, so any number returned
+    would be a fiction that a caller could compare against a real fit.
+    """
+
+    def __init__(self, message: str, *, n_matched: int):
+        super().__init__(message)
+        self.n_matched = n_matched
+
+
 def _associate(pred_valid: torch.Tensor, observed: torch.Tensor,
                max_dist: float) -> "tuple[torch.Tensor, torch.Tensor]":
     """Nearest-neighbour observed->predicted association.
@@ -97,9 +110,10 @@ def optimize_single_grain(
         Discard observed spots whose nearest predicted neighbour is
         farther than this in the angular metric (radians).
     min_matches : int
-        If fewer than this many spots are matched at any iteration,
-        return a sentinel large loss. Prevents L-BFGS from diverging
-        through a near-empty match set.
+        If fewer than this many spots are matched at any iteration, the fit is
+        not defined and the function returns a record with ``success=False``
+        and ``loss_history=[inf]`` rather than a sentinel value that could be
+        mistaken for a converged fit.
     phase1_steps, phase2_steps, phase3_steps : int
         Outer-loop step counts per phase. Each step is one L-BFGS call
         with up to ``lbfgs_max_iter`` inner iterations.
@@ -118,6 +132,10 @@ def optimize_single_grain(
                          when no ground truth is supplied; see
                          :func:`evaluate_recovery` for ground-truth eval.
         ``loss_history`` -- list of phase-final loss values.
+        ``success``   -- False if too few spots ever matched; the parameters
+                         are then the caller's own seed, ``loss_history`` is
+                         ``[inf]``, and ``failure_reason`` says why.
+        ``n_matched`` -- spots matched at the last evaluated state.
     """
     if loss is None:
         loss = SpotMatchingLoss(metric="l2")
@@ -130,6 +148,7 @@ def optimize_single_grain(
     opt_euler = init_euler.clone().requires_grad_(True)
     opt_latc = init_lattice.clone().requires_grad_(False)
     loss_history: list = []
+    _last_n_matched = [0]
 
     def make_closure(params):
         def closure():
@@ -143,17 +162,31 @@ def optimize_single_grain(
             pred_flat = coords.squeeze().reshape(-1, 3)
             valid_flat = valid.squeeze().reshape(-1)
             pred_valid = pred_flat[valid_flat > 0.5]
+            # These two early-outs used to return torch.tensor(1e6,
+            # requires_grad=True) -- a FRESH LEAF, disconnected from opt_euler,
+            # on which .backward() was never called. So opt_euler.grad stayed
+            # None, L-BFGS took no step, and optimize_single_grain returned
+            # loss_history = [1e6, 1e6, 1e6] with the parameters moved exactly
+            # 0.0 and no error raised. A constant is the worst possible signal
+            # here: it looks like a fit that plateaued.
+            #
+            # There is nothing to restore -- with no matched spots there are no
+            # observations and hence no gradient. Raise instead, and let the
+            # caller record the failure.
             if pred_valid.shape[0] == 0:
-                return torch.tensor(
-                    1e6, dtype=opt_euler.dtype, requires_grad=True,
+                raise NoMatchError(
+                    "no valid predicted spots at the current state", n_matched=0
                 )
             pred_match, obs_match = _associate(
                 pred_valid, observed_spots, max_match_distance
             )
             if pred_match.shape[0] < min_matches:
-                return torch.tensor(
-                    1e6, dtype=opt_euler.dtype, requires_grad=True,
+                raise NoMatchError(
+                    f"only {pred_match.shape[0]} spots matched, need "
+                    f"{min_matches}",
+                    n_matched=int(pred_match.shape[0]),
                 )
+            _last_n_matched[0] = int(pred_match.shape[0])
             l = loss(pred_match, obs_match)
             l.backward()
             return l
@@ -172,6 +205,25 @@ def optimize_single_grain(
         lat_err = (opt_latc[:3] - init_lattice[:3]).abs().max().item()
         print(f"{step:5d}  {l.item():12.6e}  {misori:12.6f}  {lat_err:10.6f}")
 
+    def _failed(exc: "NoMatchError") -> Dict[str, Any]:
+        """Report a fit that never had enough data to be defined.
+
+        ``loss_history`` is inf rather than a finite constant so a caller
+        cannot compare it against a real fit and prefer it, and ``success`` is
+        False so "the parameters did not move" is distinguishable from "the fit
+        converged".
+        """
+        return {
+            "euler_rad": opt_euler.detach().clone(),
+            "euler_deg": opt_euler.detach().clone() * RAD2DEG,
+            "lattice": opt_latc.detach().clone(),
+            "misori_deg": current_misori_deg(),
+            "loss_history": [float("inf")],
+            "success": False,
+            "n_matched": exc.n_matched,
+            "failure_reason": str(exc),
+        }
+
     if verbose:
         print(f"{'Step':>5}  {'Loss':>12}  {'dMisori(deg)':>12}  {'dLat':>10}")
         print("-" * 55)
@@ -181,11 +233,14 @@ def optimize_single_grain(
         [opt_euler], lr=1.0, max_iter=lbfgs_max_iter,
         line_search_fn="strong_wolfe",
     )
-    for step in range(phase1_steps):
-        l = optimizer.step(make_closure([opt_euler]))
-        log(step, l)
-        if current_misori_deg() < convergence_misori_deg and step > 0:
-            break
+    try:
+        for step in range(phase1_steps):
+            l = optimizer.step(make_closure([opt_euler]))
+            log(step, l)
+            if current_misori_deg() < convergence_misori_deg and step > 0:
+                break
+    except NoMatchError as exc:
+        return _failed(exc)
     loss_history.append(float(l.detach()))
 
     if verbose:
@@ -196,12 +251,15 @@ def optimize_single_grain(
         [opt_latc], lr=1.0, max_iter=lbfgs_max_iter,
         line_search_fn="strong_wolfe",
     )
-    for step in range(phase2_steps):
-        l = optimizer.step(make_closure([opt_latc]))
-        log(step + phase1_steps, l)
-        if (opt_latc[:3] - init_lattice[:3]).abs().max().item() > 0 \
-                and abs(float(l.detach()) - loss_history[-1]) < convergence_lattice_err:
-            break
+    try:
+        for step in range(phase2_steps):
+            l = optimizer.step(make_closure([opt_latc]))
+            log(step + phase1_steps, l)
+            if (opt_latc[:3] - init_lattice[:3]).abs().max().item() > 0 \
+                    and abs(float(l.detach()) - loss_history[-1]) < convergence_lattice_err:
+                break
+    except NoMatchError as exc:
+        return _failed(exc)
     loss_history.append(float(l.detach()))
 
     if verbose:
@@ -212,9 +270,12 @@ def optimize_single_grain(
         [opt_euler, opt_latc], lr=0.5, max_iter=lbfgs_max_iter,
         line_search_fn="strong_wolfe",
     )
-    for step in range(phase3_steps):
-        l = optimizer.step(make_closure([opt_euler, opt_latc]))
-        log(step + phase1_steps + phase2_steps, l)
+    try:
+        for step in range(phase3_steps):
+            l = optimizer.step(make_closure([opt_euler, opt_latc]))
+            log(step + phase1_steps + phase2_steps, l)
+    except NoMatchError as exc:
+        return _failed(exc)
     loss_history.append(float(l.detach()))
 
     return {
@@ -223,6 +284,8 @@ def optimize_single_grain(
         "lattice": opt_latc.detach().clone(),
         "misori_deg": current_misori_deg(),
         "loss_history": loss_history,
+        "success": True,
+        "n_matched": _last_n_matched[0],
     }
 
 

{midas_diffract-0.6.0 → midas_diffract-0.8.0}/midas_diffract.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: midas-diffract
-Version: 0.6.0
+Version: 0.8.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.6.0 → midas_diffract-0.8.0}/midas_diffract.egg-info/SOURCES.txt

@@ -18,6 +18,9 @@ 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_omega_solver_singularities.py
+tests/test_optimize_no_match.py
 tests/test_strain_tensor.py
 tests/test_tilts.py
 tests/test_wedge.py

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "midas-diffract"
-version = "0.6.0"
+version = "0.8.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.6.0 → midas_diffract-0.8.0}/tests/test_distortion_layer.py

@@ -19,7 +19,7 @@ _GD = dict(
     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).
+# A representative calibrated v2 coefficient set (datasetD 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]
 

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

midas_diffract-0.8.0/tests/test_omega_solver_singularities.py

@@ -0,0 +1,205 @@
+"""The omega solver: accuracy and gradients where the old form failed.
+
+The solver satisfies  -Gx cos(w) + Gy sin(w) = v,  so the residual of THAT
+identity is the arbiter -- not agreement with the previous implementation,
+which was itself biased. The closed form
+
+    rho = |(Gx, Gy)|,  phi = atan2(Gy, -Gx),  w = phi +- acos(v/rho)
+
+replaced a Gy^2-divided quadratic whose ``+ 1e-7`` denominators perturbed every
+spot, and whose ``safe_arccos`` froze both value and gradient in a +-0.0256 deg
+band around omega = 0.
+
+C parity is covered separately by tests/test_c_comparison.py (FF and NF); note
+its tolerance is 0.5 deg, which confirms no regression but cannot resolve the
+1e-2 deg effects tested here.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+import torch
+
+from midas_diffract.forward import HEDMForwardModel, HEDMGeometry, solve_omega
+
+
+@pytest.fixture(autouse=True)
+def _float64_default():
+    """float64 for this module only.
+
+    Setting the default dtype at import time leaks into every test module
+    imported afterwards -- it broke 49 tests in test_forward.py, which expects
+    the float32 default.
+    """
+    prev = torch.get_default_dtype()
+    torch.set_default_dtype(torch.float64)
+    yield
+    torch.set_default_dtype(prev)
+
+
+def _solve(Gx, Gy, v):
+    """The SHIPPED solver -- not a copy of it.
+
+    An earlier draft of this file reimplemented the closed form here, which
+    tested the copy rather than the code and silently dropped the tangency
+    guard. ``solve_omega`` is module-level precisely so the test can call it.
+    """
+    wp, wn, _ = solve_omega(Gx, Gy, v)
+    return wp, wn
+
+
+def _residual(w, Gx, Gy, v):
+    return torch.abs(-Gx * torch.cos(w) + Gy * torch.sin(w) - v)
+
+
+def _coeffs(n, gy_scale, seed):
+    rng = np.random.default_rng(seed)
+    Gx = torch.tensor(rng.normal(0, 0.35, n))
+    Gy = torch.tensor(rng.uniform(-gy_scale, gy_scale, n)) if gy_scale < 0.1 \
+        else torch.tensor(rng.normal(0, gy_scale, n))
+    v = torch.tensor(rng.normal(0, 0.20, n))
+    keep = (v * v) <= (Gx * Gx + Gy * Gy)
+    return Gx[keep], Gy[keep], v[keep]
+
+
+@pytest.mark.parametrize("gy_scale,label", [
+    (0.35, "ordinary"), (3e-4, "small |Gy|"), (1e-8, "tiny |Gy|"),
+])
+def test_constraint_residual_is_at_machine_precision(gy_scale, label):
+    """The old form gave median 1.2e-07 (ordinary) to 3.0e-04 (tiny |Gy|)."""
+    Gx, Gy, v = _coeffs(50_000, gy_scale, seed=0)
+    wp, wn = _solve(Gx, Gy, v)
+    r = torch.minimum(_residual(wp, Gx, Gy, v), _residual(wn, Gx, Gy, v))
+    assert float(r.median()) < 1e-15, f"{label}: median residual {float(r.median()):.3e}"
+    assert float(r.max()) < 1e-12, f"{label}: max residual {float(r.max()):.3e}"
+
+
+def test_gy_exactly_zero_needs_no_special_branch():
+    """The old code carried a separate |Gy| < 1e-12 branch. Gy = 0 is now
+    just phi = atan2(0, -Gx), which is 0 or pi. Nothing special about it."""
+    rng = np.random.default_rng(1)
+    Gx = torch.tensor(rng.normal(0, 0.35, 20_000))
+    Gy = torch.zeros_like(Gx)
+    v = torch.tensor(rng.normal(0, 0.20, 20_000))
+    keep = (v * v) <= (Gx * Gx)
+    Gx, Gy, v = Gx[keep], Gy[keep], v[keep]
+    wp, wn = _solve(Gx, Gy, v)
+    r = torch.minimum(_residual(wp, Gx, Gy, v), _residual(wn, Gx, Gy, v))
+    assert float(r.max()) < 1e-12
+
+
+def test_no_frozen_band_near_omega_zero():
+    """Recover a KNOWN omega inside the band the old clamp pinned.
+
+    Old: median error 1.56e-02 deg, i.e. the answer was the band edge rather
+    than the spot's actual omega.
+    """
+    rng = np.random.default_rng(2)
+    n = 50_000
+    Gx = torch.tensor(rng.normal(0, 0.35, n))
+    Gy = torch.tensor(rng.normal(0, 0.35, n))
+    w_true = torch.tensor(rng.uniform(-0.02, 0.02, n)) * math.pi / 180.0
+    v = -Gx * torch.cos(w_true) + Gy * torch.sin(w_true)   # exact by construction
+    wp, wn = _solve(Gx, Gy, v)
+    # Either branch may carry the true root; take the nearer.
+    err = torch.minimum((wp - w_true).abs(), (wn - w_true).abs())
+    err_deg = err * 180.0 / math.pi
+    assert float(err_deg.median()) < 1e-10, (
+        f"median error {float(err_deg.median()):.3e} deg "
+        f"(old form: 1.56e-02 deg -- the frozen band)")
+    assert float(np.percentile(err_deg.numpy(), 99.9)) < 1e-6
+
+
+def test_no_epsilon_bias_at_moderate_gy():
+    """``y2 + 1e-7`` perturbed EVERY coefficient, not only small-|Gy| spots.
+
+    At |Gy| ~ 0.35 the old relative perturbation was ~1e-6, which is what put
+    a ~1e-7 residual on ordinary spots. Nothing here may depend on a tolerance
+    that loose.
+    """
+    Gx, Gy, v = _coeffs(20_000, 0.35, seed=3)
+    wp, wn = _solve(Gx, Gy, v)
+    r = torch.minimum(_residual(wp, Gx, Gy, v), _residual(wn, Gx, Gy, v))
+    assert float((r > 1e-14).double().mean()) < 1e-3
+
+
+# ------------------------------------------------------------- gradients
+
+def test_gradient_finite_and_fd_matching_at_omega_zero():
+    """omega = 0 is an ORDINARY spot. Its gradient must be right, not merely
+    finite -- the old clamp made it exactly 0."""
+    GX, GY = 0.3, 0.4
+    # w = 0  =>  -Gx cos 0 + Gy sin 0 = -Gx, so v = -Gx puts a root at omega=0.
+    # It lands on the MINUS branch here (phi + acos(v/rho) is the other root).
+    v = torch.tensor(-GX)
+
+    def f(gx, gy):
+        return _solve(gx, gy, v)[1]
+
+    Gx = torch.tensor(GX, requires_grad=True)
+    Gy = torch.tensor(GY, requires_grad=True)
+    out = f(Gx, Gy)
+    assert abs(float(out.detach())) < 1e-12, "setup should put omega at 0"
+    g = torch.autograd.grad(out, [Gx, Gy])
+    assert all(torch.isfinite(x).all() for x in g)
+    h = 1e-7
+    for i, (val, gi) in enumerate(zip((GX, GY), g)):
+        a = [torch.tensor(GX), torch.tensor(GY)]
+        a[i] = torch.tensor(val + h); up = float(f(*a))
+        a[i] = torch.tensor(val - h); dn = float(f(*a))
+        fd = (up - dn) / (2 * h)
+        assert abs(float(gi) - fd) < 1e-4, f"param {i}: {float(gi)} vs FD {fd}"
+    # The old clamp made this exactly 0. At least one component must be alive.
+    assert max(abs(float(x)) for x in g) > 1e-6, "gradient is dead at omega = 0"
+
+
+def test_one_tangential_spot_does_not_poison_the_batch():
+    """acos'(+-1) is infinite at exact tangency; torch sums gradients, so an
+    unguarded singular spot would NaN every other spot's gradient."""
+    Gx = torch.tensor([0.3, 0.3, 0.3], requires_grad=True)
+    Gy = torch.tensor([0.4, 0.4, 0.4], requires_grad=True)
+    v = torch.tensor([0.5, 0.1, -0.2])           # v[0] = rho exactly: tangency
+    wp, _ = _solve(Gx, Gy, v)
+    g = torch.autograd.grad(wp.sum(), [Gx, Gy])
+    assert all(torch.isfinite(x).all() for x in g), "tangency poisoned the batch"
+
+
+# ------------------------------------------------------- end to end
+
+def _ff_model():
+    geom = HEDMGeometry(
+        Lsd=1_000_000.0, y_BC=1024.0, z_BC=1024.0, px=200.0,
+        omega_start=0.0, omega_step=0.25, n_frames=1440,
+        n_pixels_y=2048, n_pixels_z=2048, min_eta=6.0, wavelength=0.295,
+    )
+    a, wl = 2.87, 0.295
+    hkls_int = torch.tensor([[1, 1, 0], [2, 0, 0], [2, 1, 1], [2, 2, 0]])
+    hkls_cart = (torch.eye(3) / a @ hkls_int.double().T).T
+    thetas = torch.asin(wl / (2.0 * (1.0 / torch.norm(hkls_cart, dim=-1))))
+    return HEDMForwardModel(hkls=hkls_cart, thetas=thetas, geometry=geom,
+                            hkls_int=hkls_int.double(),
+                            device=torch.device("cpu"))
+
+
+def test_model_forward_is_finite_and_differentiable():
+    model = _ff_model()
+    eul = torch.tensor([[0.4, 0.6, 0.8]], requires_grad=True)
+    pos = torch.zeros(1, 3)
+    out = model(eul.unsqueeze(0), pos.unsqueeze(0))
+    assert torch.isfinite(out.omega).all()
+    assert torch.isfinite(out.eta).all()
+    g = torch.autograd.grad(
+        (out.omega * out.valid).sum(), eul, allow_unused=True)[0]
+    assert g is not None and torch.isfinite(g).all()
+
+
+def test_model_omega_within_180():
+    model = _ff_model()
+    rng = np.random.default_rng(0)
+    eul = torch.tensor(rng.uniform(0, 2 * np.pi, size=(64, 3)))
+    out = model(eul.unsqueeze(0), torch.zeros(64, 3).unsqueeze(0))
+    om = out.omega[out.valid > 0.5]
+    assert float(om.abs().max()) <= 180.0 + 1e-9

midas_diffract-0.8.0/tests/test_optimize_no_match.py

@@ -0,0 +1,92 @@
+"""A single-grain fit with too few matches must not look like a converged one.
+
+Both early-outs in ``optimize_single_grain``'s closure used to return
+``torch.tensor(1e6, requires_grad=True)`` -- a FRESH LEAF, disconnected from
+``opt_euler``, on which ``.backward()`` was never called. So ``opt_euler.grad``
+stayed None, L-BFGS took no step, and the function returned
+``loss_history = [1e6, 1e6, 1e6]`` with the parameters moved exactly 0.0 and no
+error raised. A constant sentinel is the worst possible signal: it reads as a
+fit that plateaued.
+
+Nothing can be restored -- with no matched spots there are no observations and
+so no gradient. The contract (spec_autograd_classB_classC.md, C2) is that the
+failure is legible instead.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+import torch
+
+from midas_diffract import NoMatchError, optimize_single_grain
+from midas_diffract.forward import HEDMForwardModel, HEDMGeometry
+
+
+def _model():
+    geom = HEDMGeometry(
+        Lsd=1_000_000.0, y_BC=1024.0, z_BC=1024.0, px=200.0,
+        omega_start=0.0, omega_step=0.25, n_frames=1440,
+        n_pixels_y=2048, n_pixels_z=2048, min_eta=6.0, wavelength=0.295,
+    )
+    a, wl = 2.87, 0.295
+    hkls_int = torch.tensor([[1, 1, 0], [2, 0, 0], [2, 1, 1], [2, 2, 0]],
+                            dtype=torch.float32)
+    hkls_cart = (torch.eye(3) / a @ hkls_int.T).T
+    thetas = torch.asin(torch.tensor(wl) / (2.0 * (1.0 / torch.norm(hkls_cart, dim=-1))))
+    return HEDMForwardModel(hkls=hkls_cart, thetas=thetas, geometry=geom,
+                            hkls_int=hkls_int, device=torch.device("cpu"))
+
+
+def _run_with_unmatchable_observations():
+    model = _model()
+    init_euler = torch.tensor([0.4, 0.6, 0.8])
+    init_lattice = torch.tensor([2.87, 2.87, 2.87, 90.0, 90.0, 90.0])
+    # Observations nowhere near any prediction, and a tolerance too tight for
+    # anything to associate.
+    obs = torch.tensor([[3.0, 3.0, 3.0], [3.1, 3.1, 3.1]])
+    return optimize_single_grain(
+        model, observed_spots=obs, init_euler=init_euler,
+        init_lattice=init_lattice, position=torch.zeros(3),
+        max_match_distance=1e-12, min_matches=50,
+        phase1_steps=2, phase2_steps=2, phase3_steps=2,
+    )
+
+
+def test_failure_is_reported_not_disguised_as_a_plateau():
+    r = _run_with_unmatchable_observations()
+    assert r["success"] is False
+    assert "failure_reason" in r
+    assert all(math.isinf(x) for x in r["loss_history"]), (
+        f"loss_history={r['loss_history']} -- a finite constant is "
+        f"indistinguishable from a fit that stopped improving"
+    )
+
+
+def test_failed_fit_never_beats_a_real_one_on_loss():
+    r = _run_with_unmatchable_observations()
+    assert min(r["loss_history"]) == float("inf")
+    assert not any(x <= 1e6 for x in r["loss_history"])
+
+
+def test_parameters_are_returned_as_the_caller_s_own_seed():
+    """Not a defect -- there is nothing to move toward. Pinned so it is
+    explicit that unmoved parameters come labelled as a failure."""
+    init_euler = torch.tensor([0.4, 0.6, 0.8])
+    r = _run_with_unmatchable_observations()
+    assert torch.allclose(r["euler_rad"], init_euler, atol=1e-12)
+    assert r["success"] is False
+
+
+def test_n_matched_is_reported():
+    r = _run_with_unmatchable_observations()
+    assert "n_matched" in r
+    assert r["n_matched"] < 50
+
+
+def test_no_match_error_is_public_and_carries_the_count():
+    assert issubclass(NoMatchError, RuntimeError)
+    e = NoMatchError("x", n_matched=3)
+    assert e.n_matched == 3