diff --git a/ROADMAP.md b/ROADMAP.md index 4345820..1613a9f 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -227,7 +227,7 @@ Acceptance (met) Composition/space/volume coverage maps; energy and per-element force distributions with outlier flags; extrapolation grade; redundancy report. -### 🚧 Chunk 3 — Validation (`validation/`) *(parity + quality report done)* +### 🚧 Chunk 3 — Validation (`validation/`) *(parity, NVE probe, spectra, learning curves done)* - ✅ `validation/` — parity metrics (RMSE/MAE for E [total + per-atom], forces [overall + per-element], stress/dipole/polarizability) against reference @@ -240,15 +240,35 @@ with outlier flags; extrapolation grade; redundancy report. training holdout `model/valid.xyz` (never trained on), overridable via `test_file`. Slurm executor maps it to `traincraft-mlip.sif`. Failed frames skip-and-record (same philosophy as labeling). -- 🔜 Learning curves (error vs dataset size), NVE/MD stability, EOS/phonons. -- 🔜 **IR and Raman spectra reconstruction** from MLIP-driven MD vs - DFT/experiment. +- ✅ NVE stability probe (`[validation.stability]`): Langevin equilibration → + Velocity-Verlet production with an explosion guard; `nve_aborted`, + `nve_drift_mev_per_atom_ps` etc. fold into the quality report as + thresholdable metrics. An exploding model is a *result*, not a crash. +- ✅ **Spectra reconstruction from MLIP-driven MD** (`[validation.spectra]`, + TCF route, clean-room from published theory): VDOS (mass-weighted, + per-element), IR (dipole-derivative ACF), Raman (iso/aniso invariants → + VV/VH/unpolarized/depolarization, laser + Bose prefactors, selectable QCF), + and SFG (Morita–Hynes cross-TCF, signed Im χ⁽²⁾). Dielectric model loaded + via `MaceCalc.model_type`; artifacts under `validation/spectra/`; optional + reference files → `_spectral_error` metrics (the AL convergence hook). + SERS = documented recipe (chemical enhancement via adsorbed-system training + data + optional `enhancement_file` EM profile; TrainCraft computes no EM + fields). +- ✅ Learning curves: `traincraft learning-curve` retrains on nested random + subsets against one fixed holdout → `validation/learning_curve.json` + (diagnostic command, not a stage). +- 🔜 EOS/phonons; harmonic (Hessian-based) spectra path. Acceptance - ✅ A model is validated per property with recorded pass/fail vs thresholds (covered by `test_validation`, EMT standing in for the MLIP). +- ✅ The physics probes are exercised end-to-end scheduler-free + (`test_spectra` against synthetic signals with known answers; + `test_dynamics`/`test_validation_science`/`test_learning_curve` with EMT and + an analytic dielectric stand-in). - 🔜 A fine-tuned model trains on a seed set and produces a quality report - covering every property, including reconstructed IR/Raman spectra. + covering every property, including reconstructed IR/Raman spectra + (needs real MACE + machine access). --- diff --git a/docs/how-to/validation-science.md b/docs/how-to/validation-science.md new file mode 100644 index 0000000..4bcc12d --- /dev/null +++ b/docs/how-to/validation-science.md @@ -0,0 +1,153 @@ +# Validation Science: NVE Probe, Spectra, Learning Curves + +Parity metrics (RMSE/MAE against reference labels) tell you the model +interpolates; they don't tell you it will **survive dynamics** or **reproduce +observables**. Three tools close that gap, and all of them write into the same +`validation/quality_report.json`, so their numbers threshold exactly like +parity metrics — including as convergence criteria for the future +active-learning loop. + +## NVE stability probe + +Add a `[validation.stability]` section and the `validate` stage runs short MD +with the model under test: Langevin equilibration to `temperature`, then +**NVE (Velocity-Verlet) production** — no thermostat to hide a broken +potential-energy surface. + +```toml +[validation.stability] +steps = 2000 # NVE production steps +timestep = 1.0 # fs +temperature = 300.0 +equilibration_steps = 500 +n_structures = 1 # probe the first N reference frames; worst case reported + +[validation] +thresholds = { nve_aborted = 0.0, nve_drift_mev_per_atom_ps = 1.0 } +``` + +Metrics folded into the report (thresholdable, `value <= limit`): + +| Metric | Meaning | +|---|---| +| `nve_aborted` | `1.0` if any probe trajectory exploded (force/energy guard) — threshold at `0.0` | +| `nve_drift_mev_per_atom_ps` | \|slope\| of total energy vs time — the standard NVE health number | +| `nve_energy_std_mev_per_atom` | detrended energy fluctuation | +| `nve_temp_mean` | mean production temperature (K) | + +An exploding trajectory is a **result** (`nve_aborted = 1`), not a crash: the +run completes and the report shows why the model failed. + +## Spectra from MLIP-driven MD (VDOS, IR, Raman, SFG) + +A `[validation.spectra]` section reconstructs vibrational spectra from one +production trajectory via time-correlation functions: velocities → VDOS, +dipole → IR, polarizability → Raman (VV/VH/depolarization), and the +dipole–polarizability cross-correlation → the resonant SFG response. + +```toml +[validation.spectra] +which = ["vdos", "ir", "raman"] # any of vdos | ir | raman | sfg +temperature = 300.0 +timestep = 0.5 # fs — resolve the fastest mode (H stretches!) +equilibration_steps = 2000 +steps = 20000 # production; longer = sharper lines +laser_wavelength = 532.0 # nm, Raman prefactor +qcf = "harmonic" # quantum correction: harmonic|standard|schofield|none + +# the dielectric model evaluating dipole/polarizability along the trajectory +[validation.spectra.calculator] +type = "mace" +model_path = "runs/dielectric/model/mace_model.model" +model_type = "DipoleMACE" # MACE calculator class for dielectric models + +# optional: reference spectra -> _spectral_error metrics +[validation.spectra.reference] +ir = "experiment_ir.dat" # two columns: freq_cm1 intensity +raman = "dft_raman.dat" +``` + +What lands in `runs//validation/spectra/`: one two-column `.dat` per +spectrum (`vdos.dat` + per-element `vdos_.dat`, `ir.dat`, `raman_vv.dat`, +`raman_vh.dat`, `raman_unpolarized.dat`, `raman_depolarization.dat`, +`raman_iso.dat`, `raman_aniso.dat`, `sfg_re.dat`, `sfg_im.dat`, +`sfg_abs2.dat`), plus `spectra.json` (parameters, peak tables, metrics). Set +`save_series = true` to also keep the raw time series (`series.npz`). + +Practical guidance: + +- **Which model computes what.** The model under test drives the MD (energy/ + forces). Dipole/polarizability come from `[validation.spectra.calculator]` — + typically a separately trained dielectric MACE (`AtomicDipolesMACE` / + `AtomicDielectricMACE` heads, loaded with `model_type`). If omitted, the + model under test is asked for those properties directly (works only for + models with dielectric heads). VDOS needs neither. +- **MD length sets resolution.** The frequency resolution is roughly + `1 / (correlation window)`: 20 000 steps × 0.5 fs → 10 ps → ~3 cm⁻¹ grid + before broadening. Aim for at least ~10× the period of your slowest mode. +- **Timestep sets the ceiling.** The Nyquist limit is `1/(2·dt)`: 0.5 fs + covers ~33 000 cm⁻¹ — fine; 4 fs caps at ~4 200 cm⁻¹ — too close to O–H + stretches for comfort. +- **`qcf`** applies a quantum correction on top of the derivative-ACF spectra + (which already correspond to the *harmonic* correction, hence + `harmonic` = default and a no-op). Change it only if you know which + convention your reference uses. +- **SFG** is surface-specific: pick the χ⁽²⁾ component with `sfg_component` + (default `"xxz"` = the ssp geometry with the surface normal along z) and + keep your slab's normal along z. The sign of `sfg_im.dat` carries the net + orientation of each mode. +- **References** are compared after Gaussian broadening (`broadening`, cm⁻¹) + and normalization on a common grid over `freq_range`; each contributes + `_spectral_error = 1 − cosine similarity` (0 = same shape) — this is + the "spectral error" convergence criterion. Stick spectra (e.g. harmonic + DFT output) work as references: both sides go through the same broadening. + +Slurm note: the spectra MD runs *inside* the `validate` job +(`traincraft-mlip.sif`, 1 GPU by default) — size `stages.validate.time` +accordingly. + +## SERS (surface-enhanced Raman): what is and isn't computable + +SERS enhancement has two parts, and it's important to be honest about which +one an MD/TCF tool can give you: + +- **Chemical enhancement** (adsorption-induced changes of the polarizability) + is captured *implicitly*: train the dielectric model on DFT + dipoles/polarizabilities of the **molecule + metal cluster/slab** system, + then run `raman` spectra of the adsorbed geometry. The difference to the + gas-phase spectrum *is* the chemical enhancement. +- **Electromagnetic enhancement** (the plasmonic 10⁴–10⁸ factor) is classical + electrodynamics of the nanostructure (Mie/BEM/FDTD) — TrainCraft does not + compute it. If you have an enhancement profile from such a tool (or an + empirical one), point `enhancement_file` at a two-column + `freq_cm1 enhancement` file: it is interpolated onto the Raman grid and + multiplied onto the VV/VH/unpolarized spectra, written as + `raman_*_enhanced.dat`. + +```toml +[validation.spectra] +which = ["raman"] +enhancement_file = "em_enhancement.dat" # from your EM solver, or empirical +``` + +## Learning curve: error vs training-set size + +How data-efficient is your setup — and is more data still buying accuracy? + +```bash +pixi run -e mace traincraft learning-curve my_run.toml --fractions 0.125,0.25,0.5,1.0 +``` + +The command retrains the `[training]` model on **nested** random subsets of +the training frames (smaller ⊂ larger — one shuffle, seeded by `[run].seed`) +and validates every point against **one fixed holdout** (split off once with +the trainer's `valid_fraction`; never trained on at any fraction). Each point +validates its own trained model. Results: +`runs//validation/learning_curve.json` with +`points: [{fraction, n_train, metrics{...}}, ...]`, plus per-point model dirs +under `validation/learning_curve/`. + +It retrains once per fraction — run it where training runs (a GPU allocation); +it is a diagnostic command, not a pipeline stage, so it never touches the +stage cache. A curve still steep at 100% of your data says "label more"; a +flat one says the model, not the data, is the limit. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 728745f..f26ac63 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -105,11 +105,34 @@ stage 'train': 42 frames The `train` stage writes `runs//model/.model` and a `manifest.json`. The `validate` stage writes `runs//validation/quality_report.json` — parity metrics + threshold checks for the trained model (see -[`[validation]`](config.md#validation)). +[`[validation]`](config.md#validation)), plus the NVE stability metrics and +`validation/spectra/` artifacts when configured (see +[Validation Science](../how-to/validation-science.md)). See [Tutorial 10 · Training](../tutorials/10-training.md). --- +## `traincraft learning-curve` + +Validation error vs training-set size: retrain the `[training]` model on +nested random subsets and validate every point against one fixed holdout. + +```bash +pixi run -e mace traincraft learning-curve CONFIG [--fractions 0.125,0.25,0.5,1.0] +``` + +| Option | Description | +|---|---| +| `--fractions` | Comma-separated training-set fractions in `(0, 1]` | + +Prints a table (`n_train` vs the headline metrics) and writes +`runs//validation/learning_curve.json`. It retrains once per fraction — +run it where training runs (e.g. inside a GPU allocation). A diagnostic +command, not a pipeline stage: no events, no cache keys. See +[Validation Science](../how-to/validation-science.md#learning-curve-error-vs-training-set-size). + +--- + ## `traincraft submit` Render (and, unless `--dry-run`, submit) the workflow as dependency-chained diff --git a/docs/reference/config.md b/docs/reference/config.md index c0ac028..c1fc2ce 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -385,6 +385,7 @@ Same as `tblite`. Uses the `xtb-python` binding instead. | `model_path` | `str \| null` | `null` | Path to a local `.model` checkpoint (overrides `model`) | | `device` | `str` | `"cpu"` | `"cpu"` or `"cuda"` | | `default_dtype` | `str` | `"float32"` | `"float32"` or `"float64"` | +| `model_type` | `str \| null` | `null` | MACE calculator class for dielectric models trained with dipole/polarizability heads (e.g. `"DipoleMACE"`, `"EnergyDipoleMACE"`). Requires `model_path`; verify the name against your installed MACE version | --- @@ -618,11 +619,15 @@ what the future active-learning loop will use as its convergence criterion. | `properties` | list \| null | `null` | Which properties to check (`energy`/`forces`/`stress`/`dipole`/`polarizability`). Default: everything the reference frames carry | | `max_frames` | `int \| null` | `null` | Cap the number of frames compared (quick checks) | | `thresholds` | `dict` | `{}` | Metric name → maximum allowed value; the report records each check and an overall `passed`. Unknown metric names fail fast | +| `stability` | table \| null | `null` | NVE stability probe — see [`[validation.stability]`](#validationstability) | +| `spectra` | table \| null | `null` | MD spectra reconstruction — see [`[validation.spectra]`](#validationspectra) | Metric names for `thresholds`: `energy_rmse`, `energy_mae`, `energy_rmse_per_atom`, `energy_mae_per_atom` (eV, eV/atom), `forces_rmse`, -`forces_mae` (eV/Å), and `_rmse`/`_mae` for -`stress`/`dipole`/`polarizability` when present. +`forces_mae` (eV/Å), `_rmse`/`_mae` for +`stress`/`dipole`/`polarizability` when present, the `nve_*` metrics when +`stability` is set, and `spectra_md_aborted` / `_spectral_error` +when `spectra` is set. ```toml [validation] @@ -633,3 +638,46 @@ thresholds = { forces_rmse = 0.05, energy_mae_per_atom = 0.005 } A frame whose energy/forces prediction fails is skipped and recorded (same philosophy as `[labeling]`'s `on_error="skip"`); optional properties the model cannot produce are dropped per property without voiding the frame. + +### `[validation.stability]` + +NVE stability probe: Langevin equilibration then Velocity-Verlet (NVE) +production with the model under test; energy-conservation metrics +(`nve_aborted`, `nve_drift_mev_per_atom_ps`, `nve_energy_std_mev_per_atom`, +`nve_temp_mean`) join the quality report. An exploding trajectory sets +`nve_aborted = 1` instead of failing the stage. See +[Validation Science](../how-to/validation-science.md). + +| Field | Type | Default | Description | +|---|---|---|---| +| `steps` | `int` | `2000` | NVE production steps | +| `timestep` | `float` | `1.0` | fs | +| `temperature` | `float` | `300.0` | Init + equilibration target (K) | +| `equilibration_steps` | `int` | `500` | Langevin steps before production | +| `n_structures` | `int` | `1` | Probe the first N reference frames; worst case reported | + +### `[validation.spectra]` + +Vibrational spectra reconstructed from one MLIP-MD trajectory via +time-correlation functions: velocities → VDOS, dipole → IR, polarizability → +Raman (VV/VH/depolarization), dipole×polarizability → SFG. Artifacts land in +`validation/spectra/` (two-column `.dat` per spectrum + `spectra.json`). See +[Validation Science](../how-to/validation-science.md) for the physics knobs. + +| Field | Type | Default | Description | +|---|---|---|---| +| `which` | `list` | `["vdos", "ir", "raman"]` | Any of `vdos`/`ir`/`raman`/`sfg` | +| `temperature` | `float` | `300.0` | K | +| `timestep` | `float` | `0.5` | fs (sets the Nyquist ceiling) | +| `equilibration_steps` | `int` | `2000` | Langevin steps before production | +| `steps` | `int` | `20000` | NVE production steps (sets resolution) | +| `sample_every` | `int` | `1` | Record every Nth step | +| `calculator` | calculator config \| null | `null` | Dielectric model evaluating dipole/polarizability along the trajectory (e.g. `mace` with `model_type`). Default: the model under test | +| `qcf` | `str` | `"harmonic"` | Quantum correction: `harmonic`/`standard`/`schofield`/`none` | +| `laser_wavelength` | `float` | `532.0` | nm — Raman Stokes prefactor | +| `sfg_component` | `str` | `"xxz"` | χ⁽²⁾ indices (`"xxz"` = ssp, surface normal ∥ z) | +| `freq_range` | `[lo, hi]` | `[0.0, 4000.0]` | cm⁻¹ — written + compared range | +| `broadening` | `float` | `10.0` | cm⁻¹ Gaussian width for the similarity metric | +| `reference` | `dict` | `{}` | Spectrum name → two-column file (`freq_cm1 intensity`); each adds a `_spectral_error` metric | +| `enhancement_file` | path \| null | `null` | Frequency-dependent EM enhancement profile multiplied onto Raman spectra (the SERS recipe) | +| `save_series` | `bool` | `false` | Also keep the raw time series as `spectra/series.npz` | diff --git a/mkdocs.yml b/mkdocs.yml index 7c60f93..ca1b608 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -126,6 +126,7 @@ nav: - Write a Custom Builder: how-to/custom-builder.md - Write a Custom Sampler: how-to/custom-sampler.md - Train & Use Your Own MACE Model: how-to/own-mace-model.md + - Validation Science (NVE, Spectra, Learning Curves): how-to/validation-science.md - Manage Datasets: how-to/datasets.md - Run on HPC (Slurm + Apptainer): how-to/hpc.md - Serve Runs to Agents & UIs: how-to/agent-service.md diff --git a/skills/traincraft/SKILL.md b/skills/traincraft/SKILL.md index 53757a7..44191e7 100644 --- a/skills/traincraft/SKILL.md +++ b/skills/traincraft/SKILL.md @@ -103,6 +103,11 @@ Map the user's words to config keys (run `traincraft plugins` for the live list) optional `thresholds` → `passed`). Defaults to the model `[training]` just produced and its holdout `valid.xyz`. **Read `quality_report.json` and report the metrics + pass/fail to the user** — never claim a model is good without it. + Physics probes are opt-in sub-sections: `[validation.stability]` (NVE probe → + `nve_*` metrics; threshold `nve_aborted = 0.0`) and `[validation.spectra]` + (MD-reconstructed VDOS/IR/Raman/SFG → `validation/spectra/*.dat` + + `_spectral_error` vs reference files). See + `docs/how-to/validation-science.md`. ## ⚠️ Gotchas (learned the hard way) @@ -126,6 +131,14 @@ Map the user's words to config keys (run `traincraft plugins` for the live list) variants (stochastic builders get derived seeds); `smiles` with `n_conformers > 1` and `file` with `index = ":"` also yield several seeds. Every seed is sampled and the candidates pooled before selection. +- **IR/Raman/SFG spectra need a dielectric model.** `[validation.spectra]` with + `ir`/`raman`/`sfg` requires dipole/polarizability predictions — point + `[validation.spectra.calculator]` at a trained dielectric MACE (`model_path` + + `model_type = "DipoleMACE"` etc.). `vdos` works with any model. The MD runs + inside the `validate` job: size `stages.validate.time` for it. +- **`traincraft learning-curve` retrains once per fraction** — never launch it + casually; run it where `stage train` runs (a GPU allocation), with the `mace` + env. ## 📂 Outputs (under `runs//`) @@ -143,7 +156,12 @@ Map the user's words to config keys (run `traincraft plugins` for the live list) - `model/.model` (+ `manifest.json`, `valid.xyz` holdout) — a trained/fine-tuned MACE model. - `validation/quality_report.json` — parity metrics, threshold checks, and the - overall `passed` verdict for the validated model. + overall `passed` verdict for the validated model (plus `nve_*` and spectral + metrics when those probes are configured). +- `validation/spectra/*.dat` + `spectra.json` — reconstructed spectra + (two-column, plottable as-is) and their peak tables. +- `validation/learning_curve.json` — error vs training-set size from + `traincraft learning-curve`. ## 🖼 Visualising (headless-VM friendly) @@ -169,6 +187,7 @@ traincraft validate # validate + show resolved stages pixi run -e traincraft run # run the pipeline pixi run -e traincraft stage # one stage standalone # (geometry|sample|select|label|dataset|train|validate) +pixi run -e mace traincraft learning-curve [--fractions 0.25,0.5,1.0] # error vs data size traincraft runs [outdir] [--json] # list all runs + status (default outdir: runs/) traincraft status [--json] # per-stage status of one run (live label progress) ``` diff --git a/src/traincraft/calculators/potentials.py b/src/traincraft/calculators/potentials.py index 4dfa30c..224cc7d 100644 --- a/src/traincraft/calculators/potentials.py +++ b/src/traincraft/calculators/potentials.py @@ -45,10 +45,16 @@ def build_mace(cfg): if cfg.model_path: from mace.calculators import MACECalculator + kwargs = {} + # Dielectric model classes (DipoleMACE, ...) — runtime-gated: an + # unknown name fails inside MACE with its own list of valid types. + if getattr(cfg, "model_type", None): + kwargs["model_type"] = cfg.model_type return MACECalculator( model_paths=cfg.model_path, device=cfg.device, default_dtype=cfg.default_dtype, + **kwargs, ) if cfg.model == "mace-off23": from mace.calculators import mace_off diff --git a/src/traincraft/cli.py b/src/traincraft/cli.py index 0f3e9e6..9bb03f2 100644 --- a/src/traincraft/cli.py +++ b/src/traincraft/cli.py @@ -231,6 +231,54 @@ def submit( typer.echo(f" {job.stage}{dep}: {job.script}{jid}") +@app.command("learning-curve", rich_help_panel=_RUN) +def learning_curve_cmd( + config: Path = typer.Argument(..., exists=True, readable=True, help="the .toml to use"), + fractions: str = typer.Option( + "0.125,0.25,0.5,1.0", "--fractions", + help="comma-separated training-set fractions in (0, 1]", + ), +) -> None: + """Validation error vs training-set size (retrains once per fraction!). + + Trains the \\[training] model on nested random subsets of the training + frames, validates each against one fixed holdout, and writes + [cyan]runs//validation/learning_curve.json[/]. Run it where the + trainer runs (e.g. inside a GPU allocation). + + [b]Example[/] + + \b + [cyan]pixi run -e mace traincraft learning-curve my_run.toml --fractions 0.25,0.5,1.0[/] + """ + from .validation.learning_curve import learning_curve + + _setup_logging() + cfg = load_config(config) + try: + parsed = [float(x) for x in fractions.split(",") if x.strip()] + except ValueError: + typer.echo(f"--fractions must be comma-separated numbers, got {fractions!r}") + raise typer.Exit(code=2) from None + curve = learning_curve(cfg, fractions=parsed) + + table = Table(title="learning curve") + table.add_column("n_train", justify="right") + shown = [ + m for m in ("energy_mae_per_atom", "forces_rmse") + if any(m in p["metrics"] for p in curve["points"]) + ] + for metric in shown: + table.add_column(metric, justify="right") + for point in curve["points"]: + table.add_row( + str(point["n_train"]), + *(f"{point['metrics'].get(m, float('nan')):.6f}" for m in shown), + ) + _console.print(table) + typer.echo(f"wrote {curve['path']} (holdout: {curve['n_holdout']} frames)") + + # --- 3 · inspect --------------------------------------------------------------- @app.command(rich_help_panel=_INSPECT) diff --git a/src/traincraft/config/loader.py b/src/traincraft/config/loader.py index 8f1ff26..a49998b 100644 --- a/src/traincraft/config/loader.py +++ b/src/traincraft/config/loader.py @@ -86,30 +86,54 @@ def _resolve_slurm_profile(plain: dict) -> dict: def _resolve_input_paths(cfg: TrainCraftConfig, base: Path) -> TrainCraftConfig: """Resolve relative input file paths against *base* (the config file's dir).""" - if cfg.geometry is None: - return cfg - - updates: dict = {} - - # source.path (e.g. FileSource) - if cfg.geometry.source is not None and hasattr(cfg.geometry.source, "path"): - src = cfg.geometry.source - p = Path(src.path) - if not p.is_absolute(): - updates["source"] = src.model_copy(update={"path": str(base / p)}) - - # builder.file (e.g. SurfaceAdsorbateBuilder, SurfacePackingBuilder) - if cfg.geometry.builder is not None and hasattr(cfg.geometry.builder, "file"): - builder = cfg.geometry.builder - if builder.file is not None: - p = Path(builder.file) - if not p.is_absolute(): - updates["builder"] = builder.model_copy(update={"file": str(base / p)}) - - if not updates: - return cfg - new_geom = cfg.geometry.model_copy(update=updates) - return cfg.model_copy(update={"geometry": new_geom}) + + def _abs(value) -> str | None: + p = Path(value) + return None if p.is_absolute() else str(base / p) + + if cfg.geometry is not None: + updates: dict = {} + + # source.path (e.g. FileSource) + if cfg.geometry.source is not None and hasattr(cfg.geometry.source, "path"): + src = cfg.geometry.source + if (p := _abs(src.path)) is not None: + updates["source"] = src.model_copy(update={"path": p}) + + # builder.file (e.g. SurfaceAdsorbateBuilder, SurfacePackingBuilder) + if cfg.geometry.builder is not None and hasattr(cfg.geometry.builder, "file"): + builder = cfg.geometry.builder + if builder.file is not None and (p := _abs(builder.file)) is not None: + updates["builder"] = builder.model_copy(update={"file": p}) + + if updates: + cfg = cfg.model_copy( + update={"geometry": cfg.geometry.model_copy(update=updates)} + ) + + if cfg.validation is not None: + updates = {} + if cfg.validation.test_file and (p := _abs(cfg.validation.test_file)) is not None: + updates["test_file"] = p + if (spectra := cfg.validation.spectra) is not None: + sp_updates: dict = {} + reference = { + key: (_abs(path) or path) for key, path in spectra.reference.items() + } + if reference != spectra.reference: + sp_updates["reference"] = reference + if spectra.enhancement_file and ( + p := _abs(spectra.enhancement_file) + ) is not None: + sp_updates["enhancement_file"] = p + if sp_updates: + updates["spectra"] = spectra.model_copy(update=sp_updates) + if updates: + cfg = cfg.model_copy( + update={"validation": cfg.validation.model_copy(update=updates)} + ) + + return cfg def dump_starter_config() -> str: diff --git a/src/traincraft/config/models.py b/src/traincraft/config/models.py index 68f7917..4d273e6 100644 --- a/src/traincraft/config/models.py +++ b/src/traincraft/config/models.py @@ -581,6 +581,17 @@ class MaceCalc(TCModel): model_path: FsPath | None = None # local fine-tuned model device: str = "cpu" default_dtype: str = "float32" + # MACE calculator class for dielectric models trained with dipole/ + # polarizability heads (e.g. "DipoleMACE", "EnergyDipoleMACE"). These + # evolve faster than the energy/forces path — verify the name against your + # installed MACE, the same way the QE ph.x polarizability path is gated. + model_type: str | None = None + + @model_validator(mode="after") + def _model_type_needs_path(self) -> MaceCalc: + if self.model_type and not self.model_path: + raise ValueError("[calculator] 'model_type' requires a local 'model_path'") + return self # Properties a DFT labeler may be asked to produce beyond E/F/stress (always on). @@ -865,6 +876,85 @@ def _check(self) -> MaceTrainer: # ------------------------------------------------------------------------ validation +class StabilityConfig(TCModel): + """NVE stability probe: can the model survive its own dynamics? + + Langevin equilibration to ``temperature``, then ``steps`` of NVE + (Velocity-Verlet) production; energy conservation is folded into the + quality report as ``nve_*`` metrics (thresholdable, e.g. + ``nve_drift_mev_per_atom_ps``). Probes start from the first + ``n_structures`` reference frames; the worst drift is reported. + """ + + steps: int = Field(default=2000, ge=2) # NVE production steps + timestep: float = Field(default=1.0, gt=0) # fs + temperature: float = Field(default=300.0, gt=0) # K + equilibration_steps: int = Field(default=500, ge=0) + n_structures: int = Field(default=1, ge=1) + + +_SPECTRUM = Literal["vdos", "ir", "raman", "sfg"] + + +class SpectraConfig(TCModel): + """Vibrational spectra reconstructed from MLIP-driven MD (the TCF route). + + One production trajectory (NVT equilibration, then NVE) starting from the + first reference frame; velocities give the VDOS, the dipole time series the + IR spectrum, the polarizability the Raman spectra (VV/VH/depolarization), + and their cross-correlation the SFG response. ``calculator`` names the + dielectric model evaluating dipole/polarizability along the trajectory + (e.g. a trained ``AtomicDielectricMACE`` via ``model_type``); by default + the model under test is asked for them directly. + + ``reference`` maps spectra to two-column files (``freq_cm1 intensity``, + experimental or DFT); each adds a ``_spectral_error`` metric + (``1 - cosine similarity`` after common-grid Gaussian broadening) to the + quality report — the AL convergence hook. ``enhancement_file`` multiplies a + frequency-dependent electromagnetic enhancement profile onto the Raman + spectra (the SERS recipe; TrainCraft does not compute the EM field itself). + """ + + which: list[_SPECTRUM] = Field(default_factory=lambda: ["vdos", "ir", "raman"]) + temperature: float = Field(default=300.0, gt=0) # K + timestep: float = Field(default=0.5, gt=0) # fs + equilibration_steps: int = Field(default=2000, ge=0) + steps: int = Field(default=20000, ge=16) # NVE production steps + sample_every: int = Field(default=1, ge=1) + calculator: CalculatorConfig | None = None # dielectric model (default: model under test) + qcf: Literal["harmonic", "standard", "schofield", "none"] = "harmonic" + laser_wavelength: float = Field(default=532.0, gt=0) # nm (Raman prefactor) + sfg_component: str = "xxz" # chi^(2) indices; "xxz" = ssp with z surface normal + freq_range: tuple[float, float] = (0.0, 4000.0) # cm^-1 (written + compared) + broadening: float = Field(default=10.0, gt=0) # cm^-1 (similarity metric) + reference: dict[_SPECTRUM, FsPath] = Field(default_factory=dict) + enhancement_file: FsPath | None = None # two-column EM enhancement profile + save_series: bool = False # keep the raw time series as spectra/series.npz + + @model_validator(mode="after") + def _consistent(self) -> SpectraConfig: + if not self.which: + raise ValueError("[validation.spectra] 'which' must not be empty") + if len(self.sfg_component) != 3 or any(c not in "xyz" for c in self.sfg_component): + raise ValueError( + f"[validation.spectra] sfg_component must be three of x/y/z, " + f"got {self.sfg_component!r}" + ) + if self.freq_range[0] >= self.freq_range[1] or self.freq_range[0] < 0: + raise ValueError("[validation.spectra] freq_range must be 0 <= lo < hi") + if missing := set(self.reference) - set(self.which): + raise ValueError( + f"[validation.spectra] reference for {sorted(missing)} but 'which' " + f"is {self.which}" + ) + if self.enhancement_file and "raman" not in self.which: + raise ValueError( + "[validation.spectra] enhancement_file applies to Raman spectra; " + "add 'raman' to 'which'" + ) + return self + + class ValidationConfig(TCModel): """Measure a model against reference-labeled frames → ``quality_report.json``. @@ -884,6 +974,8 @@ class ValidationConfig(TCModel): properties: list[_HEAD] | None = None # default: whatever the references carry max_frames: int | None = None # cap for quick checks thresholds: dict[str, float] = Field(default_factory=dict) + stability: StabilityConfig | None = None # NVE probe (off by default) + spectra: SpectraConfig | None = None # MD spectra reconstruction (off by default) # --------------------------------------------------------------------- orchestration diff --git a/src/traincraft/orchestration/stages.py b/src/traincraft/orchestration/stages.py index 44d7fe5..b78d58c 100644 --- a/src/traincraft/orchestration/stages.py +++ b/src/traincraft/orchestration/stages.py @@ -46,7 +46,12 @@ from ..selection import run_funnel from ..state import log_event from ..training import run_training -from ..validation import REPORT_NAME, load_reference_frames, resolve_calculator, run_validation +from ..validation import ( + REPORT_NAME, + load_reference_frames, + resolve_calculator, + run_validation_science, +) logger = logging.getLogger(__name__) @@ -368,7 +373,9 @@ def stage_validate(config, ws, force=False): calc_cfg = resolve_calculator(config.validation, ws.root / "model" / "manifest.json") _emit(ws, "validate", "start", n_in=len(frames)) t0 = time.perf_counter() - result = run_validation(frames, calc_cfg, config.validation, out_dir=ws.root / "validation") + result = run_validation_science( + frames, calc_cfg, config.validation, out_dir=ws.root / "validation" + ) _write_key(out, key) verdict = {True: "passed", False: "FAILED", None: "no thresholds set"}[result.passed] logger.info( diff --git a/src/traincraft/validation/__init__.py b/src/traincraft/validation/__init__.py index 0808e59..952b859 100644 --- a/src/traincraft/validation/__init__.py +++ b/src/traincraft/validation/__init__.py @@ -1,19 +1,41 @@ -"""Validation: parity metrics + quality report for a trained model.""" +"""Validation: parity metrics, quality report, and spectra for a trained model.""" from __future__ import annotations +from .dynamics import ProbeSeries, run_probe_md, stability_metrics +from .learning_curve import learning_curve +from .spectra import ( + ir_spectrum, + power_spectrum, + raman_spectra, + sfg_spectrum, + spectral_similarity, + vdos, +) from .validate import ( REPORT_NAME, ValidationResult, load_reference_frames, resolve_calculator, run_validation, + run_validation_science, ) __all__ = [ "REPORT_NAME", + "ProbeSeries", "ValidationResult", + "ir_spectrum", + "learning_curve", "load_reference_frames", + "power_spectrum", + "raman_spectra", "resolve_calculator", + "run_probe_md", "run_validation", + "run_validation_science", + "sfg_spectrum", + "spectral_similarity", + "stability_metrics", + "vdos", ] diff --git a/src/traincraft/validation/dynamics.py b/src/traincraft/validation/dynamics.py new file mode 100644 index 0000000..ed32715 --- /dev/null +++ b/src/traincraft/validation/dynamics.py @@ -0,0 +1,183 @@ +"""MD probes for validation: NVE stability and property time series for spectra. + +``run_probe_md`` drives short MLIP MD the way spectroscopy needs it: Langevin +(NVT) equilibration to the target temperature, then **NVE (Velocity-Verlet) +production** — a thermostat would distort the very dynamics (line widths, +energy drift) the probe is measuring. Along production it records energies, +temperature, velocities, and optionally dipole/polarizability evaluated by a +separate property calculator (the dielectric model need not be the model +driving the MD). + +An exploding trajectory is detected by a force/energy guard and aborts the +production loop cleanly — a model that cannot survive dynamics is exactly the +result (``nve_completed = 0``), not an exception. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import numpy as np +from ase import units +from ase.md.langevin import Langevin +from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary +from ase.md.verlet import VelocityVerlet + +logger = logging.getLogger(__name__) + +#: Guard limits: a healthy MLIP trajectory stays far below these. +MAX_FORCE_EV_A = 50.0 + + +@dataclass +class ProbeSeries: + """Time series recorded along one NVE production run. + + ``dt_fs`` is the *sampling* interval (timestep x sample_every); velocities + are in Å/fs. ``dipole``/``polarizability`` are None unless captured. + """ + + dt_fs: float + time_fs: np.ndarray + energy: np.ndarray # total (potential + kinetic), eV + temperature: np.ndarray # K + velocities: np.ndarray # (n_samples, n_atoms, 3), Å/fs + masses: np.ndarray + symbols: list[str] + dipole: np.ndarray | None = None # (n_samples, 3) + polarizability: np.ndarray | None = None # (n_samples, 3, 3) + completed: bool = True + aborted_at: int | None = None # production step where the guard fired + + @property + def n_atoms(self) -> int: + return len(self.masses) + + +def _capture_properties(atoms, prop_calc, capture: tuple[str, ...]) -> dict: + """Evaluate dipole/polarizability on the current positions. + + Follows the calculator read contracts established in labeling: dipole via + ``atoms.get_dipole_moment()``, polarizability via + ``calc.results["polarizability"]`` (no ASE getter exists). A requested + property the calculator cannot produce is a hard error — a gap in the + time series would silently corrupt every spectrum built from it. + """ + snap = atoms.copy() + snap.calc = prop_calc + out: dict = {} + if "dipole" in capture: + try: + out["dipole"] = np.asarray(snap.get_dipole_moment(), dtype=float) + except Exception as exc: + raise ValueError( + f"the property calculator produced no dipole ({exc}); spectra " + "needing the dipole (ir, sfg) require a dipole-capable model" + ) from exc + if "polarizability" in capture: + if "dipole" not in capture: # trigger a calculation without the getter + try: + snap.get_potential_energy() + except Exception: + pass + pol = getattr(snap.calc, "results", {}).get("polarizability") + if pol is None: + raise ValueError( + "the property calculator exposed no results['polarizability']; " + "spectra needing it (raman, sfg) require a polarizability-capable " + "model (e.g. a trained AtomicDielectricMACE)" + ) + out["polarizability"] = np.asarray(pol, dtype=float).reshape(3, 3) + return out + + +def run_probe_md(atoms, calc, *, temperature: float, timestep: float, + equilibration_steps: int, steps: int, sample_every: int = 1, + prop_calc=None, capture: tuple[str, ...] = (), + seed: int | None = None) -> ProbeSeries: + """NVT-equilibrate then NVE-run ``atoms`` under ``calc``; sample the series. + + ``timestep`` in fs. ``prop_calc``/``capture`` request dipole and/or + polarizability evaluation each sample (see ``_capture_properties``). + """ + atoms = atoms.copy() + atoms.calc = calc + rng = np.random.default_rng(seed) + MaxwellBoltzmannDistribution(atoms, temperature_K=temperature, rng=rng) + Stationary(atoms) # no center-of-mass drift polluting the velocity ACF + + if equilibration_steps > 0: + nvt = Langevin(atoms, timestep=timestep * units.fs, + temperature_K=temperature, friction=0.02, rng=rng) + nvt.run(equilibration_steps) + + nve = VelocityVerlet(atoms, timestep=timestep * units.fs) + time_fs: list[float] = [] + energy: list[float] = [] + temp: list[float] = [] + velocities: list[np.ndarray] = [] + props: dict[str, list] = {key: [] for key in capture} + completed, aborted_at = True, None + + step = 0 + while True: + e_pot = atoms.get_potential_energy() + f_max = float(np.abs(atoms.get_forces()).max()) + if not np.isfinite(e_pot) or f_max > MAX_FORCE_EV_A: + logger.warning( + "NVE probe aborted at step %d (max |F| = %.1f eV/Å)", step, f_max + ) + completed, aborted_at = False, step + break + time_fs.append(step * timestep) + energy.append(float(e_pot + atoms.get_kinetic_energy())) + temp.append(float(atoms.get_temperature())) + velocities.append(atoms.get_velocities() * units.fs) # -> Å/fs + if prop_calc is not None and capture: + for key, value in _capture_properties(atoms, prop_calc, capture).items(): + props[key].append(value) + step += sample_every + if step > steps: + break + nve.run(sample_every) + + return ProbeSeries( + dt_fs=timestep * sample_every, + time_fs=np.asarray(time_fs), + energy=np.asarray(energy), + temperature=np.asarray(temp), + velocities=np.asarray(velocities), + masses=atoms.get_masses().copy(), + symbols=atoms.get_chemical_symbols(), + dipole=np.asarray(props["dipole"]) if props.get("dipole") else None, + polarizability=( + np.asarray(props["polarizability"]) + if props.get("polarizability") else None + ), + completed=completed, + aborted_at=aborted_at, + ) + + +def stability_metrics(series: ProbeSeries) -> dict: + """Flat NVE-conservation metrics for the quality report. + + ``nve_drift_mev_per_atom_ps`` is the |slope| of a linear fit to the total + energy (the standard NVE health number); the std is of the detrended + residuals (fluctuation, not drift). All thresholdable via ``value <= limit`` + — hence ``nve_aborted`` (0 = survived) rather than a "completed" flag. + """ + metrics = { + "nve_aborted": 0.0 if series.completed else 1.0, + "nve_temp_mean": float(series.temperature.mean()) if len(series.temperature) else 0.0, + } + if len(series.energy) >= 2: + e_mev = series.energy / series.n_atoms * 1000.0 + t_ps = series.time_fs / 1000.0 + slope, intercept = np.polyfit(t_ps, e_mev, 1) + metrics["nve_drift_mev_per_atom_ps"] = float(abs(slope)) + metrics["nve_energy_std_mev_per_atom"] = float( + np.std(e_mev - (slope * t_ps + intercept)) + ) + return metrics diff --git a/src/traincraft/validation/learning_curve.py b/src/traincraft/validation/learning_curve.py new file mode 100644 index 0000000..e9c2cfa --- /dev/null +++ b/src/traincraft/validation/learning_curve.py @@ -0,0 +1,97 @@ +"""Learning curve: validation error vs training-set size. + +The data-efficiency diagnostic: retrain the configured model on **nested** +random subsets of the training data and validate every point against one +**fixed holdout**, so the only thing varying along the curve is how much data +the model saw. Nested subsets (smaller ⊂ larger, one shuffle) isolate the +quantity effect from sampling noise; the holdout is split off once with the +trainer's own ``valid_fraction``/seed and is never trained on at any fraction. + +This is an opt-in diagnostic command (k retrains = k× the training cost), not +a pipeline stage: no events, no cache keys — just +``/validation/learning_curve.json`` and per-point model dirs under +``/validation/learning_curve/``. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +from ..config.models import ValidationConfig +from ..core import make_rng +from ..training import run_training +from ..training.mace import _split +from .validate import resolve_calculator, run_validation + +logger = logging.getLogger(__name__) + +CURVE_NAME = "learning_curve.json" +DEFAULT_FRACTIONS = (0.125, 0.25, 0.5, 1.0) + + +def learning_curve(config, ws=None, *, fractions=DEFAULT_FRACTIONS) -> dict: + """Train on nested subsets of the training frames; validate each on the + same holdout; write and return the curve. + + Each point validates **its own trained model** (from the point's + ``manifest.json``) — an explicit ``[validation.calculator]`` is ignored + here, since a fixed external model would make the curve flat by + construction. Parity metrics only: thresholds and the stability/spectra + probes don't apply per point. + """ + # Deferred: stages imports validation at module load (cycle otherwise). + from ..orchestration.stages import _training_frames, workspace_for + + if config.training is None: + raise ValueError("learning-curve needs a [training] section to retrain from") + fractions = sorted({float(f) for f in fractions}) + if not fractions or fractions[0] <= 0.0 or fractions[-1] > 1.0: + raise ValueError(f"fractions must lie in (0, 1], got {fractions}") + + ws = ws or workspace_for(config) + frames = _training_frames(config, ws) + if len(frames) < 4: + raise ValueError( + f"learning curve needs a few frames to subset ({len(frames)} found)" + ) + + seed = config.training.seed if config.training.seed is not None else (config.run.seed or 0) + train, holdout = _split(frames, config.training.valid_fraction, seed) + order = make_rng(config.run.seed or 0).permutation(len(train)) + + vcfg = (config.validation or ValidationConfig()).model_copy( + update={"calculator": None, "thresholds": {}, "stability": None, "spectra": None} + ) + + points: list[dict] = [] + for fraction in fractions: + n = max(1, round(fraction * len(train))) + if points and n == points[-1]["n_train"]: + continue # rounding collision on a small dataset + subset = [train[i] for i in order[:n]] + job = ws.job("validation", "learning_curve", f"n{n:05d}") + logger.info("learning curve: training on %d/%d frames -> %s", n, len(train), job.dir) + run_training(subset, config.training, job) + calc_cfg = resolve_calculator(vcfg, Path(job.dir) / "manifest.json") + result = run_validation(holdout, calc_cfg, vcfg, out_dir=None) + points.append({ + "fraction": fraction, + "n_train": n, + "dir": str(job.dir), + "metrics": result.metrics, + }) + + curve = { + "n_frames": len(frames), + "n_holdout": len(holdout), + "seed": seed, + "valid_fraction": config.training.valid_fraction, + "fractions": fractions, + "points": points, + } + out = ws.subdir("validation") / CURVE_NAME + out.write_text(json.dumps(curve, indent=2)) + curve["path"] = str(out) + return curve diff --git a/src/traincraft/validation/spectra.py b/src/traincraft/validation/spectra.py new file mode 100644 index 0000000..0c9d9e7 --- /dev/null +++ b/src/traincraft/validation/spectra.py @@ -0,0 +1,287 @@ +"""Vibrational spectra from MD time series: VDOS, IR, Raman, SFG (pure numpy). + +Time-correlation-function (TCF) route: a property recorded along an MD +trajectory (velocities, dipole, polarizability) is autocorrelated, windowed, +and Fourier-transformed into a spectrum. Everything here is a pure function of +numpy arrays plus the timestep — no ASE, no I/O — so the physics is testable +against synthetic signals. + +Formulas follow the standard published theory (independent implementation, +no prior code consulted): + +- IR and Raman from time-*derivative* ACFs with the ``(nu_L - nu)^4 / nu`` + Stokes prefactor and Bose factor: Thomas, Brehm, Fligg, Vöhringer & + Kirchner, Phys. Chem. Chem. Phys. 15, 6608 (2013). +- Raman VV/VH decomposition into isotropic/anisotropic invariants: + Berne & Pecora, *Dynamic Light Scattering* (Wiley, 1976), ch. 3. +- Quantum correction factors: Ramírez, López-Ciudad, Kumar & Marx, + J. Chem. Phys. 121, 3973 (2004). +- SFG resonant chi^(2) from the dipole-polarizability cross-TCF: + Morita & Hynes, J. Phys. Chem. B 106, 673 (2002). + +Conventions: time series enter with a timestep ``dt_fs`` in femtoseconds; +frequency axes are wavenumbers in cm^-1; intensities are in arbitrary units +(the physically meaningful frequency- and temperature-dependence is included, +constant dimensional prefactors are not). The derivative-ACF spectrum already +corresponds to the *harmonic* quantum correction of the position-ACF spectrum, +so ``qcf="harmonic"`` is the identity and other choices are applied as ratios. +""" + +from __future__ import annotations + +import numpy as np + +C_CM_PER_FS = 2.99792458e-5 # speed of light [cm/fs]: nu[cm^-1] = f[1/fs] / c +C2_CM_K = 1.4387768775 # second radiation constant hc/k_B [cm K]: u = C2*nu/T + +QCF_KINDS = ("harmonic", "standard", "schofield", "none") +_AXES = {"x": 0, "y": 1, "z": 2} + + +# --- correlation functions ---------------------------------------------------- + +def _flat(x) -> np.ndarray: + x = np.asarray(x, dtype=float) + if x.shape[0] < 8: + raise ValueError(f"time series too short for a spectrum ({x.shape[0]} steps)") + return x.reshape(x.shape[0], -1) + + +def autocorrelation(x, max_lag: int | None = None) -> np.ndarray: + """Unbiased ACF ``C(tau) = ``, summed over component axes. + + FFT-based (Wiener-Khinchin) with zero padding, so it is exact and O(n log n). + """ + x = _flat(x) + n = x.shape[0] + max_lag = min(max_lag or n // 2, n) + m = 1 << (2 * n - 1).bit_length() + f = np.fft.rfft(x, n=m, axis=0) + acf = np.fft.irfft(np.abs(f) ** 2, n=m, axis=0)[:max_lag].sum(axis=1) + return acf / (n - np.arange(max_lag)) + + +def cross_correlation(a, b, max_lag: int | None = None) -> np.ndarray: + """Unbiased cross-TCF ``C(tau) = `` (b lags a), summed + over component axes. ``a`` and ``b`` must share the time axis length.""" + a, b = _flat(a), _flat(b) + if a.shape != b.shape: + raise ValueError(f"mismatched series shapes {a.shape} vs {b.shape}") + n = a.shape[0] + max_lag = min(max_lag or n // 2, n) + m = 1 << (2 * n - 1).bit_length() + fa = np.fft.rfft(a, n=m, axis=0) + fb = np.fft.rfft(b, n=m, axis=0) + ccf = np.fft.irfft(fa.conj() * fb, n=m, axis=0)[:max_lag].sum(axis=1) + return ccf / (n - np.arange(max_lag)) + + +def _one_sided_ft(tcf: np.ndarray, dt_fs: float, pad: int = 4): + """Complex one-sided FT ``int_0^T C(tau) e^{+i w tau} W(tau) dtau`` of a + Hann-windowed TCF, zero-padded ``pad``x for a smooth grid. + + Returns ``(nu_cm1, ft)``; for an ACF (even in tau) the physical power + spectrum is ``ft.real`` (one-sidedness only changes a constant factor). + """ + L = len(tcf) + window = np.hanning(2 * L)[L:] # decaying half: 1 at tau=0, -> 0 at tau=L + ft = np.conj(np.fft.rfft(tcf * window, n=L * pad)) * dt_fs + nu = np.fft.rfftfreq(L * pad, d=dt_fs) / C_CM_PER_FS + return nu, ft + + +def power_spectrum(x, dt_fs: float, *, max_lag: int | None = None, pad: int = 4): + """Power spectrum of a time series: FT of its windowed ACF -> ``(nu, I)``.""" + nu, ft = _one_sided_ft(autocorrelation(x, max_lag), dt_fs, pad) + return nu, ft.real + + +# --- quantum correction factors ----------------------------------------------- + +def _qcf(u: np.ndarray, kind: str) -> np.ndarray: + """Q(u) with u = beta*hbar*omega, mapping classical to quantum lineshape.""" + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + q = { + "none": np.ones_like(u), + "harmonic": u / (1.0 - np.exp(-u)), + "standard": 2.0 / (1.0 + np.exp(-u)), + "schofield": np.exp(u / 2.0), + }[kind] + return np.where(u > 0, q, 1.0) # every QCF -> 1 as u -> 0 + + +def qcf_ratio(nu_cm1, temperature: float, kind: str = "harmonic") -> np.ndarray: + """Correction applied on top of a derivative-ACF spectrum. + + The classical derivative-ACF spectrum equals the harmonic-QCF result, so + ``kind="harmonic"`` returns 1 and other kinds return ``Q_kind/Q_harmonic`` + (Ramírez et al., J. Chem. Phys. 121, 3973 (2004)).""" + if kind not in QCF_KINDS: + raise ValueError(f"unknown qcf {kind!r} (choose from {QCF_KINDS})") + u = C2_CM_K * np.asarray(nu_cm1, dtype=float) / float(temperature) + return _qcf(u, kind) / _qcf(u, "harmonic") + + +def _bose_stokes(nu_cm1: np.ndarray, temperature: float) -> np.ndarray: + """Stokes detailed-balance factor ``n(nu)+1 = 1/(1 - e^{-u})`` (0 at nu=0).""" + u = C2_CM_K * nu_cm1 / float(temperature) + with np.errstate(divide="ignore", invalid="ignore"): + f = 1.0 / (1.0 - np.exp(-u)) + return np.where(u > 0, f, 0.0) + + +# --- spectra -------------------------------------------------------------------- + +def vdos(velocities, masses, dt_fs: float, *, symbols=None, + max_lag: int | None = None, pad: int = 4): + """Vibrational density of states: mass-weighted velocity ACF spectrum. + + ``velocities``: (n_steps, n_atoms, 3); ``masses``: (n_atoms,) in any unit. + Returns ``(nu_cm1, {"total": I, : I, ...})`` — per-element curves + only when ``symbols`` is given. + """ + v = np.asarray(velocities, dtype=float) + w = np.sqrt(np.asarray(masses, dtype=float))[None, :, None] + nu, total = power_spectrum(v * w, dt_fs, max_lag=max_lag, pad=pad) + curves = {"total": total} + if symbols is not None: + for el in sorted(set(symbols)): + mask = np.asarray([s == el for s in symbols]) + _, curves[el] = power_spectrum( + (v * w)[:, mask, :], dt_fs, max_lag=max_lag, pad=pad + ) + return nu, curves + + +def ir_spectrum(dipole, dt_fs: float, temperature: float, *, qcf: str = "harmonic", + max_lag: int | None = None, pad: int = 4): + """IR absorption lineshape from the dipole-derivative ACF (arb. units). + + ``A(nu) ∝ FT``, the harmonic-QCF form of the + absorption coefficient (Thomas et al., PCCP 15, 6608 (2013)).""" + mu_dot = np.gradient(np.asarray(dipole, dtype=float), dt_fs, axis=0) + nu, spec = power_spectrum(mu_dot, dt_fs, max_lag=max_lag, pad=pad) + return nu, spec * qcf_ratio(nu, temperature, qcf) + + +def raman_spectra(polarizability, dt_fs: float, temperature: float, *, + laser_wavelength: float = 532.0, qcf: str = "harmonic", + max_lag: int | None = None, pad: int = 4) -> dict: + """Stokes Raman spectra from the polarizability-derivative TCFs. + + The tensor is split into the isotropic invariant ``a = Tr(alpha)/3`` and + the traceless anisotropic part ``beta``; with the full-contraction spectra + ``S_iso = FT`` and ``S_beta = FT``: + + I_VV ∝ pref * (S_iso + 4/45 S_beta), I_VH ∝ pref * S_beta/15 + pref = (nu_L - nu)^4 / nu * (n(nu)+1) + + (Berne & Pecora ch. 3; prefactor as in Thomas et al., PCCP 2013.) + Returns ``{"freq", "iso", "aniso", "vv", "vh", "unpolarized", + "depolarization"}``; the depolarization ratio ``VH/VV`` -> 3/4 for purely + anisotropic modes. ``laser_wavelength`` in nm. + """ + alpha = np.asarray(polarizability, dtype=float) + if alpha.ndim != 3 or alpha.shape[1:] != (3, 3): + raise ValueError(f"polarizability series must be (n, 3, 3), got {alpha.shape}") + a = np.trace(alpha, axis1=1, axis2=2) / 3.0 + beta = alpha - a[:, None, None] * np.eye(3) + a_dot = np.gradient(a, dt_fs, axis=0) + beta_dot = np.gradient(beta, dt_fs, axis=0) + + nu, s_iso = power_spectrum(a_dot, dt_fs, max_lag=max_lag, pad=pad) + _, s_beta = power_spectrum(beta_dot, dt_fs, max_lag=max_lag, pad=pad) + + nu_laser = 1.0e7 / float(laser_wavelength) # nm -> cm^-1 + with np.errstate(divide="ignore", invalid="ignore"): + pref = (nu_laser - nu) ** 4 / nu + pref = np.where(nu > 0, pref, 0.0) + pref *= _bose_stokes(nu, temperature) * qcf_ratio(nu, temperature, qcf) + + iso, aniso = pref * s_iso, pref * s_beta + vv = iso + (4.0 / 45.0) * aniso + vh = aniso / 15.0 + with np.errstate(divide="ignore", invalid="ignore"): + depol = np.where(vv > 0, vh / vv, 0.0) + return { + "freq": nu, "iso": iso, "aniso": aniso, "vv": vv, "vh": vh, + "unpolarized": vv + vh, "depolarization": depol, + } + + +def sfg_spectrum(polarizability, dipole, dt_fs: float, temperature: float, *, + component: str = "xxz", qcf: str = "harmonic", + max_lag: int | None = None, pad: int = 4) -> dict: + """Resonant SFG susceptibility chi^(2)_ijk from the cross-TCF + ```` (velocity form of Morita & Hynes, + J. Phys. Chem. B 106, 673 (2002); dividing by nu^2 recovers the + position-form FT). + + ``component``: three axis letters, e.g. ``"xxz"`` (the ssp geometry with + the surface normal along z) — alpha_xx paired with mu_z. Returns + ``{"freq", "re", "im", "abs2"}``; the sign of ``im`` carries the net + orientation of the mode, which is the surface-specific information. + """ + if len(component) != 3 or any(c not in _AXES for c in component): + raise ValueError(f"sfg component must be three of x/y/z, got {component!r}") + i, j, k = (_AXES[c] for c in component) + alpha = np.asarray(polarizability, dtype=float)[:, i, j] + mu = np.asarray(dipole, dtype=float)[:, k] + alpha_dot = np.gradient(alpha, dt_fs, axis=0) + mu_dot = np.gradient(mu, dt_fs, axis=0) + + ccf = cross_correlation(mu_dot, alpha_dot, max_lag) + nu, ft = _one_sided_ft(ccf, dt_fs, pad) + with np.errstate(divide="ignore", invalid="ignore"): + chi = np.where(nu > 0, ft / nu**2, 0.0) + chi = chi * qcf_ratio(nu, temperature, qcf) + return {"freq": nu, "re": chi.real, "im": chi.imag, "abs2": np.abs(chi) ** 2} + + +# --- comparison ------------------------------------------------------------------ + +def broaden(freq, intensity, *, lo: float, hi: float, sigma: float) -> tuple: + """Bin a spectrum (continuous or sticks) onto a uniform grid over + ``[lo, hi]`` and convolve with a Gaussian of width ``sigma`` (cm^-1). + Returns ``(grid, curve)``. Both sides of a comparison must pass through + this same transform for the similarity metric to be fair.""" + freq = np.asarray(freq, dtype=float) + intensity = np.asarray(intensity, dtype=float) + step = sigma / 4.0 + edges = np.arange(lo, hi + step, step) + mask = (freq >= lo) & (freq <= hi) + hist, _ = np.histogram(freq[mask], bins=edges, weights=intensity[mask]) + tail = np.arange(0.0, 4.0 * sigma + step, step) + kernel = np.exp(-0.5 * (np.concatenate([-tail[:0:-1], tail]) / sigma) ** 2) + curve = np.convolve(hist, kernel, mode="same") + return (edges[:-1] + edges[1:]) / 2.0, curve + + +def spectral_similarity(freq_a, int_a, freq_b, int_b, *, lo: float = 0.0, + hi: float = 4000.0, sigma: float = 10.0) -> float: + """Cosine similarity of two Gaussian-broadened spectra on a common grid. + + 1.0 = identical shapes (scale-invariant), 0.0 = no overlap; signed spectra + (e.g. Im chi^(2)) may go negative. Report ``1 - similarity`` as an error + so ``value <= limit`` thresholds apply.""" + _, a = broaden(freq_a, int_a, lo=lo, hi=hi, sigma=sigma) + _, b = broaden(freq_b, int_b, lo=lo, hi=hi, sigma=sigma) + na, nb = np.linalg.norm(a), np.linalg.norm(b) + if na == 0.0 or nb == 0.0: + return 0.0 + return float(np.dot(a, b) / (na * nb)) + + +def find_peaks(freq, intensity, *, min_rel_height: float = 0.05) -> list[dict]: + """Local maxima above ``min_rel_height * max``, strongest first — the + peak table recorded in ``spectra.json`` (not a fitting routine).""" + freq = np.asarray(freq, dtype=float) + y = np.asarray(intensity, dtype=float) + if len(y) < 3 or y.max() <= 0: + return [] + idx = np.flatnonzero((y[1:-1] > y[:-2]) & (y[1:-1] >= y[2:])) + 1 + idx = idx[y[idx] >= min_rel_height * y.max()] + order = np.argsort(y[idx])[::-1] + return [ + {"freq_cm1": float(freq[i]), "height": float(y[i])} for i in idx[order] + ] diff --git a/src/traincraft/validation/validate.py b/src/traincraft/validation/validate.py index c9b66e0..7c71c32 100644 --- a/src/traincraft/validation/validate.py +++ b/src/traincraft/validation/validate.py @@ -251,10 +251,7 @@ def want(prop: str) -> bool: report_path = None if out_dir is not None: - out_dir = Path(out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - report_path = out_dir / REPORT_NAME - report_path.write_text(json.dumps(report, indent=2)) + report_path = _write_report(report, out_dir) return ValidationResult( metrics=metrics, @@ -265,3 +262,202 @@ def want(prop: str) -> bool: report=report, report_path=report_path, ) + + +def _write_report(report: dict, out_dir) -> Path: + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + report_path = out_dir / REPORT_NAME + report_path.write_text(json.dumps(report, indent=2)) + return report_path + + +# --- validation science: NVE stability probe + MD spectra ---------------------- + +def _stability_probes(frames, calc_cfg, scfg) -> dict: + """NVE-probe the first ``n_structures`` frames; fold to worst-case metrics.""" + from .dynamics import run_probe_md, stability_metrics + + calc = make_calculator(calc_cfg) + probes = [] + for i, s in enumerate(frames[: scfg.n_structures]): + series = run_probe_md( + s.to_ase(with_properties=False), calc, + temperature=scfg.temperature, timestep=scfg.timestep, + equilibration_steps=scfg.equilibration_steps, steps=scfg.steps, + seed=i, + ) + probes.append( + {"frame": i, "aborted_at": series.aborted_at, **stability_metrics(series)} + ) + metrics: dict = {"nve_aborted": max(p["nve_aborted"] for p in probes)} + for name in ("nve_drift_mev_per_atom_ps", "nve_energy_std_mev_per_atom"): + values = [p[name] for p in probes if name in p] + if values: # absent only when a probe died before its second sample + metrics[name] = max(values) + metrics["nve_temp_mean"] = float( + np.mean([p["nve_temp_mean"] for p in probes]) + ) + return {"config": scfg.model_dump(), "probes": probes, "metrics": metrics} + + +def _load_two_column(path, what: str) -> np.ndarray: + """A ``freq_cm1 intensity`` file (reference spectrum, enhancement profile).""" + try: + data = np.loadtxt(path, ndmin=2) + except Exception as exc: + raise ValueError(f"could not read {what} {path}: {exc}") from exc + if data.shape[1] < 2: + raise ValueError(f"{what} {path} needs two columns (freq_cm1, intensity)") + return data[:, :2] + + +def _spectra_probe(frames, calc_cfg, cfg, out_dir) -> dict: + """Production MD from the first reference frame -> reconstructed spectra. + + Writes each spectrum as a two-column ``.dat`` under ``out_dir`` plus one + ``spectra.json`` (parameters, peak tables, metrics) — O(1) files. Returns + ``{"detail", "metrics"}``; a reference-compared spectrum contributes a + ``_spectral_error`` metric. + """ + from . import spectra as sp + from .dynamics import run_probe_md + + which = set(cfg.which) + capture = tuple( + prop for prop, needed_by in + (("dipole", {"ir", "sfg"}), ("polarizability", {"raman", "sfg"})) + if needed_by & which + ) + driver = make_calculator(calc_cfg) + prop_calc = make_calculator(cfg.calculator) if cfg.calculator is not None else driver + + series = run_probe_md( + frames[0].to_ase(with_properties=False), driver, + temperature=cfg.temperature, timestep=cfg.timestep, + equilibration_steps=cfg.equilibration_steps, steps=cfg.steps, + sample_every=cfg.sample_every, + prop_calc=prop_calc if capture else None, capture=capture, seed=0, + ) + metrics: dict = {"spectra_md_aborted": 0.0 if series.completed else 1.0} + detail: dict = { + "config": cfg.model_dump(mode="json"), + "n_samples": len(series.energy), + "dt_fs": series.dt_fs, + "aborted_at": series.aborted_at, + } + lo, hi = cfg.freq_range + curves: dict[str, tuple] = {} # file stem -> (freq, intensity) + primary: dict[str, tuple] = {} # spectrum name -> curve compared to a reference + + if series.completed: + if "vdos" in which: + nu, per = sp.vdos( + series.velocities, series.masses, series.dt_fs, symbols=series.symbols + ) + primary["vdos"] = curves["vdos"] = (nu, per.pop("total")) + curves.update({f"vdos_{el}": (nu, y) for el, y in per.items()}) + if "ir" in which: + nu, y = sp.ir_spectrum(series.dipole, series.dt_fs, cfg.temperature, qcf=cfg.qcf) + primary["ir"] = curves["ir"] = (nu, y) + if "raman" in which: + raman = sp.raman_spectra( + series.polarizability, series.dt_fs, cfg.temperature, + laser_wavelength=cfg.laser_wavelength, qcf=cfg.qcf, + ) + nu = raman.pop("freq") + curves.update({f"raman_{name}": (nu, y) for name, y in raman.items()}) + primary["raman"] = curves["raman_unpolarized"] + if cfg.enhancement_file: + profile = _load_two_column(cfg.enhancement_file, "enhancement_file") + gain = np.interp(nu, profile[:, 0], profile[:, 1]) + curves.update({ + f"raman_{name}_enhanced": (nu, raman[name] * gain) + for name in ("vv", "vh", "unpolarized") + }) + if "sfg" in which: + sfg = sp.sfg_spectrum( + series.polarizability, series.dipole, series.dt_fs, cfg.temperature, + component=cfg.sfg_component, qcf=cfg.qcf, + ) + nu = sfg.pop("freq") + curves.update({f"sfg_{name}": (nu, y) for name, y in sfg.items()}) + primary["sfg"] = curves["sfg_im"] + + for name, ref_path in cfg.reference.items(): + ref = _load_two_column(ref_path, f"reference[{name}]") + nu, y = primary[name] + similarity = sp.spectral_similarity( + nu, y, ref[:, 0], ref[:, 1], lo=lo, hi=hi, sigma=cfg.broadening + ) + metrics[f"{name}_spectral_error"] = 1.0 - similarity + + detail["peaks"] = { + name: sp.find_peaks(nu[(nu >= lo) & (nu <= hi)], y[(nu >= lo) & (nu <= hi)])[:10] + for name, (nu, y) in primary.items() + } + detail["metrics"] = metrics + + if out_dir is not None: + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + for stem, (nu, y) in curves.items(): + mask = (nu >= lo) & (nu <= hi) + np.savetxt( + out_dir / f"{stem}.dat", + np.column_stack([nu[mask], y[mask]]), + header="freq_cm1 intensity", + ) + detail["files"] = sorted(f"{stem}.dat" for stem in curves) + if cfg.save_series: + arrays = { + "time_fs": series.time_fs, "energy": series.energy, + "temperature": series.temperature, "velocities": series.velocities, + "masses": series.masses, + } + for key in ("dipole", "polarizability"): + if (value := getattr(series, key)) is not None: + arrays[key] = value + np.savez_compressed(out_dir / "series.npz", **arrays) + (out_dir / "spectra.json").write_text(json.dumps(detail, indent=2)) + + return {"detail": detail, "metrics": metrics} + + +def run_validation_science(frames, calc_cfg, cfg, *, out_dir=None) -> ValidationResult: + """Parity plus the configured physics probes, folded into one report. + + Runs ``run_validation`` (parity), then the NVE stability probe + (``cfg.stability``) and the MD spectra reconstruction (``cfg.spectra``, + artifacts under ``/spectra/``); probe metrics join the flat + ``metrics`` dict and thresholds are applied once over the union, so + ``nve_*`` / ``*_spectral_error`` names threshold exactly like parity + metrics. A probe whose MD aborts is a *result* (``nve_aborted = 1``, + ``spectra_md_aborted = 1``), not an exception. + """ + parity_cfg = cfg.model_copy(update={"thresholds": {}}) if cfg.thresholds else cfg + result = run_validation(frames, calc_cfg, parity_cfg, out_dir=None) + report = result.report + + if cfg.stability is not None: + stability = _stability_probes(frames, calc_cfg, cfg.stability) + report["stability"] = stability + result.metrics.update(stability["metrics"]) + report["metrics"] = result.metrics + + if cfg.spectra is not None: + spectra = _spectra_probe( + frames, calc_cfg, cfg.spectra, + Path(out_dir) / "spectra" if out_dir is not None else None, + ) + report["spectra"] = spectra["detail"] + result.metrics.update(spectra["metrics"]) + report["metrics"] = result.metrics + + checks, passed = _apply_thresholds(result.metrics, cfg.thresholds) + report.update(thresholds=dict(cfg.thresholds), checks=checks, passed=passed) + result.checks, result.passed = checks, passed + + if out_dir is not None: + result.report_path = _write_report(report, out_dir) + return result diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py new file mode 100644 index 0000000..9b69b56 --- /dev/null +++ b/tests/test_dynamics.py @@ -0,0 +1,147 @@ +"""NVE probe MD: conservation metrics, the explosion guard, property capture.""" + +from __future__ import annotations + +import numpy as np +import pytest +from ase.calculators.calculator import Calculator, all_changes + +from traincraft import label_frames +from traincraft.config.models import ( + CrystalBuilder, + EmtCalc, + GeometryConfig, + StabilityConfig, + ValidationConfig, +) +from traincraft.geometry import build_geometry +from traincraft.validation import dynamics as dyn_mod +from traincraft.validation import run_probe_md, run_validation_science, stability_metrics + + +def _cu_atoms(): + return build_geometry( + GeometryConfig(builder=CrystalBuilder(name="Cu", crystalstructure="fcc", a=3.6, cubic=True)) + ).atoms + + +class PointChargeDielectric(Calculator): + """Test stand-in for a dielectric MLIP: analytic dipole + polarizability.""" + + implemented_properties = ["energy", "dipole"] + + def calculate(self, atoms=None, properties=None, system_changes=all_changes): + super().calculate(atoms, properties, system_changes) + pos = self.atoms.get_positions() + charges = np.linspace(-1.0, 1.0, len(self.atoms)) + self.results["energy"] = 0.0 + self.results["dipole"] = charges @ pos + self.results["polarizability"] = np.eye(3) * (1.0 + 0.1 * np.sin(pos.sum())) + + +def _probe(atoms=None, **overrides): + from ase.calculators.emt import EMT + + kwargs = dict( + temperature=300.0, timestep=1.0, equilibration_steps=20, + steps=100, sample_every=2, seed=7, + ) + kwargs.update(overrides) + return run_probe_md(atoms if atoms is not None else _cu_atoms(), EMT(), **kwargs) + + +def test_nve_probe_conserves_energy(): + series = _probe() + assert series.completed and series.aborted_at is None + assert series.dt_fs == 2.0 + assert len(series.energy) == 51 # samples at 0, 2, ..., 100 + assert series.velocities.shape == (51, series.n_atoms, 3) + + m = stability_metrics(series) + assert m["nve_aborted"] == 0.0 + assert m["nve_temp_mean"] > 0.0 + # EMT + Velocity-Verlet at 1 fs conserves energy to well under this + assert m["nve_drift_mev_per_atom_ps"] < 5.0 + assert m["nve_energy_std_mev_per_atom"] < 1.0 + + +def test_nve_probe_is_seeded(): + a = _probe(steps=20) + b = _probe(steps=20) + assert np.array_equal(a.energy, b.energy) + assert np.array_equal(a.velocities, b.velocities) + + +def test_explosion_guard_aborts_cleanly(monkeypatch): + monkeypatch.setattr(dyn_mod, "MAX_FORCE_EV_A", 1e-9) + series = _probe(equilibration_steps=5) + assert not series.completed + assert series.aborted_at == 0 # rattled cell trips a zero tolerance at once + assert stability_metrics(series)["nve_aborted"] == 1.0 + assert "nve_drift_mev_per_atom_ps" not in stability_metrics(series) + + +def test_property_capture_along_the_run(): + series = _probe( + steps=20, prop_calc=PointChargeDielectric(), + capture=("dipole", "polarizability"), + ) + assert series.dipole.shape == (11, 3) + assert series.polarizability.shape == (11, 3, 3) + assert not np.allclose(series.dipole, series.dipole[0]) # it moves + + +def test_missing_property_is_a_hard_error(): + from ase.calculators.emt import EMT + + with pytest.raises(ValueError, match="no dipole"): + _probe(steps=10, prop_calc=EMT(), capture=("dipole",)) + with pytest.raises(ValueError, match=r"results\['polarizability'\]"): + _probe(steps=10, prop_calc=EMT(), capture=("polarizability",)) + + +# --- run_validation_science: stability folded into the quality report ---------- + +def test_validation_science_reports_nve_metrics(tmp_path): + base = _cu_atoms() + frames = [] + for i in range(2): + s = build_geometry( + GeometryConfig( + builder=CrystalBuilder(name="Cu", crystalstructure="fcc", a=3.6, cubic=True) + ) + ) + s.atoms.rattle(0.02, seed=i) + frames.append(s) + frames = label_frames(frames, EmtCalc()) + assert len(base) == len(frames[0].atoms) + + cfg = ValidationConfig( + calculator=EmtCalc(), + stability=StabilityConfig(steps=40, equilibration_steps=10, n_structures=2), + thresholds={"nve_aborted": 0.0, "nve_drift_mev_per_atom_ps": 50.0}, + ) + result = run_validation_science(frames, EmtCalc(), cfg, out_dir=tmp_path) + assert result.passed is True + assert result.metrics["nve_aborted"] == 0.0 + assert "energy_rmse" in result.metrics # parity still there + + report = result.report + assert len(report["stability"]["probes"]) == 2 + assert {c["metric"] for c in report["checks"]} == { + "nve_aborted", "nve_drift_mev_per_atom_ps", + } + assert result.report_path is not None and result.report_path.exists() + + +def test_validation_science_without_probes_matches_plain(tmp_path): + frames = label_frames( + [build_geometry(GeometryConfig(builder=CrystalBuilder( + name="Cu", crystalstructure="fcc", a=3.6, cubic=True)))], + EmtCalc(), + ) + cfg = ValidationConfig(calculator=EmtCalc(), thresholds={"energy_rmse": 1.0}) + result = run_validation_science(frames, EmtCalc(), cfg, out_dir=tmp_path) + assert result.passed is True + assert "stability" not in result.report + assert result.metrics["energy_rmse"] == pytest.approx(0.0, abs=1e-10) diff --git a/tests/test_learning_curve.py b/tests/test_learning_curve.py new file mode 100644 index 0000000..8d873eb --- /dev/null +++ b/tests/test_learning_curve.py @@ -0,0 +1,140 @@ +"""Learning curve: nested subsets, fixed holdout, curve artifact, CLI.""" + +from __future__ import annotations + +import json +import sys + +import pytest +from ase.io import read as ase_read +from typer.testing import CliRunner + +from traincraft import label_frames, write_frames +from traincraft.cli import app +from traincraft.config.models import ( + CrystalBuilder, + EmtCalc, + GeometryConfig, + MaceCalc, + MaceTrainer, + RunConfig, + TrainCraftConfig, +) +from traincraft.core import Workspace +from traincraft.training.mace import TRAIN_COMMAND_ENV +from traincraft.validation import learning_curve +from traincraft.validation import validate as validate_mod + +# Stands in for mace_run_train: writes the .model file the trainer looks for. +_FAKE_TRAINER = """\ +import sys +from pathlib import Path +argv = sys.argv +model_dir = argv[argv.index("--model_dir") + 1] +name = argv[argv.index("--name") + 1] +Path(model_dir, name + ".model").write_text("dummy") +""" + + +@pytest.fixture +def fake_training(tmp_path, monkeypatch): + """Fake mace_run_train + route the resulting MaceCalc to EMT.""" + script = tmp_path / "fake_mace_run_train.py" + script.write_text(_FAKE_TRAINER) + monkeypatch.setenv(TRAIN_COMMAND_ENV, f"{sys.executable} {script}") + + real = validate_mod.make_calculator + + def fake(cfg): + if isinstance(cfg, MaceCalc): + from ase.calculators.emt import EMT + + return EMT() + return real(cfg) + + monkeypatch.setattr(validate_mod, "make_calculator", fake) + + +def _seeded_run(tmp_path, n_frames=8) -> TrainCraftConfig: + """A run dir with labeled frames already in place of the label artifact.""" + from traincraft.geometry import build_geometry + + cfg = TrainCraftConfig( + run=RunConfig(name="lc", outdir=str(tmp_path), seed=1), + geometry=GeometryConfig( + builder=CrystalBuilder(name="Cu", crystalstructure="fcc", a=3.6, cubic=True) + ), + training=MaceTrainer(valid_fraction=0.25, seed=3, max_num_epochs=1), + ) + frames = [] + for i in range(n_frames): + s = build_geometry(cfg.geometry) + s.atoms.rattle(0.03, seed=i) + frames.append(s) + ws = Workspace(tmp_path / "lc") + write_frames(ws.subdir("labeled_dft") / "labeled.extxyz", label_frames(frames, EmtCalc())) + return cfg + + +def _positions(path): + return {atoms.positions.tobytes() for atoms in ase_read(path, index=":")} + + +def test_learning_curve_nested_subsets_fixed_holdout(tmp_path, fake_training): + cfg = _seeded_run(tmp_path) + curve = learning_curve(cfg, fractions=(0.34, 1.0)) + + assert curve["n_frames"] == 8 + assert curve["n_holdout"] == 2 # valid_fraction 0.25 of 8 + assert [p["n_train"] for p in curve["points"]] == [2, 6] + + # every point scored against the same holdout with its own trained model + for point in curve["points"]: + # ~0 up to the extxyz text round-trip of the stored labels + assert point["metrics"]["energy_rmse"] == pytest.approx(0.0, abs=1e-6) + + # nested: the small subset's frames all appear in the large one + small = _positions(tmp_path / "lc/validation/learning_curve/n00002/train.xyz") + large = _positions(tmp_path / "lc/validation/learning_curve/n00006/train.xyz") + assert small < large + + on_disk = json.loads((tmp_path / "lc/validation/learning_curve.json").read_text()) + assert on_disk["points"] == curve["points"] + + +def test_learning_curve_dedups_rounding_collisions(tmp_path, fake_training): + cfg = _seeded_run(tmp_path) + curve = learning_curve(cfg, fractions=(0.1, 0.15)) # both round to n=1 + assert [p["n_train"] for p in curve["points"]] == [1] + + +def test_learning_curve_guards(): + cfg = TrainCraftConfig( + run=RunConfig(name="x", outdir="runs"), + geometry=GeometryConfig(builder=CrystalBuilder(name="Cu", crystalstructure="fcc", a=3.6)), + ) + with pytest.raises(ValueError, match=r"\[training\] section"): + learning_curve(cfg) + cfg = cfg.model_copy(update={"training": MaceTrainer()}) + with pytest.raises(ValueError, match=r"fractions must lie"): + learning_curve(cfg, fractions=(0.5, 1.5)) + with pytest.raises(ValueError, match=r"fractions must lie"): + learning_curve(cfg, fractions=()) + + +def test_learning_curve_cli(tmp_path, fake_training): + _seeded_run(tmp_path) + config = tmp_path / "lc.toml" + config.write_text( + f'[run]\nname = "lc"\noutdir = "{tmp_path}"\nseed = 1\n\n' + '[geometry.builder]\ntype = "crystal"\nname = "Cu"\n' + 'crystalstructure = "fcc"\na = 3.6\ncubic = true\n\n' + '[training]\ntype = "mace"\nvalid_fraction = 0.25\nseed = 3\n' + ) + runner = CliRunner() + result = runner.invoke(app, ["learning-curve", str(config), "--fractions", "0.5,1.0"]) + assert result.exit_code == 0, result.output + assert "wrote" in result.output and "learning_curve.json" in result.output + + bad = runner.invoke(app, ["learning-curve", str(config), "--fractions", "half"]) + assert bad.exit_code == 2 diff --git a/tests/test_spectra.py b/tests/test_spectra.py new file mode 100644 index 0000000..eff9e32 --- /dev/null +++ b/tests/test_spectra.py @@ -0,0 +1,148 @@ +"""Spectra library: TCFs and spectra checked against synthetic signals.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from traincraft.validation.spectra import ( + C_CM_PER_FS, + autocorrelation, + broaden, + cross_correlation, + find_peaks, + ir_spectrum, + power_spectrum, + qcf_ratio, + raman_spectra, + sfg_spectrum, + spectral_similarity, + vdos, +) + +N, DT = 4096, 1.0 # steps, fs -> ~4 cm^-1 grid resolution after 4x padding +NU0 = 1000.0 # cm^-1 + + +def _cos(nu_cm1: float, n: int = N, dt: float = DT, phase: float = 0.0) -> np.ndarray: + omega = 2.0 * np.pi * C_CM_PER_FS * nu_cm1 # rad/fs + return np.cos(omega * np.arange(n) * dt + phase) + + +def _peak(freq, intensity, lo=200.0): + """Peak position above ``lo`` (the Rayleigh wing is noise-amplified).""" + mask = freq > lo + return freq[mask][np.argmax(intensity[mask])] + + +def test_autocorrelation_unbiased_on_constant(): + acf = autocorrelation(np.ones((64, 3)), max_lag=16) + assert np.allclose(acf, 3.0) # = 3 for a constant unit 3-vector + + +def test_autocorrelation_of_cosine(): + acf = autocorrelation(_cos(NU0), max_lag=200) + assert acf[0] == pytest.approx(0.5, rel=1e-2) # = 1/2 + period_fs = 1.0 / (C_CM_PER_FS * NU0) + assert acf[int(round(period_fs))] == pytest.approx(0.5, rel=0.05) + + +def test_short_series_rejected(): + with pytest.raises(ValueError, match="too short"): + autocorrelation(np.ones(4)) + with pytest.raises(ValueError, match="mismatched"): + cross_correlation(np.ones(64), np.ones((64, 2))) + + +def test_power_spectrum_peak_position(): + nu, spec = power_spectrum(_cos(NU0), DT) + assert _peak(nu, spec) == pytest.approx(NU0, abs=10.0) + + +def test_vdos_per_element_decomposition(): + v = np.zeros((N, 2, 3)) + v[:, 0, 0] = _cos(800.0) + v[:, 1, 0] = _cos(2000.0) + nu, curves = vdos(v, masses=[1.0, 1.0], dt_fs=DT, symbols=["Cu", "O"]) + assert set(curves) == {"total", "Cu", "O"} + assert _peak(nu, curves["Cu"]) == pytest.approx(800.0, abs=10.0) + assert _peak(nu, curves["O"]) == pytest.approx(2000.0, abs=10.0) + # total contains both modes + total_peaks = {round(p["freq_cm1"], -1) for p in find_peaks(nu, curves["total"])} + assert {800.0, 2000.0} <= total_peaks + + +def test_ir_peak_position(): + mu = np.zeros((N, 3)) + mu[:, 2] = _cos(NU0) + nu, spec = ir_spectrum(mu, DT, temperature=300.0) + assert _peak(nu, spec) == pytest.approx(NU0, abs=10.0) + + +def test_raman_isotropic_mode_is_polarized(): + pol = _cos(NU0)[:, None, None] * np.eye(3) + out = raman_spectra(pol, DT, temperature=300.0) + i = np.argmin(np.abs(out["freq"] - NU0)) + assert _peak(out["freq"], out["vv"]) == pytest.approx(NU0, abs=10.0) + # a purely isotropic mode has no VH intensity and depolarization -> 0 + assert out["vh"][i] < 1e-6 * out["vv"][i] + assert out["depolarization"][i] == pytest.approx(0.0, abs=1e-6) + + +def test_raman_anisotropic_mode_depolarization_is_three_quarters(): + pol = _cos(NU0)[:, None, None] * np.diag([1.0, -1.0, 0.0]) + out = raman_spectra(pol, DT, temperature=300.0) + i = np.argmin(np.abs(out["freq"] - NU0)) + assert out["depolarization"][i] == pytest.approx(0.75, abs=0.01) + assert np.allclose(out["unpolarized"], out["vv"] + out["vh"]) + + +def test_raman_rejects_bad_shape(): + with pytest.raises(ValueError, match=r"\(n, 3, 3\)"): + raman_spectra(np.zeros((100, 3)), DT, temperature=300.0) + + +def test_qcf_ratios(): + nu = np.linspace(0.0, 4000.0, 50) + assert np.allclose(qcf_ratio(nu, 300.0, "harmonic"), 1.0) + none = qcf_ratio(nu[1:], 300.0, "none") + assert np.all(none < 1.0) # classical spectrum sits below the harmonic QCF + # all corrections agree in the classical (high-T) limit + for kind in ("none", "standard", "schofield"): + assert np.allclose(qcf_ratio(nu, 1e9, kind), 1.0, atol=1e-3) + with pytest.raises(ValueError, match="unknown qcf"): + qcf_ratio(nu, 300.0, "kubo") + + +def test_sfg_peak_and_component_validation(): + pol = np.zeros((N, 3, 3)) + dip = np.zeros((N, 3)) + pol[:, 0, 0] = _cos(NU0) + dip[:, 2] = _cos(NU0) + out = sfg_spectrum(pol, dip, DT, temperature=300.0, component="xxz") + assert _peak(out["freq"], out["abs2"]) == pytest.approx(NU0, abs=10.0) + # an uncorrelated component carries no signal + quiet = sfg_spectrum(pol, dip, DT, temperature=300.0, component="yyz") + assert quiet["abs2"].max() < 1e-6 * out["abs2"].max() + with pytest.raises(ValueError, match="three of x/y/z"): + sfg_spectrum(pol, dip, DT, temperature=300.0, component="ab") + + +def test_spectral_similarity_identity_and_disjoint(): + nu, spec = ir_spectrum(_cos(NU0)[:, None] * np.array([0, 0, 1.0]), DT, 300.0) + assert spectral_similarity(nu, spec, nu, spec) == pytest.approx(1.0) + grid = np.linspace(0, 4000, 2000) + a = np.exp(-0.5 * ((grid - 1000) / 8) ** 2) + b = np.exp(-0.5 * ((grid - 3000) / 8) ** 2) + assert spectral_similarity(grid, a, grid, b) == pytest.approx(0.0, abs=1e-6) + # a stick spectrum against its broadened self still matches well + assert spectral_similarity([1000.0], [1.0], grid, a, sigma=8.0) > 0.95 + + +def test_broaden_and_find_peaks(): + grid, curve = broaden([500.0, 1500.0], [1.0, 2.0], lo=0, hi=2000, sigma=20.0) + peaks = find_peaks(grid, curve) + assert len(peaks) == 2 + assert peaks[0]["freq_cm1"] == pytest.approx(1500.0, abs=10.0) # strongest first + assert peaks[1]["freq_cm1"] == pytest.approx(500.0, abs=10.0) + assert find_peaks(grid, np.zeros_like(curve)) == [] diff --git a/tests/test_validation_science.py b/tests/test_validation_science.py new file mode 100644 index 0000000..d74be03 --- /dev/null +++ b/tests/test_validation_science.py @@ -0,0 +1,202 @@ +"""Spectra pipeline: [validation.spectra] config -> MD -> .dat artifacts + metrics.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from test_dynamics import PointChargeDielectric +from traincraft import label_frames +from traincraft.config import loads_config +from traincraft.config.loader import load_config +from traincraft.config.models import ( + CrystalBuilder, + EmtCalc, + GeometryConfig, + MaceCalc, + SpectraConfig, + ValidationConfig, +) +from traincraft.geometry import build_geometry +from traincraft.validation import run_validation_science +from traincraft.validation import validate as validate_mod + + +def _frames(n=1): + frames = [] + for i in range(n): + s = build_geometry( + GeometryConfig( + builder=CrystalBuilder(name="Cu", crystalstructure="fcc", a=3.6, cubic=True) + ) + ) + s.atoms.rattle(0.02, seed=i) + frames.append(s) + return label_frames(frames, EmtCalc()) + + +def _fast_spectra(**overrides) -> SpectraConfig: + kwargs = dict( + which=["vdos"], steps=64, equilibration_steps=5, timestep=1.0, + freq_range=(0.0, 4000.0), + ) + kwargs.update(overrides) + return SpectraConfig(**kwargs) + + +@pytest.fixture +def dielectric_stub(monkeypatch): + """Route the spectra 'dielectric model' config to the analytic test calc. + + The CalculatorConfig union is closed, so tests mark the dielectric model + with a MaceCalc placeholder and intercept its construction here — nothing + ever tries to import mace. + """ + real = validate_mod.make_calculator + + def fake(cfg): + if isinstance(cfg, MaceCalc): + return PointChargeDielectric() + return real(cfg) + + monkeypatch.setattr(validate_mod, "make_calculator", fake) + return MaceCalc(model_path="/nonexistent/dielectric.model", model_type="DipoleMACE") + + +def test_vdos_only_needs_no_dielectric_model(tmp_path): + cfg = ValidationConfig(calculator=EmtCalc(), spectra=_fast_spectra()) + result = run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path) + + assert result.metrics["spectra_md_aborted"] == 0.0 + spectra_dir = tmp_path / "spectra" + assert (spectra_dir / "vdos.dat").is_file() + assert (spectra_dir / "vdos_Cu.dat").is_file() + data = np.loadtxt(spectra_dir / "vdos.dat") + assert data.ndim == 2 and data.shape[1] == 2 + assert data[:, 0].max() <= 4000.0 # freq_range applied + + meta = json.loads((spectra_dir / "spectra.json").read_text()) + assert meta["n_samples"] == 65 + assert "vdos" in meta["peaks"] + assert meta["files"] == sorted(meta["files"]) + assert result.report["spectra"]["metrics"] == {"spectra_md_aborted": 0.0} + + +def test_full_spectra_set_with_dielectric_model(tmp_path, dielectric_stub): + cfg = ValidationConfig( + calculator=EmtCalc(), + spectra=_fast_spectra( + which=["ir", "raman", "sfg"], calculator=dielectric_stub, save_series=True + ), + ) + result = run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path) + assert result.metrics["spectra_md_aborted"] == 0.0 + + spectra_dir = tmp_path / "spectra" + for stem in ("ir", "raman_vv", "raman_vh", "raman_unpolarized", + "raman_depolarization", "raman_iso", "raman_aniso", + "sfg_re", "sfg_im", "sfg_abs2"): + assert (spectra_dir / f"{stem}.dat").is_file(), stem + + series = np.load(spectra_dir / "series.npz") + assert series["dipole"].shape == (65, 3) + assert series["polarizability"].shape == (65, 3, 3) + + +def test_reference_similarity_becomes_a_thresholdable_metric(tmp_path, dielectric_stub): + # first pass: compute an IR spectrum, then feed it back as its own reference + cfg = ValidationConfig( + calculator=EmtCalc(), + spectra=_fast_spectra(which=["ir"], calculator=dielectric_stub), + ) + run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path) + self_ref = tmp_path / "spectra" / "ir.dat" + + cfg = ValidationConfig( + calculator=EmtCalc(), + spectra=_fast_spectra( + which=["ir"], calculator=dielectric_stub, reference={"ir": str(self_ref)} + ), + thresholds={"ir_spectral_error": 0.05}, + ) + result = run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path / "second") + assert result.metrics["ir_spectral_error"] == pytest.approx(0.0, abs=1e-6) + assert result.passed is True + + # a reference peaked far away fails the same threshold + wrong = tmp_path / "wrong_ref.dat" + grid = np.linspace(0, 4000, 500) + np.savetxt(wrong, np.column_stack([grid, np.exp(-((grid - 3500) / 10) ** 2)])) + cfg = ValidationConfig( + calculator=EmtCalc(), + spectra=_fast_spectra( + which=["ir"], calculator=dielectric_stub, reference={"ir": str(wrong)} + ), + thresholds={"ir_spectral_error": 0.05}, + ) + result = run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path / "third") + assert result.metrics["ir_spectral_error"] > 0.5 + assert result.passed is False + + +def test_enhancement_profile_applies_to_raman(tmp_path, dielectric_stub): + profile = tmp_path / "enh.dat" + np.savetxt(profile, np.column_stack([[0.0, 4000.0], [0.0, 0.0]])) # kill everything + cfg = ValidationConfig( + calculator=EmtCalc(), + spectra=_fast_spectra( + which=["raman"], calculator=dielectric_stub, enhancement_file=str(profile) + ), + ) + run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path) + enhanced = np.loadtxt(tmp_path / "spectra" / "raman_vv_enhanced.dat") + plain = np.loadtxt(tmp_path / "spectra" / "raman_vv.dat") + assert np.allclose(enhanced[:, 1], 0.0) + assert plain[:, 1].max() > 0.0 + + +def test_aborted_md_is_a_metric_not_a_crash(tmp_path, monkeypatch): + from traincraft.validation import dynamics as dyn_mod + + monkeypatch.setattr(dyn_mod, "MAX_FORCE_EV_A", 1e-9) + cfg = ValidationConfig(calculator=EmtCalc(), spectra=_fast_spectra()) + result = run_validation_science(_frames(), EmtCalc(), cfg, out_dir=tmp_path) + assert result.metrics["spectra_md_aborted"] == 1.0 + assert not (tmp_path / "spectra" / "vdos.dat").exists() + assert (tmp_path / "spectra" / "spectra.json").is_file() + + +def test_spectra_config_validators(): + with pytest.raises(ValueError, match="must not be empty"): + SpectraConfig(which=[]) + with pytest.raises(ValueError, match="three of x/y/z"): + SpectraConfig(sfg_component="qq") + with pytest.raises(ValueError, match="reference for"): + SpectraConfig(which=["vdos"], reference={"ir": "ref.dat"}) + with pytest.raises(ValueError, match="enhancement_file"): + SpectraConfig(which=["ir"], enhancement_file="enh.dat") + with pytest.raises(ValueError, match="freq_range"): + SpectraConfig(freq_range=(500.0, 100.0)) + with pytest.raises(ValueError, match="model_type"): + MaceCalc(model_type="DipoleMACE") # no model_path + + +def test_loader_resolves_validation_paths(tmp_path): + for name in ("test.extxyz", "ref_ir.dat", "enh.dat"): + (tmp_path / name).touch() + cfg_file = tmp_path / "config.toml" + cfg_file.write_text( + '[run]\nname = "r"\n\n[geometry.builder]\ntype = "molecule"\nname = "H2O"\n\n' + '[validation]\ntest_file = "test.extxyz"\n\n' + '[validation.spectra]\nwhich = ["ir", "raman"]\n' + 'reference = { ir = "ref_ir.dat" }\nenhancement_file = "enh.dat"\n' + ) + cfg = load_config(cfg_file) + assert cfg.validation.test_file == str(tmp_path / "test.extxyz") + assert cfg.validation.spectra.reference["ir"] == str(tmp_path / "ref_ir.dat") + assert cfg.validation.spectra.enhancement_file == str(tmp_path / "enh.dat") + # loads_config (string form) is documented not to resolve paths + plain = loads_config(cfg_file.read_text()) + assert plain.validation.test_file == "test.extxyz"