Joint multi-measurement fitting, optimizer resume, and a runnable example

.gitignore

@@ -114,3 +114,8 @@ __pycache__/
 # Per-solver checkpoint directories written by run_rcwa.py / run_neviere.py.
 # The unsuffixed checkpoints/ baseline stays tracked; these are scratch.
 validation/*/results/checkpoints_*/
+
+# Optimizer checkpoints written by the joint measurement-fit example. The Ax
+# snapshot and the problem fingerprint both record absolute paths from the
+# machine that produced them, so they are scratch, not shareable artifacts.
+examples/optimizer/*/results/*/checkpoint/

CHANGELOG.md

@@ -2,6 +2,19 @@
 
 ## Unreleased
 
+- Added `grax_opt.optimize_to_joint_measurements` for fitting one parameter set jointly against several measured curves, each with its own measurement file and energy grid, with a configurable `joint_loss_reduction` (`mean`, `sum`, `pooled`, or `weighted`). Each trial evaluates every measurement in a single `BatchSimulationRunner` batch so trial-level `max_workers` parallelizes across measurements as well as energies.
+- Joint fits are not limited to curves that differ by grazing angle. `angle_mode`, `grazing_angle_deg`, `cff`, `diffraction_order` and `polarization` are run-level defaults that any individual `MeasurementSpec` may override, so one fit can span angles, angle modes, diffraction orders and polarizations. Numerical settings (`fourier_orders`, `solver`, `backend`, `max_workers`) stay run-level, since they describe how a curve is computed rather than what was measured.
+- `grax_opt` measurement fits now take `polarization`, accepting `s`/`p` or `TE`/`TM`. Both optimizer entrypoints previously had no such argument anywhere, so every fit silently ran the `run_simulation` default of `s` -- the same gap the theta-search workflow had. Default stays `s`.
+- `optimize_to_joint_measurements` accepts `solver` and `solver_options`, matching `optimize_to_measurements`, and records them in `best_result.json`. Both are part of the resume fingerprint, alongside each measurement's resolved conditions, so a resumed run cannot silently switch the physics it is fitting.
+- Added `examples/optimizer/optimizer_joint/`, the first runnable joint-fit workflow: it simulates a four-condition measurement set from a known grating, fits it, resumes the fit to extend the trial budget, and reports how close each parameter came to the value the data was generated from. It is also the first optimizer example with a `--solver` flag; the measurements are always generated with `rcwa`, so fitting them with `--solver neviere` recovers the same parameters to within a few parts in ten thousand.
+- Fixed an interrupted optimizer run over-running its trial budget on resume. The trial-record log is appended every trial but the Ax snapshot is only rewritten every `checkpoint_interval` trials, so an interruption left the log ahead of the snapshot; the recovered history then counted trials Ax was about to generate again, duplicating a row in `trial_history.csv` and running past `total_trials`. Records at or beyond the reconciled cursor are now discarded, since Ax is authoritative for what was issued. Only reachable with `checkpoint_interval > 1`, which is why the default of `1` hid it.
+- The single-measurement trial evaluation now raises on a missing batch result instead of reading uninitialized memory. It allocated its efficiency array with `np.empty` and filled by result index, so a case the runner never returned left whatever was in memory as that point's efficiency; the joint path already guarded against this.
+- An unreadable measurement file now raises when the resume fingerprint is built instead of hashing to a shared `"unreadable"` sentinel, which made two different unreadable files compare equal. Both entrypoints load their measurements before fingerprinting, so this is a guard against the ordering changing rather than a bug that could be hit today.
+- Added checkpoint and resume support to `optimize_to_measurements` and `optimize_to_joint_measurements` through the new `resume`, `checkpoint_dir`, and `checkpoint_interval` spec keys. The Ax client snapshot is persisted alongside optimizer run state and an append-only trial log, so a resumed run keeps its surrogate model and `total_trials` can be raised to extend a finished run. A problem fingerprint refuses to resume into a changed search space, naming the settings that differ.
+- Fixed joint multi-angle evaluation assigning simulated efficiencies to the wrong angle/energy slot when `BatchSimulationRunner` returned results in completion order rather than input order; results are now reassembled by `CaseExecutionResult.index`.
+- Extracted the shared Ax trial loop into `grax_opt.loop.run_ax_trial_loop`, so both optimizer entrypoints share candidate generation, best-so-far tracking, and early stopping. `early_stopping_min_relative_improvement` is now honored instead of being validated and ignored.
+- Optimizer artifacts are now written atomically, the best-fit plot reuses the winning trial's cached simulated curve instead of re-simulating it after every trial, and a penalized trial logs the failing case or exception instead of failing silently.
+
 - Every example that solves now accepts `--solver rcwa|neviere`, defaulting to `rcwa`. Solver-dependent example outputs are suffixed (`*_rcwa.*`, `*_neviere.*`) so the two runs sit side by side; geometry artifacts such as `*_profile.png` stay unsuffixed because they do not depend on the solver.
 - Added three examples covering the cases where the two solvers genuinely differ, rather than duplicating existing examples whose curves would be indistinguishable: `deep_grating_limits` (the modal solver stops at 8.4 wavelengths of depth where the differential method reaches 167), `continuous_vs_staircase` (continuous z-sampling is bit-identical across every `z_resolution_nm`, while a staircase run carries 2.7e-3 of discretization error at 2 nm), and `solver_runtime` (2.4x to 3.0x faster at production resolution, with the solvers still agreeing to ~1e-11).
 - Fixed `multilayer_theta_search` and `blazed_multilayer_memory_comparison` failing with `BrokenProcessPool` on macOS. Both use `max_workers`, and a spawned worker re-imports the script by path; without a `__main__` guard each worker re-ran the whole example and recursively spawned more.

docs/api/optimization.md

@@ -12,14 +12,37 @@ This page exposes the primary public optimizer API for end users.
 .. autofunction:: grax_opt.optimize_to_measurements
 ```
 
-## Result type
+## Joint measurement fitting
+
+Fits one parameter set against several measured curves at once. Each measurement
+keeps its own energy grid and its own conditions -- grazing angle, angle mode,
+diffraction order and polarization -- inheriting whichever it does not set from
+the run.
+
+```{eval-rst}
+.. autofunction:: grax_opt.optimize_to_joint_measurements
+
+.. autoclass:: grax_opt.MeasurementSpec
+
+.. autoclass:: grax_opt.JointMeasurement
+
+.. autoclass:: grax_opt.JointMeasurementFitConfig
+
+.. autofunction:: grax_opt.reduce_joint_losses
+```
+
+## Result types
 
 ```{eval-rst}
 .. autoclass:: grax_opt.OptimizationResult
+
+.. autoclass:: grax_opt.JointOptimizationResult
 ```
 
 ## See tutorials
 
 - [Optimizer setup guide](../tutorials/optimizer.md)
 - [Laminar Grating](../tutorials/optimizer-laminar-fit.md)
 - [Blazed Grating](../tutorials/optimizer-blazed-fit.md)
+- [Joint measurement fits](../tutorials/optimizer-joint-fit.md)
+- [Resume an optimizer run](../tutorials/optimizer-resume.md)

docs/developer/module-guide.md

@@ -41,6 +41,28 @@ not part of the first-class user documentation set for the core simulation
 package. Keep optimization-specific user material separate unless the core docs
 need to mention installation of the optional `opt` extra.
 
+- `config.py`: `ParameterBounds`, the shared bounds value type
+- `data.py`: measurement loading and interpolation onto an evaluation grid
+- `evaluation.py`: normalization of evaluation energies/angles for the
+  single-measurement optimizer
+- `objective.py`: trial evaluation, loss functions, joint multi-angle
+  evaluation, and the failure-penalty path
+- `optimize.py`: Ax imports and version sniffing, candidate batching, trial
+  records, plotting, and atomic artifact writes
+- `loop.py`: the Ax ask-and-tell loop shared by both optimizer entrypoints,
+  including best-so-far tracking and early stopping
+- `dynamic.py`: `MeasurementFitConfig` and `optimize_to_measurements`, the
+  single-measurement fit
+- `joint.py`: `JointMeasurementFitConfig` and `optimize_to_joint_measurements`,
+  the multi-angle fit
+- `checkpoint.py`: checkpoint layout, problem fingerprinting, and the
+  resume session used by both entrypoints
+
+Both optimizer entrypoints share `loop.py` and `checkpoint.py`; mode-specific
+work is limited to config validation, candidate evaluation, and artifact
+persistence. Add new optimizer modes the same way rather than duplicating the
+trial loop.
+
 ## Where to make changes
 
 Add new grating shapes in `gratings.py` when they need to participate in the

docs/tutorials/optimizer-joint-fit.md

@@ -0,0 +1,173 @@
+# Joint measurement fits
+
+`grax_opt.optimize_to_joint_measurements` fits **one** parameter set against
+**several** measured curves at once. Use it when a single geometry has to explain
+all of your measurements together, rather than fitting each curve separately and
+comparing the results afterwards.
+
+The curves do not have to differ only by grazing angle. Each measurement carries
+its own conditions — angle, angle mode, diffraction order, polarization — so one
+fit can span a whole measurement campaign.
+
+This is a separate entrypoint from
+{doc}`optimize_to_measurements <optimizer>`, which fits a single measurement.
+
+A runnable end-to-end workflow lives in
+`examples/optimizer/optimizer_joint/`:
+
+```bash
+./examples/optimizer/optimizer_joint/run_all.sh
+```
+
+It generates a four-condition measurement set from a known grating, fits it,
+resumes the fit to extend it, and reports how close each parameter came to the
+value the data was built from.
+
+## Basic setup
+
+```python
+from grax_opt import optimize_to_joint_measurements
+
+result = optimize_to_joint_measurements({
+    "build_grating": build_candidate_grating,
+    "parameter_bounds": {
+        "width_to_period_ratio": (0.45, 0.80),
+        "depth_nm": (4.5, 6.5),
+        "wall_angle_deg": (1.0, 40.0),
+    },
+    "output_dir": "results/joint_fit",
+    "grazing_angle_deg": 4.0,       # run-level default
+    "diffraction_order": 1,         # run-level default
+    "polarization": "s",            # run-level default
+    "measurements": [
+        {"measurement_path": "data/alpha1.dat", "grazing_angle_deg": 1.0},
+        {"measurement_path": "data/alpha2.dat", "grazing_angle_deg": 2.0},
+        {"measurement_path": "data/order2.dat", "diffraction_order": 2},
+        {"measurement_path": "data/cff_p.dat",
+         "angle_mode": "cff", "cff": 2.25, "polarization": "p"},
+    ],
+    "fourier_orders": 15,
+    "total_trials": 200,
+    "max_workers": "auto",
+})
+
+print(result.best_loss, result.per_measurement_best_losses)
+```
+
+Every measurement keeps its **own** energy grid. The grids do not have to match
+in range or in length.
+
+## Run-level defaults and per-measurement overrides
+
+The conditions below are set once on the run and inherited by every measurement.
+Any measurement may override any of them:
+
+| Key | Meaning |
+| --- | --- |
+| `angle_mode` | `"fixed"` or `"cff"` |
+| `grazing_angle_deg` | fixed angle, used when the resolved mode is `"fixed"` |
+| `cff` | fixed-focus constant, used when the resolved mode is `"cff"` |
+| `diffraction_order` | the order the curve was recorded in |
+| `polarization` | `s`/`p`, or the equivalent `TE`/`TM` |
+
+A measurement that sets none of them inherits all of them, so the common case —
+several angles, everything else shared — stays a one-key spec per curve.
+
+Numerical settings stay run-level and cannot be overridden per measurement:
+`fourier_orders`, `solver`, `solver_options`, `backend`, `max_workers`. They
+describe how the simulation is computed, not what was measured.
+
+## Measurement keys
+
+Each entry in `measurements` accepts the five condition keys above, plus:
+
+- `measurement_path` (required): the measured two-column dataset.
+- `evaluation_energies_ev`: energies to evaluate at. When omitted, the file's
+  own energy grid is used. When given, the measured curve is interpolated onto
+  it.
+- `measurement_efficiency`: measured efficiencies to use **directly** instead of
+  interpolating the file, as described under "Pre-prepared measurements" below.
+- `weight`: relative weight for the `"weighted"` reduction. Defaults to `1.0`.
+- `label`: identifier used in artifacts. Defaults to a description of whichever
+  conditions the measurement sets, for example `alpha2deg` or
+  `cff2p25_order1_p`, falling back to the measurement file stem when it sets
+  none. Labels must be unique, so give an explicit `label` when two measurements
+  would otherwise describe themselves the same way.
+
+## Choosing the solver
+
+`solver="rcwa"` (the default) or `solver="neviere"` selects the electromagnetic
+solver for every trial, with `solver_options` carrying the differential method's
+integration settings. See {doc}`choosing-a-solver`.
+
+The solver is part of the problem fingerprint, so it cannot be swapped on a
+resume — see {doc}`optimizer-resume`.
+
+## Combining the per-measurement losses
+
+Each measurement contributes its own mean squared error. `joint_loss_reduction`
+controls how those are combined into the single value the optimizer minimizes:
+
+| Reduction | Joint loss | Use when |
+| --- | --- | --- |
+| `"mean"` (default) | mean of the per-measurement MSEs | every curve should count equally |
+| `"sum"` | sum of the per-measurement MSEs | same ranking as `mean`, larger magnitude |
+| `"pooled"` | MSE over all points pooled together | grids have **different point counts** and every measured point should count equally |
+| `"weighted"` | weighted mean using each `weight` | some curves are more trustworthy than others |
+
+`"mean"` and `"pooled"` differ only when the measurements have unequal point
+counts. If you exclude absorption edges or downsample one curve more than
+another, prefer `"pooled"` — otherwise the curve with fewer points is weighted
+more heavily per measured point.
+
+A caution when the curves differ in magnitude: a weak order-2 curve contributes
+a much smaller squared error than a strong order-1 curve at the same *relative*
+accuracy, so an unweighted reduction quietly lets the strong curve dominate. Use
+`"weighted"` when you want the weak curve to carry its share.
+
+## Pre-prepared measurements
+
+`grax_opt` deliberately contains no smoothing, downsampling, or edge-exclusion
+logic. When you preprocess the measured curves yourself, pass the resulting
+values through `measurement_efficiency` so the optimizer fits exactly the
+numbers you prepared:
+
+```python
+energies, efficiencies = my_own_preprocessing(raw_path)
+
+spec = {
+    "measurement_path": raw_path,        # kept for provenance
+    "grazing_angle_deg": 2.0,
+    "evaluation_energies_ev": energies,
+    "measurement_efficiency": efficiencies,
+}
+```
+
+When `measurement_efficiency` is omitted, the measured curve is interpolated
+from the file onto `evaluation_energies_ev`.
+
+## Parallelism
+
+Each trial evaluates **all** measurements in a single batch. With
+`max_workers="auto"`, the worker pool parallelizes across measurements and
+energies together, so a four-curve fit keeps the pool busy far better than four
+separate single-curve fits would.
+
+Because a batch runner may return results out of completion order, simulated
+efficiencies are reassembled by result index rather than by arrival order.
+
+## Written artifacts
+
+- `best_result.json`: best fit, `per_measurement_best_losses`, the resolved
+  conditions of every measurement, and run metadata including `solver`.
+- `trial_history.csv`: per-trial history with one `loss_<label>` column per
+  measurement alongside the joint `loss`.
+- `best_fit.png`: one measured-versus-simulated panel per measurement, titled
+  with that measurement's conditions.
+- `best_fit_comparison.csv`: long-form `label, angle_mode, grazing_angle_deg,
+  cff, diffraction_order, polarization, energy_ev, measured_efficiency,
+  simulated_efficiency`.
+- `optimization_loss_history.png`: joint-loss history.
+
+Joint runs support checkpointing and resume on the same terms as the
+single-measurement optimizer — see {doc}`optimizer-resume`.