midas-diffract 0.1.0__py3-none-any.whl

midas_diffract/optimize.py

@@ -0,0 +1,248 @@
+"""Single-grain orientation + lattice-parameter recovery via gradient descent.
+
+Three-phase L-BFGS schedule (orientation -> lattice -> joint), with
+nearest-neighbour spot association at every step. Designed for the
+companion-paper FF-HEDM single-grain demo, but works for any geometry
+the underlying ``HEDMForwardModel`` supports.
+
+Quick start
+-----------
+    import midas_diffract as md
+    result = md.optimize_single_grain(
+        model,
+        observed_spots=obs_angular,       # (N, 3): (2theta, eta, omega) in rad
+        init_euler=init_euler_rad,        # (3,) Bunge angles in rad
+        init_lattice=init_latc,           # (6,) [a, b, c, alpha, beta, gamma]
+        position=torch.zeros(3),          # grain centroid in lab frame (um)
+        loss=md.SpotMatchingLoss(metric="l2"),
+    )
+    print(result["misori_deg"], result["lattice_errors"])
+
+The function does not assume any particular weighting; pass a
+:class:`midas_diffract.SpotMatchingLoss` configured with whatever
+``weights=`` you need (typically derived from measurement resolution).
+"""
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, Dict, Optional
+
+import torch
+
+from .forward import HEDMForwardModel
+from .losses import SpotMatchingLoss
+
+DEG2RAD = math.pi / 180.0
+RAD2DEG = 180.0 / math.pi
+
+
+def _associate(pred_valid: torch.Tensor, observed: torch.Tensor,
+               max_dist: float) -> "tuple[torch.Tensor, torch.Tensor]":
+    """Nearest-neighbour observed->predicted association.
+
+    Returns ``(pred_matched, obs_matched)``. Both have leading dim equal to
+    the number of observed spots whose nearest predicted neighbour is
+    within ``max_dist``.
+    """
+    dists = torch.cdist(observed, pred_valid)
+    min_dists, nn_idx = dists.min(dim=1)
+    keep = min_dists < max_dist
+    return pred_valid[nn_idx[keep]], observed[keep]
+
+
+def optimize_single_grain(
+    model: HEDMForwardModel,
+    observed_spots: torch.Tensor,
+    init_euler: torch.Tensor,
+    init_lattice: torch.Tensor,
+    position: Optional[torch.Tensor] = None,
+    *,
+    loss: Optional[SpotMatchingLoss] = None,
+    max_match_distance: float = 0.5,
+    min_matches: int = 5,
+    phase1_steps: int = 15,
+    phase2_steps: int = 15,
+    phase3_steps: int = 10,
+    lbfgs_max_iter: int = 20,
+    convergence_misori_deg: float = 1e-3,
+    convergence_lattice_err: float = 1e-5,
+    verbose: bool = False,
+) -> Dict[str, Any]:
+    """Recover Bunge Euler angles and lattice parameters for a single grain.
+
+    Three-phase L-BFGS schedule:
+      1. Orientation only (Euler angles) -- eta/omega-sensitive.
+      2. Lattice parameters only -- 2theta-sensitive.
+      3. Joint refinement.
+
+    Parameters
+    ----------
+    model : HEDMForwardModel
+        Pre-built forward model. Its ``hkls`` and ``thetas`` set the
+        reflection list against which the grain is fit.
+    observed_spots : Tensor (N, 3)
+        Angular coordinates of observed spots in radians:
+        ``(2theta, eta, omega)``. Use the ``angular`` output of
+        :meth:`HEDMForwardModel.predict_spot_coords` for synthetic data.
+    init_euler : Tensor (3,)
+        Initial Bunge Euler angles in radians.
+    init_lattice : Tensor (6,)
+        Initial lattice parameters ``[a, b, c, alpha, beta, gamma]``,
+        in Angstroms / degrees.
+    position : Tensor (3,), optional
+        Grain centroid in the lab frame (microns). Defaults to the origin.
+    loss : SpotMatchingLoss, optional
+        Loss object. Defaults to ``SpotMatchingLoss(metric="l2")``.
+    max_match_distance : float
+        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.
+    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.
+    convergence_misori_deg, convergence_lattice_err : float
+        Early-exit thresholds for phases 1, 2, and 3.
+    verbose : bool
+        If True, print a per-step progress table.
+
+    Returns
+    -------
+    dict with keys:
+        ``euler_rad`` -- (3,) recovered Euler angles in radians.
+        ``euler_deg`` -- (3,) same in degrees.
+        ``lattice``   -- (6,) recovered lattice parameters.
+        ``misori_deg``-- final misorientation against ``init_euler`` (deg)
+                         when no ground truth is supplied; see
+                         :func:`evaluate_recovery` for ground-truth eval.
+        ``loss_history`` -- list of phase-final loss values.
+    """
+    if loss is None:
+        loss = SpotMatchingLoss(metric="l2")
+    if position is None:
+        position = torch.zeros(3, dtype=init_euler.dtype, device=init_euler.device)
+
+    pos = position.unsqueeze(0)
+    R_init = HEDMForwardModel.euler2mat(init_euler).detach()
+
+    opt_euler = init_euler.clone().requires_grad_(True)
+    opt_latc = init_lattice.clone().requires_grad_(False)
+    loss_history: list = []
+
+    def make_closure(params):
+        def closure():
+            for p in params:
+                if p.grad is not None:
+                    p.grad.zero_()
+            spots = model(opt_euler.unsqueeze(0), pos, lattice_params=opt_latc)
+            coords, valid = HEDMForwardModel.predict_spot_coords(
+                spots, space="angular"
+            )
+            pred_flat = coords.squeeze().reshape(-1, 3)
+            valid_flat = valid.squeeze().reshape(-1)
+            pred_valid = pred_flat[valid_flat > 0.5]
+            if pred_valid.shape[0] == 0:
+                return torch.tensor(
+                    1e6, dtype=opt_euler.dtype, requires_grad=True,
+                )
+            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,
+                )
+            l = loss(pred_match, obs_match)
+            l.backward()
+            return l
+        return closure
+
+    def current_misori_deg() -> float:
+        with torch.no_grad():
+            R_cur = HEDMForwardModel.euler2mat(opt_euler)
+            trace = torch.trace(R_init.T @ R_cur)
+            return torch.acos(torch.clamp((trace - 1) / 2, -1, 1)).item() * RAD2DEG
+
+    def log(step: int, l: torch.Tensor) -> None:
+        if not verbose:
+            return
+        misori = current_misori_deg()
+        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}")
+
+    if verbose:
+        print(f"{'Step':>5}  {'Loss':>12}  {'dMisori(deg)':>12}  {'dLat':>10}")
+        print("-" * 55)
+        print("--- Phase 1: Orientation ---")
+
+    optimizer = torch.optim.LBFGS(
+        [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
+    loss_history.append(float(l.detach()))
+
+    if verbose:
+        print("--- Phase 2: Lattice parameters ---")
+    opt_euler.requires_grad_(False)
+    opt_latc.requires_grad_(True)
+    optimizer = torch.optim.LBFGS(
+        [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
+    loss_history.append(float(l.detach()))
+
+    if verbose:
+        print("--- Phase 3: Joint refinement ---")
+    opt_euler.requires_grad_(True)
+    opt_latc.requires_grad_(True)
+    optimizer = torch.optim.LBFGS(
+        [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)
+    loss_history.append(float(l.detach()))
+
+    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": loss_history,
+    }
+
+
+def evaluate_recovery(
+    result: Dict[str, Any],
+    gt_euler: torch.Tensor,
+    gt_lattice: torch.Tensor,
+) -> Dict[str, float]:
+    """Evaluate a recovery against ground truth.
+
+    Returns misorientation (deg) and per-element lattice errors.
+    Useful in unit tests and the demo notebooks.
+    """
+    R_gt = HEDMForwardModel.euler2mat(gt_euler)
+    R_rec = HEDMForwardModel.euler2mat(result["euler_rad"])
+    trace = torch.trace(R_gt.T @ R_rec)
+    misori = torch.acos(torch.clamp((trace - 1) / 2, -1, 1)).item() * RAD2DEG
+    lat_err = (result["lattice"] - gt_lattice).abs()
+    return {
+        "misori_deg": misori,
+        "lattice_max_err": lat_err.max().item(),
+        "lattice_errors": lat_err.detach().clone(),
+    }

midas_diffract-0.1.0.dist-info/METADATA

@@ -0,0 +1,122 @@
+Metadata-Version: 2.4
+Name: midas-diffract
+Version: 0.1.0
+Summary: End-to-end differentiable forward model for High-Energy Diffraction Microscopy (FF, NF, pf-HEDM)
+Author: Simon Zhang, Nina Andrejevic, Mathew Cherukara
+Author-email: Hemant Sharma <hsharma@anl.gov>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/marinerhemant/MIDAS
+Project-URL: Documentation, https://github.com/marinerhemant/MIDAS
+Project-URL: Issues, https://github.com/marinerhemant/MIDAS/issues
+Keywords: HEDM,3DXRD,differentiable physics,PyTorch,forward model,crystallography,polycrystal,grain
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.22
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Provides-Extra: hkls
+Requires-Dist: midas-hkls>=0.1.0; extra == "hkls"
+Dynamic: license-file
+
+# midas-diffract
+
+End-to-end differentiable forward model for High-Energy Diffraction Microscopy (HEDM), covering far-field (FF), near-field (NF), and point-focused (pf-HEDM) geometries. Pixel-exact agreement with the canonical C reference simulators in [MIDAS](https://github.com/marinerhemant/MIDAS).
+
+Companion paper: Sharma, Zhang, Andrejevic & Cherukara, *An End-to-End Differentiable Forward Model for High-Energy Diffraction Microscopy*, IUCrJ (in preparation, 2026).
+
+## Installation
+
+```bash
+pip install midas-diffract           # core forward model + losses + optimizer
+pip install midas-diffract[hkls]     # also installs midas-hkls for the
+                                     # pure-Python reflection-list helper
+```
+
+Optional PyTorch CUDA or MPS back-ends are used automatically if available.
+
+## Quick start
+
+```python
+import torch
+import midas_diffract as md
+from midas_hkls import Lattice, SpaceGroup           # optional, see [hkls]
+
+# Detector + scan geometry
+geom = md.HEDMGeometry(
+    Lsd=1_000_000.0,              # um
+    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.172979,           # Angstroms
+)
+
+# Reflection list: either compute from a SpaceGroup + Lattice via the
+# midas-hkls helper, or supply (hkls_cart, thetas, hkls_int) yourself
+# (e.g. parsed from MIDAS GetHKLList output).
+sg = SpaceGroup.from_number(225)                        # FCC
+lat = Lattice.for_system("cubic", a=4.08)               # Au
+hkls_cart, thetas, hkls_int = md.hkls_for_forward_model(
+    sg, lat, wavelength_A=geom.wavelength, two_theta_max_deg=15.0,
+)
+
+model = md.HEDMForwardModel(
+    hkls=hkls_cart, thetas=thetas, geometry=geom, hkls_int=hkls_int,
+)
+
+# Forward pass: grain state -> predicted spots. All inputs are leaves
+# of the autograd graph.
+euler = torch.tensor([[45.0, 30.0, 60.0]], requires_grad=True) * (3.14159 / 180)
+pos   = torch.tensor([[0.0, 0.0, 0.0]],  requires_grad=True)
+latc  = torch.tensor([4.08, 4.08, 4.08, 90.0, 90.0, 90.0], requires_grad=True)
+spots = model(euler, pos, lattice_params=latc)
+
+# Scalar loss -> gradients w.r.t. orientation, position, lattice
+loss = ((spots.omega * spots.valid) ** 2).sum()
+loss.backward()
+```
+
+## Output modes
+
+- `md.HEDMForwardModel.predict_spot_coords(spots, space="angular")` — returns
+  `(2θ, η, ω)` in radians for each valid reflection (FF and pf-HEDM).
+- `md.HEDMForwardModel.predict_spot_coords(spots, space="detector")` — returns
+  `(y_pixel, z_pixel, frame_nr)` in fractional units (FF and pf-HEDM).
+- `md.HEDMForwardModel.predict_images(spots, ...)` — renders a differentiable
+  3D detector volume via Gaussian splatting (NF-HEDM output mode).
+
+## Validation
+
+The forward model has been validated to pixel-exact agreement against the
+canonical C simulators `ForwardSimulationCompressed` and `simulateNF` in the
+MIDAS distribution. See the companion paper and the MIDAS repository
+`fwd_sim/paper/` directory for reproducibility scripts.
+
+## Scope
+
+`midas-diffract` v0.1.0 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:
+
+- Sub-voxel grain mixtures
+- Physics-informed regularisation
+- Bayesian uncertainty quantification via HMC / variational inference
+- Temporal 4D-HEDM tracking
+- Coupling to differentiable crystal plasticity (JAX-FEM)
+- EM spot ownership for ambiguous FF patterns
+
+## Citation
+
+If you use `midas-diffract` in published work, please cite the companion paper.
+
+## Licence
+
+BSD-3-Clause.

midas_diffract-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+midas_diffract/__init__.py,sha256=YzApb7RC00sN8nQyWGnABXWR6tdsnV0EXTDfwg4acCA,1709
+midas_diffract/forward.py,sha256=S0VXDsZmUVuhxOA1WfihH7xFhZTU_aggRgoASJNkNhk,63110
+midas_diffract/hkls.py,sha256=yihZkshLhTEaSF11hjN5m1EzbZAZo5wHv9DX6nWFERE,6496
+midas_diffract/losses.py,sha256=QTpSq-I6N_BBc1i0ocom1JGe4HPcA5fbTt2vJegLHJs,16109
+midas_diffract/optimize.py,sha256=kU2lh5ekM0honOIxg1n3-GwKckcoStovU97epHzKKUQ,9406
+midas_diffract-0.1.0.dist-info/licenses/LICENSE,sha256=aWLB2q1q7h9xTi8OBqYGyjcAS_ik-mDFbRx-yoVZxeI,1603
+midas_diffract-0.1.0.dist-info/METADATA,sha256=SxdtyEnavHcNq9JTha8_SURsqNxUg-mabPSelElgsCs,4751
+midas_diffract-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+midas_diffract-0.1.0.dist-info/top_level.txt,sha256=M7QaYqJ9qjhJINmex9O9LMEAhW85zTvR3kpdwG_1oSE,15
+midas_diffract-0.1.0.dist-info/RECORD,,

midas_diffract-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

midas_diffract-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,31 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, UChicago Argonne, LLC, operator of Argonne National
+Laboratory, and the midas-diffract authors.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice,
+   this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.

midas_diffract-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+midas_diffract