Counterfactual-Structure-Generation

Generate counterfactual structures to diagnose when, where, and why a trained MLIP becomes uncertain

Pipeline

1. Train a NequIP model on your MD trajectory data (bring your own NequIP/Hydra config).
2. extract_features.py — evaluate training data with the trained model, extract per-atom node features.
3. fit_gmm.py — fit a per-element GMM on those features.
4. run_counterfactual_optimization.py — NLL-ascent, run as a SLURM job array (--step split|submit|collect), generating counterfactual structures.
5. extract_features.py --input-xyz <counterfactual xyz> --gmm-path outputs/gmm_per_type.pkl — reuse on the collected counterfactual structures to score them with NequIP and (via the --gmm-path flag) the fitted GMM, producing a per-atom NLL CSV.
6. nll_gradient_attribution.py — NLL gradient-attribution map over (phi, psi) dihedral space.

Run the steps above in order; each script's --help lists its options.

Install

pip install -r requirements.txt

Everything runs directly, e.g. python scripts/extract_features.py --help — no package install step needed.

Configure

cp config/paths.example.yaml config/paths.yaml
# edit config/paths.yaml to point at your checkpoint/data/output paths

Every script accepts --config config/paths.yaml, plus per-value CLI overrides (e.g. --model-path).

Repository layout

config/paths.example.yaml   configuration template
scripts/                    all code (flat, no package layout)
requirements.txt            dependencies

License

MIT — see LICENSE.