midas-pf-odf 0.1.0__tar.gz

midas_pf_odf-0.1.0/LICENSE

@@ -0,0 +1,53 @@
+Copyright (c) 2012, UChicago Argonne, LLC
+ 
+All Rights Reserved
+ 
+MIDAS Microstructural Imaging using Diffraction Analysis Software
+
+Materials Physics and Engineering
+Computational X-ray Science
+Advanced Photon Source
+Argonne National Laboratory
+
+Contributing Authors:
+Hemant Sharma (hsharma@anl.gov)
+
+OPEN SOURCE LICENSE
+ 
+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.  Software changes,  
+   modifications, or derivative works, should be noted with comments and  
+   the author and organization's name.
+ 
+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 names of UChicago Argonne, LLC or the Department of Energy  
+   nor the names of its contributors may be used to endorse or promote  
+   products derived from this software without specific prior written  
+   permission.
+ 
+4. The software and the end-user documentation included with the  
+   redistribution, if any, must include the following acknowledgment:
+ 
+   "This product includes software produced by UChicago Argonne, LLC  
+   under Contract No. DE-AC02-06CH11357 with the Department of Energy."
+ 
+****************************************************************************
+ 
+DISCLAIMER
+ 
+THE SOFTWARE IS SUPPLIED "AS IS" WITHOUT WARRANTY OF ANY KIND.
+ 
+Neither the United States GOVERNMENT, nor the United States Department  
+of Energy, NOR UChicago Argonne, LLC, nor any of their employees, makes  
+any warranty, express or implied, or assumes any legal liability or  
+responsibility for the accuracy, completeness, or usefulness of any  
+information, data, apparatus, product, or process disclosed, or  
+represents that its use would not infringe privately owned rights.
+ 
+****************************************************************************

midas_pf_odf-0.1.0/PKG-INFO

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: midas-pf-odf
+Version: 0.1.0
+Summary: Joint per-grain peak-shape inversion of pf-HEDM data; per-voxel orientation/strain (and Phase 2: per-voxel ODF)
+Author-email: Hemant Sharma <hsharma@anl.gov>
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20
+Requires-Dist: torch>=2.0
+Requires-Dist: h5py>=3.0
+Requires-Dist: midas-diffract>=0.1
+Requires-Dist: midas-stress>=0.1
+Requires-Dist: midas-grain-odf>=0.1
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: matplotlib>=3.5; extra == "dev"
+Dynamic: license-file
+
+# midas-pf-odf
+
+Joint per-grain peak-shape inversion of pf-HEDM (point-focused HEDM, also
+known as scanning 3DXRD) data.
+
+**Phase 1 (current):** for each grain, fit all of its voxels'
+`(R_V, ε_V)` jointly to image-MSE on measured 3D peak patches, given
+voxel→grain assignment as input. Differentiable PyTorch forward;
+voxel-summed splat per (spot, scan); closed-form per-spot intensity
+scale `c_s*`.
+
+**Phase 2 (later):** per-voxel ODF (sub-grain mosaic / GND density) on
+the same forward.
+
+## Notebooks
+
+Worked-example Jupyter notebooks live in `notebooks/`. They are **not shipped with `pip install`** — get them by cloning the [MIDAS repository](https://github.com/marinerhemant/MIDAS/tree/master/packages/midas_pf_odf/notebooks).
+
+## Status
+
+Pre-alpha; private. See [dev/RESTART.md](dev/RESTART.md) for the active
+worklog and resume point.
+
+## Where this sits relative to the literature
+
+Henningsson et al. 2020 (J. Appl. Cryst.) introduced PCR / ASR — joint
+multi-voxel intragranular strain reconstruction from scanning 3DXRD,
+using **peak centroids only**. They explicitly call out peak-shape
+inversion as the next direction (untapped at time of writing). This
+package takes that step:
+
+- Centroid → peak-shape (image-MSE on `(F, P, P)` patches).
+- Hand-derived Jacobian → autograd / PyTorch.
+- Single (R_V, ε_V) per voxel → per-voxel ODF (Phase 2).
+
+The forward model is `midas_diffract.HEDMForwardModel`; reusable
+infrastructure (sparse splatter, ODF parameterizations, held-out CV
+selector, sparse-tile chunking) comes from `midas_grain_odf`.
+
+## Quickstart
+
+```bash
+cd packages/midas_pf_odf
+pip install -e .[dev]
+python -m pytest tests/ -xvs
+```
+
+## Layout
+
+```
+midas_pf_odf/        — public library
+    simulate.py       — synthetic plant (orientation+strain gradients, multi-grain, noise)
+    forward.py        — joint per-grain forward (soft beam gate, voxel-summed splat)
+    inversion.py      — joint inversion driver (Adam / L-BFGS, identifiability knobs)
+    validation.py     — per-voxel RMSE vs plant, held-out R²
+    io.py             — I/O for real data (deferred)
+dev/                  — implementation plan, paper, notebooks, worklog
+tests/                — synthetic round-trip tests
+```
+
+## Decisions locked
+
+1. **Identifiability knob:** project ε to mean-zero per grain (default,
+   recommended) or per-voxel-lattice over-parameterized (toggle).
+2. **Smoothness regularizer:** off by default (λ_smooth = 0).
+3. **Synthetic source:** torch-native simulator inside this package.
+4. **Outlier voxels:** soft-mask by held-out R² in the loss.
+5. **Multi-grain parallelism:** one grain × one GPU each via parsl
+   (deferred until workflow integration).
+6. **Output schema:** parallel HDF5; centroid path untouched (deferred).

midas_pf_odf-0.1.0/README.md

@@ -0,0 +1,70 @@
+# midas-pf-odf
+
+Joint per-grain peak-shape inversion of pf-HEDM (point-focused HEDM, also
+known as scanning 3DXRD) data.
+
+**Phase 1 (current):** for each grain, fit all of its voxels'
+`(R_V, ε_V)` jointly to image-MSE on measured 3D peak patches, given
+voxel→grain assignment as input. Differentiable PyTorch forward;
+voxel-summed splat per (spot, scan); closed-form per-spot intensity
+scale `c_s*`.
+
+**Phase 2 (later):** per-voxel ODF (sub-grain mosaic / GND density) on
+the same forward.
+
+## Notebooks
+
+Worked-example Jupyter notebooks live in `notebooks/`. They are **not shipped with `pip install`** — get them by cloning the [MIDAS repository](https://github.com/marinerhemant/MIDAS/tree/master/packages/midas_pf_odf/notebooks).
+
+## Status
+
+Pre-alpha; private. See [dev/RESTART.md](dev/RESTART.md) for the active
+worklog and resume point.
+
+## Where this sits relative to the literature
+
+Henningsson et al. 2020 (J. Appl. Cryst.) introduced PCR / ASR — joint
+multi-voxel intragranular strain reconstruction from scanning 3DXRD,
+using **peak centroids only**. They explicitly call out peak-shape
+inversion as the next direction (untapped at time of writing). This
+package takes that step:
+
+- Centroid → peak-shape (image-MSE on `(F, P, P)` patches).
+- Hand-derived Jacobian → autograd / PyTorch.
+- Single (R_V, ε_V) per voxel → per-voxel ODF (Phase 2).
+
+The forward model is `midas_diffract.HEDMForwardModel`; reusable
+infrastructure (sparse splatter, ODF parameterizations, held-out CV
+selector, sparse-tile chunking) comes from `midas_grain_odf`.
+
+## Quickstart
+
+```bash
+cd packages/midas_pf_odf
+pip install -e .[dev]
+python -m pytest tests/ -xvs
+```
+
+## Layout
+
+```
+midas_pf_odf/        — public library
+    simulate.py       — synthetic plant (orientation+strain gradients, multi-grain, noise)
+    forward.py        — joint per-grain forward (soft beam gate, voxel-summed splat)
+    inversion.py      — joint inversion driver (Adam / L-BFGS, identifiability knobs)
+    validation.py     — per-voxel RMSE vs plant, held-out R²
+    io.py             — I/O for real data (deferred)
+dev/                  — implementation plan, paper, notebooks, worklog
+tests/                — synthetic round-trip tests
+```
+
+## Decisions locked
+
+1. **Identifiability knob:** project ε to mean-zero per grain (default,
+   recommended) or per-voxel-lattice over-parameterized (toggle).
+2. **Smoothness regularizer:** off by default (λ_smooth = 0).
+3. **Synthetic source:** torch-native simulator inside this package.
+4. **Outlier voxels:** soft-mask by held-out R² in the loss.
+5. **Multi-grain parallelism:** one grain × one GPU each via parsl
+   (deferred until workflow integration).
+6. **Output schema:** parallel HDF5; centroid path untouched (deferred).

midas_pf_odf-0.1.0/midas_pf_odf/__init__.py

@@ -0,0 +1,98 @@
+"""midas-pf-odf — joint per-grain peak-shape inversion of pf-HEDM data.
+
+Phase 1: per-voxel (R, ε) from peak shapes. Phase 2: per-voxel ODF.
+"""
+
+__version__ = "0.1.0"
+
+from midas_pf_odf.simulate import (
+    SinglePhaseGrainPlant,
+    plant_single_grain,
+    simulate_grain_patches,
+)
+from midas_pf_odf.forward import (
+    joint_grain_forward,
+    soft_beam_gate,
+)
+from midas_pf_odf.inversion import (
+    neighbor_edges_from_grid_ij,
+    fit_grain_peakshape,
+    GrainPeakFitResult,
+    IdentifiabilityMode,
+)
+from midas_pf_odf.validation import (
+    recovery_metrics,
+    holdout_score,
+)
+from midas_pf_odf.multi_grain import (
+    MultiGrainPlant,
+    plant_multi_grain,
+    split_into_grains,
+    simulate_multi_grain,
+    fit_multi_grain,
+)
+from midas_pf_odf.centroid_baseline import (
+    fit_grain_centroid_baseline,
+    measured_centroids_from_patches,
+    predicted_centroids,
+)
+from midas_pf_odf.calibrate import (
+    RawFrameCalibration,
+    calibrate_raw_frame_geometry,
+    layer_model_factory,
+    measure_patch_offsets,
+)
+from midas_pf_odf.io import (
+    PFGrainDataset,
+    distortion_from_paramstest,
+    saturation_threshold_from_paramstest,
+    load_pf_grain,
+    build_model_from_paramstest,
+    geometry_from_paramstest,
+    parse_paramstest,
+    assemble_grain_patch_data,
+    crop_patches_from_frames,
+    build_model_from_zarr,
+    geometry_from_zarr,
+    read_zarr_params,
+    ZarrFrameSource,
+)
+
+__all__ = [
+    "SinglePhaseGrainPlant",
+    "plant_single_grain",
+    "simulate_grain_patches",
+    "joint_grain_forward",
+    "soft_beam_gate",
+    "fit_grain_peakshape",
+    "GrainPeakFitResult",
+    "IdentifiabilityMode",
+    "neighbor_edges_from_grid_ij",
+    "recovery_metrics",
+    "holdout_score",
+    "MultiGrainPlant",
+    "plant_multi_grain",
+    "split_into_grains",
+    "simulate_multi_grain",
+    "fit_multi_grain",
+    "fit_grain_centroid_baseline",
+    "measured_centroids_from_patches",
+    "predicted_centroids",
+    "PFGrainDataset",
+    "load_pf_grain",
+    "build_model_from_paramstest",
+    "geometry_from_paramstest",
+    "parse_paramstest",
+    "assemble_grain_patch_data",
+    "crop_patches_from_frames",
+    "build_model_from_zarr",
+    "geometry_from_zarr",
+    "read_zarr_params",
+    "ZarrFrameSource",
+    "distortion_from_paramstest",
+    "saturation_threshold_from_paramstest",
+    "RawFrameCalibration",
+    "calibrate_raw_frame_geometry",
+    "layer_model_factory",
+    "measure_patch_offsets",
+]

midas_pf_odf-0.1.0/midas_pf_odf/calibrate.py

@@ -0,0 +1,362 @@
+"""Raw-frame geometry mini-calibration (P1-6).
+
+Empirical-Jacobian damped Gauss-Newton: bump each geometry parameter in
+the ACTUAL forward model (convention-proof — no analytic derivative can
+disagree with the model), fit the parameter steps to the measured anchor
+residuals with a robust (MAD-rejecting) LSQ, iterate, verify collapse.
+
+On the SOH wt316 campaign this collapsed anchor residuals from 7.06 to
+1.4 px RMS and cross-validated between independent loads (tx/BC/tz agree;
+ty/Lsd/ω₀ scatter with 2-ring data). Any raw-frame consumer of a promoted
+``paramstest.txt`` needs it: the powder calibration behind paramstest is
+blind to tx (~0.27° ≈ 3-4e3 µε fake strain) and its convention chain
+(ideal↔raw, flips, ω origin) is easy to get subtly wrong — this
+calibration absorbs all of it into the model actually being used.
+
+Typical use::
+
+    factory = layer_model_factory(layer_dir, ring_numbers=[1, 2, 3, 4],
+                                  n_pixels_y=2880, n_pixels_z=2880,
+                                  n_frames=1440, apply_distortion=True,
+                                  device=dev)
+    ds = load_pf_grain(layer_dir, grain_id, ..., model=factory())
+    cal = calibrate_raw_frame_geometry(ds, cached_patches,
+                                       model_factory=factory)
+    model = factory(geom_mod=cal.calibrated, omega_start=cal.calibrated["om0"])
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+
+from midas_diffract.forward import HEDMForwardModel
+
+__all__ = [
+    "RawFrameCalibration",
+    "calibrate_raw_frame_geometry",
+    "layer_model_factory",
+    "measure_patch_offsets",
+]
+
+# Default parameter set + bump sizes (the empirical-derivative step per
+# parameter; also the natural scale of one GN step component).
+DEFAULT_BUMPS: Dict[str, float] = {
+    "y_BC": 1.0,      # px
+    "z_BC": 1.0,      # px
+    "tx": 0.05,       # deg
+    "ty": 0.05,       # deg
+    "tz": 0.05,       # deg
+    "Lsd": 500.0,     # um
+    "om0": 0.05,      # deg
+}
+
+
+@dataclass
+class RawFrameCalibration:
+    """Result of :func:`calibrate_raw_frame_geometry`."""
+
+    calibrated: Dict[str, float]        # fitted parameter values
+    original: Dict[str, float]          # starting values
+    rms_before_px: float                # sqrt(mean(dy^2+dz^2)), first iter
+    rms_after_px: float                 # same, after the final iteration
+    rms_frames_after: float             # RMS omega-frame residual after
+    n_spots_used: int                   # spots with a measured centroid
+    per_iter_rms_px: list = field(default_factory=list)
+
+    @property
+    def delta(self) -> Dict[str, float]:
+        return {k: self.calibrated[k] - self.original[k]
+                for k in self.calibrated}
+
+
+def measure_patch_offsets(
+    measured: torch.Tensor,
+    *,
+    signal_threshold_frac: float = 0.2,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Per-spot measured centroid offsets from a cached patch tensor.
+
+    Parameters
+    ----------
+    measured : (S, Sigma, F, P, P) tensor
+        The cached measured patches (``assemble_grain_patch_data`` layout).
+    signal_threshold_frac : float
+        A (spot, scan) cell counts as signal when its summed intensity
+        exceeds this fraction of the median non-zero cell total.
+
+    Returns
+    -------
+    (dy, dz, df, ok) : per-spot median centroid offsets (px, px, frames)
+        relative to the patch centre, NaN where no scan had signal;
+        ``ok`` = finite mask.
+    """
+    meas = measured.detach().to(torch.float64).cpu()
+    S, _Sigma, F, P, _P2 = meas.shape
+    c = P // 2
+    ps = meas.sum(dim=2)                                  # (S, Sigma, P, P)
+    tot = ps.sum(dim=(-1, -2))                            # (S, Sigma)
+    nz = tot[tot > 0]
+    thr = float(nz.median()) * signal_threshold_frac if nz.numel() else 0.0
+    yy, zz = torch.meshgrid(
+        torch.arange(P, dtype=torch.float64),
+        torch.arange(P, dtype=torch.float64), indexing="ij",
+    )
+    cy = ((ps * yy).sum(dim=(-1, -2)) / (tot + 1e-12)).numpy() - c
+    cz = ((ps * zz).sum(dim=(-1, -2)) / (tot + 1e-12)).numpy() - c
+    fidx = torch.arange(F, dtype=torch.float64)
+    pf = meas.sum(dim=(-1, -2))                           # (S, Sigma, F)
+    cf = ((pf * fidx).sum(dim=-1) / (pf.sum(dim=-1) + 1e-12)).numpy() - F // 2
+    has = (tot.numpy() > thr)
+
+    dy = np.full(S, np.nan)
+    dz = np.full(S, np.nan)
+    df = np.full(S, np.nan)
+    for s in range(S):
+        g = np.nonzero(has[s])[0]
+        if g.size >= 1:
+            dy[s] = np.median(cy[s, g])
+            dz[s] = np.median(cz[s, g])
+            df[s] = np.median(cf[s, g])
+    ok = np.isfinite(dy)
+    return dy, dz, df, ok
+
+
+def _anchors_with_model(ds, model: HEDMForwardModel, dtype, device):
+    from . import io as _io
+    d2 = copy.copy(ds)
+    d2.model = model
+    ay, az, af, _valid, _obs, S = _io._forward_anchors(d2, dtype, device)
+    return (ay.detach().cpu().numpy(), az.detach().cpu().numpy(),
+            af.detach().cpu().numpy(), S)
+
+
+def calibrate_raw_frame_geometry(
+    ds,
+    measured: torch.Tensor,
+    *,
+    model_factory: Callable[..., HEDMForwardModel],
+    anchor_model: Optional[HEDMForwardModel] = None,
+    params_to_fit: Sequence[str] = tuple(DEFAULT_BUMPS),
+    bumps: Optional[Dict[str, float]] = None,
+    n_iters: int = 3,
+    frame_weight: float = 4.0,
+    damping: float = 1e-2,
+    max_step_bumps: float = 40.0,
+    mad_reject: float = 6.0,
+    signal_threshold_frac: float = 0.2,
+    dtype: torch.dtype = torch.float64,
+    device: torch.device | str = "cpu",
+    verbose: bool = True,
+) -> RawFrameCalibration:
+    """Fit (y_BC, z_BC, tx, ty, tz, Lsd, ω₀) to measured patch centroids.
+
+    Parameters
+    ----------
+    ds : PFGrainDataset
+        The grain whose cached patches anchor the calibration.
+    measured : (S, Sigma, F, P, P) tensor
+        Cached measured patches, in the anchor layout of ``anchor_model``.
+    model_factory : callable
+        ``model_factory(geom_mod: dict | None = None, omega_start:
+        float | None = None) -> HEDMForwardModel``. ``geom_mod`` maps
+        parameter names (see :data:`DEFAULT_BUMPS`, minus ``om0`` which is
+        passed as ``omega_start``) to absolute values. Must build the model
+        the CALIBRATION should converge (e.g. distortion ON for raw-frame
+        work); :func:`layer_model_factory` builds one from a layer dir.
+    anchor_model : optional
+        The model the cache was ASSEMBLED with (defaults to ``ds.model``).
+        Absolute measured positions are ``anchors(anchor_model) + offsets``.
+
+    Notes
+    -----
+    Derivatives are empirical: each parameter is bumped by ``bumps[name]``
+    in a rebuilt model and the anchor shift is the Jacobian column — no
+    convention (flips, ideal↔raw, ω origin) can silently disagree with
+    the model. The LSQ is Levenberg-damped and MAD-rejecting; steps are
+    clipped to ``max_step_bumps`` bump units.
+    """
+    bumps = dict(DEFAULT_BUMPS if bumps is None else bumps)
+    names = [n for n in params_to_fit if n in bumps]
+    if not names:
+        raise ValueError(f"params_to_fit {params_to_fit!r} matches no known "
+                         f"parameter (choose from {sorted(DEFAULT_BUMPS)})")
+
+    base = anchor_model if anchor_model is not None else ds.model
+    ay0, az0, af0, S = _anchors_with_model(ds, base, dtype, device)
+
+    dy, dz, df, ok = measure_patch_offsets(
+        measured, signal_threshold_frac=signal_threshold_frac)
+    if int(ok.sum()) < len(names):
+        raise ValueError(
+            f"only {int(ok.sum())} spots have a measured centroid — fewer "
+            f"than the {len(names)} parameters being fit."
+        )
+    abs_y = ay0 + dy
+    abs_z = az0 + dz
+    abs_f = af0 + df
+
+    def _getf(model, name, fallback=0.0):
+        v = getattr(model, name, fallback)
+        try:
+            v = v.detach().cpu().item()
+        except AttributeError:
+            pass
+        if isinstance(v, (list, tuple)):
+            v = v[0]
+        return float(v)
+
+    fit0 = model_factory()
+    state: Dict[str, float] = {}
+    for n in names:
+        if n == "om0":
+            state[n] = _getf(fit0, "omega_start")
+        else:
+            state[n] = _getf(fit0, n)
+    original = dict(state)
+
+    def model_at(st: Dict[str, float]) -> HEDMForwardModel:
+        gm = {k: v for k, v in st.items() if k != "om0"}
+        om0 = st.get("om0")
+        return model_factory(geom_mod=gm, omega_start=om0)
+
+    rms_before = None
+    per_iter = []
+    for it in range(n_iters):
+        m_it = model_at(state)
+        ayi, azi, afi, _ = _anchors_with_model(ds, m_it, dtype, device)
+        ry = (abs_y - ayi)[ok]
+        rz = (abs_z - azi)[ok]
+        rf = (abs_f - afi)[ok]
+        rms = math.sqrt(np.nanmean(ry ** 2 + rz ** 2))
+        per_iter.append(rms)
+        if rms_before is None:
+            rms_before = rms
+        if verbose:
+            print(f"[mini-calib iter {it}] RMS(dy,dz)={rms:.2f}px  "
+                  f"RMS(df)={math.sqrt(np.nanmean(rf ** 2)):.2f}fr",
+                  flush=True)
+        Jy, Jz, Jf = [], [], []
+        for n in names:
+            st2 = dict(state)
+            st2[n] = state[n] + bumps[n]
+            ayp, azp, afp, _ = _anchors_with_model(
+                ds, model_at(st2), dtype, device)
+            Jy.append((ayp - ayi)[ok])
+            Jz.append((azp - azi)[ok])
+            Jf.append((afp - afi)[ok])
+        Jy = np.stack(Jy, 1)
+        Jz = np.stack(Jz, 1)
+        Jf = np.stack(Jf, 1)
+        A = np.vstack([Jy, Jz, frame_weight * Jf])
+        b = np.concatenate([ry, rz, frame_weight * rf])
+        good = np.isfinite(b) & np.all(np.isfinite(A), axis=1)
+        med = np.median(b[good])
+        mad = np.median(np.abs(b[good] - med)) + 1e-9
+        good &= np.abs(b - med) < mad_reject * 1.4826 * mad
+        Ag, bg = A[good], b[good]
+        ata = Ag.T @ Ag
+        ata += damping * np.eye(len(names)) * np.trace(ata) / len(names)
+        x = np.linalg.solve(ata, Ag.T @ bg)          # step in bump units
+        x = np.clip(x, -max_step_bumps, max_step_bumps)
+        for k, n in enumerate(names):
+            state[n] += x[k] * bumps[n]
+        if verbose:
+            print("   step: " + " ".join(
+                f"d{n}={x[k] * bumps[n]:+.4f}" for k, n in enumerate(names)),
+                flush=True)
+
+    m_fin = model_at(state)
+    ayf, azf, aff, _ = _anchors_with_model(ds, m_fin, dtype, device)
+    ryf = (abs_y - ayf)[ok]
+    rzf = (abs_z - azf)[ok]
+    rff = (abs_f - aff)[ok]
+    rms_after = math.sqrt(np.nanmean(ryf ** 2 + rzf ** 2))
+    per_iter.append(rms_after)
+    return RawFrameCalibration(
+        calibrated=state,
+        original=original,
+        rms_before_px=float(rms_before),
+        rms_after_px=float(rms_after),
+        rms_frames_after=float(math.sqrt(np.nanmean(rff ** 2))),
+        n_spots_used=int(ok.sum()),
+        per_iter_rms_px=per_iter,
+    )
+
+
+def layer_model_factory(
+    layer_dir: str | Path,
+    *,
+    ring_numbers: Sequence[int],
+    n_pixels_y: int,
+    n_pixels_z: int,
+    n_frames: Optional[int] = None,
+    omega_step: Optional[float] = None,
+    apply_distortion: bool = True,
+    max_two_theta_deg: Optional[float] = None,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float64,
+) -> Callable[..., HEDMForwardModel]:
+    """Build a ``model_factory`` for :func:`calibrate_raw_frame_geometry`
+    from a MIDAS layer directory (paramstest/hkls/positions).
+
+    The returned callable accepts ``geom_mod`` (dict of absolute geometry
+    values: y_BC, z_BC, tx, ty, tz, Lsd) and ``omega_start`` overrides.
+    """
+    from midas_fit_grain.driver import _cartesian_B_matrix, _read_hkls_csv
+
+    from .io import (
+        _first_float, geometry_from_paramstest, parse_paramstest,
+        scan_config_from_positions,
+    )
+
+    layer = Path(layer_dir)
+    params = parse_paramstest(layer / "paramstest.txt")
+    lat = tuple(float(x) for x in params["LatticeParameter"][0][:6])
+    beam_size = _first_float(params, "BeamSize")
+
+    def factory(geom_mod: Optional[Dict[str, float]] = None,
+                omega_start: Optional[float] = None) -> HEDMForwardModel:
+        geom = geometry_from_paramstest(
+            params, n_pixels_y=n_pixels_y, n_pixels_z=n_pixels_z,
+            n_frames=n_frames, omega_step=omega_step,
+            apply_distortion=apply_distortion,
+        )
+        if geom_mod:
+            for k, v in geom_mod.items():
+                setattr(geom, k, float(v))
+        if omega_start is not None:
+            geom.omega_start = float(omega_start)
+        mtth = max_two_theta_deg
+        if mtth is None:
+            if "MaxRingRad" in params:
+                mrr = _first_float(params, "MaxRingRad")
+                lsd = geom.Lsd[0] if isinstance(geom.Lsd, list) else geom.Lsd
+                mtth = 2.0 * math.degrees(math.atan(mrr / lsd))
+            else:
+                mtth = 180.0
+        hkls_int, thetas_deg, _ring_nr = _read_hkls_csv(
+            layer / "hkls.csv", [int(r) for r in ring_numbers], mtth)
+        B = _cartesian_B_matrix(lat)
+        cart = (B @ hkls_int.astype(np.float64).T).T
+        scan_cfg = scan_config_from_positions(
+            layer / "positions.csv", beam_size, dtype=dtype)
+        model = HEDMForwardModel(
+            torch.from_numpy(cart),
+            torch.from_numpy(np.asarray(thetas_deg) * math.pi / 180.0),
+            geom,
+            hkls_int=torch.from_numpy(hkls_int.astype(np.float64)),
+            scan_config=scan_cfg,
+            device=device,
+        )
+        # Raw-frame convention: tilts are part of the raw prediction.
+        model.apply_tilts = True
+        return model
+
+    return factory