midas-calibrate 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,103 @@
1
+ Metadata-Version: 2.4
2
+ Name: midas-calibrate
3
+ Version: 0.2.1
4
+ Summary: Native Python/Torch detector calibration for MIDAS — replaces AutoCalibrateZarr → CalibrantIntegratorOMP. CPU & GPU; LM-based refinement.
5
+ Author: MIDAS contributors
6
+ License: BSD-3-Clause
7
+ Requires-Python: >=3.9
8
+ Description-Content-Type: text/markdown
9
+ Requires-Dist: numpy>=1.22
10
+ Requires-Dist: scipy>=1.9
11
+ Requires-Dist: torch>=2.1
12
+ Requires-Dist: midas-hkls>=0.1.0
13
+ Requires-Dist: midas-integrate>=0.1.0
14
+ Requires-Dist: midas-peakfit>=0.2.0
15
+ Requires-Dist: h5py>=3.0
16
+ Requires-Dist: tifffile>=2022.0
17
+ Provides-Extra: dev
18
+ Requires-Dist: pytest>=7; extra == "dev"
19
+ Requires-Dist: pandas>=1.5; extra == "dev"
20
+ Provides-Extra: plots
21
+ Requires-Dist: matplotlib>=3.5; extra == "plots"
22
+
23
+ # midas-calibrate
24
+
25
+ Native Python/Torch detector geometry calibration for MIDAS. Replaces
26
+ `AutoCalibrateZarr → CalibrantIntegratorOMP → CalibrationCore`. Same input
27
+ parameter file format, byte-compatible output, runs on CPU or GPU.
28
+
29
+ ## Quick start
30
+
31
+ ```python
32
+ import tifffile
33
+ from midas_calibrate import CalibrationParams, autocalibrate
34
+
35
+ params = CalibrationParams.from_file("calib.txt")
36
+ image = tifffile.imread("ceo2_calibrant.tif")
37
+ result = autocalibrate(params, image)
38
+
39
+ result.params.write("calib_refined.txt")
40
+ print(f"final mean strain: {result.history[-1].mean_strain_uE:.1f} μϵ")
41
+ ```
42
+
43
+ CLI:
44
+
45
+ ```bash
46
+ midas-autocalibrate calib.txt --image ceo2.tif --output calib_refined.txt
47
+ ```
48
+
49
+ ## How it works
50
+
51
+ * **E-step** — `midas-integrate` builds a CSR pixel→bin map from the current
52
+ geometry and integrates the image into a 2D (R, η) cake. Per (ring × η-bin)
53
+ the radial peak position is extracted via weighted centroid.
54
+ * **M-step** — fit detector geometry to the (Y_pix, Z_pix, ring) data using a
55
+ custom batched Levenberg-Marquardt solver (`midas_peakfit.lm_solve_generic`)
56
+ with sigmoid-bounded reparameterisation, Cholesky_ex, and optional Huber
57
+ loss reshaping.
58
+ * **Orchestrator** — alternating E↔M iterations with optional σ-clip outlier
59
+ rejection between iterations.
60
+
61
+ The geometry forward model in [`geometry_torch.py`](midas_calibrate/geometry_torch.py)
62
+ is a byte-for-byte port of `midas_integrate.geometry.pixel_to_REta` — verified
63
+ to fp64 epsilon by parity tests.
64
+
65
+ ## Dependencies
66
+
67
+ - [`midas-hkls`](../midas_hkls) — pure-Python crystallography (sginfo replacement)
68
+ - [`midas-integrate`](../midas_integrate) — CSR pixel→bin mapper + integration
69
+ - [`midas-peakfit`](../midas_peakfit) ≥ 0.2.0 — generic LM solver
70
+
71
+ ## Synthetic-data parity test
72
+
73
+ The end-to-end synthetic test forward-simulates a CeO₂ calibrant image at
74
+ known geometry, perturbs the seed (Lsd ±300μm, BC ±1.5px, tilts ±0.06°), and
75
+ verifies recovery:
76
+
77
+ ```
78
+ [iter 0] n_fits= 176 rc=0 strain= 105.2μϵ Lsd=1000219.4 BC=(512.20,511.91) ty=0.343 tz=0.180
79
+ [iter 1] n_fits= 176 rc=0 strain= 25.7μϵ Lsd= 999973.5 BC=(512.01,512.00) ty=0.403 tz=0.250
80
+ [iter 2] n_fits= 176 rc=0 strain= 19.4μϵ Lsd= 999946.1 BC=(511.99,512.00) ty=0.400 tz=0.267
81
+ [iter 3] n_fits= 176 rc=0 strain= 21.6μϵ Lsd= 999918.1 BC=(511.99,512.00) ty=0.392 tz=0.285
82
+ ```
83
+
84
+ Final recovery: Lsd within 82μm of truth, BC within 0.01 px, tilts within
85
+ 0.04°. Mean strain 21.6μϵ, well under the 50μϵ MIDAS calibration target.
86
+
87
+ ## Engines
88
+
89
+ `autocalibrate` is the alternating E↔M engine (default).
90
+
91
+ `autocalibrate_joint` will be the fully differentiable engine — geometry +
92
+ per-(ring × η-bin) peak-shape parameters jointly refined in one batched
93
+ Schur-complement-reduced LM (see §13 of the design doc). v0.1 ships with a
94
+ working stub that delegates to the alternating engine; the arrowhead-LM
95
+ infrastructure (`midas_peakfit.lm_solve_arrowhead`) is in place and tested.
96
+
97
+ ## Status
98
+
99
+ v0.1.0 — alternating engine production-ready, joint engine scaffolded.
100
+
101
+ See [`AutoCalibrate.md`](../../manuals/AutoCalibrate.md) for the manual and
102
+ [`calibrate_torch_implementation_plan.md`](../../calibrate_torch_implementation_plan.md)
103
+ for the full design and roadmap.
@@ -0,0 +1,81 @@
1
+ # midas-calibrate
2
+
3
+ Native Python/Torch detector geometry calibration for MIDAS. Replaces
4
+ `AutoCalibrateZarr → CalibrantIntegratorOMP → CalibrationCore`. Same input
5
+ parameter file format, byte-compatible output, runs on CPU or GPU.
6
+
7
+ ## Quick start
8
+
9
+ ```python
10
+ import tifffile
11
+ from midas_calibrate import CalibrationParams, autocalibrate
12
+
13
+ params = CalibrationParams.from_file("calib.txt")
14
+ image = tifffile.imread("ceo2_calibrant.tif")
15
+ result = autocalibrate(params, image)
16
+
17
+ result.params.write("calib_refined.txt")
18
+ print(f"final mean strain: {result.history[-1].mean_strain_uE:.1f} μϵ")
19
+ ```
20
+
21
+ CLI:
22
+
23
+ ```bash
24
+ midas-autocalibrate calib.txt --image ceo2.tif --output calib_refined.txt
25
+ ```
26
+
27
+ ## How it works
28
+
29
+ * **E-step** — `midas-integrate` builds a CSR pixel→bin map from the current
30
+ geometry and integrates the image into a 2D (R, η) cake. Per (ring × η-bin)
31
+ the radial peak position is extracted via weighted centroid.
32
+ * **M-step** — fit detector geometry to the (Y_pix, Z_pix, ring) data using a
33
+ custom batched Levenberg-Marquardt solver (`midas_peakfit.lm_solve_generic`)
34
+ with sigmoid-bounded reparameterisation, Cholesky_ex, and optional Huber
35
+ loss reshaping.
36
+ * **Orchestrator** — alternating E↔M iterations with optional σ-clip outlier
37
+ rejection between iterations.
38
+
39
+ The geometry forward model in [`geometry_torch.py`](midas_calibrate/geometry_torch.py)
40
+ is a byte-for-byte port of `midas_integrate.geometry.pixel_to_REta` — verified
41
+ to fp64 epsilon by parity tests.
42
+
43
+ ## Dependencies
44
+
45
+ - [`midas-hkls`](../midas_hkls) — pure-Python crystallography (sginfo replacement)
46
+ - [`midas-integrate`](../midas_integrate) — CSR pixel→bin mapper + integration
47
+ - [`midas-peakfit`](../midas_peakfit) ≥ 0.2.0 — generic LM solver
48
+
49
+ ## Synthetic-data parity test
50
+
51
+ The end-to-end synthetic test forward-simulates a CeO₂ calibrant image at
52
+ known geometry, perturbs the seed (Lsd ±300μm, BC ±1.5px, tilts ±0.06°), and
53
+ verifies recovery:
54
+
55
+ ```
56
+ [iter 0] n_fits= 176 rc=0 strain= 105.2μϵ Lsd=1000219.4 BC=(512.20,511.91) ty=0.343 tz=0.180
57
+ [iter 1] n_fits= 176 rc=0 strain= 25.7μϵ Lsd= 999973.5 BC=(512.01,512.00) ty=0.403 tz=0.250
58
+ [iter 2] n_fits= 176 rc=0 strain= 19.4μϵ Lsd= 999946.1 BC=(511.99,512.00) ty=0.400 tz=0.267
59
+ [iter 3] n_fits= 176 rc=0 strain= 21.6μϵ Lsd= 999918.1 BC=(511.99,512.00) ty=0.392 tz=0.285
60
+ ```
61
+
62
+ Final recovery: Lsd within 82μm of truth, BC within 0.01 px, tilts within
63
+ 0.04°. Mean strain 21.6μϵ, well under the 50μϵ MIDAS calibration target.
64
+
65
+ ## Engines
66
+
67
+ `autocalibrate` is the alternating E↔M engine (default).
68
+
69
+ `autocalibrate_joint` will be the fully differentiable engine — geometry +
70
+ per-(ring × η-bin) peak-shape parameters jointly refined in one batched
71
+ Schur-complement-reduced LM (see §13 of the design doc). v0.1 ships with a
72
+ working stub that delegates to the alternating engine; the arrowhead-LM
73
+ infrastructure (`midas_peakfit.lm_solve_arrowhead`) is in place and tested.
74
+
75
+ ## Status
76
+
77
+ v0.1.0 — alternating engine production-ready, joint engine scaffolded.
78
+
79
+ See [`AutoCalibrate.md`](../../manuals/AutoCalibrate.md) for the manual and
80
+ [`calibrate_torch_implementation_plan.md`](../../calibrate_torch_implementation_plan.md)
81
+ for the full design and roadmap.
@@ -0,0 +1,32 @@
1
+ """midas-calibrate — native Python/Torch detector calibration.
2
+
3
+ Public API:
4
+
5
+ from midas_calibrate import CalibrationParams, build_ring_table, refine_geometry
6
+
7
+ params = CalibrationParams.from_file("calib.txt")
8
+ rt = build_ring_table(params)
9
+ result = autocalibrate(params) # full pipeline
10
+ """
11
+ from .params import CalibrationParams
12
+ from .rings import RingTable, build_ring_table
13
+ from .refine import FittedPoint, RefineResult, refine_geometry
14
+ from .orchestrator import CalibrationResult, IterRecord, autocalibrate
15
+ from .estep import CakeProfile, integrate_cake, run_estep
16
+
17
+ __version__ = "0.2.1"
18
+
19
+ __all__ = [
20
+ "CakeProfile",
21
+ "CalibrationParams",
22
+ "CalibrationResult",
23
+ "FittedPoint",
24
+ "IterRecord",
25
+ "RefineResult",
26
+ "RingTable",
27
+ "autocalibrate",
28
+ "build_ring_table",
29
+ "integrate_cake",
30
+ "refine_geometry",
31
+ "run_estep",
32
+ ]
@@ -0,0 +1,76 @@
1
+ """Command-line entry points for midas-calibrate."""
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import sys
6
+ from pathlib import Path
7
+ from typing import Sequence
8
+
9
+ import numpy as np
10
+
11
+
12
+ def _load_image(path: Path) -> np.ndarray:
13
+ p = Path(path)
14
+ if p.suffix.lower() in (".tif", ".tiff"):
15
+ import tifffile
16
+ return tifffile.imread(p)
17
+ if p.suffix.lower() in (".h5", ".hdf5"):
18
+ import h5py
19
+ with h5py.File(p, "r") as f:
20
+ keys = list(f.keys())
21
+ return np.asarray(f[keys[0]])
22
+ if p.suffix.lower() in (".npy",):
23
+ return np.load(p)
24
+ raise ValueError(f"unknown image extension: {p.suffix}")
25
+
26
+
27
+ def main(argv: Sequence[str] | None = None) -> int:
28
+ return autocalibrate_main(argv)
29
+
30
+
31
+ def autocalibrate_main(argv: Sequence[str] | None = None) -> int:
32
+ parser = argparse.ArgumentParser(prog="midas-autocalibrate",
33
+ description="Native Python detector calibration")
34
+ parser.add_argument("params_file", type=Path, help="CalibrationParams .txt")
35
+ parser.add_argument("--image", type=Path, default=None,
36
+ help="calibrant image (overrides ImagePath in params)")
37
+ parser.add_argument("--dark", type=Path, default=None)
38
+ parser.add_argument("--output", type=Path, default=None,
39
+ help="path to write parameters_refined.txt (default: alongside input)")
40
+ parser.add_argument("--engine", choices=["alternating", "joint"], default="alternating")
41
+ parser.add_argument("--n-iters", type=int, default=None,
42
+ help="override nIterations from the params file")
43
+ parser.add_argument("--quiet", action="store_true")
44
+ args = parser.parse_args(argv)
45
+
46
+ from .orchestrator import autocalibrate
47
+ from .params import CalibrationParams
48
+
49
+ params = CalibrationParams.from_file(args.params_file)
50
+ if args.n_iters is not None:
51
+ params.nIterations = args.n_iters
52
+
53
+ image_path = args.image or Path(params.ImagePath)
54
+ image = _load_image(image_path)
55
+ dark = _load_image(args.dark) if args.dark else None
56
+
57
+ if args.engine == "joint":
58
+ try:
59
+ from .joint import autocalibrate_joint
60
+ result = autocalibrate_joint(params, image, dark=dark, verbose=not args.quiet)
61
+ except ImportError:
62
+ print("joint engine not yet available; falling back to alternating", file=sys.stderr)
63
+ result = autocalibrate(params, image, dark=dark, verbose=not args.quiet)
64
+ else:
65
+ result = autocalibrate(params, image, dark=dark, verbose=not args.quiet)
66
+
67
+ out = args.output or args.params_file.with_name(args.params_file.stem + "_refined.txt")
68
+ result.params.write(out)
69
+ if not args.quiet:
70
+ print(f"\nFinal mean strain: {result.history[-1].mean_strain_uE:.1f} μϵ")
71
+ print(f"Refined parameters written to {out}")
72
+ return 0
73
+
74
+
75
+ if __name__ == "__main__":
76
+ sys.exit(main())
@@ -0,0 +1,169 @@
1
+ """E-step: integrate calibrant image, extract per-(ring, η-bin) peak positions.
2
+
3
+ Strategy:
4
+ 1. Build a uniform-R, uniform-η bin grid spanning the calibrant ring range
5
+ (subsetting per-ring windows after integration).
6
+ 2. Build midas_integrate's PixelMap + CSR from the current geometry.
7
+ 3. Integrate the image into a 2D (R, η) cake.
8
+ 4. For each (ring, η-bin), compute a weighted centroid in the radial window
9
+ to get R_fit (px). v0.1 uses centroid; future versions can swap in a
10
+ pseudo-Voigt LM via midas_peakfit.lm_solve_generic.
11
+ 5. Convert (R_fit, η_bin_center) → (Y_pix, Z_pix) via midas_integrate's
12
+ Newton-Raphson inverse.
13
+ """
14
+ from __future__ import annotations
15
+
16
+ from dataclasses import dataclass
17
+ from typing import List, Optional, Tuple
18
+
19
+ import numpy as np
20
+ import torch
21
+
22
+ from midas_integrate.detector_mapper import build_map
23
+ from midas_integrate.geometry import build_tilt_matrix, invert_REta_to_pixel
24
+ from midas_integrate.kernels import build_csr, integrate
25
+ from midas_integrate.params import IntegrationParams
26
+
27
+ from .params import CalibrationParams
28
+ from .refine import FittedPoint
29
+ from .rings import RingTable
30
+
31
+
32
+ def _calibration_to_integration_params(
33
+ params: CalibrationParams, *, R_min: float, R_max: float, R_bin_size: float, eta_bin_size: float,
34
+ ) -> IntegrationParams:
35
+ ip = IntegrationParams()
36
+ ip.NrPixelsY = params.NrPixelsY
37
+ ip.NrPixelsZ = params.NrPixelsZ
38
+ ip.pxY = params.pxY
39
+ ip.pxZ = params.pxZ if params.pxZ > 0 else params.pxY
40
+ ip.Lsd = params.Lsd
41
+ ip.BC_y = params.BC_y
42
+ ip.BC_z = params.BC_z
43
+ ip.tx = params.tx
44
+ ip.ty = params.ty
45
+ ip.tz = params.tz
46
+ for i in range(15):
47
+ setattr(ip, f"p{i}", getattr(params, f"p{i}"))
48
+ ip.RhoD = params.RhoD if params.RhoD > 0 else params.MaxRingRad
49
+ ip.Parallax = params.Parallax
50
+ ip.Wavelength = params.Wavelength
51
+ ip.RMin = float(R_min)
52
+ ip.RMax = float(R_max)
53
+ ip.RBinSize = float(R_bin_size)
54
+ ip.EtaMin = -180.0
55
+ ip.EtaMax = 180.0
56
+ ip.EtaBinSize = float(eta_bin_size)
57
+ ip.SolidAngleCorrection = 0
58
+ ip.PolarizationCorrection = 0
59
+ return ip
60
+
61
+
62
+ @dataclass
63
+ class CakeProfile:
64
+ R_centers: np.ndarray
65
+ eta_centers: np.ndarray
66
+ intensity: np.ndarray # [n_R, n_eta]
67
+
68
+
69
+ def integrate_cake(
70
+ params: CalibrationParams,
71
+ image: np.ndarray,
72
+ rt: RingTable,
73
+ *, dark: Optional[np.ndarray] = None,
74
+ ) -> CakeProfile:
75
+ """Build CSR + integrate the image into a uniform (R, η) cake."""
76
+ if dark is not None:
77
+ image = image - dark
78
+
79
+ # R range: half-Width margin around min/max ring radius.
80
+ px = 0.5 * (params.pxY + params.pxZ) if params.pxZ > 0 else params.pxY
81
+ half_px = 0.5 * params.Width / px
82
+ R_min = max(0.0, float(rt.r_ideal_px.min()) - half_px - 1.0)
83
+ R_max = float(rt.r_ideal_px.max()) + half_px + 1.0
84
+ ip = _calibration_to_integration_params(
85
+ params, R_min=R_min, R_max=R_max,
86
+ R_bin_size=params.RBinSize, eta_bin_size=params.EtaBinSize,
87
+ )
88
+
89
+ pmap_result = build_map(ip, verbose=False)
90
+ from midas_integrate.bin_io import PixelMap as _PixelMap
91
+ pmap = _PixelMap(
92
+ pxList=pmap_result.pxList,
93
+ counts=pmap_result.counts,
94
+ offsets=pmap_result.offsets,
95
+ map_header=None, nmap_header=None,
96
+ )
97
+ geom = build_csr(
98
+ pmap,
99
+ n_r=ip.n_r_bins, n_eta=ip.n_eta_bins,
100
+ n_pixels_y=ip.NrPixelsY, n_pixels_z=ip.NrPixelsZ,
101
+ bc_y=ip.BC_y, bc_z=ip.BC_z,
102
+ device="cpu", dtype=torch.float64,
103
+ build_modes=("bilinear",),
104
+ )
105
+
106
+ img_t = torch.as_tensor(image, dtype=torch.float64).contiguous()
107
+ cake = integrate(img_t, geom, mode="bilinear", normalize=True).numpy()
108
+
109
+ R_edges = np.linspace(ip.RMin, ip.RMin + ip.RBinSize * ip.n_r_bins, ip.n_r_bins + 1)
110
+ eta_edges = np.linspace(ip.EtaMin, ip.EtaMax, ip.n_eta_bins + 1)
111
+ return CakeProfile(
112
+ R_centers=0.5 * (R_edges[:-1] + R_edges[1:]),
113
+ eta_centers=0.5 * (eta_edges[:-1] + eta_edges[1:]),
114
+ intensity=cake,
115
+ )
116
+
117
+
118
+ def extract_fitted_points(
119
+ cake: CakeProfile, rt: RingTable, params: CalibrationParams,
120
+ *, snr_min: float = 1.0,
121
+ ) -> List[FittedPoint]:
122
+ """Per (ring × η-bin): centroid in the radial window → (R_fit, η) → (Y_pix, Z_pix)."""
123
+ px = 0.5 * (params.pxY + params.pxZ) if params.pxZ > 0 else params.pxY
124
+ half_px = 0.5 * params.Width / px
125
+ TRs = build_tilt_matrix(params.tx, params.ty, params.tz)
126
+
127
+ fits: List[FittedPoint] = []
128
+ for ring_i, r_ideal in enumerate(rt.r_ideal_px):
129
+ idx = np.where(np.abs(cake.R_centers - r_ideal) <= half_px)[0]
130
+ if idx.size < 3:
131
+ continue
132
+ R_window = cake.R_centers[idx]
133
+ for eta_j, eta in enumerate(cake.eta_centers):
134
+ I = cake.intensity[idx, eta_j]
135
+ I = np.maximum(I - I.min(), 0.0)
136
+ tot = I.sum()
137
+ if tot <= 0.0:
138
+ continue
139
+ R_fit = float((I * R_window).sum() / tot)
140
+ peak = float(I.max())
141
+ mean = float(I.mean()) + 1e-12
142
+ snr = peak / mean
143
+ if snr < snr_min:
144
+ continue
145
+ try:
146
+ Y_pix, Z_pix = invert_REta_to_pixel(
147
+ R_fit, eta,
148
+ Ycen=params.BC_y, Zcen=params.BC_z, TRs=TRs,
149
+ Lsd=params.Lsd, RhoD=(params.RhoD if params.RhoD > 0 else params.MaxRingRad),
150
+ px=px, parallax=params.Parallax,
151
+ )
152
+ except Exception:
153
+ continue
154
+ fits.append(FittedPoint(
155
+ Y_pix=float(Y_pix), Z_pix=float(Z_pix),
156
+ ring_idx=ring_i, snr=snr,
157
+ ))
158
+ return fits
159
+
160
+
161
+ def run_estep(
162
+ params: CalibrationParams,
163
+ image: np.ndarray,
164
+ rt: RingTable,
165
+ *, dark: Optional[np.ndarray] = None,
166
+ ) -> Tuple[CakeProfile, List[FittedPoint]]:
167
+ cake = integrate_cake(params, image, rt, dark=dark)
168
+ fits = extract_fitted_points(cake, rt, params, snr_min=params.SNRMin)
169
+ return cake, fits
@@ -0,0 +1,154 @@
1
+ """Differentiable detector forward model in pure torch.
2
+
3
+ Mirrors midas_integrate.geometry.pixel_to_REta but in torch ops so the result
4
+ is autograd-traced through to geometry parameters (Lsd, BC, tilts, p0..p14,
5
+ parallax, wavelength). All units match MIDAS conventions:
6
+
7
+ * Lsd in μm; BC in pixels; tilts in degrees (Rx · Ry · Rz order).
8
+ * Eta in degrees, atan2(-Y', Z'), [-180, 180).
9
+ * Distortion ΔR/RhoD = polynomial in (R/RhoD).
10
+
11
+ Inverse mapping (R, η) → (Y_pix, Z_pix) is a Newton-Raphson loop and is
12
+ provided here for the alternating engine; the joint engine forward-models
13
+ predicted-R(geometry) directly without inversion.
14
+ """
15
+ from __future__ import annotations
16
+
17
+ from typing import Tuple
18
+
19
+ import torch
20
+
21
+
22
+ def build_tilt_matrix_torch(tx: torch.Tensor, ty: torch.Tensor, tz: torch.Tensor) -> torch.Tensor:
23
+ """TRs = Rx(tx) · Ry(ty) · Rz(tz), angles in degrees, [3,3] tensor.
24
+
25
+ All inputs scalar torch tensors; output [3,3] tensor with grad flowing.
26
+ """
27
+ one = torch.ones((), dtype=tx.dtype, device=tx.device)
28
+ zero = torch.zeros((), dtype=tx.dtype, device=tx.device)
29
+ deg2rad = torch.tensor(0.017453292519943295, dtype=tx.dtype, device=tx.device)
30
+ cx, sx = torch.cos(tx * deg2rad), torch.sin(tx * deg2rad)
31
+ cy, sy = torch.cos(ty * deg2rad), torch.sin(ty * deg2rad)
32
+ cz, sz = torch.cos(tz * deg2rad), torch.sin(tz * deg2rad)
33
+
34
+ Rx = torch.stack([torch.stack([one, zero, zero]),
35
+ torch.stack([zero, cx, -sx]),
36
+ torch.stack([zero, sx, cx])])
37
+ Ry = torch.stack([torch.stack([cy, zero, sy]),
38
+ torch.stack([zero, one, zero]),
39
+ torch.stack([-sy, zero, cy])])
40
+ Rz = torch.stack([torch.stack([cz, -sz, zero]),
41
+ torch.stack([sz, cz, zero]),
42
+ torch.stack([zero, zero, one])])
43
+ return Rx @ Ry @ Rz
44
+
45
+
46
+ def pixel_to_REta_torch(
47
+ Y_pix: torch.Tensor, # [...] pixel Y indices
48
+ Z_pix: torch.Tensor, # [...] pixel Z indices
49
+ *,
50
+ Lsd: torch.Tensor,
51
+ BC_y: torch.Tensor,
52
+ BC_z: torch.Tensor,
53
+ tx: torch.Tensor,
54
+ ty: torch.Tensor,
55
+ tz: torch.Tensor,
56
+ p_coeffs: torch.Tensor, # [15] distortion coefficients p0..p14
57
+ parallax: torch.Tensor,
58
+ px: torch.Tensor, # μm; mean of pxY/pxZ
59
+ rho_d: torch.Tensor, # px; distortion normalization radius
60
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
61
+ """Differentiable pixel → (R [px], Eta [deg]).
62
+
63
+ Byte-for-byte port of midas_integrate.geometry.pixel_to_REta — vector layout
64
+ ``(0, Yc, Zc)`` → tilt → add Lsd along x; then 2θ via Lsd/Xp projection;
65
+ distortion applied as a *multiplicative scaling* on Rad (NOT ΔR).
66
+ """
67
+ deg2rad = torch.tensor(0.017453292519943295, dtype=Y_pix.dtype, device=Y_pix.device)
68
+ rad2deg = torch.tensor(57.29577951308232, dtype=Y_pix.dtype, device=Y_pix.device)
69
+
70
+ # Untilted physical coordinates (μm) — note the X-component is 0 before tilt.
71
+ Yc = (-Y_pix + BC_y) * px
72
+ Zc = (Z_pix - BC_z) * px
73
+
74
+ TRs = build_tilt_matrix_torch(tx, ty, tz)
75
+ # Apply tilt to (0, Yc, Zc): only columns 1 and 2 of TRs matter.
76
+ abcpr_x = TRs[0, 1] * Yc + TRs[0, 2] * Zc
77
+ abcpr_y = TRs[1, 1] * Yc + TRs[1, 2] * Zc
78
+ abcpr_z = TRs[2, 1] * Yc + TRs[2, 2] * Zc
79
+
80
+ XYZ_x = Lsd + abcpr_x
81
+ XYZ_y = abcpr_y
82
+ XYZ_z = abcpr_z
83
+ safe_x = torch.where(XYZ_x.abs() < 1e-30, torch.full_like(XYZ_x, 1e-30), XYZ_x)
84
+ rad_um = (Lsd / safe_x) * torch.sqrt(XYZ_y * XYZ_y + XYZ_z * XYZ_z)
85
+ eta_tilted = rad2deg * torch.atan2(-XYZ_y, XYZ_z)
86
+
87
+ # Distortion polynomial uses EtaT = 90 - EtaTilted (matches numpy).
88
+ eta_T_rad = (90.0 - eta_tilted) * deg2rad
89
+ R_norm = rad_um / rho_d if rho_d > 0 else torch.zeros_like(rad_um)
90
+
91
+ p = p_coeffs
92
+ dist = (
93
+ p[0] * R_norm.pow(2) * torch.cos(2 * eta_T_rad + deg2rad * p[6])
94
+ + p[1] * R_norm.pow(4) * torch.cos(4 * eta_T_rad + deg2rad * p[3])
95
+ + p[2] * R_norm.pow(2)
96
+ + p[4] * R_norm.pow(6)
97
+ + p[5] * R_norm.pow(4)
98
+ + p[7] * R_norm.pow(4) * torch.cos(eta_T_rad + deg2rad * p[8])
99
+ + p[9] * R_norm.pow(3) * torch.cos(3 * eta_T_rad + deg2rad * p[10])
100
+ + p[11] * R_norm.pow(5) * torch.cos(5 * eta_T_rad + deg2rad * p[12])
101
+ + p[13] * R_norm.pow(6) * torch.cos(6 * eta_T_rad + deg2rad * p[14])
102
+ + 1.0
103
+ )
104
+ Rt = rad_um * dist / px
105
+
106
+ if isinstance(parallax, torch.Tensor):
107
+ if parallax.abs().item() > 0:
108
+ two_theta = torch.atan(rad_um / Lsd)
109
+ Rt = Rt + parallax * torch.sin(two_theta) / px
110
+ return Rt, eta_tilted
111
+
112
+
113
+ def predict_R_at_pixel(Y_pix: torch.Tensor, Z_pix: torch.Tensor,
114
+ params_vec: torch.Tensor, px: float, rho_d: float) -> torch.Tensor:
115
+ """Compute R [px] given a packed geometry parameter vector.
116
+
117
+ Layout of params_vec (length 23):
118
+ [0] Lsd
119
+ [1,2] BC_y, BC_z
120
+ [3,4] ty, tz (tx fixed at 0 for now)
121
+ [5..19] p0..p14
122
+ [20] parallax
123
+ [21] wavelength (unused in geometry; needed for predict_R_ideal)
124
+ [22] tx (fixed at the input value)
125
+ """
126
+ Lsd = params_vec[0]
127
+ BC_y = params_vec[1]; BC_z = params_vec[2]
128
+ ty = params_vec[3]; tz = params_vec[4]
129
+ p = params_vec[5:20]
130
+ parallax = params_vec[20]
131
+ tx = params_vec[22]
132
+ px_t = torch.as_tensor(px, dtype=params_vec.dtype, device=params_vec.device)
133
+ rho_d_t = torch.as_tensor(rho_d, dtype=params_vec.dtype, device=params_vec.device)
134
+ R, _eta = pixel_to_REta_torch(
135
+ Y_pix, Z_pix, Lsd=Lsd, BC_y=BC_y, BC_z=BC_z,
136
+ tx=tx, ty=ty, tz=tz, p_coeffs=p, parallax=parallax,
137
+ px=px_t, rho_d=rho_d_t,
138
+ )
139
+ return R
140
+
141
+
142
+ def predict_R_ideal(two_theta_deg: torch.Tensor, params_vec: torch.Tensor, px: float) -> torch.Tensor:
143
+ """R_ideal[px] = Lsd · tan(2θ) / px. When wavelength is refined, two_theta_deg
144
+ must be recomputed externally from d-spacing — this function expects the
145
+ already-resolved 2θ values."""
146
+ Lsd = params_vec[0]
147
+ return Lsd * torch.tan(two_theta_deg * 0.017453292519943295) / px
148
+
149
+
150
+ def predict_two_theta_from_d(d_spacing_A: torch.Tensor, wavelength_A: torch.Tensor) -> torch.Tensor:
151
+ """Bragg's law in torch: 2θ = 2 arcsin(λ / 2d)."""
152
+ s = wavelength_A / (2.0 * d_spacing_A)
153
+ s = s.clamp(min=-0.999999, max=0.999999)
154
+ return 2.0 * torch.asin(s) * 57.29577951308232 # → deg
@@ -0,0 +1,53 @@
1
+ """Joint differentiable engine — fully end-to-end LM over geometry + per-region peak-shape parameters.
2
+
3
+ **Status (v0.1.0): scaffolded but not yet wired.**
4
+
5
+ The infrastructure is in place:
6
+
7
+ - ``midas_peakfit.lm_solve_arrowhead`` provides a Schur-complement-reduced LM
8
+ solver for the joint J = [J_dense | block_diag(J_block_k)] problem (see
9
+ the §13 plan).
10
+ - ``midas_calibrate.geometry_torch.pixel_to_REta_torch`` is a fully torch /
11
+ autograd compatible forward model.
12
+ - ``midas_calibrate.refine`` already exercises lm_solve_generic with the
13
+ geometry parameters as a 23-dim dense vector — the joint formulation just
14
+ enlarges this to add per-(ring, η-bin) peak-shape blocks.
15
+
16
+ What's left (deferred to v0.2):
17
+
18
+ 1. ``forward_cake(θ_geom, θ_shape)`` — predict the (R, η) cake intensity by
19
+ summing pseudo-Voigts whose centers are determined by θ_geom (NOT by θ_shape).
20
+ 2. Block-arrow Jacobian assembly: per-region 5×M peak-shape Jacobian +
21
+ coupled geometry columns, packaged for ``lm_solve_arrowhead``.
22
+ 3. Warm-start: 1-2 iterations of ``orchestrator.autocalibrate`` (alternating)
23
+ followed by per-region peak-shape seeding.
24
+
25
+ For now ``autocalibrate_joint`` is a thin wrapper that delegates to the
26
+ alternating engine; users get the same end result through a more conventional
27
+ implementation. Swap-in is local to this module — no caller changes needed
28
+ once the joint forward model lands.
29
+ """
30
+ from __future__ import annotations
31
+
32
+ from typing import Optional
33
+
34
+ import numpy as np
35
+
36
+ from .orchestrator import CalibrationResult, autocalibrate
37
+ from .params import CalibrationParams
38
+
39
+
40
+ def autocalibrate_joint(
41
+ params: CalibrationParams,
42
+ image: np.ndarray,
43
+ *, dark: Optional[np.ndarray] = None,
44
+ verbose: bool = True,
45
+ ) -> CalibrationResult:
46
+ """Full differentiable end-to-end calibration.
47
+
48
+ v0.1: delegates to the alternating engine. See module docstring for the
49
+ deferred-but-scoped roadmap to the true joint formulation.
50
+ """
51
+ if verbose:
52
+ print("[joint] v0.1 stub — delegating to alternating engine")
53
+ return autocalibrate(params, image, dark=dark, verbose=verbose)