diff --git a/.gitignore b/.gitignore index d335922..0a87355 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ examples/reference_baselines/**/results/latest.csv .codex uv.lock examples/simulation/multilayer_theta_search/results/theta_scans/* +examples/simulation/multilayer_grating_design/results/ reticolo *.log .token diff --git a/CHANGELOG.md b/CHANGELOG.md index 496fa17..7638d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,19 @@ ## Unreleased -- Added a three-stage multilayer-grating design workflow: `grax.run_d_spacing_study`, `grax.run_gamma_study` and `grax.run_blaze_study`, driven by one frozen `grax.MultilayerOptimizationConfig`. Stage 0 derives a bilayer d-spacing from the grating geometry (grazing angle at the configured CFF, then the first-order Bragg law) and scans practical candidates with planar-multilayer reflectivity; stage 1 scans the bilayer thickness ratio `gamma`; stage 2 builds the multilayer-coated blazed grating and scans the blaze angle through `run_multilayer_theta_search_sweep`. Stages hand values forward only through `optimization_state.json` and only when a config value is the string `"auto"` -- a numeric value always wins, and no stage rewrites the config. Reflectivity for stages 0-1 comes from the new `grax.MultilayerReflectivity`, a thin wrapper over XRT's dynamical-diffraction engine; `xrt` (already a hard dependency) is imported lazily, so `import grax` still pulls in neither `xrt` nor `matplotlib.pyplot` (now covered by a test). A runnable Ru/B4C second-order example lives at `examples/simulation/multilayer_optimization_rub4c/`. +- Added a multilayer-grating design workflow: `grax.MultilayerGratingDesigner`, driven by one frozen `grax.MultilayerDesignConfig`. The rough/fine/final theta-search scan settings are a separate frozen `grax.ThetaSearchScanSettings`, held on two independent config fields -- `survey_scan_settings` and `energy_scan_settings` -- so the survey (many single-energy searches, one per grid cell) and the energy scan (fewer designs, often many energies each) can be tuned for speed or accuracy independently. `run_survey` scans a 2-D grid of bilayer d-spacing against blaze angle and, for every pair, builds the multilayer-coated blazed grating and runs a single-energy multilayer theta search through `run_multilayer_theta_search_sweep` (using `survey_scan_settings`) with `output_dir` set to that pair's own folder; the incident angle is a result of that search, not an input parameter. Each pair therefore keeps the full standard theta-search artifact bundle under `survey/runs/dnm/blazedeg/` (`multilayer_theta_search_summary.csv`, `*_all_orders.csv`, `theta_scans/`, profile/stack plots, `checkpoints/`), plus a `search_parameters.json` recording the search settings for review and tuning; each `dnm/` folder also gets an `overlay.png` of its blaze angles. From the selected results the survey reports, per d-spacing, the blaze angle with the highest efficiency (`optimal_blaze_deg = argmax_blaze efficiency`), and the headline artifacts are three plots, all under `MultilayerDesignConfig.plot_dir` (`/plots/`): optimal blaze angle versus d-spacing (labelled with efficiency), max efficiency versus d-spacing (labelled with the optimal blaze angle), and a `(d, blaze) -> efficiency` heatmap of the full `efficiency_map` with the optimal-blaze ridge overlaid. `survey/survey.csv` holds the combined per-cell table. `run_energy_scan` then sweeps chosen `(d, blaze)` designs over an explicit energy range with `run_multilayer_theta_search_sweep`, and additionally writes each design's `EnergyScanResult.titled_plot_path` into that same shared `plot_dir` -- efficiency versus energy titled with the coating (`MultilayerDesignConfig.coating_label`, defaulting to `"/"`), the diffraction order, and that design's d-spacing and blaze angle, so plot titles show the real compound (e.g. "Ru/B4C") even when a material is modelled with a stand-in optical-constants table; since every design's plot shares one folder, the filename itself carries all four (`efficiency_vs_energy__order_dnm_blazedeg.png`). `evaluate_survey` / `evaluate_energy_scan` re-derive the aggregates (rebuilt `survey.csv`, regenerated plots, collected results) from the per-run artifacts already on disk without launching any solver -- the example's `0_run_survey.py --eval` and `1_run_energy_scan.py --eval`. `run_survey` / `run_energy_scan` take optional `progress_callback` (called with a `grax.StageProgress` per item) and `run_survey` also takes `should_continue` (stops between cells, optimal blaze computed from the completed subset, `aborted=True` on the result). `import grax` still pulls in neither `xrt` nor `matplotlib.pyplot` (covered by a test). A runnable Ru/B4C second-order example lives at `examples/simulation/multilayer_grating_design/`. This replaces the earlier unreleased three-stage `run_d_spacing_study` / `run_gamma_study` / `run_blaze_study` workflow, its `grax.MultilayerOptimizationConfig`, and the `grax-web` "Multilayer study" page, all of which are removed. +- `grax.MultilayerGratingDesigner` gained `plot_energy_scan_overlay(results)`, which overlays several designs' efficiency-versus-energy curves on one axis and writes `plots/efficiency_vs_energy_comparison__order.png`; `run_energy_scan` and `evaluate_energy_scan` call it automatically whenever they produce two or more designs. `run_energy_scan` also gained a keyword-only `should_continue`, mirroring `run_survey`, so a long scan can be stopped cooperatively between designs (it returns the designs completed so far). +- Added a **Multilayer design** tab to the grax web app, driving the two-step `grax.MultilayerGratingDesigner` workflow end to end. The creation form exposes every `grax.MultilayerDesignConfig` field grouped into the dataclass's three sections (shared / survey-only / energy-scan-only), with the two `grax.ThetaSearchScanSettings` blocks and the runtime knobs behind a collapsed *Advanced*, and a live readout of how many theta searches the survey grid implies (warning past 200). Running the survey shows a live progress bar and then its three headline plots; the energy scan is then chosen with three buttons -- *Scan best* (the survey's global argmax), *Scan all d* (the optimal blaze at every d-spacing), or *Choose manually* to build a list of `(d, blaze)` cases from dropdowns that show each cell's surveyed efficiency. Ticking *scan the best design over energy automatically* on the form chains step 2 straight off the survey, so the plots render while the scan runs underneath them. Two or more scanned designs also get an overlay comparison plot. Either stage runs in a background thread against the shared run registry (so it appears in `/system/resource-status`), can be aborted with a keep-or-discard choice, and can be reset; re-running the survey marks a finished energy scan *stale*. Each study lives in `/multilayer_designs//`, holding `study.json` beside the library's own `survey/`, `plots/` and `energy_scan/` trees. +- The **Multilayer design** tab's three step-1 plots are now interactive Plotly charts instead of static PNGs: hover reads d-spacing, blaze angle and efficiency off any point, and the two line plots sit side by side at equal width rather than the old main-plus-sidebar split. **Clicking a cell of the `(d, blaze)` heatmap selects that design for step 2** -- it opens the manual picker, adds the case, and marks the cell; clicking it again removes it. The dropdown picker is unchanged and stays in sync with the map. The library still writes the matplotlib PNGs, which remain the CLI's output and the page's fallback when Plotly is not installed. +- The multilayer design survey's plots now fill in while the survey runs, instead of appearing only at the end. `MultilayerGratingDesigner.run_survey` rewrites `survey/survey.csv` after every cell -- atomically, via a temp file and `os.replace`, so a reader never catches a half-written table -- and the web page polls a new `survey-options` endpoint every few seconds to redraw the three charts from the table as it grows. The page reloads once when the stage finishes, which is what stops the polling. +- The multilayer design energy scan draws a live efficiency-versus-energy overlay while it runs, one curve per design, fed from each design's checkpoint file (which already gains a record per solved energy) through a new `energy-scan-points` endpoint. The finished stage still renders the saved plots. Its subtitle names the coating and order but not the survey energy, which is a property of the survey rather than of a scan across energies. +- A finished multilayer design stage now shows its results without a manual refresh. The stage monitor stopped polling on `completed`/`failed`/`aborted` but left the progress card on screen, so the plots -- which are rendered server-side -- never appeared on a page that had been open while the stage ran; the monitor now reloads once when it reaches a terminal state. The energy-scan result plots also share the survey's equal-width grid, which gives a single scanned design the full page width instead of half of it. +- **Aborting a multilayer design stage now kills the running solves** instead of waiting for the item in flight to finish -- a 200-energy scan used to be uninterruptible for its whole multi-hour run. `run_multilayer_theta_search_sweep` gained keyword-only `stop_event` and `on_worker_pids_changed`: once the event is set, queued energies are not submitted, the live worker processes are terminated, and the sweep returns `stopped_early=True` with the energies solved so far (`MultilayerThetaSearchSweepResult` also now reports `resolved_max_workers`). Because an in-process solve cannot be interrupted, passing a `stop_event` routes execution through worker processes even at `max_workers=1`; callers that pass none keep the in-process path. `MultilayerGratingDesigner.run_survey` and `run_energy_scan` forward both arguments -- a killed survey cell is discarded (its half-written summary CSV would break later `evaluate_survey` calls) and a half-scanned design is not reported as finished, though its checkpoint keeps every solved energy so re-running resumes from there. In the web app a stage killed on purpose is recorded as *aborted* rather than *failed*, and the run monitor reports the real worker-pool size instead of a hardcoded 1. +- A multilayer design study's parameters can be changed after it is created. *Edit parameters* on the study page reopens the creation form with the stored values filled in, so a survey grid, an energy range or a theta-search setting can be adjusted without starting a new study and losing the results. Saving a change that a finished stage depended on marks that stage *stale* rather than deleting anything -- the artifacts stay on disk until the stage is re-run or reset -- while purely cosmetic edits (the coating label, which plots get saved, checkpointing, worker count) invalidate nothing, and an edit to only the energy-scan settings leaves a finished survey alone. Editing is refused while a stage is running. +- A new multilayer design study's form starts from the most recent study's settings rather than the dataclass defaults, so the advanced theta-search parameters someone tuned for a quick exploratory run carry over to the next one. The form says which study it copied and links to it; with no studies yet it falls back to the defaults as before. +- The multilayer-design tutorial now walks through the bundled Ru/B4C example rather than just pointing at it, and opens with a *Quick start* section showing that example's own output: the code to run each step immediately followed by the real plots it produces, pulled from a full 50 x 28-cell survey and a 1000-point energy scan and checked into `docs/tutorials/images/simulation/` (the `(d, blaze)` heatmap and both headline curves for the survey, the efficiency-versus-energy curve for the best design) -- so the picture comes before the explanation rather than after it. The walkthrough below it covers what each of the example's four files is for, how to run it (and how to shrink the 1400-cell grid for a first pass), how to read the survey tree, and why the rough theta-search window has to be wide -- with the measured numbers from that real run, where the Bragg seed sits 0.79 deg above the angle the search actually settles on for the best cell. It also documents the `stop_event` stop path alongside `should_continue`. +- Every "there is no CFF input" mention is gone from the multilayer-design workflow's documentation and source -- the web app's design pages, the tutorial, the API reference page, the library's own module docstring, and the bundled example's parameter file -- in favor of stating what the workflow does (every cell gets its own theta search) rather than dwelling on an input a removed workflow used to need. +- A multilayer design study can be downloaded as a standalone script to run offline. *Download script* on the study page writes one self-contained Python file: the study's parameters as named constants under the same SHARED / SURVEY / ENERGY SCAN banners `MultilayerDesignConfig` groups them by, the `MultilayerDesignConfig` built from them, and both stages behind `--survey` / `--energy-scan` flags -- with `--best`, `--pairs "d,blaze; ..."` and `--eval` as in the bundled example, and both stages in order when no stage flag is given. It imports only `grax` and pandas, so it runs anywhere grax is installed, with no sibling parameters module and nothing from the web app. - Fixed a native crash (segmentation fault, no Python traceback) during a serial (`max_workers=1`) Nevière theta-search sweep on Linux/OpenBLAS. The differential method issues thousands of tiny dense `zgesv`/`zgemm` calls per photon-energy point, and a threaded BLAS both wastes its time on dispatch and, on some OpenBLAS builds, crashes under that pattern. `BatchSimulationRunner` already pinned `OPENBLAS_NUM_THREADS` and friends to `1` in its spawned workers, but a serial run or a direct `grax.run_simulation` call executes in the current process where those environment variables can no longer take effect. `run_simulation` now wraps the Nevière solve in `threadpoolctl.threadpool_limits(1, "blas")` (a no-op where no controllable native library is loaded, e.g. NumPy on Apple Accelerate, and redundant inside a spawned worker); the RCWA path, whose single large eigensolve benefits from threads, is unchanged. `threadpoolctl` (already an indirect dependency via SciPy) is now a direct one. Additionally, the interface-response cascade in the Nevière/RCWA shared code now raises a `ValueError` when a slab transfer matrix or a cascaded block goes non-finite, instead of passing it to `np.linalg.solve` where some LAPACK builds crash rather than raising. New `examples/simulation/neviere_grazing_stability/` sweeps a coated Mo/B4C grating from a Bragg angle down to 0.01 deg in p-polarization and confirms every point stays finite. - Internal cleanup pass; no public API or numerical behaviour changes. - Progress bars now follow the terminal width while running. The `tqdm` bars in `run_parameter_study`, `BatchSimulationRunner`, and `run_multilayer_theta_search_sweep` pass `dynamic_ncols=True`, so resizing the terminal mid-run re-fits the bar instead of keeping the width measured at start-up. The `DummyProgress` test doubles were widened to accept the extra keyword. diff --git a/docs/api/simulation.md b/docs/api/simulation.md index 8f1076c..58e7434 100644 --- a/docs/api/simulation.md +++ b/docs/api/simulation.md @@ -12,6 +12,6 @@ simulation/fixed-angle-sweep simulation/monochromator-sweep simulation/energy-angle-sweep simulation/multilayer-theta-search -simulation/multilayer-optimization +simulation/multilayer-design simulation/parameter-study ``` diff --git a/docs/api/simulation/multilayer-design.md b/docs/api/simulation/multilayer-design.md new file mode 100644 index 0000000..069f939 --- /dev/null +++ b/docs/api/simulation/multilayer-design.md @@ -0,0 +1,44 @@ +# Multilayer grating design + +APIs for the multilayer-grating design workflow: a 2-D +d-spacing / blaze-angle survey at one optimization energy, followed by per-design +energy scans. For every `(d_spacing, blaze angle)` pair the survey builds the +multilayer-coated blazed grating and runs graxPy's single-energy multilayer theta +search, which scans the incident angle and returns the angle that maximizes the +selected-order efficiency for that pair. + +```{eval-rst} +.. autoclass:: grax.MultilayerDesignConfig +``` + +```{eval-rst} +.. autoclass:: grax.ThetaSearchScanSettings +``` + +```{eval-rst} +.. autoclass:: grax.MultilayerGratingDesigner + :members: run_survey, evaluate_survey, run_energy_scan, evaluate_energy_scan, plot_energy_scan_overlay +``` + +```{eval-rst} +.. autoclass:: grax.SurveyResult +``` + +```{eval-rst} +.. autoclass:: grax.EnergyScanResult +``` + +```{eval-rst} +.. autoclass:: grax.StageProgress +``` + +## Planar-multilayer reflectivity + +The design workflow itself never touches XRT, but the planar-multilayer +reflectivity wrapper remains available for standalone Bragg-reflectivity work. +`xrt` is imported lazily, so importing `grax` does not pull it in. + +```{eval-rst} +.. autoclass:: grax.MultilayerReflectivity + :members: reflectivity_vs_energy +``` diff --git a/docs/api/simulation/multilayer-optimization.md b/docs/api/simulation/multilayer-optimization.md deleted file mode 100644 index 8b4e558..0000000 --- a/docs/api/simulation/multilayer-optimization.md +++ /dev/null @@ -1,44 +0,0 @@ -# Multilayer optimization - -APIs for the three-stage multilayer-grating design workflow: derive and scan the -bilayer d-spacing, scan the bilayer thickness ratio, then scan the blaze angle -with graxPy's internal theta search. Stages hand values forward only through -`optimization_state.json`, and only when a config value is `"auto"`. - -```{eval-rst} -.. autoclass:: grax.MultilayerOptimizationConfig -``` - -```{eval-rst} -.. autofunction:: grax.run_d_spacing_study -``` - -```{eval-rst} -.. autofunction:: grax.run_gamma_study -``` - -```{eval-rst} -.. autofunction:: grax.run_blaze_study -``` - -```{eval-rst} -.. autoclass:: grax.DSpacingStudyResult -``` - -```{eval-rst} -.. autoclass:: grax.GammaStudyResult -``` - -```{eval-rst} -.. autoclass:: grax.BlazeStudyResult -``` - -## Planar-multilayer reflectivity - -Stages 0 and 1 measure peak Bragg reflectivity with this XRT wrapper. `xrt` is -imported lazily, so importing `grax` does not pull it in. - -```{eval-rst} -.. autoclass:: grax.MultilayerReflectivity - :members: reflectivity_vs_energy -``` diff --git a/docs/developer/module-guide.md b/docs/developer/module-guide.md index 038eb96..0814472 100644 --- a/docs/developer/module-guide.md +++ b/docs/developer/module-guide.md @@ -24,7 +24,7 @@ This guide summarizes the source layout for contributors. parameter studies - `multilayer_reflectivity.py`: planar-multilayer peak Bragg reflectivity versus energy, wrapping the XRT dynamical-diffraction engine (imported lazily) -- `multilayer_optimization.py`: three-stage d-spacing / gamma / blaze design +- `multilayer_design.py`: d-spacing / blaze-angle survey plus per-design energy scans workflow built on the public API, with a JSON state file for the `"auto"` hand-off between stages - `solvers/`: the one-dimensional electromagnetic solvers diff --git a/docs/tutorials/choosing-a-solver.md b/docs/tutorials/choosing-a-solver.md index 7547a6b..f777a17 100644 --- a/docs/tutorials/choosing-a-solver.md +++ b/docs/tutorials/choosing-a-solver.md @@ -159,7 +159,7 @@ Solver selection reaches every workflow, not just one-point solves: | {class}`grax.BatchSimulationRunner` | `solver=`, `solver_options=`, or a per-case `"solver"` key | | {func}`grax.run_multilayer_theta_search` | `solver=` — used for all three scan stages | | {func}`grax.run_multilayer_theta_search_sweep` | `solver=` | -| {func}`grax.run_blaze_study` | `MultilayerOptimizationConfig.solver` (stage 2 only) | +| {class}`grax.MultilayerGratingDesigner` | `MultilayerDesignConfig.solver` | | {func}`grax.run_parameter_study` | `solver=` | | `grax_opt` measurement fits | `solver` / `solver_options` on the config | | Web UI | Solver dropdown in the run form | diff --git a/docs/tutorials/images/simulation/multilayer_design_efficiency_heatmap_d_vs_blaze.png b/docs/tutorials/images/simulation/multilayer_design_efficiency_heatmap_d_vs_blaze.png new file mode 100644 index 0000000..1182115 Binary files /dev/null and b/docs/tutorials/images/simulation/multilayer_design_efficiency_heatmap_d_vs_blaze.png differ diff --git a/docs/tutorials/images/simulation/multilayer_design_efficiency_vs_energy_best_design.png b/docs/tutorials/images/simulation/multilayer_design_efficiency_vs_energy_best_design.png new file mode 100644 index 0000000..fc2aafa Binary files /dev/null and b/docs/tutorials/images/simulation/multilayer_design_efficiency_vs_energy_best_design.png differ diff --git a/docs/tutorials/images/simulation/multilayer_design_max_efficiency_vs_d_spacing.png b/docs/tutorials/images/simulation/multilayer_design_max_efficiency_vs_d_spacing.png new file mode 100644 index 0000000..f34576d Binary files /dev/null and b/docs/tutorials/images/simulation/multilayer_design_max_efficiency_vs_d_spacing.png differ diff --git a/docs/tutorials/images/simulation/multilayer_design_optimal_blaze_vs_d_spacing.png b/docs/tutorials/images/simulation/multilayer_design_optimal_blaze_vs_d_spacing.png new file mode 100644 index 0000000..cb4b968 Binary files /dev/null and b/docs/tutorials/images/simulation/multilayer_design_optimal_blaze_vs_d_spacing.png differ diff --git a/docs/tutorials/images/simulation/multilayer_optimization_rub4c_d_spacing.png b/docs/tutorials/images/simulation/multilayer_optimization_rub4c_d_spacing.png deleted file mode 100644 index a32cd73..0000000 Binary files a/docs/tutorials/images/simulation/multilayer_optimization_rub4c_d_spacing.png and /dev/null differ diff --git a/docs/tutorials/multilayer-design.md b/docs/tutorials/multilayer-design.md new file mode 100644 index 0000000..dd34d3e --- /dev/null +++ b/docs/tutorials/multilayer-design.md @@ -0,0 +1,355 @@ +# Multilayer grating design workflow + +Use {class}`grax.MultilayerGratingDesigner` to size a periodic multilayer coating +for a blazed grating working in a chosen diffraction order. The workflow has two +steps that share one {class}`grax.MultilayerDesignConfig`: a **survey** across +bilayer d-spacing and blaze angle, then an **energy scan** of the designs you +pick off it. + +## Quick start: the Ru/B4C example + +`examples/simulation/multilayer_grating_design/` is a runnable project that +sizes a Ru/B4C coating on silicon for a 2400 l/mm blazed grating working in +inside order 2 at 9 keV. The images below are its actual output. + +### Step 1 -- survey a d-spacing / blaze-angle grid + +```console +$ cd examples/simulation/multilayer_grating_design +$ python 0_run_survey.py +``` + +```{figure} images/simulation/multilayer_design_efficiency_heatmap_d_vs_blaze.png +:alt: Peak selected-order efficiency over the bilayer d-spacing and blaze-angle grid, with the optimal-blaze ridge overlaid. +:align: center +:width: 85% + +Peak efficiency over the full `(d, blaze)` grid, with the optimal blaze angle +per d-spacing overlaid as a ridge. One glance shows both where the resonance +sits and how sharp it is -- here a narrow band around d = 3 nm, order-2 clear. +``` + +### Step 2 -- scan the best design over energy + +```console +$ python 1_run_energy_scan.py --best +``` + +```{figure} images/simulation/multilayer_design_efficiency_vs_energy_best_design.png +:alt: Selected-order efficiency versus photon energy for the survey's best design, d = 3.102 nm, blaze = 0.859 deg. +:align: center +:width: 75% + +The survey's best cell, d = 3.102 nm / blaze = 0.859 deg, swept from 3 keV to +12 keV. `theta*` is re-solved at every energy, so this plot is not an +extrapolation from the single survey point -- it is 1000 independent theta +searches. +``` + +Read on for what each step is doing, how to read the rest of the output, and +how to run this example (or your own) sensibly. + +## Step 1: the d-spacing / blaze-angle survey + +{meth}`grax.MultilayerGratingDesigner.run_survey` scans a 2-D grid of bilayer +d-spacing (`d_min_nm`..`d_max_nm`, `d_points` values) against blaze angle +(`blaze_min_deg`..`blaze_max_deg`, `blaze_points` values). For every `(d, blaze)` +pair it: + +1. builds the multilayer-coated blazed grating with the fixed `gamma`, + `n_bilayers` and `anti_blaze_angle_deg`, and +2. runs {func}`grax.run_multilayer_theta_search_sweep` at the single + `target_energy_ev`, using `config.survey_scan_settings`, with `output_dir` + set to that pair's own run folder. The search scans the incident angle + around the analytical multilayer Bragg estimate and returns the angle + `theta*` that maximizes the selected-order efficiency, plus the + precise-scan FWHM. + +The incident angle is not a parameter you choose: every pair gets its own theta +search, and `theta*` is whatever angle that search settles on. Step 2 then +tracks how `theta*` evolves with energy for the designs you pick. + +`survey_scan_settings` and `energy_scan_settings` are independent +{class}`grax.ThetaSearchScanSettings` instances on `MultilayerDesignConfig` -- +tune the survey's (many single-energy searches, one per grid cell) separately +from the energy scan's (fewer designs, but often many energies each). + +The survey writes: + +- `survey/survey.csv` -- one row per cell: `d_nm`, `blaze_deg`, + `bragg_estimate_deg`, `incidence_angle_deg`, `peak_efficiency`, + `precise_fwhm_deg`, `edge_clipped`. +- `plots/optimal_blaze_vs_d_spacing.png` -- **headline curve 1**: for each + d-spacing, the blaze angle with the highest efficiency + (`optimal_blaze_deg = argmax_blaze efficiency`), plotted d-spacing (x) against + optimal blaze angle (y), each point labelled with its peak efficiency. +- `plots/max_efficiency_vs_d_spacing.png` -- **headline curve 2**: the same + optimum plotted d-spacing (x) against max selected-order efficiency (y), each + point labelled with the blaze angle that achieved it. +- `plots/efficiency_heatmap_d_vs_blaze.png` -- **headline map**: the full + `efficiency_map`, not just the per-d optimum -- d-spacing (x), blaze angle + (y), peak selected-order efficiency as colour (z) -- with the + `optimal_blaze_deg` ridge overlaid, so you can see the whole resonance + structure the two curves above are slices of. Shown above, in *Quick start*. +- `survey/runs/` -- one folder per d-spacing (`dnm/`), each holding an + `overlay.png` (every blaze angle's incidence-angle scan at that d, with the + chosen blaze drawn bold and a star at each selected peak) and one sub-folder + per blaze angle (`blazedeg/`). Each blaze folder is the **full output of that + multilayer theta search**, exactly as {func}`grax.run_multilayer_theta_search_sweep` + writes it: `multilayer_theta_search_summary.csv` (selected angle, efficiency, + FWHM), `multilayer_theta_search_all_orders.csv` (every diffraction order), + `theta_scans/theta_scan_eV.csv` + `.png` (the incidence-angle scan), the + profile and stack plots, and `checkpoints/`. A `search_parameters.json` is + added listing everything that was fed into the search so you can review a run + for correctness and change its parameters. + +The two curves the heatmap is sliced from look like this on the real Ru/B4C +survey (a full 50 x 28 grid, d = 2-8 nm): + +```{figure} images/simulation/multilayer_design_optimal_blaze_vs_d_spacing.png +:alt: Optimal blaze angle versus d-spacing, points coloured and labelled by peak efficiency. +:align: center +:width: 85% + +Optimal blaze angle per d-spacing. The step pattern between d = 4 and 6 nm is +a real feature of this grid, not noise -- see below. +``` + +```{figure} images/simulation/multilayer_design_max_efficiency_vs_d_spacing.png +:alt: Max selected-order efficiency versus d-spacing, points coloured and labelled by optimal blaze angle. +:align: center +:width: 85% + +The same optimum, efficiency instead of angle on the y-axis. The main +resonance around d = 3 nm reaches ~0.69; efficiency collapses to near zero for +d = 4-6 nm before a much weaker secondary rise past it. +``` + +A cell whose search raises (for example an unphysical sub-grazing incidence +angle) is recorded as a `NaN` row when `on_error="continue"` (the default) and +re-raises when `on_error="fail_fast"`. + +{meth}`grax.MultilayerGratingDesigner.evaluate_survey` (the example's +`0_run_survey.py --eval`) runs **no new solves**: it re-reads every per-run +`multilayer_theta_search_summary.csv` already on disk, rebuilds `survey.csv` and +regenerates all three headline plots and every overlay. Use it to iterate on +the analysis after a long `run_survey`. + +## Step 2: per-design energy scans + +Pick one or more `(d_spacing_nm, blaze_angle_deg)` designs off the survey and +pass them to {meth}`grax.MultilayerGratingDesigner.run_energy_scan`. The example's +`1_run_energy_scan.py` selects them from `survey.csv`: no flag scans the optimal +blaze at every d-spacing, `--best` scans only the single highest-efficiency +`(d, blaze)`, and `--pairs "d,blaze; ..."` scans exactly what you list. Each +design is one fixed grating swept over +`energy_scan_min_ev`..`energy_scan_max_ev` (`energy_scan_points` values) with +{func}`grax.run_multilayer_theta_search_sweep`, giving efficiency versus energy +and `theta*(E)` -- shown for the example's best design in *Quick start* above. +Adding `--eval` ({meth}`grax.MultilayerGratingDesigner.evaluate_energy_scan`) +re-reads existing `energy_scan/` results instead of re-solving. + +Besides the sweep's own generic-titled plot (`EnergyScanResult.energy_efficiency_plot_path`, +inside that design's `energy_scan/` folder), every design also gets a plot in +the shared `plot_dir` (`EnergyScanResult.titled_plot_path`, the same folder as +the survey's three headline plots) -- the same curve, titled with the coating +and this design's d-spacing and blaze angle, e.g. *"Ru/B4C multilayer grating +(order 2): d = 3.102 nm, blaze = 0.859 deg"*, and named accordingly: +`efficiency_vs_energy_Ru-B4C_order2_d3.102nm_blaze0.859deg.png`. Because every +design's plot lands in the same folder, the filename itself carries the +coating, order, d-spacing and blaze angle -- there is nothing per-design about +the folder to disambiguate them. + +Scanning two or more designs in one call also writes an overlay comparison -- +every curve on one axis, legended by d-spacing and blaze angle, as +`plots/efficiency_vs_energy_comparison__order.png`. +{meth}`grax.MultilayerGratingDesigner.run_energy_scan` and +{meth}`grax.MultilayerGratingDesigner.evaluate_energy_scan` call +{meth}`grax.MultilayerGratingDesigner.plot_energy_scan_overlay` for you; call it +directly to overlay an arbitrary subset. + +`run_energy_scan` also accepts `should_continue`, checked before each design, so +a long scan can be stopped cooperatively; it returns the designs completed so +far. A `stop_event` stops it immediately instead -- see *Progress and abort +hooks* below. + +The coating name comes from `MultilayerDesignConfig.coating_label` when set, +otherwise from `"/"` -- set it explicitly +whenever a material is modelled with a stand-in optical-constants table (B4C +modelled with the carbon table, in the Ru/B4C example) so the title and +filename show the real compound. + +## Walking through the Ru/B4C example in depth + +### The four files + +| File | Role | +| --- | --- | +| `rub4c_design_parameters.py` | **The only file you edit.** Builds `CONFIG`; nothing ever rewrites it. | +| `0_run_survey.py` | Step 1. `--eval` re-derives the survey from disk without solving. | +| `1_run_energy_scan.py` | Step 2. `--best`, `--pairs "d,blaze; ..."`, `--eval`. | +| `run_all.sh` | Both steps in order (`PYTHON_BIN` overrides the interpreter). | + +The parameters file is split by three `# ===` banners that mirror +`MultilayerDesignConfig`'s own grouping -- `SHARED`, `SURVEY`, `ENERGY_SCAN`. +Check the banner before moving a value: editing the wrong section is the usual +cause of "I changed X and nothing happened". `SURVEY_SCAN.rough_scan_points` +only affects `0_run_survey.py`; `1_run_energy_scan.py` reads `ENERGY_SCAN_SCAN`. + +`COATING_LABEL = "Ru/B4C"` is there because B4C is modelled with the carbon +optical-constants table. Without it every plot title and filename would say +"Ru/C", the table's element name rather than the real compound. + +### Sizing the grid + +The shipped grid is `D_POINTS = 50` x `BLAZE_POINTS = 28` = 1400 theta +searches, which is an overnight run -- that is the grid the images above come +from. For a first pass set both to 2-3 and cut `rough_scan_points`/ +`fine_scan_points` to ~15 in **both** scan blocks; the whole thing then +finishes in a couple of minutes and still exercises every artifact path. + +### Reading the survey output tree + +```text +results/ + survey/survey.csv one row per (d, blaze) cell + survey/runs/d3.102nm/overlay.png every blaze angle at this d, chosen one bold + survey/runs/d3.102nm/blaze0.859deg/ the full theta search for one cell + plots/optimal_blaze_vs_d_spacing.png + plots/max_efficiency_vs_d_spacing.png + plots/efficiency_heatmap_d_vs_blaze.png +``` + +Read the **heatmap first**: the two curves are slices through it, and only the +map shows whether an optimum sits on a broad plateau or a knife-edge. Then use +the curves to pick designs, and drill into `survey/runs/` when a cell looks +wrong -- `theta_scans/theta_scan_9000eV.csv` and its PNG show the actual +incidence-angle scan, and `search_parameters.json` records exactly what was fed +in. + +The full 50 x 28 run above puts the optimum at **d = 3.102 nm, blaze 0.859 +deg, efficiency 0.69**, with a sharp collapse to near zero for d = 4-6 nm and a +much weaker secondary rise past it -- the step pattern visible in the optimal +blaze curve. A coarse d grid can step straight over a resonance this narrow. + +### Why the rough theta-search window has to be wide + +The search is seeded from the analytical multilayer Bragg estimate, and for +shallow inside orders that estimate **overshoots the true grating optimum**. +For the survey's best cell (d = 3.102 nm, blaze = 0.859 deg), `survey.csv` +records `bragg_estimate_deg = 1.315` but `incidence_angle_deg = 0.528` -- the +real peak sits **0.79 deg below** the seed. A `rough_scan_half_width_deg` of +0.2 deg would never reach it and every cell would report a spurious efficiency +from the window's edge, which is why the example ships a wide rough window +(`rough_scan_half_width_deg = 1`) narrowed only for the fine pass. + +The peak is also narrow: `precise_fwhm_deg = 0.064` for that cell. A wide +window therefore needs enough `rough_scan_points` to land on the peak at all -- +61 points over +/-1 deg is a 0.033 deg step, about two samples across the +FWHM. Widening the window without adding points is how a survey silently +misses resonances. `edge_clipped` in `survey.csv` flags cells whose selected +angle landed against a scan boundary; a column of those means the window is +wrong, not the physics. + +### Choosing designs and scanning them + +```console +$ python 1_run_energy_scan.py # optimal blaze at every d +$ python 1_run_energy_scan.py --best # only the global best cell +$ python 1_run_energy_scan.py --pairs "3.1,0.86; 3.5,0.96" +``` + +Each design gets `plots/efficiency_vs_energy_Ru-B4C_order2_dnm_blazedeg.png` +in the shared `plots/` folder alongside step 1's headline plots, and two or more +designs also produce +`plots/efficiency_vs_energy_comparison_Ru-B4C_order2.png` with every curve on +one axis -- that overlay is the point of scanning several designs. + +### Iterating without re-solving + +Both scripts take `--eval`, which re-derives everything from the artifacts +already on disk. After a long survey, `python 0_run_survey.py --eval` rebuilds +`survey.csv` and regenerates all three plots in seconds, so changing how results +are *presented* never costs a re-run. The same applies to +`1_run_energy_scan.py --eval`. + +Checkpointing (`CHECKPOINT = True`, `RESUME = True`) means an interrupted run +picks up where it stopped: re-running the same command re-solves only what is +missing. + +## Using the API directly + +The example scripts are a thin CLI over the same two calls. A smaller, +self-contained version: + +```python +from pathlib import Path + +from grax import MultilayerDesignConfig, MultilayerGratingDesigner, ThetaSearchScanSettings + +config = MultilayerDesignConfig( + output_dir=Path("examples/simulation/multilayer_grating_design/results"), + grating_density_lpermm=2400.0, + diffraction_order=2, + multilayer_bragg_order=1, + target_energy_ev=9000.0, + material_a=("Ru", 12.1), + material_b=("C", 2.52), # B4C modelled with the carbon table + substrate_material=("Si", 2.33), + n_bilayers=40, + gamma=0.5, + d_min_nm=1.5, + d_max_nm=6.0, + d_points=4, + blaze_min_deg=0.6, + blaze_max_deg=2.0, + blaze_points=6, + # Coarser/faster for the 4x6=24-cell survey; the defaults (finer) are kept + # for the energy scan since it runs far fewer searches. + survey_scan_settings=ThetaSearchScanSettings( + rough_scan_points=31, fine_scan_points=41, final_fourier_orders=15 + ), + energy_scan_min_ev=3000.0, + energy_scan_max_ev=12000.0, + energy_scan_points=15, + solver="neviere", + polarization="p", +) + +designer = MultilayerGratingDesigner(config) +survey = designer.run_survey() +for d_nm, blaze_opt, eff in zip( + survey.d_values_nm, survey.optimal_blaze_deg, survey.optimal_blaze_efficiency +): + print(f"d {d_nm:.2f} nm -> optimal blaze {blaze_opt:.2f} deg (eff {eff:.3g})") + +# Scan the optimal design for the first d-spacing over energy. +designer.run_energy_scan([(float(survey.d_values_nm[0]), float(survey.optimal_blaze_deg[0]))]) +``` + +## Progress and abort hooks + +Both `run_survey` and `run_energy_scan` accept an optional `progress_callback` +(called with a {class}`grax.StageProgress` before each item and once with +`current_label="done"`), and both accept the two stop mechanisms: + +- `should_continue` -- a callable checked **between items**. Returning `False` + stops after the cell or design in flight finishes. `run_survey` then computes + the optimal blaze per d-spacing from the completed subset and sets + `SurveyResult.aborted`. +- `stop_event` -- a {class}`threading.Event` forwarded into the theta search + itself. Setting it stops **immediately**: queued energies are not submitted and + the live worker processes are terminated. A cell killed this way is discarded + (a half-written summary CSV would break later `evaluate_survey` calls) and a + half-scanned design is not returned, but its checkpoint keeps every energy + already solved, so re-running resumes from there. + +Because an in-process solve cannot be interrupted, passing a `stop_event` routes +execution through worker processes even at `max_workers=1`. Callers that pass +none keep the cheaper in-process path. + +## Solver selection + +`solver` (`rcwa` or `neviere`), `polarization` and `backend` come from the +config and are used by both steps. See {doc}`choosing-a-solver` and +{doc}`multilayer-theta-search`. diff --git a/docs/tutorials/multilayer-optimization.md b/docs/tutorials/multilayer-optimization.md deleted file mode 100644 index 67c4c87..0000000 --- a/docs/tutorials/multilayer-optimization.md +++ /dev/null @@ -1,101 +0,0 @@ -# Multilayer optimization workflow - -Use {func}`grax.run_d_spacing_study`, {func}`grax.run_gamma_study` and -{func}`grax.run_blaze_study` to size a periodic multilayer coating for a blazed -grating monochromator working in a chosen diffraction order at fixed CFF. The -three stages share one {class}`grax.MultilayerOptimizationConfig` and run in -order: - -1. **D-spacing.** Compute the grazing angle at the target energy and CFF with - {func}`grax.monochromator_grazing_angles_deg`, convert it to a bilayer - d-spacing with the first-order Bragg law `d = λ / (2 sin θ)`, build a - practical 0.1 nm-rounded scan grid that is guaranteed to contain the rounded - geometry value, and scan every candidate with XRT planar-multilayer - reflectivity. The geometry value is stored as `d_suggested_nm`; the - numerically best d at the target energy is stored separately as a diagnostic. -2. **Gamma.** At the resolved d-spacing, scan the bilayer thickness ratio and - keep the value with the highest peak reflectivity at the target energy. -3. **Blaze.** Build the multilayer-coated blazed grating and scan the blaze - angle, running {func}`grax.run_multilayer_theta_search` per energy (diffraction - order 2, multilayer Bragg order 1, Nevière solver, p-polarization by default), - and keep the blaze angle with the highest selected-order efficiency at the - target energy. - -## State file and the `"auto"` hand-off - -The only channel between stages is `output_dir/optimization_state.json`. A stage -reads it only when the corresponding config value is the string `"auto"`: - -- `d_spacing_nm="auto"` in stages 1 and 2 resolves `d_suggested_nm` written by - stage 0. -- A numeric `d_spacing_nm` (or `gamma`) always wins; the state file is ignored. -- No stage ever rewrites the config. Stage 1's `gamma_suggested` is recorded for - traceability but **not** auto-applied -- copy it into the config yourself if - you want stage 2 to use it. - -## Two inherited conventions - -- The geometry d-spacing derivation uses `hc = 1239.841984` eV·nm while - {func}`grax.monochromator_grazing_angles_deg` uses `1239.8` internally. The - difference is far below the 0.1 nm rounding applied to the candidate grid. -- The XRT reflectivity path (stages 0-1) places `material_a` on top of a - `material_a` substrate; the graxPy {class}`grax.MultilayerStack` used in stage - 2 places `material_b` on top of the configured substrate. Both are modelling - choices carried over unchanged from the original workflow. - -## Example - -```python -from pathlib import Path - -from grax import ( - MultilayerOptimizationConfig, - run_blaze_study, - run_d_spacing_study, - run_gamma_study, -) - -config = MultilayerOptimizationConfig( - output_dir=Path("examples/simulation/multilayer_optimization_rub4c/results"), - d_spacing_nm="auto", - gamma=0.5, - blaze_angle_deg=1.1, - material_a=("Ru", 12.1), - material_b=("C", 2.52), # B4C modelled with the carbon table - substrate_material=("Si", 2.33), - n_bilayers=40, - target_energy_ev=9000.0, - grating_density_lpermm=2400.0, - diffraction_order=2, - cff=2.25, - multilayer_bragg_order=1, - solver="neviere", - polarization="p", -) - -d_result = run_d_spacing_study(config) -print(f"geometry d = {d_result.geometry_d_nm:.3f} nm -> suggested {d_result.d_suggested_nm:.1f} nm") - -gamma_result = run_gamma_study(config) # d_spacing_nm="auto" -> reads the state file -print(f"suggested gamma = {gamma_result.gamma_suggested:.3f}") - -blaze_result = run_blaze_study(config) # uses config.gamma, not the suggestion -print(f"suggested blaze = {blaze_result.blaze_suggested_deg:.4f} deg") -``` - -```{image} images/simulation/multilayer_optimization_rub4c_d_spacing.png -:alt: Ru/B4C d-spacing study, peak p-polarized reflectivity versus photon energy per candidate d-spacing -:align: center -:width: 90% -``` - -See `examples/simulation/multilayer_optimization_rub4c/` for the full runnable -workflow: `ru_b4c_parameters.py` builds the config and the three numbered -scripts run the stages. `run_all.sh` runs them in order. - -## Solver selection - -Stage 2 takes `solver` (`rcwa` or `neviere`) from the config; the example's -`2_ru_b4c_blaze_study.py` also exposes it as `--solver`. Stages 0 and 1 do not -use a graxPy solver -- they measure reflectivity with XRT. See -{doc}`choosing-a-solver` and {doc}`multilayer-theta-search`. diff --git a/docs/tutorials/sweep-recipes.md b/docs/tutorials/sweep-recipes.md index 8f163e6..c8b4588 100644 --- a/docs/tutorials/sweep-recipes.md +++ b/docs/tutorials/sweep-recipes.md @@ -11,6 +11,6 @@ roughness monochromator-sweep energy-angle-sweep multilayer-theta-search -multilayer-optimization +multilayer-design parameter-study ``` diff --git a/examples/simulation/multilayer_grating_design/0_run_survey.py b/examples/simulation/multilayer_grating_design/0_run_survey.py new file mode 100644 index 0000000..6ac6418 --- /dev/null +++ b/examples/simulation/multilayer_grating_design/0_run_survey.py @@ -0,0 +1,64 @@ +"""Step 1: run the Ru/B4C d-spacing x blaze-angle survey. + +Runs :meth:`grax.MultilayerGratingDesigner.run_survey` with the shared +``rub4c_design_parameters.CONFIG``. Writes ``results/survey/survey.csv``, three +headline plots in ``results/plots/`` (``optimal_blaze_vs_d_spacing.png``, +``max_efficiency_vs_d_spacing.png`` and ``efficiency_heatmap_d_vs_blaze.png`` -- +the full ``(d, blaze) -> efficiency`` map with the optimal-blaze ridge overlaid; +``1_run_energy_scan.py``'s titled plots land in this same folder), and a +per-run folder tree under +``results/survey/runs/``: one ``dnm/`` folder per period (with an +``overlay.png`` of its blaze angles, chosen one highlighted) and inside it one +``blazedeg/`` folder per run holding that theta search's full output +(``multilayer_theta_search_summary.csv``, ``*_all_orders.csv``, ``theta_scans/``, +profile and stack plots, ``checkpoints/``) plus a ``search_parameters.json``. + +``--eval`` runs no new solves: it re-reads the per-run summary CSVs already on +disk (:meth:`grax.MultilayerGratingDesigner.evaluate_survey`), rebuilds +``survey.csv`` and regenerates every plot. Use it after a completed run to +iterate on the analysis. + +The executable body is guarded because a later step spawns worker processes that +re-import the example package by path. +""" + +from __future__ import annotations + +import argparse + +from rub4c_design_parameters import CONFIG + +from grax import MultilayerGratingDesigner + + +def main() -> None: + """Run (or, with ``--eval``, re-evaluate) the survey.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--eval", + action="store_true", + help="Re-evaluate from existing results/ (no new solves): " + "rebuild survey.csv and every plot.", + ) + args = parser.parse_args() + + designer = MultilayerGratingDesigner(CONFIG) + result = designer.evaluate_survey() if args.eval else designer.run_survey() + + print(f"Survey CSV: {result.combined_csv_path}") + print(f"Optimal-blaze curve: {result.plot_path}") + print(f"Max-efficiency curve: {result.efficiency_plot_path}") + print(f"Efficiency heatmap: {result.heatmap_plot_path}") + print( + f"Runs: {len(result.run_dirs)} theta-search folders in " + f"{len(result.period_dirs)} period folders under {result.runs_dir}" + ) + for d_nm, blaze_opt, eff in zip( + result.d_values_nm, result.optimal_blaze_deg, result.optimal_blaze_efficiency + ): + print(f" d {d_nm:.3f} nm -> optimal blaze {blaze_opt:.3f} deg (efficiency {eff:.4g})") + + +if __name__ == "__main__": + main() diff --git a/examples/simulation/multilayer_grating_design/1_run_energy_scan.py b/examples/simulation/multilayer_grating_design/1_run_energy_scan.py new file mode 100644 index 0000000..46bf6d9 --- /dev/null +++ b/examples/simulation/multilayer_grating_design/1_run_energy_scan.py @@ -0,0 +1,143 @@ +"""Step 2: sweep chosen (d, blaze) designs over energy. + +Runs :meth:`grax.MultilayerGratingDesigner.run_energy_scan` for one or more +``(d_spacing_nm, blaze_angle_deg)`` designs read from ``results/survey/survey.csv``: + +* default -- one scan per d-spacing, at that d's optimal blaze angle; +* ``--best`` -- only the single ``(d, blaze)`` with the highest survey efficiency; +* ``--pairs`` -- exactly the designs you list. + +Besides the standard sweep artifacts under ``results/energy_scan/``, each +design gets a plot in the shared ``results/plots/`` folder (alongside +``0_run_survey.py``'s headline plots): selected efficiency versus energy, +titled with the coating (``CONFIG.coating_label``, "Ru/B4C" here), the +diffraction order and this design's d-spacing and blaze angle, named +``efficiency_vs_energy__order_dnm_blazedeg.png`` -- that +is the plot to open. Scanning two or more designs also writes an overlay +comparison, ``efficiency_vs_energy_comparison__order.png``, with +every curve on one axis. + +``--eval`` runs no new solves: it re-reads the existing +``results/energy_scan/*/multilayer_theta_search_summary.csv`` files +(:meth:`grax.MultilayerGratingDesigner.evaluate_energy_scan`) for the selected +designs, or for every design folder found when no selector is given; it also +regenerates the titled plot in ``results/plots/``. + +Examples:: + + python 1_run_energy_scan.py + python 1_run_energy_scan.py --best + python 1_run_energy_scan.py --pairs "3.0,1.1; 4.5,0.9" + python 1_run_energy_scan.py --eval + +The executable body is guarded because the theta-search sweep spawns worker +processes that re-import this file by path. +""" + +from __future__ import annotations + +import argparse + +import pandas as pd +from rub4c_design_parameters import CONFIG + +from grax import MultilayerGratingDesigner + + +def _parse_pairs(text: str) -> list[tuple[float, float]]: + """Parse ``"d,blaze; d,blaze"`` into ``(d, blaze)`` float pairs.""" + + pairs: list[tuple[float, float]] = [] + for chunk in text.replace("\n", ";").split(";"): + chunk = chunk.strip() + if not chunk: + continue + d_text, blaze_text = chunk.split(",") + pairs.append((float(d_text), float(blaze_text))) + if not pairs: + raise ValueError(f"no (d, blaze) pairs parsed from {text!r}") + return pairs + + +def _survey_table() -> pd.DataFrame: + """Return the usable rows of ``results/survey/survey.csv``.""" + + csv_path = CONFIG.survey_dir / "survey.csv" + if not csv_path.is_file(): + raise SystemExit( + f"{csv_path} not found -- run 0_run_survey.py first or pass --pairs." + ) + table = pd.read_csv(csv_path).dropna(subset=["peak_efficiency"]) + if table.empty: + raise SystemExit("survey.csv has no usable rows; pass --pairs explicitly.") + return table + + +def _optimal_pairs_from_survey() -> list[tuple[float, float]]: + """Return one ``(d, optimal_blaze)`` pair per d-spacing from the survey CSV.""" + + table = _survey_table() + best = table.loc[table.groupby("d_nm")["peak_efficiency"].idxmax()] + return [(float(row.d_nm), float(row.blaze_deg)) for row in best.itertuples()] + + +def _best_pair_from_survey() -> list[tuple[float, float]]: + """Return only the single ``(d, blaze)`` with the highest survey efficiency.""" + + row = _survey_table().loc[lambda frame: frame["peak_efficiency"].idxmax()] + print( + f"Best survey design: d = {float(row.d_nm):.3f} nm, " + f"blaze = {float(row.blaze_deg):.3f} deg (efficiency {float(row.peak_efficiency):.4g})" + ) + return [(float(row.d_nm), float(row.blaze_deg))] + + +def main() -> None: + """Parse the design selector and run the energy scan.""" + + parser = argparse.ArgumentParser(description=__doc__) + selector = parser.add_mutually_exclusive_group() + selector.add_argument( + "--pairs", + default=None, + help='Semicolon-separated "d_nm,blaze_deg" designs to scan.', + ) + selector.add_argument( + "--best", + action="store_true", + help="Scan only the single (d, blaze) with the highest survey efficiency.", + ) + parser.add_argument( + "--eval", + action="store_true", + help="Re-read existing energy_scan/ results (no new solves). With no " + "selector, evaluates every design folder found.", + ) + args = parser.parse_args() + + designer = MultilayerGratingDesigner(CONFIG) + if args.eval and not (args.pairs or args.best): + results = designer.evaluate_energy_scan() + else: + if args.pairs: + pairs = _parse_pairs(args.pairs) + elif args.best: + pairs = _best_pair_from_survey() + else: + pairs = _optimal_pairs_from_survey() + results = ( + designer.evaluate_energy_scan(pairs) + if args.eval + else designer.run_energy_scan(pairs) + ) + + for scan in results: + print( + f"d = {scan.d_spacing_nm:.3f} nm, blaze = {scan.blaze_angle_deg:.3f} deg " + f"-> {scan.summary_csv_path}" + ) + print(f" plot: {scan.titled_plot_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/simulation/multilayer_grating_design/rub4c_design_parameters.py b/examples/simulation/multilayer_grating_design/rub4c_design_parameters.py new file mode 100644 index 0000000..3d83140 --- /dev/null +++ b/examples/simulation/multilayer_grating_design/rub4c_design_parameters.py @@ -0,0 +1,202 @@ +"""Ru/B4C configuration for the multilayer-grating design example. + +Edit the values here, then run ``0_run_survey.py`` followed by +``1_run_energy_scan.py`` (or ``run_all.sh``). The scripts never rewrite this +file. + +Workflow: + +* ``0_run_survey.py`` scans a 2-D grid of bilayer d-spacing x blaze angle. For + every pair it builds the multilayer-coated blazed grating and runs graxPy's + single-energy multilayer theta search at ``target_energy_ev`` -- the search + scans the incident angle and returns the angle that maximizes the + selected-order efficiency. It writes ``survey/survey.csv``, three headline + plots (``plot/optimal_blaze_vs_d_spacing.png``, + ``plot/max_efficiency_vs_d_spacing.png``, and + ``plot/efficiency_heatmap_d_vs_blaze.png`` -- the full ``(d, blaze)`` + efficiency map with the optimal-blaze ridge overlaid), and a per-run folder + tree under ``survey/runs/``: ``dnm/overlay.png`` per period, and + ``dnm/blazedeg/`` per run with that theta search's full output + (summary CSV, all-orders CSV, ``theta_scans/``, profile/stack plots, + checkpoint) plus a ``search_parameters.json``. ``--eval`` re-derives all of + this from the runs already on disk without solving anything. +* ``1_run_energy_scan.py`` sweeps ``(d, blaze)`` designs from the survey over + ``[energy_scan_min_ev, energy_scan_max_ev]``: the optimal blaze per d-spacing + by default, only the single best with ``--best``, or exactly ``--pairs + "d,blaze; ..."``. ``--eval`` re-reads existing ``energy_scan/`` results + instead of re-solving. + +``gamma`` is held fixed; the anti-blaze angle defaults to 0 (a plain sawtooth). + +The config below is split into three sections, marked ``# ===`` below, matching +:class:`grax.MultilayerDesignConfig`'s own grouping -- check that comment before +moving a value, since editing the wrong section is a common way to get "I +changed X but nothing happened" (e.g. ``rough_scan_points`` inside +``SURVEY_SCAN`` only affects ``0_run_survey.py``; ``1_run_energy_scan.py`` reads +``ENERGY_SCAN_SCAN`` instead): + +1. ``SHARED`` -- grating geometry, materials, solver/runtime. Used by both + scripts. +2. ``SURVEY`` -- read only by ``0_run_survey.py`` (grid, ``SURVEY_SCAN`` + settings, ``on_error``). +3. ``ENERGY_SCAN`` -- read only by ``1_run_energy_scan.py`` (energy range, + ``ENERGY_SCAN_SCAN`` settings, worker/progress/tracking controls). + +For a fast smoke run: shrink ``d_points``/``blaze_points`` to 2 in ``SURVEY``, +``rough_scan_points``/``fine_scan_points`` to ~15 and the Fourier orders to +3/5/7 in *both* ``SURVEY_SCAN`` and ``ENERGY_SCAN_SCAN``, the rough/fine/final +resolutions to 2.0/1.0/1.0 nm, and ``energy_scan_points`` to 3 in +``ENERGY_SCAN``. +""" + +from __future__ import annotations + +from pathlib import Path + +from grax import MultilayerDesignConfig, ThetaSearchScanSettings + +# ========================================================================== # +# SHARED -- grating geometry, materials, solver/runtime. # +# Read by both 0_run_survey.py and 1_run_energy_scan.py. # +# ========================================================================== # +OUTPUT_DIR = Path(__file__).resolve().parent / "results" + +GRATING_DENSITY_LPERMM = 2400.0 +DIFFRACTION_ORDER = 2 +MULTILAYER_BRAGG_ORDER = 1 + +# Ru/B4C on a silicon substrate; B4C is modelled with the carbon table, so +# COATING_LABEL overrides what plot titles call it -- otherwise they would say +# "Ru/C", the table's element name, not the real compound. +MATERIAL_A = ("Ru", 12.1) +MATERIAL_B = ("C", 2.52) +COATING_LABEL = "Ru/B4C" +SUBSTRATE_MATERIAL = ("Si", 2.33) +N_BILAYERS = 40 +GAMMA = 0.5 +ANTI_BLAZE_ANGLE_DEG = 0.0 + +X_RESOLUTION_NM = 0.5 +Z_RESOLUTION_NM = 0.5 + +BACKEND = "numba" +SOLVER = "neviere" +POLARIZATION = "p" +CHECKPOINT = True +RESUME = True +SAVE_PROFILE_PLOT = True +SAVE_STACK_PLOT = True + +# ========================================================================== # +# SURVEY -- read only by 0_run_survey.py (run_survey / evaluate_survey). # +# ========================================================================== # +TARGET_ENERGY_EV = 9000.0 + +# d-spacing x blaze-angle grid. +D_MIN_NM = 2.0 +D_MAX_NM = 8.0 +D_POINTS = 50 # 10 +BLAZE_MIN_DEG = 0.6 +BLAZE_MAX_DEG = 2.0 +BLAZE_POINTS = 28 + +# Theta-search scan settings for the survey: one single-energy search per +# (d, blaze) grid cell (many cheap searches -- 1400 with the grid above), so +# this is usually the one to keep fast. The rough half-width is wide because +# the scan is seeded from the multilayer Bragg estimate, which overshoots the +# true grating optimum for shallow inside orders; the window has to reach +# well below it across the whole d grid. +SURVEY_SCAN = ThetaSearchScanSettings( + rough_scan_half_width_deg=1, + rough_scan_points=61, + rough_fourier_orders=5, + rough_x_resolution_nm=1.0, + rough_z_resolution_nm=1.0, + fine_scan_half_width_deg=0.2, + fine_scan_points=81, + fine_fourier_orders=15, + fine_x_resolution_nm=0.5, + fine_z_resolution_nm=0.5, + final_fourier_orders=25, + final_x_resolution_nm=0.2, + final_z_resolution_nm=0.2, + precise_peak_selection_mode="max", +) + +ON_ERROR = "continue" + +# ========================================================================== # +# ENERGY_SCAN -- read only by 1_run_energy_scan.py # +# (run_energy_scan / evaluate_energy_scan). # +# ========================================================================== # +ENERGY_SCAN_MIN_EV = 3000.0 +ENERGY_SCAN_MAX_EV = 12000.0 +ENERGY_SCAN_POINTS = 1000 + +# Theta-search scan settings for the energy scan: one search per energy, for +# each chosen (d, blaze) design -- few designs, but often many energies each, +# so this is usually the one to make more precise. +ENERGY_SCAN_SCAN = ThetaSearchScanSettings( + rough_scan_half_width_deg=1, + rough_scan_points=61, + rough_fourier_orders=5, + rough_x_resolution_nm=1.0, + rough_z_resolution_nm=1.0, + fine_scan_half_width_deg=0.2, + fine_scan_points=81, + fine_fourier_orders=15, + fine_x_resolution_nm=0.5, + fine_z_resolution_nm=0.5, + final_fourier_orders=25, + final_x_resolution_nm=0.2, + final_z_resolution_nm=0.2, + precise_peak_selection_mode="max", +) + +MAX_WORKERS = "auto" +SHOW_PROGRESS = True +THETA_TRACKING_MODE = "auto" +MAX_TRACKING_ENERGY_STEP_EV = None + +CONFIG = MultilayerDesignConfig( + output_dir=OUTPUT_DIR, + # -- Shared -- + grating_density_lpermm=GRATING_DENSITY_LPERMM, + diffraction_order=DIFFRACTION_ORDER, + multilayer_bragg_order=MULTILAYER_BRAGG_ORDER, + material_a=MATERIAL_A, + material_b=MATERIAL_B, + coating_label=COATING_LABEL, + substrate_material=SUBSTRATE_MATERIAL, + n_bilayers=N_BILAYERS, + gamma=GAMMA, + anti_blaze_angle_deg=ANTI_BLAZE_ANGLE_DEG, + x_resolution_nm=X_RESOLUTION_NM, + z_resolution_nm=Z_RESOLUTION_NM, + backend=BACKEND, + solver=SOLVER, + polarization=POLARIZATION, + checkpoint=CHECKPOINT, + resume=RESUME, + save_profile_plot=SAVE_PROFILE_PLOT, + save_stack_plot=SAVE_STACK_PLOT, + # -- Survey only -- + target_energy_ev=TARGET_ENERGY_EV, + d_min_nm=D_MIN_NM, + d_max_nm=D_MAX_NM, + d_points=D_POINTS, + blaze_min_deg=BLAZE_MIN_DEG, + blaze_max_deg=BLAZE_MAX_DEG, + blaze_points=BLAZE_POINTS, + survey_scan_settings=SURVEY_SCAN, + on_error=ON_ERROR, + # -- Energy scan only -- + energy_scan_min_ev=ENERGY_SCAN_MIN_EV, + energy_scan_max_ev=ENERGY_SCAN_MAX_EV, + energy_scan_points=ENERGY_SCAN_POINTS, + energy_scan_settings=ENERGY_SCAN_SCAN, + max_workers=MAX_WORKERS, + show_progress=SHOW_PROGRESS, + theta_tracking_mode=THETA_TRACKING_MODE, + max_tracking_energy_step_ev=MAX_TRACKING_ENERGY_STEP_EV, +) diff --git a/examples/simulation/multilayer_grating_design/run_all.sh b/examples/simulation/multilayer_grating_design/run_all.sh new file mode 100755 index 0000000..705ddce --- /dev/null +++ b/examples/simulation/multilayer_grating_design/run_all.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Run the multilayer-grating design survey, then the energy scan of its ridge. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python}" + +echo "==> Running 0_run_survey.py" +"${PYTHON_BIN}" "${SCRIPT_DIR}/0_run_survey.py" + +echo "==> Running 1_run_energy_scan.py (survey ridge)" +"${PYTHON_BIN}" "${SCRIPT_DIR}/1_run_energy_scan.py" diff --git a/examples/simulation/multilayer_optimization_rub4c/0_ru_b4c_d_spacing_study.py b/examples/simulation/multilayer_optimization_rub4c/0_ru_b4c_d_spacing_study.py deleted file mode 100644 index 3ab6909..0000000 --- a/examples/simulation/multilayer_optimization_rub4c/0_ru_b4c_d_spacing_study.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Stage 0: derive and scan the Ru/B4C multilayer d-spacing. - -Runs :func:`grax.run_d_spacing_study` with the shared ``ru_b4c_parameters.CONFIG`` -and writes ``results/optimization_state.json`` with ``d_suggested_nm`` for the -next two stages. -""" - -from __future__ import annotations - -from ru_b4c_parameters import CONFIG - -from grax import run_d_spacing_study - -if __name__ == "__main__": - run_d_spacing_study(CONFIG) diff --git a/examples/simulation/multilayer_optimization_rub4c/1_ru_b4c_gamma_study.py b/examples/simulation/multilayer_optimization_rub4c/1_ru_b4c_gamma_study.py deleted file mode 100644 index 467bc29..0000000 --- a/examples/simulation/multilayer_optimization_rub4c/1_ru_b4c_gamma_study.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Stage 1: scan the Ru/B4C bilayer thickness ratio at the selected d-spacing. - -Runs :func:`grax.run_gamma_study`. With ``D_SPACING_NM = "auto"`` this reads -``d_suggested_nm`` from ``results/optimization_state.json`` written by stage 0. -The suggested gamma is recorded in the state file for traceability; it is not -applied automatically -- copy it into ``ru_b4c_parameters.CONFIG`` if you want -stage 2 to use it. -""" - -from __future__ import annotations - -from ru_b4c_parameters import CONFIG - -from grax import run_gamma_study - -if __name__ == "__main__": - run_gamma_study(CONFIG) diff --git a/examples/simulation/multilayer_optimization_rub4c/2_ru_b4c_blaze_study.py b/examples/simulation/multilayer_optimization_rub4c/2_ru_b4c_blaze_study.py deleted file mode 100644 index 48450c8..0000000 --- a/examples/simulation/multilayer_optimization_rub4c/2_ru_b4c_blaze_study.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Stage 2: scan the blaze angle of the multilayer-coated Ru/B4C grating. - -Runs :func:`grax.run_blaze_study`, which builds the multilayer-coated blazed -grating and runs graxPy's internal theta search per energy for each blaze angle, -then records ``blaze_suggested_deg`` in ``results/optimization_state.json``. - -``--solver`` overrides ``ru_b4c_parameters.CONFIG.solver`` for this run. The -executable body is guarded because the theta-search sweep spawns worker -processes that re-import this file by path. -""" - -from __future__ import annotations - -import argparse -import dataclasses - -from ru_b4c_parameters import CONFIG - -from grax import run_blaze_study - - -def main() -> None: - """Parse ``--solver`` and run the blaze study.""" - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--solver", - choices=("rcwa", "neviere"), - default=CONFIG.solver, - help="Electromagnetic solver to run. Both compute every diffraction order; " - "they differ only in how each layer is crossed in z.", - ) - args = parser.parse_args() - run_blaze_study(dataclasses.replace(CONFIG, solver=args.solver)) - - -if __name__ == "__main__": - main() diff --git a/examples/simulation/multilayer_optimization_rub4c/ru_b4c_parameters.py b/examples/simulation/multilayer_optimization_rub4c/ru_b4c_parameters.py deleted file mode 100644 index 33689bf..0000000 --- a/examples/simulation/multilayer_optimization_rub4c/ru_b4c_parameters.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Ru/B4C configuration for the three-stage multilayer optimization example. - -Edit the values here, then run ``0_ru_b4c_d_spacing_study.py``, -``1_ru_b4c_gamma_study.py`` and ``2_ru_b4c_blaze_study.py`` in order (or -``run_all.sh``). ``D_SPACING_NM`` is left as ``"auto"`` so stage 0's geometry -suggestion flows into stages 1 and 2 through ``results/optimization_state.json``; -set it (and ``GAMMA``) to a number to pin it instead. The scripts never rewrite -this file. - -For a fast smoke run, shrink the grids: set ``d_spacing_points`` to ``5``, -``blaze_angle_points`` to ``2``, ``blaze_energy_points`` to ``3`` and widen the -energy steps. -""" - -from __future__ import annotations - -from pathlib import Path - -from grax import MultilayerOptimizationConfig - -CONFIG = MultilayerOptimizationConfig( - output_dir=Path(__file__).resolve().parent / "results", - # Selected values ("auto" consumes the previous stage's suggestion). - d_spacing_nm="auto", - gamma=0.5, - blaze_angle_deg=1.1, - # Ru/B4C on a silicon substrate; B4C is modelled with the carbon table. - material_a=("Ru", 12.1), - material_b=("C", 2.52), - substrate_material=("Si", 2.33), - n_bilayers=40, - # Target and grating geometry. - target_energy_ev=9000.0, - grating_density_lpermm=2400.0, - diffraction_order=2, - cff=2.25, - multilayer_bragg_order=1, - # Per-stage energy grids. - d_spacing_energy_min_ev=500.0, - d_spacing_energy_max_ev=12000.0, - d_spacing_energy_step_ev=100.0, - gamma_energy_min_ev=500.0, - gamma_energy_max_ev=12000.0, - gamma_energy_step_ev=100.0, - blaze_energy_min_ev=3000.0, - blaze_energy_max_ev=12000.0, - blaze_energy_points=15, - # D-spacing scan. - bragg_angle_min_deg=0.5, - bragg_angle_max_deg=2.0, - d_spacing_relative_range=0.25, - d_spacing_min_practical_nm=2.0, - d_spacing_max_practical_nm=8.0, - d_spacing_points=21, - # Gamma scan. - gamma_min=0.3, - gamma_max=0.8, - gamma_step=0.1, - # Blaze scan. - blaze_angle_half_range_deg=0.3, - blaze_angle_points=4, - anti_blaze_angle_deg=0.0, - # XRT reflectivity settings. - xrt_window_deg=0.2, - xrt_angle_points=2001, - xrt_min_angle_deg=0.0, - # graxPy theta-search settings for stage 2. - grax_x_resolution_nm=0.5, - grax_z_resolution_nm=0.5, - rough_scan_half_width_deg=0.5, - rough_scan_points=61, - rough_fourier_orders=5, - fine_scan_half_width_deg=0.2, - fine_scan_points=81, - fine_fourier_orders=15, - final_fourier_orders=25, - final_x_resolution_nm=0.2, - final_z_resolution_nm=0.2, - backend="numba", - solver="neviere", - polarization="p", - # Runtime controls. - max_workers="auto", - on_error="fail_fast", - resume=True, - theta_tracking_mode="auto", -) diff --git a/examples/simulation/multilayer_optimization_rub4c/run_all.sh b/examples/simulation/multilayer_optimization_rub4c/run_all.sh deleted file mode 100755 index 5678481..0000000 --- a/examples/simulation/multilayer_optimization_rub4c/run_all.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env bash -# Run the three Ru/B4C multilayer optimization stages in order. -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PYTHON_BIN="${PYTHON_BIN:-python}" -SOLVER="${SOLVER:-neviere}" - -echo "==> Running 0_ru_b4c_d_spacing_study.py" -"${PYTHON_BIN}" "${SCRIPT_DIR}/0_ru_b4c_d_spacing_study.py" - -echo "==> Running 1_ru_b4c_gamma_study.py" -"${PYTHON_BIN}" "${SCRIPT_DIR}/1_ru_b4c_gamma_study.py" - -echo "==> Running 2_ru_b4c_blaze_study.py (--solver ${SOLVER})" -"${PYTHON_BIN}" "${SCRIPT_DIR}/2_ru_b4c_blaze_study.py" --solver "${SOLVER}" diff --git a/src/grax/__init__.py b/src/grax/__init__.py index 87d4eaf..116bcf2 100644 --- a/src/grax/__init__.py +++ b/src/grax/__init__.py @@ -62,14 +62,13 @@ write_all_orders_csv, ) from .multilayer_reflectivity import MultilayerReflectivity -from .multilayer_optimization import ( - BlazeStudyResult, - DSpacingStudyResult, - GammaStudyResult, - MultilayerOptimizationConfig, - run_blaze_study, - run_d_spacing_study, - run_gamma_study, +from .multilayer_design import ( + EnergyScanResult, + MultilayerDesignConfig, + MultilayerGratingDesigner, + StageProgress, + SurveyResult, + ThetaSearchScanSettings, ) from .slag import SlagConfig, default_example_slag_config, run_example_slag, simulate_single_energy @@ -79,16 +78,15 @@ "BaseGrating", "BaseStack", "BatchSimulationRunner", - "BlazeStudyResult", "BlazedGrating", "CaseExecutionResult", "CustomStack", - "DSpacingStudyResult", - "GammaStudyResult", + "EnergyScanResult", "LayerSpec", "MaterialSpec", "LaminarGrating", - "MultilayerOptimizationConfig", + "MultilayerDesignConfig", + "MultilayerGratingDesigner", "MultilayerReflectivity", "MultilayerThetaSearchSweepResult", "MultilayerStack", @@ -101,7 +99,10 @@ "SingleLayerStack", "SingleSimulationResult", "SlagConfig", + "StageProgress", + "SurveyResult", "ThetaSearchDiagnostics", + "ThetaSearchScanSettings", "assemble_custom_stack", "available_material_symbols", "build_multilayer_stack", @@ -125,10 +126,7 @@ "res1", "res2", "res2_dm", - "run_blaze_study", - "run_d_spacing_study", "run_example_slag", - "run_gamma_study", "run_parameter_study", "run_multilayer_theta_search", "run_multilayer_theta_search_sweep", diff --git a/src/grax/multilayer_design.py b/src/grax/multilayer_design.py new file mode 100644 index 0000000..595e034 --- /dev/null +++ b/src/grax/multilayer_design.py @@ -0,0 +1,1688 @@ +"""Multilayer-grating design workflow: a d-spacing / blaze-angle survey plus energy scans. + +The workflow sizes a periodic multilayer coating for a blazed grating working in a +chosen diffraction order, in two steps driven by a single +:class:`MultilayerDesignConfig`: + +1. :meth:`MultilayerGratingDesigner.run_survey` -- for every ``(d_spacing, blaze + angle)`` pair on a 2-D grid, build the multilayer-coated blazed grating and run + the multilayer theta search at the optimization energy through + :func:`grax.run_multilayer_theta_search_sweep`, with its ``output_dir`` set to + that pair's own folder. Each run therefore saves the full standard artifact + bundle under ``survey/runs/dnm/blazedeg/`` -- the selected-result summary + CSV, the all-orders CSV, the incidence-angle scan CSV/PNG, the profile and + stack plots, a checkpoint, and a ``search_parameters.json`` recording what was + fed in so the search can be reviewed and tuned. From every pair's selected + result the survey reports, per d-spacing, the blaze angle with the highest + efficiency (``argmax_blaze``); the headline artifacts are two curves -- optimal + blaze angle versus d-spacing, and max efficiency versus d-spacing. Each + ``dnm/`` folder also gets an ``overlay.png`` of its blaze angles' scans with + the chosen one highlighted. + +2. :meth:`MultilayerGratingDesigner.run_energy_scan` -- for one or more + ``(d_spacing, blaze angle)`` pairs picked off the survey, sweep a fixed coated + grating over an explicit energy range with + :func:`grax.run_multilayer_theta_search_sweep`, yielding efficiency versus + energy and ``theta*(E)`` for each design. + +``gamma`` is held fixed and the anti-blaze angle defaults to ``0`` (a plain +sawtooth); both are later optimizations. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import threading +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd + +from .gratings import BlazedGrating +from .materials import MaterialSpec +from .simulation import ( + estimate_multilayer_bragg_angle_deg, + run_multilayer_theta_search_sweep, +) +from .stacks import MultilayerStack + +__all__ = [ + "EnergyScanResult", + "MultilayerDesignConfig", + "MultilayerGratingDesigner", + "StageProgress", + "SurveyResult", + "ThetaSearchScanSettings", +] + +HC_EV_NM = 1239.841984 + + +def _write_csv_atomic(frame: pd.DataFrame, path: Path) -> None: + """Write ``frame`` to ``path`` so a concurrent reader never sees a half file. + + The survey table is rewritten after every cell while the web app polls it for + the live plots; a plain ``to_csv`` truncates first, so a read landing in that + window gets a partial or empty file. + """ + + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f"{path.name}.tmp") + frame.to_csv(temporary, index=False) + os.replace(temporary, path) + + +class _CellAbortedError(RuntimeError): + """Raised when a ``stop_event`` killed a survey cell's theta search mid-search.""" + + +@dataclass(frozen=True) +class StageProgress: + """Progress report emitted before each scanned item by a workflow step. + + Attributes: + stage: ``"survey"`` or ``"energy_scan"``. + completed: Items finished so far. + total: Total items in the scan. + current_label: Human-readable label of the item about to run, or + ``"done"`` on the final call once the scan has finished. + """ + + stage: str + completed: int + total: int + current_label: str + + +@dataclass(frozen=True) +class ThetaSearchScanSettings: + """Rough/fine/final incidence-angle scan settings for one theta search. + + :class:`MultilayerDesignConfig` carries two independent instances -- + ``survey_scan_settings`` and ``energy_scan_settings`` -- so the (d, blaze) + survey (many single-energy searches, one per grid cell) and the per-design + energy scan (few designs, many energies each) can be tuned for speed or + accuracy independently. + + Attributes: + rough_scan_half_width_deg: Half-width of the coarse first pass around + the seed angle. + rough_scan_points: Number of angles sampled in the rough pass. + rough_fourier_orders: Fourier truncation order used during the rough pass. + rough_x_resolution_nm / rough_z_resolution_nm: Grating mesh override + during the rough pass. + fine_scan_half_width_deg: Half-width of the precise pass around the + rough-pass maximum. + fine_scan_points: Number of angles sampled in the precise pass. + fine_fourier_orders: Fourier truncation order used during the precise pass. + fine_x_resolution_nm / fine_z_resolution_nm: Grating mesh override + during the precise pass. + final_fourier_orders: Fourier truncation order for the final solve at + the selected angle. + final_x_resolution_nm / final_z_resolution_nm: Grating mesh override + for the final solve. + precise_peak_selection_mode: ``"max"``, ``"gauss"`` or ``"voigt"``. + roughness_sigma_nm: Optional rms interface roughness in nm. + """ + + rough_scan_half_width_deg: float = 1.5 + rough_scan_points: int = 61 + rough_fourier_orders: int = 5 + rough_x_resolution_nm: float = 1.0 + rough_z_resolution_nm: float = 1.0 + fine_scan_half_width_deg: float = 0.2 + fine_scan_points: int = 81 + fine_fourier_orders: int = 15 + fine_x_resolution_nm: float = 0.5 + fine_z_resolution_nm: float = 0.5 + final_fourier_orders: int = 25 + final_x_resolution_nm: float = 0.2 + final_z_resolution_nm: float = 0.2 + precise_peak_selection_mode: str = "max" + roughness_sigma_nm: float | None = None + + def __post_init__(self) -> None: + """Validate the peak-selection mode.""" + + if self.precise_peak_selection_mode not in {"max", "gauss", "voigt"}: + raise ValueError( + "precise_peak_selection_mode must be 'max', 'gauss' or 'voigt', got " + f"{self.precise_peak_selection_mode!r}" + ) + + +@dataclass(frozen=True) +class MultilayerDesignConfig: + """Every knob for the multilayer-grating design survey and energy scans. + + Fields fall into three groups, matching the section comments below and in + the example ``rub4c_design_parameters.py``: + + * **Shared** -- grating geometry, materials, solver/runtime. Read by both + :meth:`MultilayerGratingDesigner.run_survey` and + :meth:`MultilayerGratingDesigner.run_energy_scan`. + * **Survey only** -- read only by ``run_survey`` / + :meth:`MultilayerGratingDesigner.evaluate_survey`. + * **Energy scan only** -- read only by ``run_energy_scan`` / + :meth:`MultilayerGratingDesigner.evaluate_energy_scan`. + + ``material_a`` / ``material_b`` / ``substrate_material`` are + ``(name, density_g_cm3)`` pairs (a :class:`grax.MaterialSpec` is also + accepted). + + Attributes: + output_dir: Root directory for all generated artifacts. + grating_density_lpermm: Groove density in lines/mm. + diffraction_order: Grating diffraction order to optimize. + multilayer_bragg_order: Multilayer Bragg order (distinct from the grating + diffraction order). + material_a: Incident-side / top bilayer material. + material_b: Second bilayer material. + substrate_material: Grating substrate material. + n_bilayers: Number of bilayer periods. + gamma: Bilayer thickness ratio (``material_a`` fraction), 0 < gamma < 1; + held fixed for both steps. + anti_blaze_angle_deg: Anti-blaze facet angle; ``0`` (default) builds a + plain sawtooth. + coating_label: Display label for the coating in plot titles, e.g. + ``"Ru/B4C"``. Defaults to ``"/"`` + when ``None`` -- set this explicitly whenever a material is modelled + with a stand-in optical-constants table (for example B4C modelled + with the carbon table), so titles show the real compound. + x_resolution_nm / z_resolution_nm: Grating profile discretization. + backend / solver / polarization: Solver selection. + checkpoint / resume: Checkpointing for both steps. + save_profile_plot / save_stack_plot: Per-run diagnostic plots. + target_energy_ev: Photon energy the survey optimizes at. + d_min_nm / d_max_nm / d_points: Inclusive d-spacing grid in nm. + blaze_min_deg / blaze_max_deg / blaze_points: Inclusive blaze-angle grid. + survey_scan_settings: :class:`ThetaSearchScanSettings` for the survey -- + one single-energy search per ``(d, blaze)`` grid cell. + on_error: ``"continue"`` records a NaN row for a failing survey cell and + keeps going; ``"fail_fast"`` re-raises. + energy_scan_min_ev / energy_scan_max_ev / energy_scan_points: Explicit + energy grid. + energy_scan_settings: :class:`ThetaSearchScanSettings` for the energy + scan -- one search per energy, for each chosen design. + max_workers: Batch worker count (``"auto"`` calibrates). + show_progress: Whether a progress bar is shown. + theta_tracking_mode / max_tracking_energy_step_ev: How the seed angle + for each energy is chosen (the survey always seeds from the Bragg + estimate instead). + """ + + output_dir: Path + + # ================================================================== # + # Shared -- grating geometry, materials, solver/runtime. # + # Read by both run_survey and run_energy_scan. # + # ================================================================== # + grating_density_lpermm: float = 2400.0 + diffraction_order: int = 2 + multilayer_bragg_order: int = 1 + + material_a: Any = ("Ru", 12.1) + material_b: Any = ("C", 2.52) + substrate_material: Any = ("Si", 2.33) + n_bilayers: int = 40 + gamma: float = 0.5 + anti_blaze_angle_deg: float = 0.0 + coating_label: str | None = None + + x_resolution_nm: float = 0.5 + z_resolution_nm: float = 0.5 + + backend: str = "numba" + solver: str = "neviere" + polarization: str = "p" + checkpoint: bool = True + resume: bool = True + save_profile_plot: bool = True + save_stack_plot: bool = True + + # ================================================================== # + # Survey only -- read by run_survey / evaluate_survey. # + # ================================================================== # + target_energy_ev: float = 9000.0 + + d_min_nm: float = 1.5 + d_max_nm: float = 6.0 + d_points: int = 4 + blaze_min_deg: float = 0.6 + blaze_max_deg: float = 2.0 + blaze_points: int = 6 + + # The rough half-width is wide because the scan is seeded from the + # multilayer Bragg estimate, which overshoots the true grating optimum for + # shallow inside orders; the window has to reach well below it across the + # whole d grid. + survey_scan_settings: ThetaSearchScanSettings = field( + default_factory=ThetaSearchScanSettings + ) + on_error: str = "continue" + + # ================================================================== # + # Energy scan only -- read by run_energy_scan / evaluate_energy_scan. # + # ================================================================== # + energy_scan_min_ev: float = 3000.0 + energy_scan_max_ev: float = 12000.0 + energy_scan_points: int = 15 + energy_scan_settings: ThetaSearchScanSettings = field( + default_factory=ThetaSearchScanSettings + ) + max_workers: int | str | None = "auto" + show_progress: bool = True + theta_tracking_mode: str = "auto" + max_tracking_energy_step_ev: float | None = None + + def __post_init__(self) -> None: + """Validate cross-field constraints.""" + + object.__setattr__(self, "output_dir", Path(self.output_dir)) + if not (0.0 < float(self.gamma) < 1.0): + raise ValueError(f"gamma must satisfy 0 < gamma < 1, got {self.gamma!r}") + if self.target_energy_ev <= 0.0: + raise ValueError("target_energy_ev must be positive") + if self.n_bilayers < 1: + raise ValueError("n_bilayers must be at least 1") + if self.multilayer_bragg_order < 1: + raise ValueError("multilayer_bragg_order must be at least 1") + if not (0.0 < float(self.d_min_nm) < float(self.d_max_nm)): + raise ValueError( + f"d grid must satisfy 0 < d_min_nm < d_max_nm, got " + f"{self.d_min_nm!r}, {self.d_max_nm!r}" + ) + if self.d_points < 2: + raise ValueError("d_points must be at least 2") + if not (0.0 < float(self.blaze_min_deg) <= float(self.blaze_max_deg)): + raise ValueError( + f"blaze grid must satisfy 0 < blaze_min_deg <= blaze_max_deg, got " + f"{self.blaze_min_deg!r}, {self.blaze_max_deg!r}" + ) + if self.blaze_points < 1: + raise ValueError("blaze_points must be at least 1") + if self.anti_blaze_angle_deg < 0.0: + raise ValueError("anti_blaze_angle_deg must be non-negative") + if not (0.0 < float(self.energy_scan_min_ev) < float(self.energy_scan_max_ev)): + raise ValueError( + "energy scan must satisfy 0 < energy_scan_min_ev < energy_scan_max_ev" + ) + if self.energy_scan_points < 2: + raise ValueError("energy_scan_points must be at least 2") + if self.solver not in {"rcwa", "neviere"}: + raise ValueError(f"solver must be 'rcwa' or 'neviere', got {self.solver!r}") + if self.on_error not in {"continue", "fail_fast"}: + raise ValueError(f"on_error must be 'continue' or 'fail_fast', got {self.on_error!r}") + + @property + def survey_dir(self) -> Path: + """Directory for the survey CSV.""" + + return self.output_dir / "survey" + + @property + def plot_dir(self) -> Path: + """Directory for the survey headline plots and the titled energy-scan plots.""" + + return self.output_dir / "plots" + + @property + def runs_dir(self) -> Path: + """Root of the survey-run folder tree (``dnm/blazedeg/``).""" + + return self.output_dir / "survey" / "runs" + + def period_dir(self, d_spacing_nm: float) -> Path: + """Return one d-spacing's folder, holding its per-blaze run folders.""" + + return self.runs_dir / f"d{float(d_spacing_nm):.3f}nm" + + def run_dir(self, d_spacing_nm: float, blaze_angle_deg: float) -> Path: + """Return one ``(d, blaze)`` run's folder, nested under its period folder.""" + + return self.period_dir(d_spacing_nm) / f"blaze{float(blaze_angle_deg):.3f}deg" + + @property + def energy_scan_dir(self) -> Path: + """Root directory for per-design energy scans.""" + + return self.output_dir / "energy_scan" + + def d_grid_nm(self) -> np.ndarray: + """Return the inclusive d-spacing grid in nm.""" + + return np.linspace(float(self.d_min_nm), float(self.d_max_nm), int(self.d_points)) + + def blaze_grid_deg(self) -> np.ndarray: + """Return the inclusive blaze-angle grid in degrees.""" + + return np.linspace( + float(self.blaze_min_deg), float(self.blaze_max_deg), int(self.blaze_points) + ) + + def energy_grid_ev(self) -> np.ndarray: + """Return the inclusive energy grid for the energy scan.""" + + return np.linspace( + float(self.energy_scan_min_ev), + float(self.energy_scan_max_ev), + int(self.energy_scan_points), + ) + + +@dataclass +class SurveyResult: + """Outcome of :meth:`MultilayerGratingDesigner.run_survey`. + + The ``*_map`` arrays are shaped ``(len(d_values_nm), len(blaze_values_deg))``. + Cells that failed (or were never reached on an early stop) hold ``NaN``. + + Attributes: + d_values_nm: The d-spacing grid. + blaze_values_deg: The blaze-angle grid. + efficiency_map: Peak selected-order efficiency per ``(d, blaze)``. + incidence_angle_map_deg: ``theta*`` per ``(d, blaze)``. + fwhm_map_deg: Precise-scan FWHM per ``(d, blaze)``. + bragg_estimate_map_deg: Analytical Bragg-angle seed per ``(d, blaze)``. + optimal_blaze_deg: Per d-spacing, the blaze angle with the largest + efficiency (``NaN`` where the whole row failed). + optimal_blaze_efficiency: Efficiency along ``optimal_blaze_deg``. + optimal_blaze_incidence_deg: ``theta*`` along ``optimal_blaze_deg``. + combined_csv_path: Long-format per-cell table. + plot_path: The optimal blaze angle versus d-spacing curve. + efficiency_plot_path: The max efficiency versus d-spacing curve, each + point labelled with the blaze angle that achieved it. + heatmap_plot_path: 2-D heatmap of ``efficiency_map`` -- d-spacing (x) + against blaze angle (y), colour is peak efficiency (z) -- with the + ``optimal_blaze_deg`` ridge overlaid. + runs_dir: Root of the survey-run folder tree. + period_dirs: One folder per d-spacing that solved at least one cell + (``runs/dnm/``), each holding an ``overlay.png`` of every blaze + angle's theta search plus one sub-folder per blaze angle. In d-grid + order. + run_dirs: One leaf folder per solved cell + (``runs/dnm/blazedeg/``), each holding that theta search's full + :func:`grax.run_multilayer_theta_search_sweep` output + (``multilayer_theta_search_summary.csv``, ``*_all_orders.csv``, + ``theta_scans/``, profile/stack plots, ``checkpoints/``) plus a + ``search_parameters.json``. In scan order. + aborted: Whether the scan stopped early on a ``should_continue`` signal. + results: The per-cell table. + """ + + d_values_nm: np.ndarray + blaze_values_deg: np.ndarray + efficiency_map: np.ndarray + incidence_angle_map_deg: np.ndarray + fwhm_map_deg: np.ndarray + bragg_estimate_map_deg: np.ndarray + optimal_blaze_deg: np.ndarray + optimal_blaze_efficiency: np.ndarray + optimal_blaze_incidence_deg: np.ndarray + combined_csv_path: Path + plot_path: Path + efficiency_plot_path: Path + heatmap_plot_path: Path + runs_dir: Path + period_dirs: list[Path] + run_dirs: list[Path] + aborted: bool + results: pd.DataFrame = field(repr=False) + + +@dataclass +class EnergyScanResult: + """Outcome of one design's energy scan. + + Attributes: + d_spacing_nm: Bilayer period the grating was built with. + blaze_angle_deg: Blaze angle the grating was built with. + output_dir: Directory holding this design's sweep artifacts. + summary_csv_path: One selected theta/efficiency per energy. + all_orders_csv_path: All reflected orders per energy. + energy_efficiency_plot_path: Selected efficiency versus energy, with + graxPy's generic title (as written by + :func:`grax.run_multilayer_theta_search_sweep`). + titled_plot_path: The same curve, titled with the coating + (:attr:`MultilayerDesignConfig.coating_label`) and this design's + d-spacing and blaze angle -- the plot to actually look at. Lives in + :attr:`MultilayerDesignConfig.plot_dir` (shared with the survey's + headline plots) as + ``efficiency_vs_energy__order_dnm_blazedeg.png``, + since every design's plot lands in the same folder. + results: The summary table. + """ + + d_spacing_nm: float + blaze_angle_deg: float + output_dir: Path + summary_csv_path: Path + all_orders_csv_path: Path + energy_efficiency_plot_path: Path + titled_plot_path: Path + results: pd.DataFrame = field(repr=False) + + +def _material_pair(material: Any) -> tuple[str, float]: + """Return a ``(name, density_g_cm3)`` pair for a material specification.""" + + if isinstance(material, MaterialSpec): + if material.density_g_cm3 is None: + raise ValueError(f"MaterialSpec {material.name!r} needs a density for this workflow") + return str(material.name), float(material.density_g_cm3) + if isinstance(material, (tuple, list)) and len(material) == 2: + return str(material[0]), float(material[1]) + raise ValueError(f"material must be a (name, density) pair or MaterialSpec, got {material!r}") + + +def _coating_label(config: MultilayerDesignConfig) -> str: + """Return the coating's display label for plot titles. + + Uses ``config.coating_label`` when set; otherwise derives + ``"/"`` from the configured materials. + """ + + if config.coating_label: + return str(config.coating_label) + name_a, _ = _material_pair(config.material_a) + name_b, _ = _material_pair(config.material_b) + return f"{name_a}/{name_b}" + + +def _energy_scan_title( + config: MultilayerDesignConfig, d_spacing_nm: float, blaze_angle_deg: float +) -> str: + """Return the title for one design's efficiency-versus-energy plot.""" + + return ( + f"{_coating_label(config)} multilayer grating (order {config.diffraction_order}): " + f"d = {d_spacing_nm:.3f} nm, blaze = {blaze_angle_deg:.3f} deg" + ) + + +def _filename_slug(text: str) -> str: + """Return ``text`` with filesystem-unsafe characters collapsed to ``-``.""" + + slug = re.sub(r"[^A-Za-z0-9.+_-]+", "-", text.strip()) + return slug.strip("-") or "coating" + + +def _energy_scan_plot_filename( + config: MultilayerDesignConfig, d_spacing_nm: float, blaze_angle_deg: float +) -> str: + """Return the filename for one design's titled efficiency-versus-energy plot. + + All designs share :attr:`MultilayerDesignConfig.plot_dir`, so the filename + itself carries the coating, diffraction order, d-spacing and blaze angle: + ``efficiency_vs_energy__order_dnm_blazedeg.png``. + """ + + materials = _filename_slug(_coating_label(config)) + return ( + f"efficiency_vs_energy_{materials}_order{int(config.diffraction_order)}_" + f"d{d_spacing_nm:.3f}nm_blaze{blaze_angle_deg:.3f}deg.png" + ) + + +def _energy_scan_overlay_filename(config: MultilayerDesignConfig) -> str: + """Return the filename for the multi-design efficiency-versus-energy overlay.""" + + materials = _filename_slug(_coating_label(config)) + return ( + f"efficiency_vs_energy_comparison_{materials}_" + f"order{int(config.diffraction_order)}.png" + ) + + +def _emit_progress( + callback: Callable[[StageProgress], None] | None, + *, + stage: str, + completed: int, + total: int, + current_label: str, +) -> None: + """Report progress through ``callback`` when one was supplied.""" + + if callback is not None: + callback( + StageProgress( + stage=stage, completed=completed, total=total, current_label=current_label + ) + ) + + +class MultilayerGratingDesigner: + """Run the multilayer-grating design survey and per-design energy scans.""" + + def __init__(self, config: MultilayerDesignConfig) -> None: + """Store the shared configuration.""" + + self.config = config + + # ------------------------------------------------------------------ # + # Grating construction # + # ------------------------------------------------------------------ # + def _build_grating(self, d_spacing_nm: float, blaze_angle_deg: float) -> BlazedGrating: + """Build the multilayer-coated blazed grating for one ``(d, blaze)`` pair.""" + + config = self.config + material_a_spec = MaterialSpec(*_material_pair(config.material_a)) + material_b_spec = MaterialSpec(*_material_pair(config.material_b)) + substrate_spec = MaterialSpec(*_material_pair(config.substrate_material)) + stack = MultilayerStack( + substrate_material=substrate_spec, + material_a=material_a_spec, + material_b=material_b_spec, + d_period_nm=float(d_spacing_nm), + gamma=float(config.gamma), + n_bilayers=int(config.n_bilayers), + top_material=material_b_spec, + ) + anti_blaze_kwargs = ( + {} + if config.anti_blaze_angle_deg == 0.0 + else {"anti_blaze_angle_deg": float(config.anti_blaze_angle_deg)} + ) + return BlazedGrating( + period_lpermm=float(config.grating_density_lpermm), + blaze_angle_deg=float(blaze_angle_deg), + coating_stack=stack, + substrate_material=substrate_spec, + x_resolution_nm=float(config.x_resolution_nm), + z_resolution_nm=float(config.z_resolution_nm), + **anti_blaze_kwargs, + ) + + # ------------------------------------------------------------------ # + # Step 1: (d, blaze) survey # + # ------------------------------------------------------------------ # + def _survey_cell( + self, + d_spacing_nm: float, + blaze_angle_deg: float, + *, + stop_event: threading.Event | None = None, + on_worker_pids_changed: Callable[[set[int]], None] | None = None, + ) -> tuple[dict[str, float], Path]: + """Run one ``(d, blaze)`` multilayer theta search and save its full output. + + The search runs through :func:`grax.run_multilayer_theta_search_sweep` at + the single optimization energy, with ``output_dir`` set to this run's + folder, so it drops the standard artifact bundle there (selected-result + summary CSV, all-orders CSV, incidence-angle scan CSV/PNG, profile and + stack plots, checkpoint). A ``search_parameters.json`` is added so the + search settings can be reviewed and tuned. Returns the scalar summary and + the run folder. + """ + + config = self.config + grating = self._build_grating(d_spacing_nm, blaze_angle_deg) + bragg_estimate_deg = estimate_multilayer_bragg_angle_deg( + grating=grating, + energy_ev=float(config.target_energy_ev), + multilayer_bragg_order=int(config.multilayer_bragg_order), + ) + run_dir = config.run_dir(d_spacing_nm, blaze_angle_deg) + run_dir.mkdir(parents=True, exist_ok=True) + scan = config.survey_scan_settings + + sweep = run_multilayer_theta_search_sweep( + grating=grating, + energies_ev=[float(config.target_energy_ev)], + output_dir=run_dir, + diffraction_order=int(config.diffraction_order), + multilayer_bragg_order=int(config.multilayer_bragg_order), + rough_scan_half_width_deg=float(scan.rough_scan_half_width_deg), + rough_scan_points=int(scan.rough_scan_points), + rough_fourier_orders=int(scan.rough_fourier_orders), + rough_x_resolution_nm=float(scan.rough_x_resolution_nm), + rough_z_resolution_nm=float(scan.rough_z_resolution_nm), + fine_scan_half_width_deg=float(scan.fine_scan_half_width_deg), + fine_scan_points=int(scan.fine_scan_points), + fine_fourier_orders=int(scan.fine_fourier_orders), + fine_x_resolution_nm=float(scan.fine_x_resolution_nm), + fine_z_resolution_nm=float(scan.fine_z_resolution_nm), + final_fourier_orders=int(scan.final_fourier_orders), + final_x_resolution_nm=float(scan.final_x_resolution_nm), + final_z_resolution_nm=float(scan.final_z_resolution_nm), + roughness_sigma_nm=scan.roughness_sigma_nm, + precise_peak_selection_mode=str(scan.precise_peak_selection_mode), + max_workers=1, + stop_event=stop_event, + on_worker_pids_changed=on_worker_pids_changed, + show_progress=False, + on_error="fail_fast", + checkpoint_dir=(run_dir / "checkpoints") if config.checkpoint else None, + checkpoint_interval=1, + resume=bool(config.resume and config.checkpoint), + theta_tracking_mode="bragg", + save_profile_plot=bool(config.save_profile_plot), + save_stack_plot=bool(config.save_stack_plot), + backend=str(config.backend), + solver=str(config.solver), + polarization=str(config.polarization), + ) + if sweep.stopped_early: + # A killed cell leaves a header-only summary CSV, which would make + # every later evaluate_survey() raise on `.iloc[0]`. The cell is a + # single energy, so there is nothing worth keeping. + shutil.rmtree(run_dir, ignore_errors=True) + raise _CellAbortedError( + f"survey cell d={d_spacing_nm:.3f} nm blaze={blaze_angle_deg:.3f} deg was aborted" + ) + (run_dir / "search_parameters.json").write_text( + json.dumps( + self._search_parameters_record( + d_spacing_nm=float(d_spacing_nm), + blaze_angle_deg=float(blaze_angle_deg), + bragg_estimate_deg=float(bragg_estimate_deg), + ), + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + + row = pd.read_csv(sweep.summary_csv_path).iloc[0] + selected_angle = float(row["selected_grazing_angle_deg"]) + fwhm_value = row.get("precise_fwhm_deg", float("nan")) + summary = { + "bragg_estimate_deg": float(bragg_estimate_deg), + "incidence_angle_deg": selected_angle, + "peak_efficiency": float(row["selected_efficiency"]), + "precise_fwhm_deg": float(fwhm_value) if pd.notna(fwhm_value) else float("nan"), + "edge_clipped": _edge_clipped_from_scan(sweep.theta_scan_directory, selected_angle), + } + return summary, run_dir + + def run_survey( + self, + *, + progress_callback: Callable[[StageProgress], None] | None = None, + should_continue: Callable[[], bool] | None = None, + stop_event: threading.Event | None = None, + on_worker_pids_changed: Callable[[set[int]], None] | None = None, + ) -> SurveyResult: + """Run the ``(d, blaze)`` survey at the optimization energy. + + Args: + progress_callback: Optional callable invoked with a + :class:`StageProgress` before each cell and once more when the + scan finishes. + should_continue: Optional callable checked before each cell; when it + returns ``False`` the scan stops early and the optimal blaze per + d-spacing is computed from the completed cells (``aborted`` is set + on the result). + stop_event: Optional stop signal forwarded to each cell's theta + search. Unlike ``should_continue``, which is only consulted + between cells, setting it kills the search in flight; the + half-finished cell is discarded and the survey is marked aborted. + on_worker_pids_changed: Optional callback receiving the theta + search's current worker process IDs. + + Returns: + A :class:`SurveyResult` with the maps, the per-d optimal blaze + angle, and the per-run folder tree under ``survey/runs/``. + + Raises: + RuntimeError: If ``should_continue`` stops the scan before any cell + has been evaluated. + """ + + config = self.config + d_values = config.d_grid_nm() + blaze_values = config.blaze_grid_deg() + n_d = d_values.size + n_blaze = blaze_values.size + total = n_d * n_blaze + + efficiency_map = np.full((n_d, n_blaze), np.nan) + incidence_map = np.full((n_d, n_blaze), np.nan) + fwhm_map = np.full((n_d, n_blaze), np.nan) + bragg_map = np.full((n_d, n_blaze), np.nan) + run_dir_by_cell: dict[tuple[int, int], Path] = {} + rows: list[dict[str, float]] = [] + completed = 0 + aborted = False + + for i_d, d_spacing in enumerate(d_values): + if aborted: + break + for i_b, blaze in enumerate(blaze_values): + if should_continue is not None and not should_continue(): + aborted = True + break + _emit_progress( + progress_callback, + stage="survey", + completed=completed, + total=total, + current_label=f"d = {d_spacing:.2f} nm, blaze = {blaze:.3f} deg", + ) + print( + f"Multilayer theta search: d = {d_spacing:.2f} nm, blaze = {blaze:.3f} deg" + ) + try: + summary, run_dir = self._survey_cell( + float(d_spacing), + float(blaze), + stop_event=stop_event, + on_worker_pids_changed=on_worker_pids_changed, + ) + except _CellAbortedError: + # Must precede the generic handler below, or on_error + # "fail_fast" would re-raise an abort as a failure. + aborted = True + break + except Exception as error: # noqa: BLE001 - per-cell isolation + if config.on_error == "fail_fast": + raise + print( + f" cell failed (d = {d_spacing:.2f} nm, blaze = {blaze:.3f} deg): {error}" + ) + summary = { + "bragg_estimate_deg": float("nan"), + "incidence_angle_deg": float("nan"), + "peak_efficiency": float("nan"), + "precise_fwhm_deg": float("nan"), + "edge_clipped": float("nan"), + } + run_dir = None + if run_dir is not None: + run_dir_by_cell[(i_d, i_b)] = run_dir + efficiency_map[i_d, i_b] = summary["peak_efficiency"] + incidence_map[i_d, i_b] = summary["incidence_angle_deg"] + fwhm_map[i_d, i_b] = summary["precise_fwhm_deg"] + bragg_map[i_d, i_b] = summary["bragg_estimate_deg"] + rows.append( + { + "d_nm": float(d_spacing), + "blaze_deg": float(blaze), + "bragg_estimate_deg": summary["bragg_estimate_deg"], + "incidence_angle_deg": summary["incidence_angle_deg"], + "peak_efficiency": summary["peak_efficiency"], + "precise_fwhm_deg": summary["precise_fwhm_deg"], + "edge_clipped": summary["edge_clipped"], + } + ) + completed += 1 + # Rewrite the table after every cell so a watcher -- the web + # app's live plots -- can read the survey as it fills in. The + # grid is small (hundreds of rows) next to a theta search. + _write_csv_atomic(pd.DataFrame(rows), config.survey_dir / "survey.csv") + + if not rows: + raise RuntimeError("multilayer design survey aborted before any result") + _emit_progress( + progress_callback, + stage="survey", + completed=completed, + total=total, + current_label="done", + ) + return self._finalize_survey( + d_values=d_values, + blaze_values=blaze_values, + efficiency_map=efficiency_map, + incidence_map=incidence_map, + fwhm_map=fwhm_map, + bragg_map=bragg_map, + run_dir_by_cell=run_dir_by_cell, + rows=rows, + aborted=aborted, + ) + + def evaluate_survey(self) -> SurveyResult: + """Re-derive the survey outputs from the runs already on disk. + + Runs no new solves: reads every + ``survey/runs/dnm/blazedeg/multilayer_theta_search_summary.csv`` + that exists, rebuilds ``survey/survey.csv``, the per-period overlays and + the two headline curves, and returns a fresh :class:`SurveyResult`. Use + it to regenerate the analysis after a completed :meth:`run_survey`. + + Raises: + RuntimeError: If no completed run folder is found under + ``config.runs_dir``. + """ + + config = self.config + found: dict[tuple[float, float], Path] = {} + if config.runs_dir.is_dir(): + for run_dir in sorted(config.runs_dir.glob("d*nm/blaze*deg")): + if not (run_dir / "multilayer_theta_search_summary.csv").is_file(): + continue + try: + d_nm = float(run_dir.parent.name[1:-2]) + blaze = float(run_dir.name[5:-3]) + except ValueError: + continue + found[(d_nm, blaze)] = run_dir + if not found: + raise RuntimeError( + f"no completed survey runs found under {config.runs_dir}; run run_survey() first" + ) + + d_values = np.array(sorted({d for d, _ in found}), dtype=float) + blaze_values = np.array(sorted({b for _, b in found}), dtype=float) + n_d, n_blaze = d_values.size, blaze_values.size + efficiency_map = np.full((n_d, n_blaze), np.nan) + incidence_map = np.full((n_d, n_blaze), np.nan) + fwhm_map = np.full((n_d, n_blaze), np.nan) + bragg_map = np.full((n_d, n_blaze), np.nan) + run_dir_by_cell: dict[tuple[int, int], Path] = {} + rows: list[dict[str, float]] = [] + for (d_nm, blaze), run_dir in sorted(found.items()): + i_d = int(np.argmin(np.abs(d_values - d_nm))) + i_b = int(np.argmin(np.abs(blaze_values - blaze))) + row = pd.read_csv(run_dir / "multilayer_theta_search_summary.csv").iloc[0] + selected_angle = float(row["selected_grazing_angle_deg"]) + fwhm_value = row.get("precise_fwhm_deg", float("nan")) + bragg_estimate = float("nan") + params_path = run_dir / "search_parameters.json" + if params_path.is_file(): + bragg_estimate = float( + json.loads(params_path.read_text()).get("bragg_estimate_deg", float("nan")) + ) + summary = { + "bragg_estimate_deg": bragg_estimate, + "incidence_angle_deg": selected_angle, + "peak_efficiency": float(row["selected_efficiency"]), + "precise_fwhm_deg": float(fwhm_value) if pd.notna(fwhm_value) else float("nan"), + "edge_clipped": _edge_clipped_from_scan(run_dir / "theta_scans", selected_angle), + } + run_dir_by_cell[(i_d, i_b)] = run_dir + efficiency_map[i_d, i_b] = summary["peak_efficiency"] + incidence_map[i_d, i_b] = summary["incidence_angle_deg"] + fwhm_map[i_d, i_b] = summary["precise_fwhm_deg"] + bragg_map[i_d, i_b] = summary["bragg_estimate_deg"] + rows.append( + { + "d_nm": float(d_nm), + "blaze_deg": float(blaze), + "bragg_estimate_deg": summary["bragg_estimate_deg"], + "incidence_angle_deg": summary["incidence_angle_deg"], + "peak_efficiency": summary["peak_efficiency"], + "precise_fwhm_deg": summary["precise_fwhm_deg"], + "edge_clipped": summary["edge_clipped"], + } + ) + print(f"Re-evaluating {len(found)} completed runs under {config.runs_dir}") + return self._finalize_survey( + d_values=d_values, + blaze_values=blaze_values, + efficiency_map=efficiency_map, + incidence_map=incidence_map, + fwhm_map=fwhm_map, + bragg_map=bragg_map, + run_dir_by_cell=run_dir_by_cell, + rows=rows, + aborted=False, + ) + + def _finalize_survey( + self, + *, + d_values: np.ndarray, + blaze_values: np.ndarray, + efficiency_map: np.ndarray, + incidence_map: np.ndarray, + fwhm_map: np.ndarray, + bragg_map: np.ndarray, + run_dir_by_cell: dict[tuple[int, int], Path], + rows: list[dict[str, float]], + aborted: bool, + ) -> SurveyResult: + """Aggregate the per-cell results into ``survey.csv``, the plots and a result.""" + + config = self.config + n_blaze = blaze_values.size + optimal_blaze, optimal_eff, optimal_inc = _optimal_blaze_per_d( + blaze_values, efficiency_map, incidence_map + ) + combined = pd.DataFrame(rows) + config.survey_dir.mkdir(parents=True, exist_ok=True) + csv_path = config.survey_dir / "survey.csv" + _write_csv_atomic(combined, csv_path) + + period_dirs: list[Path] = [] + run_dirs: list[Path] = [] + for i_d, d_spacing in enumerate(d_values): + d_spacing = float(d_spacing) + cells = [ + (float(blaze_values[i_b]), run_dir_by_cell[(i_d, i_b)]) + for i_b in range(n_blaze) + if (i_d, i_b) in run_dir_by_cell + ] + if not cells: + continue + run_dirs.extend(run_dir for _, run_dir in cells) + period_dir = config.period_dir(d_spacing) + self._plot_period_overlay( + d_spacing_nm=d_spacing, + cells=cells, + chosen_blaze_deg=( + float(optimal_blaze[i_d]) if np.isfinite(optimal_blaze[i_d]) else None + ), + output_path=period_dir / "overlay.png", + ) + period_dirs.append(period_dir) + + config.plot_dir.mkdir(parents=True, exist_ok=True) + plot_path = config.plot_dir / "optimal_blaze_vs_d_spacing.png" + self._plot_optimal_blaze( + d_values=d_values, + optimal_blaze=optimal_blaze, + optimal_eff=optimal_eff, + output_path=plot_path, + ) + efficiency_plot_path = config.plot_dir / "max_efficiency_vs_d_spacing.png" + self._plot_max_efficiency( + d_values=d_values, + optimal_blaze=optimal_blaze, + optimal_eff=optimal_eff, + output_path=efficiency_plot_path, + ) + heatmap_plot_path = config.plot_dir / "efficiency_heatmap_d_vs_blaze.png" + self._plot_efficiency_heatmap( + d_values=d_values, + blaze_values=blaze_values, + efficiency_map=efficiency_map, + optimal_blaze=optimal_blaze, + output_path=heatmap_plot_path, + ) + + print(f"Survey CSV: {csv_path}") + print(f"Per-run results: {len(run_dirs)} folders under {config.runs_dir}") + for d_spacing, blaze_opt, eff in zip(d_values, optimal_blaze, optimal_eff): + print( + f" d {d_spacing:.3f} nm -> optimal blaze {blaze_opt:.3f} deg " + f"(efficiency {eff:.4g})" + ) + return SurveyResult( + d_values_nm=d_values, + blaze_values_deg=blaze_values, + efficiency_map=efficiency_map, + incidence_angle_map_deg=incidence_map, + fwhm_map_deg=fwhm_map, + bragg_estimate_map_deg=bragg_map, + optimal_blaze_deg=optimal_blaze, + optimal_blaze_efficiency=optimal_eff, + optimal_blaze_incidence_deg=optimal_inc, + combined_csv_path=csv_path, + plot_path=plot_path, + efficiency_plot_path=efficiency_plot_path, + heatmap_plot_path=heatmap_plot_path, + runs_dir=config.runs_dir, + period_dirs=period_dirs, + run_dirs=run_dirs, + aborted=aborted, + results=combined, + ) + + # ------------------------------------------------------------------ # + # Step 2: per-design energy scan # + # ------------------------------------------------------------------ # + def run_energy_scan( + self, + pairs: Iterable[tuple[float, float]], + *, + progress_callback: Callable[[StageProgress], None] | None = None, + should_continue: Callable[[], bool] | None = None, + stop_event: threading.Event | None = None, + on_worker_pids_changed: Callable[[set[int]], None] | None = None, + ) -> list[EnergyScanResult]: + """Sweep chosen ``(d_spacing_nm, blaze_angle_deg)`` designs over energy. + + Args: + pairs: ``(d_spacing_nm, blaze_angle_deg)`` designs, typically read + off the survey. + progress_callback: Optional callable invoked with a + :class:`StageProgress` before each design and once when finished. + should_continue: Optional callable checked before each design; when + it returns ``False`` the scan stops early and returns the + designs completed so far (an empty list if it stopped before + the first one). + stop_event: Optional stop signal forwarded to the sweep. Unlike + ``should_continue``, which is only consulted between designs, + setting it kills the energies in flight. The half-scanned design + is not returned, but its checkpoint keeps every energy already + solved, so re-running resumes from there. + on_worker_pids_changed: Optional callback receiving the sweep's + current worker process IDs. + + Returns: + One :class:`EnergyScanResult` per completed design, in input order. + + Raises: + ValueError: If ``pairs`` is empty. + """ + + config = self.config + scan = config.energy_scan_settings + designs = [(float(d), float(b)) for d, b in pairs] + if not designs: + raise ValueError("run_energy_scan requires at least one (d_spacing, blaze) pair") + energies = config.energy_grid_ev() + total = len(designs) + results: list[EnergyScanResult] = [] + for index, (d_spacing, blaze) in enumerate(designs): + if should_continue is not None and not should_continue(): + break + _emit_progress( + progress_callback, + stage="energy_scan", + completed=index, + total=total, + current_label=f"d = {d_spacing:.3f} nm, blaze = {blaze:.3f} deg", + ) + design_dir = config.energy_scan_dir / f"d{d_spacing:.3f}nm_blaze{blaze:.3f}deg" + design_dir.mkdir(parents=True, exist_ok=True) + sweep = run_multilayer_theta_search_sweep( + grating=self._build_grating(d_spacing, blaze), + energies_ev=energies, + output_dir=design_dir, + diffraction_order=int(config.diffraction_order), + multilayer_bragg_order=int(config.multilayer_bragg_order), + rough_scan_half_width_deg=float(scan.rough_scan_half_width_deg), + rough_scan_points=int(scan.rough_scan_points), + rough_fourier_orders=int(scan.rough_fourier_orders), + rough_x_resolution_nm=float(scan.rough_x_resolution_nm), + rough_z_resolution_nm=float(scan.rough_z_resolution_nm), + fine_scan_half_width_deg=float(scan.fine_scan_half_width_deg), + fine_scan_points=int(scan.fine_scan_points), + fine_fourier_orders=int(scan.fine_fourier_orders), + fine_x_resolution_nm=float(scan.fine_x_resolution_nm), + fine_z_resolution_nm=float(scan.fine_z_resolution_nm), + final_fourier_orders=int(scan.final_fourier_orders), + final_x_resolution_nm=float(scan.final_x_resolution_nm), + final_z_resolution_nm=float(scan.final_z_resolution_nm), + roughness_sigma_nm=scan.roughness_sigma_nm, + precise_peak_selection_mode=str(scan.precise_peak_selection_mode), + max_workers=config.max_workers, + stop_event=stop_event, + on_worker_pids_changed=on_worker_pids_changed, + show_progress=bool(config.show_progress), + on_error="fail_fast", + checkpoint_dir=(design_dir / "checkpoints") if config.checkpoint else None, + checkpoint_interval=1, + resume=bool(config.resume and config.checkpoint), + theta_tracking_mode=str(config.theta_tracking_mode), + max_tracking_energy_step_ev=config.max_tracking_energy_step_ev, + save_profile_plot=bool(config.save_profile_plot), + save_stack_plot=bool(config.save_stack_plot), + backend=str(config.backend), + solver=str(config.solver), + polarization=str(config.polarization), + ) + if sweep.stopped_early: + # Do not report a half-scanned design as done: its checkpoint + # holds every solved energy, so re-running resumes from there. + break + scan_results = pd.read_csv(sweep.summary_csv_path) + config.plot_dir.mkdir(parents=True, exist_ok=True) + titled_plot_path = config.plot_dir / _energy_scan_plot_filename( + config, d_spacing, blaze + ) + self._plot_energy_scan( + d_spacing_nm=d_spacing, + blaze_angle_deg=blaze, + results=scan_results, + output_path=titled_plot_path, + ) + results.append( + EnergyScanResult( + d_spacing_nm=d_spacing, + blaze_angle_deg=blaze, + output_dir=design_dir, + summary_csv_path=sweep.summary_csv_path, + all_orders_csv_path=sweep.all_orders_csv_path, + energy_efficiency_plot_path=sweep.energy_efficiency_plot_path, + titled_plot_path=titled_plot_path, + results=scan_results, + ) + ) + _emit_progress( + progress_callback, + stage="energy_scan", + completed=total, + total=total, + current_label="done", + ) + if len(results) >= 2: + print(f"Comparison plot: {self.plot_energy_scan_overlay(results)}") + return results + + def evaluate_energy_scan( + self, pairs: Iterable[tuple[float, float]] | None = None + ) -> list[EnergyScanResult]: + """Collect energy-scan results already on disk, running no new solves. + + Reads + ``energy_scan/dnm_blazedeg/multilayer_theta_search_summary.csv`` for + each requested design (or, when ``pairs`` is ``None``, for every design + folder found) and returns the matching :class:`EnergyScanResult` objects. + + Raises: + RuntimeError: If no requested/found design folder holds a summary CSV. + """ + + config = self.config + if pairs is None: + designs = [] + if config.energy_scan_dir.is_dir(): + for design_dir in sorted(config.energy_scan_dir.glob("d*nm_blaze*deg")): + stem = design_dir.name + try: + d_text, blaze_text = stem.split("_blaze") + designs.append((float(d_text[1:-2]), float(blaze_text[:-3]))) + except ValueError: + continue + else: + designs = [(float(d), float(b)) for d, b in pairs] + + results: list[EnergyScanResult] = [] + for d_spacing, blaze in designs: + design_dir = config.energy_scan_dir / f"d{d_spacing:.3f}nm_blaze{blaze:.3f}deg" + summary_csv = design_dir / "multilayer_theta_search_summary.csv" + if not summary_csv.is_file(): + print(f" skipping {design_dir} (no summary CSV)") + continue + scan_results = pd.read_csv(summary_csv) + config.plot_dir.mkdir(parents=True, exist_ok=True) + titled_plot_path = config.plot_dir / _energy_scan_plot_filename( + config, d_spacing, blaze + ) + self._plot_energy_scan( + d_spacing_nm=d_spacing, + blaze_angle_deg=blaze, + results=scan_results, + output_path=titled_plot_path, + ) + results.append( + EnergyScanResult( + d_spacing_nm=d_spacing, + blaze_angle_deg=blaze, + output_dir=design_dir, + summary_csv_path=summary_csv, + all_orders_csv_path=design_dir + / "multilayer_theta_search_all_orders.csv", + energy_efficiency_plot_path=design_dir + / "multilayer_theta_search_energy_vs_efficiency.png", + titled_plot_path=titled_plot_path, + results=scan_results, + ) + ) + if not results: + raise RuntimeError( + f"no completed energy scans found under {config.energy_scan_dir}" + ) + print(f"Re-evaluated {len(results)} energy scans under {config.energy_scan_dir}") + if len(results) >= 2: + print(f"Comparison plot: {self.plot_energy_scan_overlay(results)}") + return results + + # ------------------------------------------------------------------ # + # Plotting # + # ------------------------------------------------------------------ # + def _plot_optimal_blaze( + self, + *, + d_values: np.ndarray, + optimal_blaze: np.ndarray, + optimal_eff: np.ndarray, + output_path: Path, + ) -> None: + """Write the headline curve: optimal blaze angle versus d-spacing. + + The line is coloured by peak efficiency; a sparse set of points is + labelled with their efficiency so the text stays readable on fine grids. + """ + + import matplotlib.pyplot as plt + + config = self.config + finite = np.isfinite(optimal_blaze) + figure, axis = plt.subplots(figsize=(9, 6)) + axis.plot(d_values[finite], optimal_blaze[finite], "-", color="0.6", zorder=1) + scatter = axis.scatter( + d_values[finite], + optimal_blaze[finite], + c=optimal_eff[finite], + cmap="viridis", + zorder=2, + ) + figure.colorbar(scatter, ax=axis, label="peak efficiency") + for i in _label_indices(optimal_eff): + axis.annotate( + f"{optimal_eff[i]:.3g}", + xy=(d_values[i], optimal_blaze[i]), + xytext=(4, 4), + textcoords="offset points", + fontsize=8, + ) + axis.set_xlabel("Bilayer d-spacing (nm)") + axis.set_ylabel("Optimal blaze angle (deg)") + axis.set_title( + f"Optimal blaze angle versus d-spacing at {config.target_energy_ev:g} eV " + f"(order {config.diffraction_order}; colour/labels: peak efficiency)" + ) + axis.grid(True, alpha=0.3) + figure.tight_layout() + figure.savefig(output_path, dpi=150) + plt.close(figure) + + def _plot_max_efficiency( + self, + *, + d_values: np.ndarray, + optimal_blaze: np.ndarray, + optimal_eff: np.ndarray, + output_path: Path, + ) -> None: + """Write max selected-order efficiency versus d-spacing. + + The line is coloured by the optimal blaze angle; a sparse set of points + (the peak and where the optimal blaze changes) is labelled with the blaze + angle so the text stays readable on fine grids. + """ + + import matplotlib.pyplot as plt + + config = self.config + finite = np.isfinite(optimal_eff) + figure, axis = plt.subplots(figsize=(9, 6)) + axis.plot(d_values[finite], optimal_eff[finite], "-", color="0.6", zorder=1) + scatter = axis.scatter( + d_values[finite], + optimal_eff[finite], + c=optimal_blaze[finite], + cmap="plasma", + zorder=2, + ) + figure.colorbar(scatter, ax=axis, label="optimal blaze angle (deg)") + for i in _label_indices(optimal_eff, change_series=optimal_blaze): + axis.annotate( + f"blaze {optimal_blaze[i]:.3f} deg", + xy=(d_values[i], optimal_eff[i]), + xytext=(4, 4), + textcoords="offset points", + fontsize=8, + ) + axis.set_xlabel("Bilayer d-spacing (nm)") + axis.set_ylabel(f"Max selected-order efficiency ({config.polarization}-pol)") + axis.set_title( + f"Max efficiency versus d-spacing at {config.target_energy_ev:g} eV " + f"(order {config.diffraction_order}; colour/labels: optimal blaze angle)" + ) + axis.grid(True, alpha=0.3) + figure.tight_layout() + figure.savefig(output_path, dpi=150) + plt.close(figure) + + def _plot_efficiency_heatmap( + self, + *, + d_values: np.ndarray, + blaze_values: np.ndarray, + efficiency_map: np.ndarray, + optimal_blaze: np.ndarray, + output_path: Path, + ) -> None: + """Write the full ``(d, blaze) -> efficiency`` heatmap. + + x is d-spacing, y is blaze angle, colour is peak selected-order + efficiency -- the whole ``efficiency_map``, not just the per-d optimum. + The ``optimal_blaze_deg`` ridge is overlaid so the two headline curves + and this map read together. + """ + + import matplotlib.pyplot as plt + + config = self.config + figure, axis = plt.subplots(figsize=(9, 6)) + mesh = axis.pcolormesh( + d_values, + blaze_values, + efficiency_map.T, + shading="nearest", + cmap="viridis", + ) + figure.colorbar( + mesh, ax=axis, label=f"peak selected-order efficiency ({config.polarization}-pol)" + ) + finite = np.isfinite(optimal_blaze) + axis.plot( + d_values[finite], + optimal_blaze[finite], + "o-", + color="white", + markeredgecolor="black", + linewidth=1.6, + label="optimal blaze per d", + ) + axis.set_xlabel("Bilayer d-spacing (nm)") + axis.set_ylabel("Blaze angle (deg)") + axis.set_title( + f"Peak efficiency over (d, blaze) at {config.target_energy_ev:g} eV " + f"(order {config.diffraction_order})" + ) + axis.legend(loc="best", fontsize=8) + figure.tight_layout() + figure.savefig(output_path, dpi=150) + plt.close(figure) + + def _plot_energy_scan( + self, + *, + d_spacing_nm: float, + blaze_angle_deg: float, + results: pd.DataFrame, + output_path: Path, + ) -> None: + """Write one design's efficiency-versus-energy curve, titled with the design. + + The title names the coating (:func:`_coating_label`) and this design's + d-spacing and blaze angle, so the plot is self-describing without + needing its folder name. + """ + + import matplotlib.pyplot as plt + + config = self.config + figure, axis = plt.subplots(figsize=(9, 6)) + axis.plot( + results["energy_ev"], + results["selected_efficiency"], + "-", + color="tab:blue", + ) + axis.set_xlabel("Photon energy (eV)") + axis.set_ylabel(f"Selected-order efficiency ({config.polarization}-pol)") + axis.set_title(_energy_scan_title(config, d_spacing_nm, blaze_angle_deg)) + axis.grid(True, alpha=0.3) + figure.tight_layout() + figure.savefig(output_path, dpi=150) + plt.close(figure) + + def plot_energy_scan_overlay(self, results: Sequence[EnergyScanResult]) -> Path: + """Overlay several designs' efficiency-versus-energy curves on one axis. + + Written to :attr:`MultilayerDesignConfig.plot_dir` as + ``efficiency_vs_energy_comparison__order.png``, so the + designs scanned in one batch can be compared at a glance. Called + automatically by :meth:`run_energy_scan` and + :meth:`evaluate_energy_scan` whenever they produce two or more designs. + + Args: + results: The designs to overlay, in legend order. + + Returns: + The path the overlay was written to. + + Raises: + ValueError: If ``results`` is empty. + """ + + import matplotlib.pyplot as plt + + if not results: + raise ValueError("plot_energy_scan_overlay requires at least one result") + config = self.config + config.plot_dir.mkdir(parents=True, exist_ok=True) + output_path = config.plot_dir / _energy_scan_overlay_filename(config) + + figure, axis = plt.subplots(figsize=(9, 6)) + for scan in results: + axis.plot( + scan.results["energy_ev"], + scan.results["selected_efficiency"], + "-", + linewidth=1.6, + label=f"d = {scan.d_spacing_nm:.3f} nm, blaze = {scan.blaze_angle_deg:.3f} deg", + ) + axis.set_xlabel("Photon energy (eV)") + axis.set_ylabel(f"Selected-order efficiency ({config.polarization}-pol)") + axis.set_title( + f"{_coating_label(config)} multilayer grating " + f"(order {config.diffraction_order}): {len(results)} designs" + ) + axis.grid(True, alpha=0.3) + axis.legend(loc="best", fontsize=8) + figure.tight_layout() + figure.savefig(output_path, dpi=150) + plt.close(figure) + return output_path + + def _search_parameters_record( + self, *, d_spacing_nm: float, blaze_angle_deg: float, bragg_estimate_deg: float + ) -> dict[str, object]: + """Build ``search_parameters.json`` for one run so the search can be tuned. + + The run's numerical *results* are already in the sweep's own artifacts + (``multilayer_theta_search_summary.csv``, ``*_all_orders.csv``, + ``theta_scans/``); this records only what was fed into the search. + """ + + config = self.config + scan = config.survey_scan_settings + return { + "d_spacing_nm": float(d_spacing_nm), + "blaze_angle_deg": float(blaze_angle_deg), + "target_energy_ev": float(config.target_energy_ev), + "diffraction_order": int(config.diffraction_order), + "multilayer_bragg_order": int(config.multilayer_bragg_order), + "bragg_estimate_deg": float(bragg_estimate_deg), + "theta_seed": "multilayer Bragg estimate (theta_tracking_mode='bragg')", + "rough_scan_half_width_deg": float(scan.rough_scan_half_width_deg), + "rough_scan_points": int(scan.rough_scan_points), + "rough_fourier_orders": int(scan.rough_fourier_orders), + "rough_x_resolution_nm": float(scan.rough_x_resolution_nm), + "rough_z_resolution_nm": float(scan.rough_z_resolution_nm), + "fine_scan_half_width_deg": float(scan.fine_scan_half_width_deg), + "fine_scan_points": int(scan.fine_scan_points), + "fine_fourier_orders": int(scan.fine_fourier_orders), + "fine_x_resolution_nm": float(scan.fine_x_resolution_nm), + "fine_z_resolution_nm": float(scan.fine_z_resolution_nm), + "final_fourier_orders": int(scan.final_fourier_orders), + "final_x_resolution_nm": float(scan.final_x_resolution_nm), + "final_z_resolution_nm": float(scan.final_z_resolution_nm), + "precise_peak_selection_mode": str(scan.precise_peak_selection_mode), + "roughness_sigma_nm": ( + None if scan.roughness_sigma_nm is None else float(scan.roughness_sigma_nm) + ), + "grating_x_resolution_nm": float(config.x_resolution_nm), + "grating_z_resolution_nm": float(config.z_resolution_nm), + "solver": str(config.solver), + "polarization": str(config.polarization), + "backend": str(config.backend), + } + + def _plot_period_overlay( + self, + *, + d_spacing_nm: float, + cells: list[tuple[float, Path]], + chosen_blaze_deg: float | None, + output_path: Path, + ) -> None: + """Overlay every blaze angle's theta scan at one d-spacing. + + The blaze angle that won (highest peak efficiency, ``chosen_blaze_deg``) + is drawn bold with a ``(chosen)`` legend entry, a dashed vertical line at + its selected incidence angle and a large star at its peak; the other + blaze angles are drawn light with a small star at each selected peak. + """ + + import matplotlib.pyplot as plt + + config = self.config + figure, axis = plt.subplots(figsize=(9, 6)) + color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"] + for order, (blaze, run_dir) in enumerate(cells): + frame = _read_theta_scan_frame(run_dir) + if frame is None: + continue + is_chosen = chosen_blaze_deg is not None and bool( + np.isclose(blaze, chosen_blaze_deg, rtol=0.0, atol=1.0e-9) + ) + color = color_cycle[order % len(color_cycle)] + rough = frame[frame["scan_type"] == "rough"] + precise = frame[frame["scan_type"] == "precise"] + axis.plot( + rough["grazing_angle_deg"], + rough["selected_efficiency"], + ":", + color=color, + linewidth=0.9, + alpha=0.6 if is_chosen else 0.35, + ) + axis.plot( + precise["grazing_angle_deg"], + precise["selected_efficiency"], + "-", + color=color, + linewidth=2.8 if is_chosen else 1.3, + alpha=1.0 if is_chosen else 0.75, + label=f"blaze {blaze:.3f} deg" + (" (chosen)" if is_chosen else ""), + zorder=5 if is_chosen else 3, + ) + peak = frame[frame["is_selected_peak"] == 1] + if not peak.empty: + peak_angle = float(peak["grazing_angle_deg"].iloc[0]) + peak_eff = float(peak["selected_efficiency"].iloc[0]) + axis.plot( + [peak_angle], + [peak_eff], + "*", + color=color, + markersize=20 if is_chosen else 9, + markeredgecolor="black", + zorder=6 if is_chosen else 4, + ) + if is_chosen: + axis.axvline( + peak_angle, color=color, linestyle="--", linewidth=1.2, alpha=0.7 + ) + axis.text( + 0.02, + 0.72, + f"chosen: blaze {blaze:.3f} deg\n" + f"theta* = {peak_angle:.3f} deg\n" + f"efficiency = {peak_eff:.4g}", + transform=axis.transAxes, + va="top", + fontsize=9, + bbox={"boxstyle": "round", "fc": "white", "ec": color, "alpha": 0.9}, + ) + axis.set_xlabel("Incidence (grazing) angle (deg)") + axis.set_ylabel(f"Selected-order efficiency ({config.polarization}-pol)") + axis.set_title( + f"Multilayer theta search at d = {d_spacing_nm:.3f} nm, " + f"{config.target_energy_ev:g} eV (dotted: rough, solid: precise, *: selected peak)" + ) + axis.grid(True, alpha=0.3) + axis.legend(loc="best", fontsize=8) + figure.tight_layout() + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path, dpi=150) + plt.close(figure) + + +def _read_theta_scan_frame(run_dir: Path) -> pd.DataFrame | None: + """Return the rough/precise scan table a sweep wrote under ``run_dir``.""" + + scan_files = sorted((run_dir / "theta_scans").glob("theta_scan_*.csv")) + if not scan_files: + return None + return pd.read_csv(scan_files[0]) + + +def _edge_clipped_from_scan(theta_scan_directory: Path, selected_angle_deg: float) -> float: + """Flag whether the selected angle sits within 5% of a precise-scan edge.""" + + scan_files = sorted(Path(theta_scan_directory).glob("theta_scan_*.csv")) + if not scan_files: + return float("nan") + frame = pd.read_csv(scan_files[0]) + precise = np.asarray( + frame.loc[frame["scan_type"] == "precise", "grazing_angle_deg"], dtype=float + ) + if precise.size == 0: + return float("nan") + span = float(precise.max() - precise.min()) + margin = 0.05 * span if span > 0.0 else 0.0 + return float( + selected_angle_deg <= precise.min() + margin + or selected_angle_deg >= precise.max() - margin + ) + + +def _label_indices( + values: np.ndarray, + *, + change_series: np.ndarray | None = None, + max_labels: int = 14, +) -> list[int]: + """Pick a readable subset of indices to annotate on a per-d headline curve. + + Always keeps the first and last finite points and the efficiency peak; adds + every point where ``change_series`` steps to a new value, then tops up (or + thins) to about ``max_labels`` evenly spaced points. + """ + + finite = np.where(np.isfinite(values))[0] + if finite.size == 0: + return [] + peak = int(finite[int(np.nanargmax(values[finite]))]) + chosen = {int(finite[0]), int(finite[-1]), peak} + if change_series is not None: + previous: float | None = None + for i in finite: + current = float(change_series[i]) + if not np.isfinite(current): + continue + if previous is None or not np.isclose(current, previous): + chosen.add(int(i)) + previous = current + if len(chosen) < max_labels: + step = max(1, finite.size // max_labels) + chosen.update(int(i) for i in finite[::step]) + ordered = sorted(chosen) + if len(ordered) > max_labels: + keep = {ordered[0], ordered[-1], peak} + rest = [i for i in ordered if i not in keep] + step = max(1, len(rest) // max(1, max_labels - len(keep))) + keep.update(rest[::step]) + ordered = sorted(keep) + return ordered + + +def _optimal_blaze_per_d( + blaze_values: np.ndarray, + efficiency_map: np.ndarray, + incidence_map: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return ``(blaze_opt, efficiency, theta*)`` per d-spacing row of the maps.""" + + n_d = efficiency_map.shape[0] + blaze_opt = np.full(n_d, np.nan) + eff_opt = np.full(n_d, np.nan) + inc_opt = np.full(n_d, np.nan) + for i_d in range(n_d): + row = efficiency_map[i_d, :] + if not np.any(np.isfinite(row)): + continue + i_b = int(np.nanargmax(row)) + blaze_opt[i_d] = float(blaze_values[i_b]) + eff_opt[i_d] = float(row[i_b]) + inc_opt[i_d] = float(incidence_map[i_d, i_b]) + return blaze_opt, eff_opt, inc_opt diff --git a/src/grax/multilayer_optimization.py b/src/grax/multilayer_optimization.py deleted file mode 100644 index 34a054a..0000000 --- a/src/grax/multilayer_optimization.py +++ /dev/null @@ -1,1207 +0,0 @@ -"""Three-stage multilayer-grating design workflow. - -The workflow sizes a periodic multilayer coating for a blazed grating -monochromator working in a chosen diffraction order at fixed CFF, in three -stages driven by a single :class:`MultilayerOptimizationConfig`: - -1. :func:`run_d_spacing_study` -- derive a bilayer d-spacing from the grating - geometry (grazing angle at the configured CFF, then the first-order Bragg - law) and scan practical d-spacing candidates with XRT planar-multilayer - reflectivity. -2. :func:`run_gamma_study` -- at the selected d-spacing, scan the bilayer - thickness ratio ``gamma`` and pick the value with the highest peak - reflectivity at the target energy. -3. :func:`run_blaze_study` -- build the multilayer-coated blazed grating and - scan the blaze angle, using graxPy's internal theta search per energy, then - pick the blaze angle with the highest selected-order efficiency at the target - energy. - -Stages hand values forward only through ``optimization_state.json`` and only when -a config value is the string ``"auto"``. A numeric config value always wins over -anything in the state file, and no stage ever rewrites the config. Stage 2 reads -``config.gamma`` directly -- the gamma suggestion from stage 1 is recorded for -traceability but is not auto-applied. - -Two numerical conventions are inherited from the original workflow and kept -deliberately: the geometry d-spacing derivation uses ``HC_EV_NM = 1239.841984`` -while :func:`grax.monochromator_grazing_angles_deg` uses ``1239.8`` internally -(immaterial at the 0.1 nm rounding used here), and the XRT reflectivity path puts -``material_a`` on top of a ``material_a`` substrate whereas the graxPy -:class:`grax.MultilayerStack` puts ``material_b`` on top of the configured -substrate. -""" - -from __future__ import annotations - -import json -from collections.abc import Iterable, Mapping -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Literal - -import numpy as np -import pandas as pd - -from .gratings import BlazedGrating -from .materials import MaterialSpec -from .multilayer_reflectivity import MultilayerReflectivity -from .simulation import monochromator_grazing_angles_deg, run_multilayer_theta_search_sweep -from .simulation.core import normalize_polarization -from .stacks import MultilayerStack - -__all__ = [ - "BlazeStudyResult", - "DSpacingStudyResult", - "GammaStudyResult", - "MultilayerOptimizationConfig", - "d_spacing_bounds_from_bragg_angles", - "energy_to_wavelength_nm", - "ensure_target_energy", - "intersect_search_bounds", - "resolve_configured_value", - "run_blaze_study", - "run_d_spacing_study", - "run_gamma_study", - "select_target_energy_optimum", - "update_optimization_state", -] - -HC_EV_NM = 1239.841984 - - -@dataclass(frozen=True) -class MultilayerOptimizationConfig: - """Every knob for the three multilayer-optimization stages. - - The defaults reproduce the Ru/B4C second-order study the workflow was ported - from. ``material_a`` / ``material_b`` / ``substrate_material`` are - ``(name, density_g_cm3)`` pairs (a :class:`grax.MaterialSpec` is also - accepted). ``d_spacing_nm`` and, downstream, ``gamma`` may be the string - ``"auto"`` to consume the previous stage's suggestion from the state file. - - Attributes: - output_dir: Root directory for all generated artifacts. - d_spacing_nm: Bilayer period in nm, or ``"auto"``. - gamma: Bilayer thickness ratio (``material_a`` fraction), 0 < gamma < 1. - blaze_angle_deg: Center blaze angle for the stage-2 scan. - material_a: Incident-side / top bilayer material. - material_b: Second bilayer material. - substrate_material: Grating substrate material. - n_bilayers: Number of bilayer periods. - target_energy_ev: Photon energy the stages optimize at. - grating_density_lpermm: Groove density in lines/mm. - diffraction_order: Grating diffraction order to optimize. - cff: Fixed-focus constant used for the geometry grazing angle. - multilayer_bragg_order: Multilayer Bragg order (distinct from the - grating diffraction order). - """ - - output_dir: Path - - # Selected values (numeric, or "auto" where noted). - d_spacing_nm: float | Literal["auto"] = "auto" - gamma: float = 0.5 - blaze_angle_deg: float = 1.1 - - # Materials. - material_a: Any = ("Ru", 12.1) - material_b: Any = ("C", 2.52) - substrate_material: Any = ("Si", 2.33) - n_bilayers: int = 40 - - # Target and grating geometry. - target_energy_ev: float = 9000.0 - grating_density_lpermm: float = 2400.0 - diffraction_order: int = 2 - cff: float = 2.25 - multilayer_bragg_order: int = 1 - - # Per-stage energy grids. - d_spacing_energy_min_ev: float = 500.0 - d_spacing_energy_max_ev: float = 12000.0 - d_spacing_energy_step_ev: float = 100.0 - d_spacing_energy_quick_step_ev: float = 250.0 - gamma_energy_min_ev: float = 500.0 - gamma_energy_max_ev: float = 12000.0 - gamma_energy_step_ev: float = 100.0 - gamma_energy_quick_step_ev: float = 250.0 - blaze_energy_min_ev: float = 3000.0 - blaze_energy_max_ev: float = 12000.0 - blaze_energy_points: int = 15 - blaze_energy_quick_points: int = 5 - - # D-spacing geometry and scan settings. - bragg_angle_min_deg: float = 0.5 - bragg_angle_max_deg: float = 2.0 - d_spacing_relative_range: float = 0.25 - d_spacing_min_practical_nm: float = 2.0 - d_spacing_max_practical_nm: float = 8.0 - d_spacing_points: int = 21 - - # Gamma scan settings. - gamma_min: float = 0.3 - gamma_max: float = 0.8 - gamma_step: float = 0.1 - - # Blaze scan settings. - blaze_angle_half_range_deg: float = 0.3 - blaze_angle_points: int = 4 - anti_blaze_angle_deg: float = 0.0 - - # XRT reflectivity settings. - xrt_window_deg: float = 0.2 - xrt_angle_points: int = 2001 - xrt_min_angle_deg: float = 0.0 - xrt_individuals: bool = False - - # graxPy theta-search settings. - grax_x_resolution_nm: float = 0.5 - grax_z_resolution_nm: float = 0.5 - rough_scan_half_width_deg: float = 0.5 - rough_scan_points: int = 61 - rough_fourier_orders: int = 5 - rough_x_resolution_nm: float = 1.0 - rough_z_resolution_nm: float = 1.0 - fine_scan_half_width_deg: float = 0.2 - fine_scan_points: int = 81 - fine_fourier_orders: int = 15 - fine_x_resolution_nm: float = 0.5 - fine_z_resolution_nm: float = 0.5 - final_fourier_orders: int = 25 - final_x_resolution_nm: float = 0.2 - final_z_resolution_nm: float = 0.2 - roughness_sigma_nm: float | None = None - precise_peak_selection_mode: str = "max" - retry_on_selected_efficiency_zero: bool = True - retry_selected_efficiency_threshold: float = 1.0e-4 - max_zero_efficiency_retries: int = 3 - backend: str = "numba" - solver: str = "neviere" - polarization: str = "p" - - # Runtime controls. - quick: bool = False - max_workers: int | str | None = "auto" - show_progress: bool = True - live_plot: bool = False - on_error: str = "fail_fast" - checkpoint_interval: int = 1 - resume: bool = True - theta_tracking_mode: str = "auto" - max_tracking_energy_step_ev: float | None = None - save_profile_plot: bool = True - save_stack_plot: bool = True - - def __post_init__(self) -> None: - """Validate cross-field constraints.""" - - object.__setattr__(self, "output_dir", Path(self.output_dir)) - if not (0.0 < float(self.gamma) < 1.0): - raise ValueError(f"gamma must satisfy 0 < gamma < 1, got {self.gamma!r}") - if isinstance(self.d_spacing_nm, str): - if self.d_spacing_nm.strip().lower() != "auto": - raise ValueError( - f"d_spacing_nm must be numeric or 'auto', got {self.d_spacing_nm!r}" - ) - elif not np.isfinite(float(self.d_spacing_nm)) or float(self.d_spacing_nm) <= 0.0: - raise ValueError(f"d_spacing_nm must be finite and positive, got {self.d_spacing_nm!r}") - if self.target_energy_ev <= 0.0: - raise ValueError("target_energy_ev must be positive") - if self.n_bilayers < 1: - raise ValueError("n_bilayers must be at least 1") - if self.d_spacing_points < 2: - raise ValueError("d_spacing_points must be at least 2") - if self.blaze_angle_points < 1: - raise ValueError("blaze_angle_points must be at least 1") - if self.blaze_angle_half_range_deg < 0.0: - raise ValueError("blaze_angle_half_range_deg must be non-negative") - if self.blaze_energy_points < 2: - raise ValueError("blaze_energy_points must be at least 2") - if self.solver not in {"rcwa", "neviere"}: - raise ValueError(f"solver must be 'rcwa' or 'neviere', got {self.solver!r}") - - @property - def plot_dir(self) -> Path: - """Directory for the per-stage summary plots.""" - - return self.output_dir / "plot" - - @property - def state_path(self) -> Path: - """Path to the cross-stage ``optimization_state.json`` file.""" - - return self.output_dir / "optimization_state.json" - - @property - def d_spacing_results_dir(self) -> Path: - """Directory for stage-0 (d-spacing) artifacts.""" - - return self.output_dir / "0_d_spacing" - - @property - def gamma_results_dir(self) -> Path: - """Directory for stage-1 (gamma) artifacts.""" - - return self.output_dir / "1_gamma" - - @property - def blaze_results_dir(self) -> Path: - """Directory for stage-2 (blaze) artifacts.""" - - return self.output_dir / "2_blaze" - - -@dataclass -class DSpacingStudyResult: - """Outcome of :func:`run_d_spacing_study`. - - Attributes: - geometry_grazing_angle_deg: Grazing angle at the target energy and CFF. - geometry_d_nm: d-spacing from the first-order Bragg law at that angle. - d_suggested_nm: ``geometry_d_nm`` rounded to 0.1 nm; handed to stages 1-2. - d_suggested_peak_rp: Peak reflectivity at ``d_suggested_nm`` and target. - d_reflectivity_best_nm: Numerically best d at the target (diagnostic). - d_reflectivity_best_peak_rp: Peak reflectivity at ``d_reflectivity_best_nm``. - search_min_nm: Lower edge of the resolved d-spacing search interval. - search_max_nm: Upper edge of the resolved d-spacing search interval. - combined_csv_path: Combined per-d reflectivity table. - plot_path: Reflectivity-versus-energy summary plot. - state_path: The updated state file. - results: The combined reflectivity table. - """ - - geometry_grazing_angle_deg: float - geometry_d_nm: float - d_suggested_nm: float - d_suggested_peak_rp: float - d_reflectivity_best_nm: float - d_reflectivity_best_peak_rp: float - search_min_nm: float - search_max_nm: float - combined_csv_path: Path - plot_path: Path - state_path: Path - results: pd.DataFrame = field(repr=False) - - -@dataclass -class GammaStudyResult: - """Outcome of :func:`run_gamma_study`. - - Attributes: - d_spacing_nm: Resolved d-spacing the scan ran at. - gamma_suggested: Gamma with the highest peak reflectivity at the target. - gamma_suggested_peak_rp: Peak reflectivity at ``gamma_suggested``. - combined_csv_path: Combined per-gamma reflectivity table. - plot_path: Reflectivity-versus-energy summary plot. - state_path: The updated state file. - results: The combined reflectivity table. - """ - - d_spacing_nm: float - gamma_suggested: float - gamma_suggested_peak_rp: float - combined_csv_path: Path - plot_path: Path - state_path: Path - results: pd.DataFrame = field(repr=False) - - -@dataclass -class BlazeStudyResult: - """Outcome of :func:`run_blaze_study`. - - Attributes: - d_spacing_nm: Resolved d-spacing the gratings were built with. - gamma: Gamma the gratings were built with (``config.gamma``). - blaze_suggested_deg: Blaze angle with the highest selected-order - efficiency at the target energy. - blaze_suggested_efficiency: Selected-order efficiency at that blaze. - combined_csv_path: Combined per-blaze theta-search summary table. - plot_path: Efficiency-versus-energy summary plot. - state_path: The updated state file. - results: The combined theta-search summary table. - """ - - d_spacing_nm: float - gamma: float - blaze_suggested_deg: float - blaze_suggested_efficiency: float - combined_csv_path: Path - plot_path: Path - state_path: Path - results: pd.DataFrame = field(repr=False) - - -def energy_to_wavelength_nm(energy_ev: float | Iterable[float]) -> float | np.ndarray: - """Convert photon energy in eV to wavelength in nm. - - Args: - energy_ev: One energy or an iterable of energies in eV. - - Returns: - The wavelength in nm, scalar for scalar input. - - Raises: - ValueError: If any energy is non-finite or non-positive. - """ - - energy = np.asarray(energy_ev, dtype=float) - if np.any(~np.isfinite(energy)) or np.any(energy <= 0.0): - raise ValueError(f"Photon energy must be finite and positive, got {energy_ev!r}") - wavelength = HC_EV_NM / energy - return float(wavelength) if wavelength.ndim == 0 else wavelength - - -def d_spacing_bounds_from_bragg_angles( - energy_ev: float, - angle_min_deg: float, - angle_max_deg: float, - *, - bragg_order: int = 1, -) -> tuple[float, float]: - """Return the d-spacing range whose Bragg grazing angle lies in an interval. - - Args: - energy_ev: Photon energy in eV. - angle_min_deg: Minimum grazing Bragg angle in degrees. - angle_max_deg: Maximum grazing Bragg angle in degrees. - bragg_order: Positive Bragg order. - - Returns: - ``(d_min_nm, d_max_nm)``; the larger angle maps to the smaller d. - - Raises: - ValueError: If the energy, order or angle interval is invalid. - """ - - energy = float(energy_ev) - angle_min = float(angle_min_deg) - angle_max = float(angle_max_deg) - order = int(bragg_order) - if not np.isfinite(energy) or energy <= 0.0: - raise ValueError(f"target energy must be finite and positive, got {energy_ev!r}") - if order <= 0: - raise ValueError(f"Bragg order must be positive, got {bragg_order!r}") - if not (0.0 < angle_min < angle_max < 90.0): - raise ValueError( - "Bragg grazing-angle limits must satisfy 0 < min < max < 90 deg; " - f"got {angle_min_deg!r}, {angle_max_deg!r}" - ) - wavelength_nm = float(energy_to_wavelength_nm(energy)) - d_min = order * wavelength_nm / (2.0 * np.sin(np.deg2rad(angle_max))) - d_max = order * wavelength_nm / (2.0 * np.sin(np.deg2rad(angle_min))) - return float(d_min), float(d_max) - - -def intersect_search_bounds( - derived_min_nm: float, - derived_max_nm: float, - practical_min_nm: float | None = None, - practical_max_nm: float | None = None, -) -> tuple[float, float]: - """Intersect a derived d-spacing interval with optional practical limits. - - Args: - derived_min_nm: Lower derived bound in nm. - derived_max_nm: Upper derived bound in nm. - practical_min_nm: Optional practical lower clamp in nm. - practical_max_nm: Optional practical upper clamp in nm. - - Returns: - The intersected ``(lower_nm, upper_nm)`` interval. - - Raises: - ValueError: If any interval is empty or non-positive, or the - intersection is empty. - """ - - derived_min = float(derived_min_nm) - derived_max = float(derived_max_nm) - practical_min = None if practical_min_nm is None else float(practical_min_nm) - practical_max = None if practical_max_nm is None else float(practical_max_nm) - if not (np.isfinite(derived_min) and np.isfinite(derived_max)) or derived_min <= 0.0: - raise ValueError( - f"Derived d-spacing bounds must be finite and positive: " - f"{derived_min_nm!r}, {derived_max_nm!r}" - ) - if derived_min >= derived_max: - raise ValueError( - f"Derived d-spacing interval is empty: ({derived_min:.8g}, {derived_max:.8g}) nm" - ) - if practical_min is not None and (not np.isfinite(practical_min) or practical_min <= 0.0): - raise ValueError(f"Practical d-spacing minimum must be positive, got {practical_min_nm!r}") - if practical_max is not None and (not np.isfinite(practical_max) or practical_max <= 0.0): - raise ValueError(f"Practical d-spacing maximum must be positive, got {practical_max_nm!r}") - if practical_min is not None and practical_max is not None and practical_min >= practical_max: - raise ValueError( - "Practical d-spacing interval is empty: " - f"({practical_min_nm!r}, {practical_max_nm!r}) nm" - ) - lower = max(derived_min, practical_min) if practical_min is not None else derived_min - upper = min(derived_max, practical_max) if practical_max is not None else derived_max - if lower >= upper: - raise ValueError( - "No usable d-spacing interval after intersecting bounds: " - f"derived=({derived_min:.8g}, {derived_max:.8g}) nm, " - f"practical=({practical_min_nm!r}, {practical_max_nm!r}) nm" - ) - return float(lower), float(upper) - - -def ensure_target_energy(energies: Iterable[float], target_energy_ev: float) -> np.ndarray: - """Return a sorted energy grid that contains the target energy exactly once. - - Args: - energies: Candidate energy grid in eV. - target_energy_ev: Energy that must appear in the returned grid. - - Returns: - The sorted, de-duplicated grid including ``target_energy_ev``. - - Raises: - ValueError: If the grid or the target is not finite and positive. - """ - - values = np.asarray(list(energies), dtype=float).reshape(-1) - target = float(target_energy_ev) - if values.size == 0 or np.any(~np.isfinite(values)) or np.any(values <= 0.0): - raise ValueError("Energy grid must contain finite, positive values") - if not np.isfinite(target) or target <= 0.0: - raise ValueError(f"Target energy must be finite and positive, got {target_energy_ev!r}") - close = np.isclose(values, target, rtol=0.0, atol=1.0e-12) - if np.any(close): - values[close] = target - else: - values = np.append(values, target) - return np.unique(np.sort(values)) - - -def select_target_energy_optimum( - results: pd.DataFrame, - *, - parameter_column: str, - metric_column: str, - target_energy_ev: float, - energy_column: str, -) -> tuple[float, float]: - """Return the tested parameter with the largest metric at the target energy. - - Args: - results: Long-format table with parameter, metric and energy columns. - parameter_column: Name of the swept-parameter column. - metric_column: Name of the metric column to maximize. - target_energy_ev: Energy at which to compare. - energy_column: Name of the energy column. - - Returns: - ``(parameter_value, metric_value)`` for the best row at the target. - - Raises: - KeyError: If a required column is missing. - ValueError: If no usable target-energy row exists. - """ - - missing = {energy_column, parameter_column, metric_column} - set(results.columns) - if missing: - raise KeyError(f"Results are missing required columns: {sorted(missing)}") - target = float(target_energy_ev) - target_rows = results[ - np.isclose(results[energy_column].astype(float), target, rtol=0.0, atol=1.0e-9) - ] - target_rows = target_rows.dropna(subset=[parameter_column, metric_column]) - if target_rows.empty: - raise ValueError(f"Results contain no usable target-energy row for {target:g} eV") - best = target_rows.loc[target_rows[metric_column].astype(float).idxmax()] - return float(best[parameter_column]), float(best[metric_column]) - - -def update_optimization_state(state_path: str | Path, updates: Mapping[str, Any]) -> dict[str, Any]: - """Merge ``updates`` into the JSON state file and return the full state. - - Args: - state_path: Path to ``optimization_state.json``. Created if absent. - updates: Keys to merge in. NumPy scalars and nested mappings/sequences - are converted to plain JSON types. - - Returns: - The merged state dictionary. - - Raises: - ValueError: If the existing file does not hold a JSON object. - """ - - path = Path(state_path) - state: dict[str, Any] = {} - if path.exists(): - with path.open(encoding="utf-8") as stream: - loaded = json.load(stream) - if not isinstance(loaded, dict): - raise ValueError(f"Optimization state must contain a JSON object: {path}") - state.update(loaded) - - def json_value(value: Any) -> Any: - if isinstance(value, np.generic): - return value.item() - if isinstance(value, Mapping): - return {str(key): json_value(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [json_value(item) for item in value] - return value - - state.update({str(key): json_value(value) for key, value in updates.items()}) - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w", encoding="utf-8") as stream: - json.dump(state, stream, indent=2, sort_keys=True) - stream.write("\n") - return state - - -def resolve_configured_value( - configured_value: float | str, - *, - state_path: str | Path, - suggestion_key: str, - parameter_name: str, -) -> float: - """Resolve a numeric configuration value or the ``"auto"`` sentinel. - - Args: - configured_value: A number, or ``"auto"`` to read the state file. - state_path: Path to ``optimization_state.json``. - suggestion_key: State key holding the upstream suggestion. - parameter_name: Name used in error messages. - - Returns: - The resolved positive, finite value. - - Raises: - ValueError: If a string other than ``"auto"`` is given, the state file - or key is missing, or the resolved value is not finite and positive. - """ - - if isinstance(configured_value, str): - if configured_value.strip().lower() != "auto": - raise ValueError( - f"{parameter_name} must be numeric or 'auto', got {configured_value!r}" - ) - path = Path(state_path) - if not path.exists(): - raise ValueError( - f"{parameter_name} is 'auto', but optimization state does not exist: {path}" - ) - with path.open(encoding="utf-8") as stream: - state = json.load(stream) - if suggestion_key not in state: - raise ValueError( - f"{parameter_name} is 'auto', but state has no {suggestion_key!r}; " - "run the preceding optimization stage first" - ) - value: Any = state[suggestion_key] - else: - value = configured_value - try: - resolved = float(value) - except (TypeError, ValueError) as error: - raise ValueError(f"Resolved {parameter_name} must be numeric, got {value!r}") from error - if not np.isfinite(resolved) or resolved <= 0.0: - raise ValueError(f"Resolved {parameter_name} must be finite and positive, got {resolved!r}") - return resolved - - -def _material_pair(material: Any) -> tuple[str, float]: - """Return a ``(name, density_g_cm3)`` pair for a material specification.""" - - if isinstance(material, MaterialSpec): - if material.density_g_cm3 is None: - raise ValueError(f"MaterialSpec {material.name!r} needs a density for this workflow") - return str(material.name), float(material.density_g_cm3) - if isinstance(material, (tuple, list)) and len(material) == 2: - return str(material[0]), float(material[1]) - raise ValueError(f"material must be a (name, density) pair or MaterialSpec, got {material!r}") - - -def _stage_energy_grid(config: MultilayerOptimizationConfig, stage: str) -> np.ndarray: - """Build the energy grid for one stage, honoring ``config.quick``.""" - - minimum = float(getattr(config, f"{stage}_energy_min_ev")) - maximum = float(getattr(config, f"{stage}_energy_max_ev")) - if minimum <= 0.0 or maximum < minimum: - raise ValueError(f"Invalid {stage} energy range: {minimum}, {maximum} eV") - if stage == "blaze": - points = int(config.blaze_energy_points) - energies = np.linspace(minimum, maximum, points) - if config.quick: - energies = np.linspace(minimum, maximum, max(2, int(config.blaze_energy_quick_points))) - else: - step = float(getattr(config, f"{stage}_energy_step_ev")) - if step <= 0.0: - raise ValueError(f"{stage}_energy_step_ev must be positive") - if config.quick: - step = float(getattr(config, f"{stage}_energy_quick_step_ev")) - if step <= 0.0: - raise ValueError(f"{stage}_energy_quick_step_ev must be positive") - energies = np.arange(minimum, maximum + 0.5 * step, step) - return ensure_target_energy(energies, config.target_energy_ev) - - -def _rounded_d_grid( - lower_nm: float, upper_nm: float, points: int, required_nm: float -) -> np.ndarray: - """Return a 0.1 nm-rounded d-spacing grid that includes the geometry value. - - Args: - lower_nm: Lower edge of the search interval in nm. - upper_nm: Upper edge of the search interval in nm. - points: Number of ``linspace`` samples before rounding and de-duplication. - required_nm: Geometry d-spacing that must appear in the grid. - - Returns: - The sorted, de-duplicated candidate grid rounded to 0.1 nm. - - Raises: - ValueError: If the rounded interval cannot contain ``required_nm`` or - fewer than two unique candidates survive. - """ - - if points < 2: - raise ValueError(f"d_spacing_points must be at least 2, got {points!r}") - rounded_min = np.ceil(lower_nm * 10.0) / 10.0 - rounded_max = np.floor(upper_nm * 10.0) / 10.0 - required = round(float(required_nm), 1) - if rounded_min > rounded_max or not (rounded_min <= required <= rounded_max): - raise ValueError( - "The automatic d-spacing interval cannot contain the one-decimal " - f"geometry value {required:.1f} nm: interval=({lower_nm:.8g}, {upper_nm:.8g}) nm" - ) - grid = np.unique(np.round(np.linspace(rounded_min, rounded_max, int(points)), decimals=1)) - if not np.any(np.isclose(grid, required, rtol=0.0, atol=1.0e-9)): - grid[np.argmin(np.abs(grid - required))] = required - grid = np.unique(np.sort(grid)) - if grid.size < 2: - raise ValueError( - "The automatic d-spacing interval provides fewer than two unique 0.1 nm candidates" - ) - return grid - - -def _reflectivity_metric(config: MultilayerOptimizationConfig) -> str: - """Return the reflectivity column matching the configured polarization.""" - - return "peak_rp" if normalize_polarization(config.polarization) == "p" else "peak_rs" - - -def _reflectivity_curve( - config: MultilayerOptimizationConfig, - d_spacing_nm: float, - gamma: float, - output_dir: Path, - energies_ev: np.ndarray, -) -> pd.DataFrame: - """Run the XRT reflectivity engine for one ``(d, gamma)`` over the grid.""" - - output_dir.mkdir(parents=True, exist_ok=True) - engine = MultilayerReflectivity( - config.material_a, - d_spacing_nm * gamma, - config.material_b, - d_spacing_nm * (1.0 - gamma), - config.n_bilayers, - save_recap=output_dir, - individuals=config.xrt_individuals, - ) - return engine.reflectivity_vs_energy( - energies_ev, - bragg_order=config.multilayer_bragg_order, - window_deg=config.xrt_window_deg, - angle_range=None, - angle_points=config.xrt_angle_points, - min_angle_deg=config.xrt_min_angle_deg, - ) - - -def _target_metric_for_parameter( - results: pd.DataFrame, - *, - parameter_column: str, - parameter_value: float, - metric_column: str, - target_energy_ev: float, - energy_column: str, -) -> float: - """Return the metric for one parameter value at the target energy.""" - - energy_match = np.isclose( - results[energy_column].astype(float), float(target_energy_ev), rtol=0.0, atol=1.0e-9 - ) - parameter_match = np.isclose( - results[parameter_column].astype(float), float(parameter_value), rtol=0.0, atol=1.0e-9 - ) - rows = results[energy_match & parameter_match] - rows = rows.dropna(subset=[metric_column]) - if rows.empty: - raise ValueError( - f"No {metric_column} result for {parameter_value:g} at {target_energy_ev:g} eV" - ) - return float(rows.iloc[0][metric_column]) - - -def _plot_reflectivity( - config: MultilayerOptimizationConfig, - results: pd.DataFrame, - *, - parameter_column: str, - metric_column: str, - suggested_value: float, - suggested_metric: float, - energy_column: str, - label_format: str, - title: str, - output_path: Path, -) -> None: - """Plot the per-parameter metric versus energy, highlighting the suggestion.""" - - import matplotlib.pyplot as plt - - target_energy = float(config.target_energy_ev) - figure, axis = plt.subplots(figsize=(10, 6)) - for value, curve in results.groupby(parameter_column, sort=True): - selected = bool(np.isclose(float(value), suggested_value, rtol=0.0, atol=1.0e-9)) - axis.plot( - curve[energy_column], - curve[metric_column], - marker="o", - markersize=2, - linewidth=2.6 if selected else 0.9, - alpha=1.0 if selected else 0.7, - label=label_format.format(value=float(value)) + (" (suggested)" if selected else ""), - zorder=3 if selected else 2, - ) - axis.axvline(target_energy, color="black", linestyle="--", linewidth=1) - axis.plot( - [target_energy], [suggested_metric], marker="o", markersize=7, color="black", zorder=5 - ) - axis.annotate( - f"{target_energy:g} eV\n{label_format.format(value=suggested_value)}\n" - f"{metric_column} = {suggested_metric:.4g}", - xy=(target_energy, suggested_metric), - xytext=(10, 12), - textcoords="offset points", - fontsize=9, - bbox={"boxstyle": "round", "fc": "white", "alpha": 0.85}, - ) - axis.set_xlabel("Photon energy (eV)") - axis.set_ylabel(f"Peak reflectivity ({normalize_polarization(config.polarization)}-polarized)") - axis.set_title(title) - axis.grid(True, alpha=0.3) - axis.legend(ncol=2, fontsize=8) - figure.tight_layout() - output_path.parent.mkdir(parents=True, exist_ok=True) - figure.savefig(output_path, dpi=150) - plt.close(figure) - - -def _build_grating( - config: MultilayerOptimizationConfig, - d_spacing_nm: float, - gamma: float, - blaze_angle_deg: float, -) -> BlazedGrating: - """Build the multilayer-coated blazed grating for one stage-2 case.""" - - material_a_spec = MaterialSpec(*_material_pair(config.material_a)) - material_b_spec = MaterialSpec(*_material_pair(config.material_b)) - substrate_spec = MaterialSpec(*_material_pair(config.substrate_material)) - stack = MultilayerStack( - substrate_material=substrate_spec, - material_a=material_a_spec, - material_b=material_b_spec, - d_period_nm=d_spacing_nm, - gamma=gamma, - n_bilayers=int(config.n_bilayers), - top_material=material_b_spec, - ) - anti_blaze_kwargs = ( - {} - if config.anti_blaze_angle_deg == 0.0 - else {"anti_blaze_angle_deg": float(config.anti_blaze_angle_deg)} - ) - return BlazedGrating( - period_lpermm=float(config.grating_density_lpermm), - blaze_angle_deg=float(blaze_angle_deg), - coating_stack=stack, - substrate_material=substrate_spec, - x_resolution_nm=float(config.grax_x_resolution_nm), - z_resolution_nm=float(config.grax_z_resolution_nm), - **anti_blaze_kwargs, - ) - - -def _run_blaze_case( - config: MultilayerOptimizationConfig, - d_spacing_nm: float, - gamma: float, - blaze_angle_deg: float, - energies_ev: np.ndarray, - output_dir: Path, -) -> pd.DataFrame: - """Run graxPy's theta-search sweep for one blaze angle and return its summary.""" - - output_dir.mkdir(parents=True, exist_ok=True) - sweep = run_multilayer_theta_search_sweep( - grating=_build_grating(config, d_spacing_nm, gamma, blaze_angle_deg), - energies_ev=energies_ev, - output_dir=output_dir, - diffraction_order=int(config.diffraction_order), - multilayer_bragg_order=int(config.multilayer_bragg_order), - rough_scan_half_width_deg=float(config.rough_scan_half_width_deg), - rough_scan_points=int(config.rough_scan_points), - rough_fourier_orders=int(config.rough_fourier_orders), - rough_x_resolution_nm=float(config.rough_x_resolution_nm), - rough_z_resolution_nm=float(config.rough_z_resolution_nm), - fine_scan_half_width_deg=float(config.fine_scan_half_width_deg), - fine_scan_points=int(config.fine_scan_points), - fine_fourier_orders=int(config.fine_fourier_orders), - fine_x_resolution_nm=float(config.fine_x_resolution_nm), - fine_z_resolution_nm=float(config.fine_z_resolution_nm), - final_fourier_orders=int(config.final_fourier_orders), - final_x_resolution_nm=float(config.final_x_resolution_nm), - final_z_resolution_nm=float(config.final_z_resolution_nm), - roughness_sigma_nm=config.roughness_sigma_nm, - precise_peak_selection_mode=str(config.precise_peak_selection_mode), - retry_on_selected_efficiency_zero=bool(config.retry_on_selected_efficiency_zero), - retry_selected_efficiency_threshold=float(config.retry_selected_efficiency_threshold), - max_zero_efficiency_retries=int(config.max_zero_efficiency_retries), - max_workers=config.max_workers, - show_progress=bool(config.show_progress), - live_plot=bool(config.live_plot), - on_error=str(config.on_error), - checkpoint_dir=output_dir / "checkpoints", - checkpoint_interval=int(config.checkpoint_interval), - resume=bool(config.resume), - theta_tracking_mode=str(config.theta_tracking_mode), - max_tracking_energy_step_ev=config.max_tracking_energy_step_ev, - save_profile_plot=bool(config.save_profile_plot), - save_stack_plot=bool(config.save_stack_plot), - backend=str(config.backend), - solver=str(config.solver), - polarization=str(config.polarization), - ) - return pd.read_csv(sweep.summary_csv_path) - - -def run_d_spacing_study(config: MultilayerOptimizationConfig) -> DSpacingStudyResult: - """Run stage 0: derive and scan the bilayer d-spacing. - - Derives the grazing angle at the target energy and CFF, converts it to a - d-spacing with the first-order Bragg law, builds a practical 0.1 nm-rounded - scan grid that includes the geometry value, runs XRT reflectivity over the - energy grid for every candidate, and writes ``d_suggested_nm`` (the geometry - value) plus the numerically best d (diagnostic) to the state file. - - Args: - config: The workflow configuration. - - Returns: - A :class:`DSpacingStudyResult` with the suggestion, diagnostics and - artifact paths. - """ - - target_energy = float(config.target_energy_ev) - wavelength_nm = float(energy_to_wavelength_nm(target_energy)) - geometry_angle = float( - monochromator_grazing_angles_deg( - [target_energy], - period_lpermm=float(config.grating_density_lpermm), - diffraction_order=int(config.diffraction_order), - cff=float(config.cff), - )[0] - ) - d_geometry = wavelength_nm / (2.0 * np.sin(np.deg2rad(geometry_angle))) - geometry_min = d_geometry * (1.0 - float(config.d_spacing_relative_range)) - geometry_max = d_geometry * (1.0 + float(config.d_spacing_relative_range)) - angle_derived_min, angle_derived_max = d_spacing_bounds_from_bragg_angles( - target_energy, - float(config.bragg_angle_min_deg), - float(config.bragg_angle_max_deg), - bragg_order=int(config.multilayer_bragg_order), - ) - lower, upper = intersect_search_bounds( - geometry_min, geometry_max, angle_derived_min, angle_derived_max - ) - lower, upper = intersect_search_bounds( - lower, - upper, - float(config.d_spacing_min_practical_nm), - float(config.d_spacing_max_practical_nm), - ) - d_suggested = round(d_geometry, 1) - d_values = _rounded_d_grid(lower, upper, int(config.d_spacing_points), d_suggested) - energies = _stage_energy_grid(config, "d_spacing") - - results_dir = config.d_spacing_results_dir - curves = [] - for d_spacing in d_values: - print(f"Calculating multilayer reflectivity, d = {d_spacing:.1f} nm") - curve = _reflectivity_curve( - config, - float(d_spacing), - float(config.gamma), - results_dir / f"d_{d_spacing:.1f}nm", - energies, - ) - curve.insert(0, "d_spacing_nm", float(d_spacing)) - curves.append(curve) - combined = pd.concat(curves, ignore_index=True) - - metric = _reflectivity_metric(config) - geometry_metric = _target_metric_for_parameter( - combined, - parameter_column="d_spacing_nm", - parameter_value=d_suggested, - metric_column=metric, - target_energy_ev=target_energy, - energy_column="energy_ev", - ) - numerical_best_d, numerical_best_metric = select_target_energy_optimum( - combined, - parameter_column="d_spacing_nm", - metric_column=metric, - target_energy_ev=target_energy, - energy_column="energy_ev", - ) - update_optimization_state( - config.state_path, - { - "target_energy_eV": target_energy, - "wavelength_nm": wavelength_nm, - "grating_grazing_angle_deg": geometry_angle, - "d_geometry_estimate_nm": d_geometry, - "d_geometry_search_min_nm": geometry_min, - "d_geometry_search_max_nm": geometry_max, - "d_search_min_nm": lower, - "d_search_max_nm": upper, - "d_suggested_nm": d_suggested, - "d_suggested_peak_rp": geometry_metric, - "d_reflectivity_best_nm": numerical_best_d, - "d_reflectivity_best_peak_rp": numerical_best_metric, - }, - ) - csv_path = results_dir / "d_spacing_study.csv" - combined.to_csv(csv_path, index=False) - plot_path = config.plot_dir / "0_d_spacing_study.png" - _plot_reflectivity( - config, - combined, - parameter_column="d_spacing_nm", - metric_column=metric, - suggested_value=d_suggested, - suggested_metric=geometry_metric, - energy_column="energy_ev", - label_format="d = {value:.1f} nm", - title=f"Multilayer d-spacing study (geometry d = {d_geometry:.3f} nm)", - output_path=plot_path, - ) - print( - f"Target {target_energy:g} eV: grazing angle = {geometry_angle:.6f} deg, " - f"geometry d = {d_geometry:.6f} nm, suggested d = {d_suggested:.1f} nm" - ) - print(f"Reflectivity at suggested d: {geometry_metric:.6g}") - print(f"Numerical best d (diagnostic): {numerical_best_d:.1f} nm ({numerical_best_metric:.6g})") - print(f"Results CSV: {csv_path}") - print(f"Optimization state: {config.state_path}") - return DSpacingStudyResult( - geometry_grazing_angle_deg=geometry_angle, - geometry_d_nm=float(d_geometry), - d_suggested_nm=float(d_suggested), - d_suggested_peak_rp=float(geometry_metric), - d_reflectivity_best_nm=float(numerical_best_d), - d_reflectivity_best_peak_rp=float(numerical_best_metric), - search_min_nm=float(lower), - search_max_nm=float(upper), - combined_csv_path=csv_path, - plot_path=plot_path, - state_path=config.state_path, - results=combined, - ) - - -def run_gamma_study(config: MultilayerOptimizationConfig) -> GammaStudyResult: - """Run stage 1: scan the bilayer thickness ratio at the selected d-spacing. - - Resolves ``config.d_spacing_nm`` (numeric, or ``"auto"`` from the state - file), scans ``gamma`` over ``[gamma_min, gamma_max]`` in ``gamma_step`` - increments, and records the gamma with the highest peak reflectivity at the - target energy. The config is never modified. - - Args: - config: The workflow configuration. - - Returns: - A :class:`GammaStudyResult` with the suggested gamma and artifact paths. - """ - - target_energy = float(config.target_energy_ev) - d_spacing = resolve_configured_value( - config.d_spacing_nm, - state_path=config.state_path, - suggestion_key="d_suggested_nm", - parameter_name="d_spacing_nm", - ) - gamma_values = np.round( - np.arange( - float(config.gamma_min), - float(config.gamma_max) + 0.5 * float(config.gamma_step), - float(config.gamma_step), - ), - 3, - ) - energies = _stage_energy_grid(config, "gamma") - - results_dir = config.gamma_results_dir - curves = [] - for gamma in gamma_values: - print(f"Calculating multilayer reflectivity, gamma = {gamma:.3f}") - curve = _reflectivity_curve( - config, d_spacing, float(gamma), results_dir / f"gamma_{gamma:.3f}", energies - ) - curve.insert(0, "gamma", float(gamma)) - curves.append(curve) - combined = pd.concat(curves, ignore_index=True) - - metric = _reflectivity_metric(config) - suggested_gamma, suggested_metric = select_target_energy_optimum( - combined, - parameter_column="gamma", - metric_column=metric, - target_energy_ev=target_energy, - energy_column="energy_ev", - ) - update_optimization_state( - config.state_path, - {"gamma_suggested": suggested_gamma, "gamma_suggested_peak_rp": suggested_metric}, - ) - csv_path = results_dir / "gamma_study.csv" - combined.to_csv(csv_path, index=False) - plot_path = config.plot_dir / "1_gamma_study.png" - _plot_reflectivity( - config, - combined, - parameter_column="gamma", - metric_column=metric, - suggested_value=suggested_gamma, - suggested_metric=suggested_metric, - energy_column="energy_ev", - label_format="gamma = {value:.3f}", - title=f"Multilayer gamma study (d = {d_spacing:.3f} nm)", - output_path=plot_path, - ) - print( - f"Target {target_energy:g} eV, d = {d_spacing:.6f} nm: suggested gamma = " - f"{suggested_gamma:.3f} ({metric} = {suggested_metric:.6g})" - ) - print(f"Results CSV: {csv_path}") - print(f"Optimization state: {config.state_path}") - return GammaStudyResult( - d_spacing_nm=float(d_spacing), - gamma_suggested=float(suggested_gamma), - gamma_suggested_peak_rp=float(suggested_metric), - combined_csv_path=csv_path, - plot_path=plot_path, - state_path=config.state_path, - results=combined, - ) - - -def run_blaze_study(config: MultilayerOptimizationConfig) -> BlazeStudyResult: - """Run stage 2: scan the blaze angle with graxPy's theta search. - - Resolves ``config.d_spacing_nm`` (numeric, or ``"auto"`` from the state - file), uses ``config.gamma`` directly, builds the multilayer-coated blazed - grating for each blaze angle in the scan, runs - :func:`grax.run_multilayer_theta_search_sweep` per blaze angle, and records - the blaze angle with the highest selected-order efficiency at the target - energy. - - Args: - config: The workflow configuration. - - Returns: - A :class:`BlazeStudyResult` with the suggested blaze angle and artifact - paths. - """ - - target_energy = float(config.target_energy_ev) - d_spacing = resolve_configured_value( - config.d_spacing_nm, - state_path=config.state_path, - suggestion_key="d_suggested_nm", - parameter_name="d_spacing_nm", - ) - gamma = float(config.gamma) - blaze_center = float(config.blaze_angle_deg) - blaze_half_range = float(config.blaze_angle_half_range_deg) - blaze_points = int(config.blaze_angle_points) - blaze_values = np.linspace( - blaze_center - blaze_half_range, blaze_center + blaze_half_range, blaze_points - ) - if np.any(blaze_values <= 0.0): - raise ValueError(f"Blaze-angle scan must be positive: {blaze_values.tolist()}") - energies = _stage_energy_grid(config, "blaze") - - results_dir = config.blaze_results_dir - curves = [] - for blaze in blaze_values: - print(f"Running theta search, blaze = {blaze:.4f} deg") - curve = _run_blaze_case( - config, d_spacing, gamma, float(blaze), energies, results_dir / f"blaze_{blaze:.4f}deg" - ) - curve.insert(0, "blaze_angle_deg", float(blaze)) - curves.append(curve) - combined = pd.concat(curves, ignore_index=True) - - suggested_blaze, suggested_efficiency = select_target_energy_optimum( - combined, - parameter_column="blaze_angle_deg", - metric_column="selected_efficiency", - target_energy_ev=target_energy, - energy_column="energy_ev", - ) - update_optimization_state( - config.state_path, - { - "blaze_suggested_deg": suggested_blaze, - "blaze_suggested_efficiency": suggested_efficiency, - }, - ) - csv_path = results_dir / "blaze_study.csv" - combined.to_csv(csv_path, index=False) - plot_path = config.plot_dir / "2_blaze_study.png" - _plot_reflectivity( - config, - combined, - parameter_column="blaze_angle_deg", - metric_column="selected_efficiency", - suggested_value=suggested_blaze, - suggested_metric=suggested_efficiency, - energy_column="energy_ev", - label_format="blaze = {value:.4f} deg", - title=f"Blaze study (d = {d_spacing:.3f} nm, gamma = {gamma:.3f})", - output_path=plot_path, - ) - print( - f"Target {target_energy:g} eV, d = {d_spacing:.6f} nm, gamma = {gamma:.3f}: " - f"suggested blaze = {suggested_blaze:.4f} deg (efficiency = {suggested_efficiency:.6g})" - ) - print(f"Results CSV: {csv_path}") - print(f"Optimization state: {config.state_path}") - return BlazeStudyResult( - d_spacing_nm=float(d_spacing), - gamma=gamma, - blaze_suggested_deg=float(suggested_blaze), - blaze_suggested_efficiency=float(suggested_efficiency), - combined_csv_path=csv_path, - plot_path=plot_path, - state_path=config.state_path, - results=combined, - ) diff --git a/src/grax/simulation/models.py b/src/grax/simulation/models.py index 0a8b645..f64a4cd 100644 --- a/src/grax/simulation/models.py +++ b/src/grax/simulation/models.py @@ -340,6 +340,9 @@ class MultilayerThetaSearchSweepResult: total_elapsed_seconds: Total accumulated sweep runtime in seconds, including prior resumed runs when checkpoint metadata is available. current_run_elapsed_seconds: Runtime spent in the current process invocation. + stopped_early: Whether a ``stop_event`` cut the sweep short, leaving some + requested energies unsolved. + resolved_max_workers: Worker count the sweep actually resolved and used. """ batch_result: BatchSimulationResult @@ -352,3 +355,5 @@ class MultilayerThetaSearchSweepResult: stack_plot_path: Path | None = None total_elapsed_seconds: float = 0.0 current_run_elapsed_seconds: float = 0.0 + stopped_early: bool = False + resolved_max_workers: int = 1 diff --git a/src/grax/simulation/theta_search_sweep.py b/src/grax/simulation/theta_search_sweep.py index 9a4aed7..33ad1e3 100644 --- a/src/grax/simulation/theta_search_sweep.py +++ b/src/grax/simulation/theta_search_sweep.py @@ -8,8 +8,9 @@ import json import logging import multiprocessing as mp +import threading import time -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -83,6 +84,37 @@ def _load_checkpoint_case_results_for_ids( loaded[case_result.case_id] = case_result return loaded + +def _terminate_worker_pool( + executor: concurrent.futures.ProcessPoolExecutor, + futures: dict[concurrent.futures.Future, object], +) -> None: + """Cancel queued cases and stop every live worker process immediately. + + Terminating the workers breaks the pool, which is fine because the caller + abandons it: the point is that a solve already in flight dies now rather than + running to completion. Results reached before this call are already on disk + when checkpointing is on. + """ + + processes = [ + process + for process in getattr(executor, "_processes", {}).values() + if process is not None + ] + for future in futures: + future.cancel() + executor.shutdown(wait=False, cancel_futures=True) + for process in processes: + if process.is_alive(): + process.terminate() + for process in processes: + process.join(timeout=2.0) + if process.is_alive(): + process.kill() + process.join(timeout=2.0) + + def _adaptive_scan_half_widths( *, energy_ev: float, @@ -216,6 +248,8 @@ def run_multilayer_theta_search_sweep( retry_selected_efficiency_threshold: float = 1e-4, max_zero_efficiency_retries: int = 3, max_workers: MaxWorkers = None, + stop_event: threading.Event | None = None, + on_worker_pids_changed: Callable[[set[int]], None] | None = None, show_progress: bool = True, live_plot: bool = False, on_error: ErrorPolicy = "fail_fast", @@ -266,6 +300,13 @@ def run_multilayer_theta_search_sweep( max_workers: Optional batch worker count. ``"auto"`` calibrates from one completed theta-search case and available system memory before launching the remaining parallel work. + stop_event: Optional stop signal. Once it is set, queued energies are no + longer submitted, the live worker processes are terminated, and the + sweep returns with ``stopped_early=True`` and the energies solved so + far. Passing one also forces worker-process execution even at + ``max_workers=1``, so a single-energy sweep stays interruptible. + on_worker_pids_changed: Optional callback receiving the current worker + process IDs, so a caller can track or signal them. show_progress: Whether to show a progress bar during execution. live_plot: Whether to update the standard batch live plot. on_error: Per-case error policy forwarded to the batch runner. ``continue`` @@ -315,6 +356,14 @@ def run_multilayer_theta_search_sweep( raise ValueError("retry_selected_efficiency_threshold must be finite and >= 0.0.") simulation_api = _simulation_api() effective_workers = simulation_api._resolve_max_workers(max_workers) + + def _stop_requested() -> bool: + return stop_event is not None and stop_event.is_set() + + def _emit_worker_pids(worker_pids: set[int]) -> None: + if on_worker_pids_changed is not None: + on_worker_pids_changed(set(worker_pids)) + output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) theta_scan_directory = output_path / "theta_scans" @@ -893,6 +942,10 @@ def _finalize_tracked_result( previous_theta_deg=tracking_previous_theta_deg, ) ) + if should_try_bragg_fallback and _stop_requested(): + # The fallback re-solve runs in this process and cannot be killed; + # skip it rather than make an abort wait out a whole extra search. + should_try_bragg_fallback = False if should_try_bragg_fallback: bragg_fallback_triggered = True fallback_case = dict(prepared_case) @@ -1060,6 +1113,7 @@ def _write_single_theta_scan_artifacts(case: CaseExecutionResult, *, skip_if_exi previous_elapsed_seconds, ) + stopped_early = False try: resumed_artifacts_written = 0 resumed_artifacts_skipped = 0 @@ -1097,7 +1151,7 @@ def _write_single_theta_scan_artifacts(case: CaseExecutionResult, *, skip_if_exi axis=live_axis, successful_cases=[case for case in cases_result if case.status == "ok"], ) - if max_workers == "auto" and pending_order: + if max_workers == "auto" and pending_order and not _stop_requested(): calibration_energy_ev, calibration_case_id = pending_order[0] resumed_successful_cases = [case for case in resumed_cases_sorted if case.status == "ok"] previous_successful_case = ( @@ -1160,7 +1214,11 @@ def _write_single_theta_scan_artifacts(case: CaseExecutionResult, *, skip_if_exi axis=live_axis, successful_cases=[case for case in cases_result if case.status == "ok"], ) - if effective_workers == 1: + # A caller that supplies a stop event always goes through worker processes, + # even at one worker: an in-process solve cannot be interrupted, but a + # child process can be terminated. Callers without one keep the cheaper + # in-process path. + if effective_workers == 1 and stop_event is None: resumed_successful_cases = [case for case in resumed_cases_sorted if case.status == "ok"] previous_successful_case: CaseExecutionResult | None = ( max(resumed_successful_cases, key=lambda case: case.energy_ev) @@ -1168,6 +1226,9 @@ def _write_single_theta_scan_artifacts(case: CaseExecutionResult, *, skip_if_exi else None ) for serial_index, (energy_ev, case_id) in enumerate(pending_order): + if _stop_requested(): + stopped_early = True + break _set_progress_postfix(active=1, queued=max(len(pending_order) - serial_index - 1, 0), completed=len(cases_result)) case = pending_by_case_id[case_id] tracked_case, single = _run_with_tracking( @@ -1240,7 +1301,11 @@ def _latest_available_lower_success(energy_ev: float) -> CaseExecutionResult | N return max(candidates, key=lambda existing: existing.energy_ev) while pending_cursor < len(pending_order) or futures: - while pending_cursor < len(pending_order) and len(futures) < effective_workers: + while ( + pending_cursor < len(pending_order) + and len(futures) < effective_workers + and not _stop_requested() + ): energy_ev, case_id = pending_order[pending_cursor] case = pending_by_case_id[case_id] previous_successful_case = _latest_available_lower_success(energy_ev) @@ -1259,6 +1324,16 @@ def _latest_available_lower_success(energy_ev: float) -> CaseExecutionResult | N auto_classification, ) pending_cursor += 1 + if on_worker_pids_changed is not None: + _emit_worker_pids( + { + process.pid + for process in getattr( + executor, "_processes", {} + ).values() + if process is not None and process.pid is not None + } + ) _set_progress_postfix( active=len(futures), queued=len(pending_order) - pending_cursor, @@ -1266,8 +1341,29 @@ def _latest_available_lower_success(energy_ev: float) -> CaseExecutionResult | N ) if not futures: + if _stop_requested(): + stopped_early = True + break continue - completed_future = next(concurrent.futures.as_completed(futures)) + if stop_event is None: + completed_future = next(concurrent.futures.as_completed(futures)) + else: + # Poll rather than block: a blocking wait cannot notice the + # stop event, so an abort would still have to sit out the + # case currently in flight. + done, _ = concurrent.futures.wait( + futures, + timeout=0.25, + return_when=concurrent.futures.FIRST_COMPLETED, + ) + if _stop_requested(): + stopped_early = True + _terminate_worker_pool(executor, futures) + futures.clear() + break + if not done: + continue + completed_future = next(iter(done)) index, case_id, energy_ev, prepared_case, center_mode, auto_classification = futures.pop( completed_future ) @@ -1355,6 +1451,7 @@ def _latest_available_lower_success(energy_ev: float) -> CaseExecutionResult | N successful_cases=[case for case in cases_result if case.status == "ok"], ) finally: + _emit_worker_pids(set()) current_run_elapsed_seconds = float(time.perf_counter() - run_started_monotonic) total_elapsed_seconds = float(previous_elapsed_seconds + current_run_elapsed_seconds) if checkpoint_metadata_file is not None: @@ -1613,4 +1710,6 @@ def _latest_available_lower_success(energy_ev: float) -> CaseExecutionResult | N stack_plot_path=stack_plot_path, total_elapsed_seconds=total_elapsed_seconds, current_run_elapsed_seconds=current_run_elapsed_seconds, + stopped_early=stopped_early, + resolved_max_workers=int(effective_workers), ) diff --git a/src/grax/web/app.py b/src/grax/web/app.py index b77656b..0a698a9 100644 --- a/src/grax/web/app.py +++ b/src/grax/web/app.py @@ -36,6 +36,24 @@ from grax.materials import available_material_symbols, material_density_catalog, material_density_g_cm3 from grax.simulation.core import normalize_polarization +from .multilayer_design_studies import ( + STAGE_LABELS, + STAGES, + SURVEY_CELL_WARNING_THRESHOLD, + MultilayerDesignStudyStore, + build_design_config, + downstream_stages, + flatten_config_values, + offline_script_filename, + parse_study_config, + render_offline_script, + resolve_designs, + stages_invalidated_by, + study_config_defaults, + study_form_sections, + survey_cell_count, + survey_design_options, +) from .persistence import GratingStore, build_grating_from_spec from .runs import RunStore @@ -121,6 +139,10 @@ class ActiveRunState: stop_event: threading.Event = field(default_factory=threading.Event) worker_thread: threading.Thread | None = None simulation_pids: set[int] = field(default_factory=set) + # Points already on disk when the run started (a resumed checkpoint). They + # count towards progress but not towards this run's rate, so the ETA has to + # subtract them. + resumed_points: int = 0 def _active_runs(app: Any) -> dict[str, ActiveRunState]: @@ -147,6 +169,7 @@ def create_app(*, data_dir: str | Path | None = None): try: from flask import ( Flask, + Response, abort, jsonify, redirect, @@ -605,6 +628,281 @@ def plot_delete(plot_id: str): return redirect(url_for("plot_index")) return render_template("plot_delete.html", plot=manifest) + # ------------------------------------------------------------------ # + # Multilayer-design studies # + # ------------------------------------------------------------------ # + def _design_store() -> MultilayerDesignStudyStore: + return MultilayerDesignStudyStore(app.config["GRAx_DATA_DIR"] / "multilayer_designs") + + def _design_study_or_404(study_id: str) -> dict[str, Any]: + try: + return _design_store().load(study_id) + except (ValueError, FileNotFoundError, OSError): + abort(404) + + def _design_view_model(manifest: dict[str, Any]) -> dict[str, Any]: + store = _design_store() + study_dir = store.study_dir(str(manifest["id"])) + options = survey_design_options(study_dir) + stages = [] + for stage in STAGES: + state = dict(manifest["stages"][stage]) + state["stage"] = stage + state["label"] = STAGE_LABELS[stage] + stages.append(state) + config = manifest["config"] + + def _material_name(key: str) -> str: + # Materials are stored as [name, density] pairs. + value = config.get(key) + return str(value[0]) if isinstance(value, list) and value else "" + + plot_meta = { + "coating_label": ( + config.get("coating_label") + or f"{_material_name('material_a')}/{_material_name('material_b')}" + ), + "target_energy_ev": config.get("target_energy_ev"), + "diffraction_order": config.get("diffraction_order"), + } + return { + "study": manifest, + "stages": stages, + "design_options": options, + "design_options_json": json.dumps(options), + "survey_plot_meta_json": json.dumps(plot_meta), + # Interactive survey plots need plotly.js inlined; without plotly the + # template falls back to the PNGs the library already wrote. + "plotly_bundle": _plotly_bundle_text() if get_plotlyjs is not None else None, + "survey_cells": survey_cell_count(manifest["config"]), + } + + @app.get("/multilayer-design") + def multilayer_design_index(): + return render_template( + "multilayer_design_index.html", + studies=_design_store().list(), + stage_labels=STAGE_LABELS, + stages=STAGES, + ) + + @app.post("/multilayer-design") + def multilayer_design_create(): + if request.form.get("action") == "delete": + _design_store().delete_many(request.form.getlist("delete_study_id")) + return redirect(url_for("multilayer_design_index")) + display_name = request.form.get("display_name", "").strip() or "Multilayer design" + try: + config = parse_study_config(request.form) + build_design_config( + app.config["GRAx_DATA_DIR"] / "multilayer_designs" / "_validate", config + ) + except (TypeError, ValueError) as error: + abort(400, str(error)) + study = _design_store().create( + display_name=display_name, + config=config, + auto_energy_scan="best" if request.form.get("auto_energy_scan") else "none", + ) + return redirect(url_for("multilayer_design_detail", study_id=study["id"])) + + @app.get("/multilayer-design/new") + def multilayer_design_new(): + # Seed from the newest study: the settings someone last tuned -- the + # advanced scan parameters above all -- are a far better starting point + # than the dataclass defaults. + recent = _design_store().list() + seed = recent[0] if recent else None + return render_template( + "multilayer_design_form.html", + seeded_from=seed, + defaults=flatten_config_values( + seed["config"] if seed else study_config_defaults() + ), + basic_sections=study_form_sections(advanced=False), + advanced_sections=study_form_sections(advanced=True), + materials=available_material_symbols(), + material_density_map=dict(material_density_catalog()), + cell_warning_threshold=SURVEY_CELL_WARNING_THRESHOLD, + ) + + @app.get("/multilayer-design//edit") + def multilayer_design_edit(study_id: str): + manifest = _design_study_or_404(study_id) + return render_template( + "multilayer_design_form.html", + study=manifest, + defaults=flatten_config_values(manifest["config"]), + basic_sections=study_form_sections(advanced=False), + advanced_sections=study_form_sections(advanced=True), + materials=available_material_symbols(), + material_density_map=dict(material_density_catalog()), + cell_warning_threshold=SURVEY_CELL_WARNING_THRESHOLD, + ) + + @app.post("/multilayer-design//edit") + def multilayer_design_update(study_id: str): + store = _design_store() + manifest = _design_study_or_404(study_id) + if any( + _is_run_active(app, _multilayer_design_job_key(study_id, stage)) for stage in STAGES + ): + abort(409, "Stop the running stage before changing the parameters.") + try: + config = parse_study_config(request.form, base=manifest["config"]) + build_design_config(store.study_dir(study_id), config) + except (TypeError, ValueError) as error: + abort(400, str(error)) + # Results computed under the old parameters are kept but flagged, so it + # is obvious which of them the edit left behind. + for stage in stages_invalidated_by(manifest["config"], config): + if manifest["stages"][stage]["status"] in {"completed", "aborted"}: + manifest["stages"][stage]["status"] = "stale" + manifest["config"] = config + manifest["display_name"] = ( + request.form.get("display_name", "").strip() or manifest["display_name"] + ) + manifest["auto_energy_scan"] = ( + "best" if request.form.get("auto_energy_scan") else "none" + ) + store.save(manifest) + return redirect(url_for("multilayer_design_detail", study_id=study_id)) + + @app.get("/multilayer-design/") + def multilayer_design_detail(study_id: str): + manifest = _design_study_or_404(study_id) + return render_template( + "multilayer_design_detail.html", + **_design_view_model(manifest), + ) + + @app.post("/multilayer-design//delete") + def multilayer_design_delete(study_id: str): + _design_store().delete_many([study_id]) + return redirect(url_for("multilayer_design_index")) + + @app.post("/multilayer-design//stages//run") + def multilayer_design_run_stage(study_id: str, stage: str): + if stage not in STAGES: + abort(404) + store = _design_store() + manifest = _design_study_or_404(study_id) + data_dir = app.config["GRAx_DATA_DIR"] + if any( + _is_run_active(app, _multilayer_design_job_key(study_id, other)) for other in STAGES + ): + abort(409, "Another stage of this study is still running.") + + designs: list[list[float]] = [] + if stage == "energy_scan": + if manifest["stages"]["survey"]["status"] != "completed": + abort(409, "Run the survey before scanning designs over energy.") + try: + designs = resolve_designs( + request.form.get("mode", "best"), + survey_design_options(store.study_dir(study_id)), + request.form.getlist("design"), + ) + except ValueError as error: + abort(400, str(error)) + + manifest["stages"][stage]["status"] = "queued" + manifest["stages"][stage]["error_text"] = "" + manifest["stages"][stage]["designs"] = designs + for later in downstream_stages(stage): + if manifest["stages"][later]["status"] in {"completed", "aborted"}: + manifest["stages"][later]["status"] = "stale" + store.save(manifest) + _start_multilayer_design_worker( + app=app, data_dir=data_dir, study_id=study_id, stage=stage + ) + return redirect(url_for("multilayer_design_detail", study_id=study_id)) + + @app.get("/multilayer-design//script") + def multilayer_design_script(study_id: str): + """Download this study as one standalone script to run outside the web app.""" + + manifest = _design_study_or_404(study_id) + filename = offline_script_filename(manifest) + return Response( + render_offline_script(manifest), + mimetype="text/x-python", + headers={"Content-Disposition": f'attachment; filename="{filename}"'}, + ) + + @app.get("/multilayer-design//survey-options") + def multilayer_design_survey_options(study_id: str): + """Return the survey grid as it stands, so the plots can follow a live run.""" + + _design_study_or_404(study_id) + return jsonify(survey_design_options(_design_store().study_dir(study_id))) + + @app.get("/multilayer-design//energy-scan-points") + def multilayer_design_energy_scan_points(study_id: str): + """Return each design's solved energies, so the scan plot can follow a live run.""" + + manifest = _design_study_or_404(study_id) + return jsonify( + { + "designs": _energy_scan_checkpoint_series( + study_dir=_design_store().study_dir(study_id), + designs=manifest["stages"]["energy_scan"].get("designs") or [], + ) + } + ) + + @app.get("/multilayer-design//stages//status") + def multilayer_design_stage_status(study_id: str, stage: str): + if stage not in STAGES: + abort(404) + _design_study_or_404(study_id) + return jsonify( + _multilayer_design_stage_status_payload( + app=app, + data_dir=app.config["GRAx_DATA_DIR"], + study_id=study_id, + stage=stage, + ) + ) + + @app.get("/multilayer-design//stages//abort") + def multilayer_design_stage_abort_dialog(study_id: str, stage: str): + if stage not in STAGES: + abort(404) + manifest = _design_study_or_404(study_id) + return render_template( + "multilayer_design_stage_abort.html", + study=manifest, + stage=stage, + stage_label=STAGE_LABELS[stage], + ) + + @app.post("/multilayer-design//stages//abort") + def multilayer_design_stage_abort(study_id: str, stage: str): + if stage not in STAGES: + abort(404) + _design_study_or_404(study_id) + _abort_multilayer_design_stage( + app=app, + store=_design_store(), + study_id=study_id, + stage=stage, + discard=request.form.get("disposition") == "discard", + ) + return redirect(url_for("multilayer_design_detail", study_id=study_id)) + + @app.post("/multilayer-design//stages//reset") + def multilayer_design_reset_stage(study_id: str, stage: str): + if stage not in STAGES: + abort(404) + _design_study_or_404(study_id) + if _is_run_active(app, _multilayer_design_job_key(study_id, stage)): + abort(409, "This stage is still running.") + _reset_multilayer_design_stage( + store=_design_store(), study_id=study_id, stage=stage + ) + return redirect(url_for("multilayer_design_detail", study_id=study_id)) + @app.get("/_data/") def data_file(filename: str): return send_from_directory(app.config["GRAx_DATA_DIR"], filename) @@ -1225,6 +1523,498 @@ def _execute_run_job( release_workers(run_id) +# --------------------------------------------------------------------------- # +# Multilayer-design studies # +# --------------------------------------------------------------------------- # +def _multilayer_design_job_key(study_id: str, stage: str) -> str: + """Registry key for one running design-study stage.""" + + return f"mldesign:{study_id}:{stage}" + + +def _multilayer_design_stage_total(manifest: dict[str, Any], stage: str) -> int: + """Best-effort count of scan items for one stage, for the progress bar.""" + + if stage == "survey": + return survey_cell_count(manifest["config"]) + return len(manifest["stages"][stage].get("designs") or []) + + +def _start_multilayer_design_worker( + *, app: Any, data_dir: Path, study_id: str, stage: str +) -> None: + """Register an active-run entry for a design stage and start its worker thread.""" + + store = MultilayerDesignStudyStore(data_dir / "multilayer_designs") + manifest = store.load(study_id) + + key = _multilayer_design_job_key(study_id, stage) + active_state = ActiveRunState( + run_id=key, + workflow=f"multilayer_design_{stage}", + total_points=_multilayer_design_stage_total(manifest, stage), + worker_mode="auto", + requested_workers=None, + resolved_workers=None, + ) + if stage == "energy_scan": + active_state.resumed_points = _energy_scan_checkpoint_progress( + study_dir=store.study_dir(study_id), + designs=manifest["stages"][stage].get("designs") or [], + energy_points=int(manifest["config"].get("energy_scan_points") or 0), + )[0] + with _active_runs_lock(app): + _active_runs(app)[key] = active_state + + worker = threading.Thread( + target=_execute_multilayer_design_job, + kwargs={"app": app, "data_dir": data_dir, "study_id": study_id, "stage": stage}, + daemon=True, + name=f"grax-mldesign-{study_id}-{stage}", + ) + active_state.worker_thread = worker + worker.start() + + +def _execute_multilayer_design_job( + *, app: Any, data_dir: Path, study_id: str, stage: str +) -> None: + """Run one design-study stage in a background thread and record the outcome.""" + + from grax.multilayer_design import MultilayerGratingDesigner + + from .resource_manager import allocate_workers, release_workers + + key = _multilayer_design_job_key(study_id, stage) + store = MultilayerDesignStudyStore(data_dir / "multilayer_designs") + + def _set_stage(**updates: Any) -> None: + manifest = store.load(study_id) + manifest["stages"][stage].update(updates) + store.save(manifest) + + allocate_workers(key) + _update_active_run(app, key, state="running", started=True) + _set_stage(status="running", error_text="", aborted=False) + auto_chain = False + with _active_runs_lock(app): + entry = _active_runs(app).get(key) + stop_event = entry.stop_event if entry is not None else threading.Event() + try: + manifest = store.load(study_id) + config = build_design_config(store.study_dir(study_id), manifest["config"]) + designer = MultilayerGratingDesigner(config) + + def _progress(report: Any) -> None: + _update_active_run(app, key, completed_points=report.completed) + with _active_runs_lock(app): + entry = _active_runs(app).get(key) + if entry is not None and report.total: + entry.total_points = report.total + + def _keep_going() -> bool: + with _active_runs_lock(app): + entry = _active_runs(app).get(key) + return entry is None or not entry.stop_event.is_set() + + def _worker_pids_changed(worker_pids: set[int]) -> None: + # The live pool size is the only honest worker count -- the config + # only ever says "auto". + _update_active_run( + app, + key, + simulation_pids=worker_pids, + resolved_workers=len(worker_pids) or None, + ) + + if stage == "survey": + result = designer.run_survey( + progress_callback=_progress, + should_continue=_keep_going, + stop_event=stop_event, + on_worker_pids_changed=_worker_pids_changed, + ) + aborted = bool(result.aborted) + artifacts = { + "optimal_blaze_plot": _relative_to(result.plot_path, store.study_dir(study_id)), + "max_efficiency_plot": _relative_to( + result.efficiency_plot_path, store.study_dir(study_id) + ), + "heatmap_plot": _relative_to( + result.heatmap_plot_path, store.study_dir(study_id) + ), + "survey_csv": _relative_to( + result.combined_csv_path, store.study_dir(study_id) + ), + } + auto_chain = ( + not aborted + and manifest.get("auto_energy_scan") == "best" + and manifest["stages"]["energy_scan"]["status"] != "running" + ) + else: + designs = [tuple(pair) for pair in manifest["stages"][stage].get("designs") or []] + scans = designer.run_energy_scan( + designs, + progress_callback=_progress, + should_continue=_keep_going, + stop_event=stop_event, + on_worker_pids_changed=_worker_pids_changed, + ) + aborted = len(scans) < len(designs) + artifacts = { + "designs": [ + { + "d_spacing_nm": scan.d_spacing_nm, + "blaze_angle_deg": scan.blaze_angle_deg, + "plot": _relative_to(scan.titled_plot_path, store.study_dir(study_id)), + "summary_csv": _relative_to( + scan.summary_csv_path, store.study_dir(study_id) + ), + } + for scan in scans + ], + } + if len(scans) >= 2: + artifacts["overlay_plot"] = _relative_to( + designer.plot_energy_scan_overlay(scans), store.study_dir(study_id) + ) + + _set_stage( + status="aborted" if aborted else "completed", + ran_at=datetime.now().isoformat(timespec="seconds"), + aborted=aborted, + artifacts=artifacts, + error_text="", + ) + _finish_active_run(app, key, state="aborted" if aborted else "completed") + except Exception as error: # noqa: BLE001 - surfaced to the study page + if stop_event.is_set(): + # Killing the workers raises out of the stage; that is an abort, not + # a failure, and the user does not need the traceback for it. + _set_stage( + status="aborted", + ran_at=datetime.now().isoformat(timespec="seconds"), + aborted=True, + error_text="", + ) + _finish_active_run(app, key, state="aborted") + else: + _set_stage(status="failed", error_text=str(error)) + _finish_active_run(app, key, state="failed", error_text=str(error)) + finally: + release_workers(key) + + if auto_chain: + _start_auto_energy_scan(app=app, data_dir=data_dir, study_id=study_id) + + +def _start_auto_energy_scan(*, app: Any, data_dir: Path, study_id: str) -> None: + """Queue the energy scan on the survey's best design, for ``auto_energy_scan``.""" + + store = MultilayerDesignStudyStore(data_dir / "multilayer_designs") + try: + options = survey_design_options(store.study_dir(study_id)) + designs = resolve_designs("best", options, []) + except (ValueError, OSError): + return + manifest = store.load(study_id) + manifest["stages"]["energy_scan"]["status"] = "queued" + manifest["stages"]["energy_scan"]["error_text"] = "" + manifest["stages"]["energy_scan"]["designs"] = designs + store.save(manifest) + _start_multilayer_design_worker( + app=app, data_dir=data_dir, study_id=study_id, stage="energy_scan" + ) + + +def _relative_to(path: Path, root: Path) -> str: + """Return ``path`` relative to ``root`` as a POSIX string, or its name.""" + + try: + return Path(path).relative_to(root).as_posix() + except ValueError: + return Path(path).name + + +def _energy_scan_checkpoint_path(study_dir: Path, pair: Any) -> Path | None: + """Return one design's checkpoint file, or ``None`` if the pair is malformed.""" + + try: + d_spacing, blaze = float(pair[0]), float(pair[1]) + except (TypeError, ValueError, IndexError): + return None + return ( + study_dir + / "energy_scan" + / f"d{d_spacing:.3f}nm_blaze{blaze:.3f}deg" + / "checkpoints" + / "results.jsonl" + ) + + +def _energy_scan_checkpoint_series(*, study_dir: Path, designs: list[Any]) -> list[dict[str, Any]]: + """Return each design's solved ``(energy, efficiency)`` points, energy-sorted. + + The summary CSV only lands when a design finishes, so the live plot reads the + checkpoint instead -- one JSON record per solved energy, written as the sweep + goes. Energies come back in completion order, hence the sort. + """ + + series: list[dict[str, Any]] = [] + for pair in designs: + results = _energy_scan_checkpoint_path(study_dir, pair) + if results is None: + continue + points: list[tuple[float, float]] = [] + if results.is_file(): + try: + with results.open("r", encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + try: + record = json.loads(line) + except ValueError: + continue # a line still being written + if record.get("status") != "ok": + continue + energy = record.get("energy_ev") + efficiency = record.get("selected_efficiency") + if energy is None or efficiency is None: + continue + points.append((float(energy), float(efficiency))) + except OSError: + points = [] + points.sort() + series.append( + { + "d_spacing_nm": float(pair[0]), + "blaze_angle_deg": float(pair[1]), + "energies_ev": [energy for energy, _ in points], + "efficiencies": [efficiency for _, efficiency in points], + } + ) + return series + + +def _energy_scan_checkpoint_progress( + *, study_dir: Path, designs: list[Any], energy_points: int +) -> tuple[int, int]: + """Return ``(completed, total)`` energy points across an energy scan's designs. + + ``run_energy_scan`` reports progress once per *design*, so a single-design + scan would otherwise sit at ``0 / 1`` for its whole run -- indistinguishable + from a job that never started. Each design's sweep checkpoints one line per + solved energy, so counting those gives a bar that actually moves. + """ + + total = max(0, int(energy_points)) * len(designs) + completed = 0 + for pair in designs: + results = _energy_scan_checkpoint_path(study_dir, pair) + if results is None or not results.is_file(): + continue + try: + with results.open("r", encoding="utf-8") as handle: + completed += sum(1 for line in handle if line.strip()) + except OSError: + continue + return min(completed, total) if total else completed, total + + +def _multilayer_design_stage_status_payload( + *, app: Any, data_dir: Path, study_id: str, stage: str +) -> dict[str, Any]: + """Return a status payload for one design stage in the shape ``initRunMonitor`` reads.""" + + _cleanup_finished_runs(app) + key = _multilayer_design_job_key(study_id, stage) + manifest = MultilayerDesignStudyStore(data_dir / "multilayer_designs").load(study_id) + stage_state = manifest["stages"][stage] + + with _active_runs_lock(app): + active = _active_runs(app).get(key) + if active is not None and _is_active_run_entry_live(active): + completed = active.completed_points + total = active.total_points + elapsed = _elapsed_seconds(active) + eta = _eta_seconds(active) + label = "" + if stage == "energy_scan": + # Count solved energies, not finished designs, so the bar moves. + designs = stage_state.get("designs") or [] + checkpoint_completed, checkpoint_total = _energy_scan_checkpoint_progress( + study_dir=data_dir / "multilayer_designs" / study_id, + designs=designs, + energy_points=int(manifest["config"].get("energy_scan_points") or 0), + ) + if checkpoint_total: + completed, total = checkpoint_completed, checkpoint_total + # Only energies solved by *this* run measure its rate; + # resumed checkpoints cost it no time. + solved_here = completed - active.resumed_points + eta = ( + elapsed * (total - completed) / solved_here + if solved_here > 0 and elapsed and total > completed + else None + ) + if designs: + index = min(active.completed_points, len(designs) - 1) + d_spacing, blaze = designs[index] + label = ( + f"design {index + 1} of {len(designs)}: " + f"d = {float(d_spacing):.3f} nm, blaze = {float(blaze):.3f} deg" + ) + return { + "state": active.state, + "completed_points": completed, + "total_points": total, + "remaining_points": max(total - completed, 0), + "elapsed_seconds": elapsed, + "eta_seconds": eta, + "current_label": ( + "stopping workers…" if active.state == "aborting" else label + ), + "worker_mode": "auto", + "requested_workers": None, + "resolved_workers": active.resolved_workers, + "plot_url": None, + "plot_token": "", + "error_text": active.error_text, + "can_abort": active.state in {"queued", "running"} + and not active.abort_requested, + } + + status = stage_state.get("status", "not_run") + normalized = {"queued": "running", "aborting": "running"}.get(status, status) + if normalized not in {"completed", "failed", "aborted", "stale", "not_run", "running"}: + normalized = "not_run" + artifacts = stage_state.get("artifacts") or {} + plot_rel = artifacts.get("heatmap_plot") if stage == "survey" else artifacts.get("overlay_plot") + plot_url = None + plot_token = "" + if plot_rel: + plot_path = data_dir / "multilayer_designs" / study_id / plot_rel + if plot_path.exists(): + plot_token = _file_token(plot_path) + plot_url = f"/_data/multilayer_designs/{study_id}/{plot_rel}?v={plot_token}" + return { + "state": normalized if normalized != "stale" else "completed", + "completed_points": 0, + "total_points": 0, + "remaining_points": 0, + "elapsed_seconds": None, + "eta_seconds": None, + "current_label": "", + "worker_mode": "auto", + "requested_workers": None, + "resolved_workers": None, + "plot_url": plot_url, + "plot_token": plot_token, + "error_text": stage_state.get("error_text", ""), + "can_abort": False, + } + + +def _abort_multilayer_design_stage( + *, app: Any, store: Any, study_id: str, stage: str, discard: bool +) -> None: + """Request a cooperative stop for a running stage and wait for it to finish.""" + + key = _multilayer_design_job_key(study_id, stage) + with _active_runs_lock(app): + active = _active_runs(app).get(key) + if active is not None and active.state in {"queued", "running"}: + active.abort_requested = True + active.stop_event.set() + active.state = "aborting" + # The sweep kills its own pool from the inside; this is the backstop for a + # worker stuck somewhere the stop check does not reach. It also matters for + # `discard`, which deletes the stage directory the worker is writing into. + _terminate_run_processes(app, key) + _wait_for_run_shutdown(app, key) + if discard: + _reset_multilayer_design_stage(store=store, study_id=study_id, stage=stage) + + +def _terminate_run_processes(app: Any, run_id: str, *, grace_seconds: float = 3.0) -> None: + """Stop one active run's simulation subprocesses, children first.""" + + with _active_runs_lock(app): + active = _active_runs(app).get(run_id) + pids = set() if active is None else set(active.simulation_pids) + pids.discard(os.getpid()) + if not pids or not hasattr(psutil, "Process"): + return + targets: list[Any] = [] + for pid in pids: + try: + process = psutil.Process(pid) + except Exception: + continue + try: + targets.extend(process.children(recursive=True)) + except Exception: + pass + targets.append(process) + for process in targets: + try: + process.terminate() + except Exception: + continue + try: + _, alive = psutil.wait_procs(targets, timeout=grace_seconds) + except Exception: + return + for process in alive: + try: + process.kill() + except Exception: + continue + + +def _reset_multilayer_design_stage(*, store: Any, study_id: str, stage: str) -> None: + """Delete one stage's outputs and blank its manifest entry.""" + + study_dir = store.study_dir(study_id) + manifest = store.load(study_id) + artifacts = manifest["stages"][stage].get("artifacts") or {} + + if stage == "survey": + shutil.rmtree(study_dir / "survey", ignore_errors=True) + for key in ("optimal_blaze_plot", "max_efficiency_plot", "heatmap_plot"): + _unlink_relative(study_dir, artifacts.get(key)) + else: + shutil.rmtree(study_dir / "energy_scan", ignore_errors=True) + _unlink_relative(study_dir, artifacts.get("overlay_plot")) + for design in artifacts.get("designs") or []: + _unlink_relative(study_dir, design.get("plot")) + + manifest["stages"][stage] = { + "status": "not_run", + "ran_at": None, + "error_text": "", + "aborted": False, + "designs": [], + "artifacts": {}, + } + for later in downstream_stages(stage): + if manifest["stages"][later]["status"] in {"completed", "aborted"}: + manifest["stages"][later]["status"] = "stale" + store.save(manifest) + + +def _unlink_relative(root: Path, relative: str | None) -> None: + """Delete ``root / relative`` when it exists and stays inside ``root``.""" + + if not relative: + return + candidate = (root / relative).resolve() + if root.resolve() in candidate.parents and candidate.exists(): + candidate.unlink() + + def _worker_settings_from_form(form: Any) -> tuple[str, str | int, int | None]: """Return worker-mode metadata and the runner max_workers setting.""" diff --git a/src/grax/web/multilayer_design_studies.py b/src/grax/web/multilayer_design_studies.py new file mode 100644 index 0000000..b9b619c --- /dev/null +++ b/src/grax/web/multilayer_design_studies.py @@ -0,0 +1,942 @@ +"""Persistence and form helpers for web multilayer-design studies. + +One *study* is a directory under ``/multilayer_designs//``. +:class:`grax.MultilayerDesignConfig`'s ``output_dir`` points straight at that +directory, so the library's own layout (``survey/``, ``plots/``, +``energy_scan/``) lands inside it. This module adds exactly one extra file, +``study.json``, holding the config the user entered plus per-stage status. + +The workflow has two stages, matching the library: + +* ``survey`` -- :meth:`grax.MultilayerGratingDesigner.run_survey`, one + single-energy theta search per ``(d, blaze)`` grid cell. +* ``energy_scan`` -- :meth:`grax.MultilayerGratingDesigner.run_energy_scan` + over the ``(d, blaze)`` designs picked off the survey. +""" + +from __future__ import annotations + +import json +import shutil +from dataclasses import dataclass, fields +from datetime import datetime +from pathlib import Path +from typing import Any + +from grax.multilayer_design import MultilayerDesignConfig, ThetaSearchScanSettings + +from .persistence import _slugify + +STAGES: tuple[str, ...] = ("survey", "energy_scan") + +STAGE_LABELS: dict[str, str] = { + "survey": "1. d-spacing x blaze-angle survey", + "energy_scan": "2. Energy scan of the chosen designs", +} + +#: Selection modes the energy-scan stage accepts. +DESIGN_MODES: tuple[str, ...] = ("best", "per_d", "manual") + +_TERMINAL_STAGE_STATES = {"completed", "aborted"} + +#: Survey cells above this count get a warning on the creation form. +SURVEY_CELL_WARNING_THRESHOLD = 200 + + +def downstream_stages(stage: str) -> tuple[str, ...]: + """Return the stages that depend on ``stage``.""" + + return STAGES[STAGES.index(stage) + 1 :] + + +# --------------------------------------------------------------------------- # +# Form field specs # +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class FieldSpec: + """One editable config field on the study form. + + Attributes: + name: :class:`grax.MultilayerDesignConfig` field name. A dotted name + such as ``"survey_scan_settings.rough_scan_points"`` addresses a + field on one of the nested :class:`grax.ThetaSearchScanSettings`. + kind: ``number`` / ``int`` / ``text`` / ``select`` / ``checkbox`` / + ``material`` (a name + density pair). + label: Human-readable label. + section: Fieldset heading it belongs to. + advanced: Rendered inside the collapsible "Advanced" section. + choices: Options for ``select`` fields. + row: Optional sub-group inside ``section``. Fields sharing a ``row`` + render on one line under that heading; ``""`` means "no sub-group", + and those fields flow in the fieldset as usual. + """ + + name: str + kind: str + label: str + section: str + advanced: bool = False + choices: tuple[str, ...] = () + row: str = "" + + @property + def html_name(self) -> str: + """Return the form input name (dots are legal in HTML field names).""" + + return self.name + + +def _scan_fields(prefix: str, section: str) -> tuple[FieldSpec, ...]: + """Return the advanced field specs for one nested scan-settings block. + + Grouped into one line per theta-search pass -- rough, then fine, then the + final solve -- followed by the two settings that apply to the search as a + whole. The pass name lives in the row heading, so the field labels drop it. + """ + + return ( + FieldSpec(f"{prefix}.rough_scan_half_width_deg", "number", + "Half-width, deg", section, advanced=True, row="Rough pass"), + FieldSpec(f"{prefix}.rough_scan_points", "int", + "Points", section, advanced=True, row="Rough pass"), + FieldSpec(f"{prefix}.rough_fourier_orders", "int", + "Fourier orders", section, advanced=True, row="Rough pass"), + FieldSpec(f"{prefix}.rough_x_resolution_nm", "number", + "x resolution, nm", section, advanced=True, row="Rough pass"), + FieldSpec(f"{prefix}.rough_z_resolution_nm", "number", + "z resolution, nm", section, advanced=True, row="Rough pass"), + FieldSpec(f"{prefix}.fine_scan_half_width_deg", "number", + "Half-width, deg", section, advanced=True, row="Fine pass"), + FieldSpec(f"{prefix}.fine_scan_points", "int", + "Points", section, advanced=True, row="Fine pass"), + FieldSpec(f"{prefix}.fine_fourier_orders", "int", + "Fourier orders", section, advanced=True, row="Fine pass"), + FieldSpec(f"{prefix}.fine_x_resolution_nm", "number", + "x resolution, nm", section, advanced=True, row="Fine pass"), + FieldSpec(f"{prefix}.fine_z_resolution_nm", "number", + "z resolution, nm", section, advanced=True, row="Fine pass"), + FieldSpec(f"{prefix}.final_fourier_orders", "int", + "Fourier orders", section, advanced=True, row="Final solve"), + FieldSpec(f"{prefix}.final_x_resolution_nm", "number", + "x resolution, nm", section, advanced=True, row="Final solve"), + FieldSpec(f"{prefix}.final_z_resolution_nm", "number", + "z resolution, nm", section, advanced=True, row="Final solve"), + FieldSpec(f"{prefix}.precise_peak_selection_mode", "select", + "Peak selection", section, advanced=True, row="Peak & roughness", + choices=("max", "gauss", "voigt")), + FieldSpec(f"{prefix}.roughness_sigma_nm", "number", + "Roughness sigma, nm (blank = none)", section, advanced=True, + row="Peak & roughness"), + ) + + +#: Every editable field, in render order. Sections mirror the three groups in +#: ``MultilayerDesignConfig``: shared, survey-only, energy-scan-only. +STUDY_FIELDS: tuple[FieldSpec, ...] = ( + # -- Shared --------------------------------------------------------- # + FieldSpec("grating_density_lpermm", "number", "Line density, l/mm", + "Shared - geometry & order"), + FieldSpec("diffraction_order", "int", "Diffraction order", + "Shared - geometry & order"), + FieldSpec("multilayer_bragg_order", "int", "Multilayer Bragg order", + "Shared - geometry & order"), + FieldSpec("anti_blaze_angle_deg", "number", "Anti-blaze angle, deg (0 = sawtooth)", + "Shared - geometry & order"), + FieldSpec("material_a", "material", "Material A (top)", "Shared - materials"), + FieldSpec("material_b", "material", "Material B", "Shared - materials"), + FieldSpec("substrate_material", "material", "Substrate", "Shared - materials"), + FieldSpec("n_bilayers", "int", "Bilayers", "Shared - materials"), + FieldSpec("gamma", "number", "Gamma (material A fraction)", "Shared - materials"), + FieldSpec("coating_label", "text", "Coating label for plots (e.g. Ru/B4C)", + "Shared - materials"), + FieldSpec("solver", "select", "Solver", "Shared - numerics", + choices=("neviere", "rcwa")), + FieldSpec("polarization", "select", "Polarization", "Shared - numerics", + choices=("p", "s")), + FieldSpec("backend", "select", "Backend", "Shared - numerics", + choices=("numba", "numpy")), + FieldSpec("x_resolution_nm", "number", "Grating x resolution, nm", + "Shared - numerics"), + FieldSpec("z_resolution_nm", "number", "Grating z resolution, nm", + "Shared - numerics"), + FieldSpec("checkpoint", "checkbox", "Write checkpoints", "Shared - runtime", + advanced=True), + FieldSpec("resume", "checkbox", "Resume from checkpoints", "Shared - runtime", + advanced=True), + FieldSpec("save_profile_plot", "checkbox", "Save profile plots", "Shared - runtime", + advanced=True), + FieldSpec("save_stack_plot", "checkbox", "Save stack plots", "Shared - runtime", + advanced=True), + # -- Survey only ---------------------------------------------------- # + FieldSpec("target_energy_ev", "number", "Survey energy, eV", + "Survey - target & grids"), + FieldSpec("d_min_nm", "number", "d-spacing min, nm", "Survey - target & grids"), + FieldSpec("d_max_nm", "number", "d-spacing max, nm", "Survey - target & grids"), + FieldSpec("d_points", "int", "d-spacing points", "Survey - target & grids"), + FieldSpec("blaze_min_deg", "number", "Blaze min, deg", "Survey - target & grids"), + FieldSpec("blaze_max_deg", "number", "Blaze max, deg", "Survey - target & grids"), + FieldSpec("blaze_points", "int", "Blaze points", "Survey - target & grids"), + FieldSpec("on_error", "select", "On a failing cell", "Survey - target & grids", + choices=("continue", "fail_fast")), + *_scan_fields("survey_scan_settings", "Survey - theta-search settings"), + # -- Energy scan only ----------------------------------------------- # + FieldSpec("energy_scan_min_ev", "number", "Energy min, eV", + "Energy scan - energy grid"), + FieldSpec("energy_scan_max_ev", "number", "Energy max, eV", + "Energy scan - energy grid"), + FieldSpec("energy_scan_points", "int", "Energy points", + "Energy scan - energy grid"), + FieldSpec("max_workers", "text", "Max workers (\"auto\" or a number)", + "Energy scan - runtime", advanced=True), + FieldSpec("theta_tracking_mode", "select", "Theta tracking", + "Energy scan - runtime", advanced=True, + choices=("auto", "bragg", "previous")), + FieldSpec("max_tracking_energy_step_ev", "number", + "Max tracking energy step, eV (blank = none)", + "Energy scan - runtime", advanced=True), + *_scan_fields("energy_scan_settings", "Energy scan - theta-search settings"), +) + +#: Fields whose blank form value means ``None`` rather than "leave unchanged". +_NULLABLE_FIELDS = frozenset( + { + "coating_label", + "max_tracking_energy_step_ev", + "survey_scan_settings.roughness_sigma_nm", + "energy_scan_settings.roughness_sigma_nm", + } +) + +_MATERIAL_FIELDS = frozenset({"material_a", "material_b", "substrate_material"}) + + +def _field_by_name() -> dict[str, FieldSpec]: + return {spec.name: spec for spec in STUDY_FIELDS} + + +def study_form_sections(advanced: bool) -> list[tuple[str, list[tuple[str, list[FieldSpec]]]]]: + """Return ``(section, rows)`` groups for the study form. + + Each section holds ``(row_label, fields)`` pairs. A row labelled ``""`` + is an ungrouped run of fields that flows in the fieldset as usual; a named + row renders on its own line under that heading (see + :attr:`FieldSpec.row`). + + Args: + advanced: ``True`` for the advanced (collapsible) fields, ``False`` for + the always-visible ones. + + Returns: + Section groups preserving :data:`STUDY_FIELDS` order. + """ + + sections: dict[str, list[tuple[str, list[FieldSpec]]]] = {} + for spec in STUDY_FIELDS: + if bool(spec.advanced) != advanced: + continue + rows = sections.setdefault(spec.section, []) + if rows and rows[-1][0] == spec.row: + rows[-1][1].append(spec) + else: + rows.append((spec.row, [spec])) + return list(sections.items()) + + +def _nested_defaults() -> dict[str, Any]: + """Return the default values of both nested scan-settings blocks.""" + + return { + name: spec.default + for name, spec in ThetaSearchScanSettings.__dataclass_fields__.items() + } + + +def study_config_defaults() -> dict[str, Any]: + """Return the JSON-safe config dict from :class:`grax.MultilayerDesignConfig` defaults.""" + + dataclass_defaults = {f.name: f.default for f in fields(MultilayerDesignConfig)} + scan_defaults = _nested_defaults() + config: dict[str, Any] = { + "survey_scan_settings": dict(scan_defaults), + "energy_scan_settings": dict(scan_defaults), + } + for spec in STUDY_FIELDS: + if "." in spec.name: + continue + default = dataclass_defaults.get(spec.name) + if spec.kind == "material": + name, density = default if isinstance(default, (tuple, list)) else ("", None) + config[spec.name] = [str(name), float(density)] + elif spec.kind == "checkbox": + config[spec.name] = bool(default) + else: + config[spec.name] = default + return config + + +def _coerce_field(spec: FieldSpec, raw: str) -> Any: + """Coerce one raw form value for ``spec`` into its JSON-safe type.""" + + text = raw.strip() + if text == "" and spec.name in _NULLABLE_FIELDS: + return None + if spec.name == "max_workers": + return text if text.lower() == "auto" else int(float(text)) + if spec.kind == "int": + return int(float(text)) + if spec.kind == "number": + return float(text) + return text + + +def _set_nested(config: dict[str, Any], dotted: str, value: Any) -> None: + """Assign ``value`` at a ``a.b`` path inside ``config``.""" + + head, _, tail = dotted.partition(".") + if not tail: + config[dotted] = value + return + block = dict(config.get(head) or {}) + block[tail] = value + config[head] = block + + +def parse_study_config(form: Any, base: dict[str, Any] | None = None) -> dict[str, Any]: + """Overlay a form's values onto ``base`` (or the defaults) and return a config dict.""" + + config = dict(base or study_config_defaults()) + for spec in STUDY_FIELDS: + if spec.kind == "material": + name = form.get(f"{spec.name}_name") + density = form.get(f"{spec.name}_density") + if name is not None and density not in (None, ""): + config[spec.name] = [str(name).strip(), float(density)] + continue + if spec.kind == "checkbox": + # The form renders a hidden "0" before the box, so a checked box + # submits both and the last value is the real one. + values = form.getlist(spec.name) if hasattr(form, "getlist") else None + if values: + config[spec.name] = str(values[-1]) not in ("", "0", "false") + elif any(key == spec.name for key in form): + config[spec.name] = form.get(spec.name) not in (None, "", "0", "false") + elif base is None: + config[spec.name] = False + continue + if spec.name not in form: + continue + raw = str(form.get(spec.name)) + if raw.strip() == "" and spec.name not in _NULLABLE_FIELDS: + continue + _set_nested(config, spec.name, _coerce_field(spec, raw)) + return config + + +def build_design_config(study_dir: Path, config: dict[str, Any]) -> MultilayerDesignConfig: + """Build a :class:`grax.MultilayerDesignConfig` for ``study_dir`` from a config dict.""" + + kwargs: dict[str, Any] = {} + for key, value in config.items(): + if key in _MATERIAL_FIELDS and isinstance(value, (list, tuple)): + kwargs[key] = (str(value[0]), float(value[1])) + elif key in {"survey_scan_settings", "energy_scan_settings"}: + kwargs[key] = ThetaSearchScanSettings(**dict(value or {})) + else: + kwargs[key] = value + return MultilayerDesignConfig(output_dir=study_dir, **kwargs) + + +def flatten_config_values(config: dict[str, Any]) -> dict[str, Any]: + """Return ``config`` keyed by :class:`FieldSpec` names, dotted paths included. + + The templates render one input per :data:`STUDY_FIELDS` entry, so they need + a flat lookup in which ``"survey_scan_settings.rough_scan_points"`` resolves + without walking into the nested block. + """ + + flat: dict[str, Any] = {} + for key, value in config.items(): + if isinstance(value, dict): + for inner_key, inner_value in value.items(): + flat[f"{key}.{inner_key}"] = inner_value + else: + flat[key] = value + return flat + + +#: Fields that change nothing a stage computed -- only plot labelling, which +#: artifacts get written, or how the work is scheduled. Editing one of these +#: leaves existing results valid. +_COSMETIC_FIELDS = frozenset( + { + "coating_label", + "save_profile_plot", + "save_stack_plot", + "checkpoint", + "resume", + "max_workers", + } +) + + +def stages_invalidated_by(old: dict[str, Any], new: dict[str, Any]) -> tuple[str, ...]: + """Return the stages whose stored results no longer match an edited config. + + A field's :attr:`FieldSpec.section` says which stage reads it: the shared + section feeds both, so changing one invalidates the survey and everything + downstream of it. + + Args: + old: The config as stored before the edit. + new: The config after the edit. + + Returns: + Stage names in :data:`STAGES` order, empty when nothing that affects a + result changed. + """ + + before = flatten_config_values(old) + after = flatten_config_values(new) + invalidated: set[str] = set() + for spec in STUDY_FIELDS: + if spec.name in _COSMETIC_FIELDS or before.get(spec.name) == after.get(spec.name): + continue + if spec.section.startswith("Energy scan"): + invalidated.add("energy_scan") + else: + invalidated.update(("survey", *downstream_stages("survey"))) + return tuple(stage for stage in STAGES if stage in invalidated) + + +def survey_cell_count(config: dict[str, Any]) -> int: + """Return how many theta searches the survey grid implies.""" + + try: + return int(config.get("d_points", 0)) * int(config.get("blaze_points", 0)) + except (TypeError, ValueError): + return 0 + + + +# --------------------------------------------------------------------------- # +# Offline script generation # +# --------------------------------------------------------------------------- # +#: Section prefix -> (banner title, banner note) for the generated script. +_SCRIPT_SECTIONS: tuple[tuple[str, str, str], ...] = ( + ( + "Shared", + "SHARED -- grating geometry, materials, solver/runtime.", + "Read by both stages.", + ), + ( + "Survey", + "SURVEY -- read only by --survey (run_survey / evaluate_survey).", + "", + ), + ( + "Energy scan", + "ENERGY SCAN -- read only by --energy-scan.", + "", + ), +) + +#: Nested scan-settings blocks, as ``config key -> generated constant``. +_SCRIPT_SCAN_BLOCKS = { + "survey_scan_settings": "SURVEY_SCAN", + "energy_scan_settings": "ENERGY_SCAN_SCAN", +} + + +def _script_literal(name: str, value: Any) -> str: + """Return ``value`` as Python source, as a tuple for the material pairs.""" + + if name in _MATERIAL_FIELDS and isinstance(value, (list, tuple)): + return f"({str(value[0])!r}, {float(value[1])!r})" + return repr(value) + + +def _script_banner(title: str, note: str) -> str: + """Return a boxed section banner in the examples' style.""" + + width = 75 + lines = [f"# {'=' * width} #", f"# {title.ljust(width)} #"] + if note: + lines.append(f"# {note.ljust(width)} #") + lines.append(f"# {'=' * width} #") + return "\n".join(lines) + + +_OFFLINE_SCRIPT_TEMPLATE = '''"""{display_name} -- standalone multilayer-grating design run. + +Generated by the grax web app from study {study_id} on {generated_at}. +Everything this needs is in this one file: edit the parameters below, then run +the stages you want. It imports nothing but ``grax`` and pandas, so it can be +copied to another machine and run there. + +Stage 1 (``--survey``) scans a 2-D grid of bilayer d-spacing x blaze angle. For +every pair it builds the multilayer-coated blazed grating and runs a +single-energy multilayer theta search at TARGET_ENERGY_EV -- the search scans +the incident angle and returns the one that maximizes the selected-order +efficiency, so there is no CFF input. It writes ``results/survey/survey.csv``, +three headline plots under ``results/plots/`` and a per-run folder tree under +``results/survey/runs/``. + +Stage 2 (``--energy-scan``) sweeps ``(d, blaze)`` designs over +[ENERGY_SCAN_MIN_EV, ENERGY_SCAN_MAX_EV]: the optimal blaze per d-spacing by +default, only the best cell with ``--best``, or exactly ``--pairs``. Two or +more designs also get an overlay comparison plot. + +With no stage flag both stages run in order. ``--eval`` re-derives results from +what is already on disk without solving anything. + +Examples:: + + python {script_name} --survey + python {script_name} --energy-scan --best + python {script_name} --energy-scan --pairs "3.0,0.8; 3.3,0.9" + python {script_name} # survey, then every optimal blaze + python {script_name} --survey --eval # rebuild plots from existing runs + +The executable body is guarded because the theta-search sweep spawns worker +processes that re-import this file by path. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import pandas as pd + +from grax import ( + MultilayerDesignConfig, + MultilayerGratingDesigner, + ThetaSearchScanSettings, +) + +{parameters} + +CONFIG = MultilayerDesignConfig( +{constructor} +) + + +def _parse_pairs(text: str) -> list[tuple[float, float]]: + """Parse ``"d,blaze; d,blaze"`` into ``(d, blaze)`` float pairs.""" + + pairs: list[tuple[float, float]] = [] + for chunk in text.replace("\\\\n", ";").split(";"): + chunk = chunk.strip() + if not chunk: + continue + d_text, blaze_text = chunk.split(",") + pairs.append((float(d_text), float(blaze_text))) + if not pairs: + raise ValueError(f"no (d, blaze) pairs parsed from {{text!r}}") + return pairs + + +def _survey_table() -> pd.DataFrame: + """Return the usable rows of ``results/survey/survey.csv``.""" + + csv_path = CONFIG.survey_dir / "survey.csv" + if not csv_path.is_file(): + raise SystemExit(f"{{csv_path}} not found -- run --survey first, or pass --pairs.") + table = pd.read_csv(csv_path).dropna(subset=["peak_efficiency"]) + if table.empty: + raise SystemExit("survey.csv has no usable rows; pass --pairs explicitly.") + return table + + +def _optimal_pairs_from_survey() -> list[tuple[float, float]]: + """Return one ``(d, optimal_blaze)`` pair per d-spacing.""" + + table = _survey_table() + best = table.loc[table.groupby("d_nm")["peak_efficiency"].idxmax()] + return [(float(row.d_nm), float(row.blaze_deg)) for row in best.itertuples()] + + +def _best_pair_from_survey() -> list[tuple[float, float]]: + """Return only the single ``(d, blaze)`` with the highest survey efficiency.""" + + row = _survey_table().loc[lambda frame: frame["peak_efficiency"].idxmax()] + print( + f"Best survey design: d = {{float(row.d_nm):.3f}} nm, " + f"blaze = {{float(row.blaze_deg):.3f}} deg " + f"(efficiency {{float(row.peak_efficiency):.4g}})" + ) + return [(float(row.d_nm), float(row.blaze_deg))] + + +def run_survey(designer: MultilayerGratingDesigner, *, evaluate: bool) -> None: + """Run stage 1, or rebuild it from the runs already on disk.""" + + result = designer.evaluate_survey() if evaluate else designer.run_survey() + print(f"Survey table: {{result.combined_csv_path}}") + print(f" optimal blaze vs d: {{result.plot_path}}") + print(f" max efficiency vs d: {{result.efficiency_plot_path}}") + print(f" (d, blaze) heatmap: {{result.heatmap_plot_path}}") + + +def run_energy_scan( + designer: MultilayerGratingDesigner, + *, + evaluate: bool, + pairs_text: str | None, + best_only: bool, +) -> None: + """Run stage 2 for the selected designs, or re-read existing results.""" + + if evaluate and not (pairs_text or best_only): + results = designer.evaluate_energy_scan() + else: + if pairs_text: + pairs = _parse_pairs(pairs_text) + elif best_only: + pairs = _best_pair_from_survey() + else: + pairs = _optimal_pairs_from_survey() + results = ( + designer.evaluate_energy_scan(pairs) + if evaluate + else designer.run_energy_scan(pairs) + ) + + for scan in results: + print( + f"d = {{scan.d_spacing_nm:.3f}} nm, blaze = {{scan.blaze_angle_deg:.3f}} deg " + f"-> {{scan.summary_csv_path}}" + ) + print(f" plot: {{scan.titled_plot_path}}") + + +def main() -> None: + """Parse the stage flags and run what was asked for.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--survey", action="store_true", help="Run stage 1.") + parser.add_argument( + "--energy-scan", action="store_true", dest="energy_scan", help="Run stage 2." + ) + selector = parser.add_mutually_exclusive_group() + selector.add_argument( + "--pairs", default=None, help='Semicolon-separated "d_nm,blaze_deg" designs.' + ) + selector.add_argument( + "--best", action="store_true", help="Scan only the best surveyed design." + ) + parser.add_argument( + "--eval", + action="store_true", + dest="evaluate", + help="Re-derive results from disk instead of solving.", + ) + args = parser.parse_args() + + # No stage flag means the whole workflow, in order. + survey = args.survey or not (args.survey or args.energy_scan) + energy_scan = args.energy_scan or not (args.survey or args.energy_scan) + + designer = MultilayerGratingDesigner(CONFIG) + if survey: + run_survey(designer, evaluate=args.evaluate) + if energy_scan: + run_energy_scan( + designer, + evaluate=args.evaluate, + pairs_text=args.pairs, + best_only=args.best, + ) + + +if __name__ == "__main__": + main() +''' + + +def render_offline_script(manifest: dict[str, Any]) -> str: + """Return a standalone script that reproduces one study outside the web app. + + The result is a single file: the study's parameters as module constants, + grouped the way :class:`grax.MultilayerDesignConfig` groups them, then both + stages behind command-line flags. It depends on nothing but ``grax`` -- the + point is that it can be copied to another machine and run there. + """ + + config = manifest["config"] + flat = flatten_config_values(config) + by_section: dict[str, list[str]] = {prefix: [] for prefix, _, _ in _SCRIPT_SECTIONS} + constructor: list[str] = [" output_dir=OUTPUT_DIR,"] + + for spec in STUDY_FIELDS: + prefix = next( + (key for key, _, _ in _SCRIPT_SECTIONS if spec.section.startswith(key)), + "Shared", + ) + block, _, leaf = spec.name.partition(".") + if leaf: + continue # the nested blocks are emitted whole, below + constant = spec.name.upper() + by_section[prefix].append( + f"{constant} = {_script_literal(spec.name, flat.get(spec.name))}" + ) + constructor.append(f" {spec.name}={constant},") + + for key, constant in _SCRIPT_SCAN_BLOCKS.items(): + prefix = "Survey" if key.startswith("survey") else "Energy scan" + settings = dict(config.get(key) or {}) + lines = [f"{constant} = ThetaSearchScanSettings("] + # Dataclass order (rough, then fine, then final, then peak/roughness), + # not whatever order the stored JSON happens to be in. + lines += [ + f" {name}={settings[name]!r}," + for name in ThetaSearchScanSettings.__dataclass_fields__ + if name in settings + ] + lines.append(")") + by_section[prefix].append("\n".join(lines)) + constructor.append(f" {key}={constant},") + + sections: list[str] = [] + for prefix, title, note in _SCRIPT_SECTIONS: + sections.append(_script_banner(title, note)) + if prefix == "Shared": + sections.append('OUTPUT_DIR = Path(__file__).resolve().parent / "results"') + # Scalar assignments run together; a multi-line block gets air around it. + body: list[str] = [] + for piece in by_section[prefix]: + if body and ("\n" in piece or "\n" in body[-1]): + body.append("") + body.append(piece) + sections.append("\n".join(body)) + parameters = "\n\n".join(sections) + + return _OFFLINE_SCRIPT_TEMPLATE.format( + display_name=manifest.get("display_name", "multilayer design"), + study_id=manifest.get("id", "unknown"), + script_name=offline_script_filename(manifest), + generated_at=datetime.now().isoformat(timespec="seconds"), + parameters=parameters, + constructor="\n".join(constructor), + ) + + +def offline_script_filename(manifest: dict[str, Any]) -> str: + """Return the filename to offer the generated script under.""" + + return f"{_slugify(str(manifest.get('id', 'multilayer_design')))}.py".replace("-", "_") + + +# --------------------------------------------------------------------------- # +# Survey-derived design choices # +# --------------------------------------------------------------------------- # +def survey_design_options(study_dir: Path) -> dict[str, Any]: + """Return the ``(d, blaze)`` choices offered after a completed survey. + + Reads ``survey/survey.csv`` and mirrors ``1_run_energy_scan.py``'s selection + logic, so the web offers exactly the designs the CLI would. + + Returns: + A dict with ``d_values``/``blaze_values`` (sorted grids for the + dropdowns), ``efficiency`` (``{"d,blaze": value}`` for the JS hint), + ``best`` (the single global-argmax pair, or ``None``) and ``per_d`` + (one pair per d-spacing). Empty lists when the CSV is missing or holds + no usable rows. + """ + + import pandas as pd + + empty: dict[str, Any] = { + "d_values": [], + "blaze_values": [], + "efficiency": {}, + "best": None, + "per_d": [], + } + csv_path = Path(study_dir) / "survey" / "survey.csv" + if not csv_path.is_file(): + return empty + try: + table = pd.read_csv(csv_path).dropna(subset=["peak_efficiency"]) + except (OSError, ValueError, KeyError): + # The survey rewrites this file as it runs; a read that lands badly is + # answered with "nothing yet" and the next poll picks it up. + return empty + if table.empty: + return empty + + best_row = table.loc[table["peak_efficiency"].idxmax()] + per_d_rows = table.loc[table.groupby("d_nm")["peak_efficiency"].idxmax()] + return { + "d_values": sorted({round(float(v), 3) for v in table["d_nm"]}), + "blaze_values": sorted({round(float(v), 3) for v in table["blaze_deg"]}), + "efficiency": { + f"{float(row.d_nm):.3f},{float(row.blaze_deg):.3f}": float(row.peak_efficiency) + for row in table.itertuples() + }, + "best": [float(best_row.d_nm), float(best_row.blaze_deg)], + "per_d": [ + [float(row.d_nm), float(row.blaze_deg)] for row in per_d_rows.itertuples() + ], + } + + +def resolve_designs( + mode: str, options: dict[str, Any], raw_designs: list[str] +) -> list[list[float]]: + """Resolve a selection ``mode`` into the ``(d, blaze)`` pairs to scan. + + Args: + mode: One of :data:`DESIGN_MODES`. + options: The dict from :func:`survey_design_options`. + raw_designs: ``"d,blaze"`` strings, used only when ``mode`` is + ``"manual"``. + + Returns: + The pairs to hand to + :meth:`grax.MultilayerGratingDesigner.run_energy_scan`. + + Raises: + ValueError: If the mode is unknown, the survey offers nothing, or a + manual pair is malformed or off the survey grid. + """ + + if mode not in DESIGN_MODES: + raise ValueError(f"Unknown design mode {mode!r}.") + if mode == "best": + if not options.get("best"): + raise ValueError("The survey has no usable cell to pick a best design from.") + return [list(options["best"])] + if mode == "per_d": + if not options.get("per_d"): + raise ValueError("The survey has no usable cells to pick designs from.") + return [list(pair) for pair in options["per_d"]] + + d_values = {round(float(v), 3) for v in options.get("d_values", ())} + blaze_values = {round(float(v), 3) for v in options.get("blaze_values", ())} + designs: list[list[float]] = [] + seen: set[tuple[float, float]] = set() + for raw in raw_designs: + text = str(raw).strip() + if not text: + continue + try: + d_text, blaze_text = text.split(",") + pair = (round(float(d_text), 3), round(float(blaze_text), 3)) + except ValueError as error: + raise ValueError(f"Could not read the design {text!r}.") from error + if pair[0] not in d_values or pair[1] not in blaze_values: + raise ValueError(f"The design {text!r} is not on the survey grid.") + if pair in seen: + continue + seen.add(pair) + designs.append([pair[0], pair[1]]) + if not designs: + raise ValueError("Pick at least one (d, blaze) design to scan.") + return designs + + +# --------------------------------------------------------------------------- # +# Study store # +# --------------------------------------------------------------------------- # +class MultilayerDesignStudyStore: + """Store multilayer-design study manifests in a filesystem directory.""" + + def __init__(self, directory: str | Path) -> None: + """Initialise the store rooted at ``directory``.""" + + self.directory = Path(directory) + + def list(self) -> list[dict[str, Any]]: + """Return study manifests, newest first.""" + + if not self.directory.exists(): + return [] + studies = [self.load(path.parent.name) for path in self.directory.glob("*/study.json")] + return sorted( + studies, + key=lambda study: (str(study.get("created_at", "")), str(study.get("id", ""))), + reverse=True, + ) + + def load(self, study_id: str) -> dict[str, Any]: + """Load one study manifest by id.""" + + path = self._study_dir(study_id) / "study.json" + with path.open("r", encoding="utf-8") as handle: + payload = json.load(handle) + payload.setdefault("id", study_id) + return payload + + def save(self, manifest: dict[str, Any]) -> dict[str, Any]: + """Persist a study manifest atomically and return it.""" + + payload = dict(manifest) + payload["updated_at"] = datetime.now().isoformat(timespec="seconds") + study_dir = self._study_dir(str(payload["id"])) + study_dir.mkdir(parents=True, exist_ok=True) + path = study_dir / "study.json" + temp_path = path.with_name(f"study.json.{datetime.now().timestamp():.9f}.tmp") + with temp_path.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + temp_path.replace(path) + return payload + + def create( + self, + *, + display_name: str, + config: dict[str, Any], + auto_energy_scan: str = "none", + ) -> dict[str, Any]: + """Create a new study directory + manifest and return it.""" + + slug = _slugify(display_name) or "study" + study_id = f"{datetime.now():%Y%m%d-%H%M%S}-{slug}" + candidate = study_id + suffix = 2 + while (self._study_dir(candidate) / "study.json").exists(): + candidate = f"{study_id}-{suffix}" + suffix += 1 + manifest = { + "id": candidate, + "created_at": datetime.now().isoformat(timespec="seconds"), + "display_name": display_name.strip() or candidate, + "comment": "", + "config": config, + "auto_energy_scan": ( + auto_energy_scan if auto_energy_scan in {"best", "none"} else "none" + ), + "stages": {stage: _blank_stage() for stage in STAGES}, + } + return self.save(manifest) + + def delete_many(self, study_ids: list[str]) -> None: + """Delete several study directories.""" + + for study_id in study_ids: + study_dir = self._study_dir(study_id) + if study_dir.exists(): + shutil.rmtree(study_dir) + + def study_dir(self, study_id: str) -> Path: + """Return the directory for one study id (validated).""" + + return self._study_dir(study_id) + + def _study_dir(self, study_id: str) -> Path: + if _slugify(study_id) != study_id: + raise ValueError("Invalid study id.") + return self.directory / study_id + + +def _blank_stage() -> dict[str, Any]: + return { + "status": "not_run", + "ran_at": None, + "error_text": "", + "aborted": False, + "designs": [], + "artifacts": {}, + } diff --git a/src/grax/web/static/web.css b/src/grax/web/static/web.css index 16dfabe..f42abdb 100644 --- a/src/grax/web/static/web.css +++ b/src/grax/web/static/web.css @@ -137,6 +137,15 @@ p { gap: 28px; } +/* Peer panels of equal width -- unlike `.grid.two`, which is a main pane plus a + narrower sidebar. At most two per row (a third 45% track cannot fit), and a + lone item takes the whole row instead of sitting in a half-width column. */ +.grid.halves { + display: grid; + grid-template-columns: repeat(auto-fit, minmax(45%, 1fr)); + gap: 28px; +} + .workspace-grid { display: grid; grid-template-columns: minmax(0, 1.2fr) minmax(320px, 0.8fr); @@ -233,6 +242,25 @@ legend { font-weight: 700; } +/* A named sub-group inside a fieldset: heading and fields each claim a full + row of the fieldset grid, so the group reads as one line. */ +.field-row-label { + grid-column: 1 / -1; + margin: 4px 0 0; + color: var(--muted); + font-size: 0.85rem; + font-weight: 700; + letter-spacing: 0.02em; + text-transform: uppercase; +} + +.field-row { + grid-column: 1 / -1; + display: grid; + grid-template-columns: repeat(auto-fit, minmax(170px, 1fr)); + gap: 14px; +} + .plot-run { grid-template-columns: 1fr; } @@ -299,6 +327,10 @@ select { margin-bottom: 12px; } +.plotly-stage.is-half { + min-height: 360px; +} + .plotly-stage { min-height: 420px; width: 100%; @@ -422,6 +454,54 @@ select { font-size: 0.85rem; } +.stage-card { + margin-bottom: 24px; + padding: 18px; + border: 1px solid var(--line); + border-radius: 8px; + background: #fcfcfc; +} + +.stage-card .form { + margin-top: 12px; +} + +.status-pill { + display: inline-flex; + align-items: center; + min-height: 24px; + padding: 2px 10px; + border: 1px solid var(--line); + border-radius: 999px; + font-size: 0.82rem; + color: var(--muted); + background: #fff; +} + +.status-pill.is-completed { + border-color: var(--accent); + color: var(--accent-dark); +} + +.status-pill.is-running, +.status-pill.is-queued, +.status-pill.is-aborting { + border-color: #b5860b; + color: #8a6508; +} + +.status-pill.is-stale { + border-color: #b5860b; + color: #8a6508; + background: #fff8e6; +} + +.status-pill.is-failed, +.status-pill.is-aborted { + border-color: #b42318; + color: #b42318; +} + .export-dialog { width: min(820px, 92vw); border: 1px solid var(--line); @@ -492,7 +572,8 @@ td { flex-direction: column; } - .grid.two { + .grid.two, + .grid.halves { grid-template-columns: 1fr; } diff --git a/src/grax/web/static/web.js b/src/grax/web/static/web.js index 525a12e..184cbbc 100644 --- a/src/grax/web/static/web.js +++ b/src/grax/web/static/web.js @@ -170,6 +170,325 @@ function renderPlotlyFigure(container, figureJson) { }); } +// The survey figures and the (d, blaze) picker live in different sections of the +// design detail page, so they talk through document-level events instead of +// reaching into each other: a click on the map asks the picker to toggle a case, +// and the picker announces its selection so the map can mark it. +const SURVEY_CELL_PICKED = "grax:survey-cell-picked"; +const DESIGN_SELECTION_CHANGED = "grax:design-selection-changed"; + +function formatDesignValue(value) { + return Number(value).toFixed(3); +} + +function designPairKey(d, blaze) { + return `${formatDesignValue(d)},${formatDesignValue(blaze)}`; +} + +function nearestGridValue(values, target) { + let best = null; + let bestDistance = Infinity; + values.forEach((value) => { + const distance = Math.abs(Number(value) - Number(target)); + if (distance < bestDistance) { + bestDistance = distance; + best = Number(value); + } + }); + return best; +} + +function surveyFigureLayout(meta, xTitle, yTitle, title, options) { + // The survey energy is a property of the survey, so it belongs on the survey + // plots only -- naming it on an energy scan would be actively misleading. + const parts = (options || {}).withEnergy === false + ? [meta.coating_label, `order ${meta.diffraction_order}`] + : [meta.coating_label, `${meta.target_energy_ev} eV`, `order ${meta.diffraction_order}`]; + const subtitle = parts + .filter((part) => part !== undefined && part !== null && part !== "") + .join(" · "); + return { + template: "plotly_white", + title: {text: subtitle ? `${title}
${subtitle}` : title, x: 0.02}, + margin: {l: 64, r: 24, t: 72, b: 56}, + xaxis: {title: {text: xTitle}}, + yaxis: {title: {text: yTitle}}, + showlegend: false, + }; +} + +function buildSurveyFigures(options, meta, selectedKeys) { + const dValues = (options.d_values || []).map(Number); + const blazeValues = (options.blaze_values || []).map(Number); + const efficiency = options.efficiency || {}; + const perD = options.per_d || []; + + const ridgeD = perD.map((pair) => Number(pair[0])); + const ridgeBlaze = perD.map((pair) => Number(pair[1])); + const ridgeEfficiency = perD.map((pair) => efficiency[designPairKey(pair[0], pair[1])] ?? null); + + const optimalBlaze = { + data: [ + { + type: "scatter", + mode: "lines+markers", + x: ridgeD, + y: ridgeBlaze, + line: {color: "#b0b7bd", width: 1.4}, + marker: { + size: 11, + color: ridgeEfficiency, + colorscale: "Viridis", + colorbar: {title: {text: "peak efficiency"}}, + }, + customdata: ridgeEfficiency, + hovertemplate: + "d = %{x:.3f} nm
blaze = %{y:.3f} deg
efficiency = %{customdata:.4f}", + }, + ], + layout: surveyFigureLayout( + meta, + "Bilayer d-spacing (nm)", + "Optimal blaze angle (deg)", + "Optimal blaze angle versus d-spacing", + ), + }; + + const maxEfficiency = { + data: [ + { + type: "scatter", + mode: "lines+markers", + x: ridgeD, + y: ridgeEfficiency, + line: {color: "#b0b7bd", width: 1.4}, + marker: { + size: 11, + color: ridgeBlaze, + colorscale: "Plasma", + colorbar: {title: {text: "optimal blaze (deg)"}}, + }, + customdata: ridgeBlaze, + hovertemplate: + "d = %{x:.3f} nm
efficiency = %{y:.4f}
blaze = %{customdata:.3f} deg", + }, + ], + layout: surveyFigureLayout( + meta, + "Bilayer d-spacing (nm)", + "Max selected-order efficiency", + "Max efficiency versus d-spacing", + ), + }; + + // z is indexed [blaze][d]; a cell the survey never solved stays null so Plotly + // leaves a gap instead of drawing a misleading zero. + const z = blazeValues.map((blaze) => + dValues.map((d) => { + const value = efficiency[designPairKey(d, blaze)]; + return value === undefined ? null : Number(value); + }), + ); + const selected = Array.from(selectedKeys || []).map((key) => key.split(",").map(Number)); + const heatmapLayout = surveyFigureLayout( + meta, + "Bilayer d-spacing (nm)", + "Blaze angle (deg)", + "Peak efficiency over (d, blaze) — click a cell to scan it", + ); + heatmapLayout.showlegend = true; + // Inside the axes, like the matplotlib version: a legend above the plot would + // land on the title's subtitle line. + heatmapLayout.legend = { + orientation: "h", + x: 0.02, + xanchor: "left", + y: 0.98, + yanchor: "top", + bgcolor: "rgba(255, 255, 255, 0.78)", + }; + const heatmap = { + data: [ + { + type: "heatmap", + x: dValues, + y: blazeValues, + z, + colorscale: "Viridis", + colorbar: {title: {text: "peak efficiency"}}, + hovertemplate: + "d = %{x:.3f} nm
blaze = %{y:.3f} deg
efficiency = %{z:.4f}", + }, + { + type: "scatter", + mode: "lines+markers", + name: "optimal blaze per d", + x: ridgeD, + y: ridgeBlaze, + line: {color: "#ffffff", width: 1.4}, + marker: {color: "#ffffff", size: 6}, + hoverinfo: "skip", + }, + { + type: "scatter", + mode: "markers", + name: "selected designs", + x: selected.map((pair) => pair[0]), + y: selected.map((pair) => pair[1]), + // square-open takes its stroke from marker.color, not marker.line. + marker: {symbol: "square-open", size: 16, color: "#ff3b30", line: {width: 3}}, + hoverinfo: "skip", + }, + ], + layout: heatmapLayout, + }; + + return {optimal_blaze: optimalBlaze, max_efficiency: maxEfficiency, heatmap}; +} + +function initEnergyScanFigure(root) { + const stage = root.querySelector("[data-energy-scan-figure]"); + const liveUrl = root.dataset.energyScanLiveUrl; + if (!stage || !liveUrl || !window.Plotly) { + return; + } + let meta; + try { + meta = JSON.parse(root.dataset.surveyPlotMeta || "{}"); + } catch (error) { + meta = {}; + } + + function draw(designs) { + const layout = surveyFigureLayout( + meta, + "Photon energy (eV)", + "Selected-order efficiency", + "Efficiency versus energy", + {withEnergy: false}, + ); + // Below the axes: one entry per design, and the curves themselves leave no + // reliable free corner. + layout.showlegend = true; + layout.legend = {orientation: "h", x: 0, y: -0.22, yanchor: "top"}; + layout.margin = {l: 64, r: 24, t: 72, b: 96}; + renderPlotlyFigure(stage, { + data: designs.map((design) => ({ + type: "scatter", + mode: "lines+markers", + name: `d = ${formatDesignValue(design.d_spacing_nm)} nm, blaze = ${formatDesignValue(design.blaze_angle_deg)} deg`, + x: design.energies_ev || [], + y: design.efficiencies || [], + marker: {size: 5}, + line: {width: 1.6}, + hovertemplate: "%{x:.1f} eV
efficiency = %{y:.4f}%{fullData.name}", + })), + layout, + }); + } + + async function poll() { + try { + const response = await window.fetch(liveUrl, {cache: "no-store"}); + if (!response.ok) { + return; + } + const payload = await response.json(); + if ((payload.designs || []).some((design) => (design.energies_ev || []).length)) { + draw(payload.designs); + } + } catch (error) { + // A dropped poll is not worth surfacing; the next one will catch up. + } + } + + poll(); + window.setInterval(poll, 4000); +} + +function initSurveyFigures(root) { + let options; + let meta; + try { + options = JSON.parse(root.dataset.designOptions || "{}"); + meta = JSON.parse(root.dataset.surveyPlotMeta || "{}"); + } catch (error) { + return; + } + const stages = new Map(); + root.querySelectorAll("[data-survey-figure]").forEach((node) => { + stages.set(node.dataset.surveyFigure, node); + }); + if (stages.size === 0 || !window.Plotly) { + return; + } + const dValues = (options.d_values || []).map(Number); + const blazeValues = (options.blaze_values || []).map(Number); + let selectedKeys = new Set(); + + function draw() { + const figures = buildSurveyFigures(options, meta, selectedKeys); + stages.forEach((node, name) => { + if (figures[name]) { + renderPlotlyFigure(node, figures[name]); + } + }); + } + + draw(); + + // While the survey runs its table is rewritten after every cell, so re-reading + // it shows the grid filling in. The monitor reloads the page when the stage + // finishes, which is what ends this polling. + const liveUrl = root.dataset.surveyLiveUrl; + if (liveUrl) { + window.setInterval(async () => { + try { + const response = await window.fetch(liveUrl, {cache: "no-store"}); + if (!response.ok) { + return; + } + const fresh = await response.json(); + if ((fresh.per_d || []).length === 0) { + return; + } + options = fresh; + dValues.splice(0, dValues.length, ...(fresh.d_values || []).map(Number)); + blazeValues.splice(0, blazeValues.length, ...(fresh.blaze_values || []).map(Number)); + draw(); + } catch (error) { + // A dropped poll is not worth surfacing; the next one will catch up. + } + }, 4000); + } + + const heatmapNode = stages.get("heatmap"); + if (heatmapNode && typeof heatmapNode.on === "function") { + heatmapNode.on("plotly_click", (event) => { + const point = (event.points || [])[0]; + if (!point) { + return; + } + const d = nearestGridValue(dValues, point.x); + const blaze = nearestGridValue(blazeValues, point.y); + if (d === null || blaze === null) { + return; + } + if ((options.efficiency || {})[designPairKey(d, blaze)] === undefined) { + return; // an unsolved cell is not a design anyone can scan + } + document.dispatchEvent( + new CustomEvent(SURVEY_CELL_PICKED, {detail: {d, blaze}}), + ); + }); + } + + document.addEventListener(DESIGN_SELECTION_CHANGED, (event) => { + selectedKeys = new Set((event.detail && event.detail.pairs) || []); + draw(); + }); +} + function initPlotWorkspace(form) { const previewUrl = form.dataset.previewUrl; const picker = form.querySelector("[data-run-picker]"); @@ -362,7 +681,12 @@ function initRunMonitor(container) { const memoryNode = container.querySelector("[data-run-memory]"); const errorNode = container.querySelector("[data-run-error]"); const progressBar = container.querySelector("[data-run-progress-bar]"); - const abortButton = document.querySelector("[data-run-abort-action]"); + // Only some monitors render a "what is running right now" line; the shared + // run-detail monitor does not, so every use below is null-guarded. + const currentLabelNode = container.querySelector("[data-run-current-label]"); + // Scoped to this monitor: a page can host several (one per workflow stage), + // and a document-wide lookup would let every monitor drive the first button. + const abortButton = container.querySelector("[data-run-abort-action]"); let latestPlotToken = ""; let latestPlotUrl = plotImage.getAttribute("src") || ""; let statusTimerId = null; @@ -398,6 +722,9 @@ function initRunMonitor(container) { const percent = payload.total_points > 0 ? (payload.completed_points / payload.total_points) * 100 : 0; progressBar.style.width = `${percent}%`; errorNode.textContent = payload.error_text || ""; + if (currentLabelNode) { + currentLabelNode.textContent = payload.current_label || ""; + } if (abortButton) { if (payload.can_abort) { abortButton.removeAttribute("aria-disabled"); @@ -415,6 +742,13 @@ function initRunMonitor(container) { if (memoryTimerId !== null) { window.clearInterval(memoryTimerId); } + // Results are rendered server-side, so a page loaded while the stage was + // running has no way to show them. Monitors that opt in reload once; the + // reloaded page renders the results instead of this monitor, so there is + // no loop. + if (container.hasAttribute("data-run-reload-on-finish")) { + window.location.reload(); + } } } @@ -440,6 +774,165 @@ function initRunMonitor(container) { memoryTimerId = window.setInterval(pollMemory, 500); } +function initSurveyCellCounter(form) { + const dPoints = form.querySelector("[data-survey-d-points]"); + const blazePoints = form.querySelector("[data-survey-blaze-points]"); + const readout = form.querySelector("[data-survey-cell-readout]"); + if (!dPoints || !blazePoints || !readout) { + return; + } + const threshold = Number(form.dataset.cellWarningThreshold || 200); + + function update() { + const cells = Math.max(0, Number(dPoints.value) || 0) * Math.max(0, Number(blazePoints.value) || 0); + if (cells > threshold) { + readout.textContent = + `This survey will run ${cells} theta searches — well above ${threshold}, so expect it to take a long time. ` + + `You can abort it from the study page once it starts.`; + readout.classList.add("notice-error"); + } else { + readout.textContent = `This survey will run ${cells} theta searches.`; + readout.classList.remove("notice-error"); + } + } + + dPoints.addEventListener("input", update); + blazePoints.addEventListener("input", update); + update(); +} + +function initDesignPicker(form) { + let options; + try { + options = JSON.parse(form.dataset.designOptions || "{}"); + } catch (error) { + return; + } + const toggle = form.querySelector("[data-design-toggle]"); + const panel = form.querySelector("[data-design-manual]"); + const rows = form.querySelector("[data-design-rows]"); + const addButton = form.querySelector("[data-design-add]"); + if (!toggle || !panel || !rows || !addButton) { + return; + } + const dValues = options.d_values || []; + const blazeValues = options.blaze_values || []; + const efficiency = options.efficiency || {}; + + const formatValue = formatDesignValue; + + function announceSelection() { + const pairs = Array.from(rows.querySelectorAll('input[name="design"]')).map( + (input) => input.value, + ); + document.dispatchEvent( + new CustomEvent(DESIGN_SELECTION_CHANGED, {detail: {pairs}}), + ); + } + + function buildSelect(values, initial) { + const select = document.createElement("select"); + values.forEach((value) => { + const option = document.createElement("option"); + option.value = formatValue(value); + option.textContent = formatValue(value); + select.appendChild(option); + }); + if (initial !== undefined) { + select.value = formatValue(initial); + } + return select; + } + + function addRow(pair) { + const best = pair || options.best || [dValues[0], blazeValues[0]]; + const row = document.createElement("div"); + row.className = "row"; + + const dSelect = buildSelect(dValues, best[0]); + const blazeSelect = buildSelect(blazeValues, best[1]); + const hidden = document.createElement("input"); + hidden.type = "hidden"; + hidden.name = "design"; + + const hint = document.createElement("span"); + hint.className = "subtle"; + + function sync() { + const pair = `${dSelect.value},${blazeSelect.value}`; + hidden.value = pair; + const value = efficiency[pair]; + hint.textContent = value === undefined ? "not surveyed" : `efficiency ${Number(value).toPrecision(4)}`; + announceSelection(); + } + + dSelect.addEventListener("change", sync); + blazeSelect.addEventListener("change", sync); + sync(); + + const remove = document.createElement("button"); + remove.type = "button"; + remove.className = "button danger"; + remove.textContent = "Remove"; + remove.addEventListener("click", () => { + row.remove(); + announceSelection(); + }); + + const dLabel = document.createElement("label"); + dLabel.append("d-spacing, nm", dSelect); + const blazeLabel = document.createElement("label"); + blazeLabel.append("Blaze, deg", blazeSelect); + + row.append(dLabel, blazeLabel, hint, hidden, remove); + rows.appendChild(row); + // Only now is the hidden input part of the list announceSelection reads. + announceSelection(); + } + + toggle.addEventListener("click", () => { + panel.classList.toggle("is-hidden"); + if (!panel.classList.contains("is-hidden") && rows.children.length === 0) { + addRow(); + } + }); + addButton.addEventListener("click", () => addRow()); + + // Clicking a heatmap cell toggles that design: a second click on an already + // chosen cell removes its row, so the map doubles as the selection list. + document.addEventListener(SURVEY_CELL_PICKED, (event) => { + const {d, blaze} = event.detail || {}; + if (d === undefined || blaze === undefined) { + return; + } + const key = designPairKey(d, blaze); + const existing = Array.from(rows.querySelectorAll('input[name="design"]')).find( + (input) => input.value === key, + ); + if (existing) { + existing.closest(".row").remove(); + announceSelection(); + return; + } + panel.classList.remove("is-hidden"); + addRow([d, blaze]); + }); + + form.addEventListener("submit", (event) => { + const submitter = event.submitter; + if (!submitter || submitter.value !== "manual") { + return; + } + const chosen = new Set( + Array.from(rows.querySelectorAll('input[name="design"]')).map((input) => input.value), + ); + if (chosen.size === 0) { + event.preventDefault(); + window.alert("Add at least one (d, blaze) design to scan."); + } + }); +} + document.addEventListener("DOMContentLoaded", () => { document.querySelectorAll("[data-grating-type]").forEach((select) => { syncGratingSections(select); @@ -484,8 +977,23 @@ document.addEventListener("DOMContentLoaded", () => { initSavedPlotFigure(container); }); - const runMonitor = document.querySelector("[data-live-run-monitor]"); - if (runMonitor) { + document.querySelectorAll("[data-survey-cell-counter]").forEach((form) => { + initSurveyCellCounter(form); + }); + + document.querySelectorAll("[data-design-picker]").forEach((form) => { + initDesignPicker(form); + }); + + document.querySelectorAll("[data-survey-figures]").forEach((root) => { + initSurveyFigures(root); + }); + + document.querySelectorAll("[data-energy-scan-figure-root]").forEach((root) => { + initEnergyScanFigure(root); + }); + + document.querySelectorAll("[data-live-run-monitor]").forEach((runMonitor) => { initRunMonitor(runMonitor); - } + }); }); diff --git a/src/grax/web/templates/_multilayer_design_macros.html b/src/grax/web/templates/_multilayer_design_macros.html new file mode 100644 index 0000000..a34426b --- /dev/null +++ b/src/grax/web/templates/_multilayer_design_macros.html @@ -0,0 +1,76 @@ +{% macro render_field(spec, values) %} + {% set value = values.get(spec.name) %} + {% if spec.kind == "material" %} + {% set pair = value if value is iterable and value is not string else ["", ""] %} + + + {% elif spec.kind == "select" %} + + {% elif spec.kind == "checkbox" %} + {# The hidden "0" always submits, so an unchecked box reads as False rather + than as "field absent" -- which on an edit would keep the old value. #} + + {% else %} + + {% endif %} +{% endmacro %} + +{% macro material_datalist(materials, density_map) %} + + {% for material in materials %} + + {% endfor %} + +{% endmacro %} + +{% macro stage_monitor(study_id, stage) %} +
+ +

Status: queued — 0 / 0

+
+

+
+
Completed0
+
Remaining0
+
Elapsed—
+
ETA—
+
Workersauto
+
Memory—
+
+

+
+ Abort stage +
+
+{% endmacro %} diff --git a/src/grax/web/templates/base.html b/src/grax/web/templates/base.html index 8d80831..6046a85 100644 --- a/src/grax/web/templates/base.html +++ b/src/grax/web/templates/base.html @@ -15,6 +15,7 @@ Grax Web