battery-materials-ml

A single-file pipeline that assembles a clean, de-duplicated crystal-structure dataset for battery-relevant chemistries from two public DFT databases, featurises it, turns it into graphs, and trains a baseline property regressor.

Materials Project  +  Alexandria (OPTIMADE)
      │
      ├─ merge + structural de-duplication      pymatgen StructureMatcher
      ├─ symmetry / validity sanitisation       conventional cell, Niggli, min-distance filter
      ├─ features                               matminer density + packing, dscribe SOAP
      ├─ E(3)-ready graphs                      Z nodes, PBC-aware edge vectors
      └─ baseline model                         RandomForest + GroupKFold CV

The group key for cross-validation is the structural-duplicate cluster, so a material that appears in both databases can never straddle the train/test split — the usual cause of over-optimistic scores in materials ML.

What you get

Running the full pipeline writes:

file	contents
datasets/merged_sanitized.parquet	de-duplicated, sanitised structures + labels
datasets/features_soap.parquet	the above + density/packing + SOAP columns
datasets/graphs_preview_e3.npz	first 100 structures as equivariant-ready graphs
artifacts/rf_<target>.joblib	trained RandomForest + feature list

Default regression target is band_gap_eV; E_form_per_atom and e_above_hull are also available.

Setup

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

cp .env.example .env          # then edit .env
# MP_API_KEY=...   (free key from https://materialsproject.org/api)

dscribe, scikit-learn, python-dotenv and m3gnet are optional — the script checks for them at runtime and skips the corresponding step with a message rather than crashing.

Run

# quick offline + online sanity checks, then build + train
python build_dataset.py

Switches are near the bottom of the file:

RUN_SMOKE = True           # offline (NaCl) + online (tiny MP/Alexandria) checks
RUN_MAIN  = True           # full build + baseline training
TARGET    = "band_gap_eV"  # or "E_form_per_atom", "e_above_hull"

The chemistry scope is the DEFAULT_ELEMENTS_ANY list near the top — alkali + common TM + chalcogen/halogen + a few framework cations. Alkali–X binaries are enumerated first for broad coverage without huge pulls; extend to ternaries by adding chemsys combinations in mp_fetch_slice.

Key functions

function	role
mp_fetch_slice / alex_fetch_slice	pull summary + structure records from each source
merge_and_dedup	concat, bucket by (formula, nsites), cluster with StructureMatcher, pick a representative
symmetry_sanitize	conventional standard cell → Niggli reduction → wrap sites → reject if any pair closer than min_dist Å
add_matminer_light	DensityFeatures + MaximumPackingEfficiency
make_soap_descriptor	per-structure averaged SOAP vector; tolerates both rcut and r_cut dscribe signatures
structure_to_graph	nodes = atomic number, edges within a cutoff carrying the minimum-image displacement vector and the lattice matrix — ready for an E(3)-equivariant GNN
train_baseline_rf	GroupKFold CV on cluster id, then a final fit on all data; saves the bundle
predict_from_structure	featurise one new Structure and predict with the saved model

Notes

• Data comes from Materials Project and the Alexandria library; cite them if you use the assembled dataset. Respect each database's terms of use.
• This is a research scaffold, not a packaged library — one file, edit freely.
• The graph export is deliberately framework-agnostic (plain NumPy arrays) so it can be adapted to PyG, e3nn, or JAX without a hard dependency.

License

MIT — see LICENSE. Applies to this pipeline code, not to the upstream database contents.