feat(adastra): port ChargE3Net fine-tuning to AMD MI250X on CINES Adastra

boa_ft/README.md

@@ -0,0 +1,97 @@
+# boa_ft: BOA on LeMat-Rho
+
+Trains BOA (Basis Overlap Architecture, github.com/sciai-lab/boa) from scratch
+on LeMat-Rho r2SCAN charge densities (15x15x15 grids, periodic solids).
+
+BOA learns per atom-pair Gaussian-orbital coefficients and is supervised
+directly on grid density values (no grid-to-basis projection step). We add:
+
+- `preprocess.py`: LeMat-Rho parquet -> sharded LMDB of scdp `AtomicData`.
+- `transforms.py`: `PBCRadiusEdgeIndex`, a periodic (minimum-image) neighbor
+  graph that replaces BOA's open-boundary `radius_graph` and removes the
+  `torch_cluster` dependency from the training path.
+- `configs/`: Hydra overlays wiring the datamodule to our LMDB and transform,
+  and turning the periodic decoder on (`model.pbc=true`, `model.orb_cutoff=3.0`).
+
+## Setup
+
+BOA and its vendored `scdp` and `sciai-dft` (`mldft`) packages are not on PyPI.
+Clone BOA as a sibling of this repo and editable-install the three packages into
+the same environment that has this repo's dependencies:
+
+```bash
+git clone https://github.com/sciai-lab/boa ../boa
+uv pip install -e ../boa/scdp -e ../boa/sciai-dft -e ../boa
+```
+
+`preprocess.py` also imports the shared parquet decoder from `charge3net_ft`,
+which requires the charge3net repo cloned as a sibling (`../charge3net`), exactly
+as the deepdft and charge3net arms already need.
+
+## Preprocess
+
+```bash
+python -m boa_ft.preprocess \
+    --parquet-dir /path/to/lemat_rho_chunks \
+    --out-dir $BOA_DATA/lematrho \
+    --num-shards 16
+```
+
+This writes `$BOA_DATA/lematrho/data/*.lmdb`, `$BOA_DATA/lematrho/datasplits.json`
+(charge3net split convention: seed 42, 5% val, 5% test), and
+`$BOA_DATA/lematrho/atomic_numbers.json`. The element list is the union over the
+*train* split only, because BOA builds one orbital per training element and
+aligns the basis positionally against the train split. This means the train
+split must contain every element you want a basis for (a 90% split of a large
+dataset does; a tiny subset may not). Use `--limit N` for a quick smoke subset.
+
+## Train
+
+BOA reads three environment variables for its paths (`PROJECT_ROOT` = the boa
+clone, `BOA_DATA` = the dataset root that holds `lematrho/`, `BOA_MODELS` = the
+run/checkpoint output root). BOA's Hydra root lives in the boa clone, so add our
+config directory to the search path and select the experiment:
+
+```bash
+export PROJECT_ROOT=/path/to/boa
+export BOA_DATA=/path/to/data_root
+export BOA_MODELS=/path/to/model_root
+
+ATOMIC_NUMBERS=$(python -c "import json; print(json.load(open('$BOA_DATA/lematrho/atomic_numbers.json')))")
+
+python "$PROJECT_ROOT/boa/train.py" \
+    hydra.searchpath="[file://$(pwd)/boa_ft/configs]" \
+    experiment=lematrho \
+    data.basis_info.atomic_numbers="$ATOMIC_NUMBERS"
+```
+
+The `experiment=lematrho` overlay sets `model.pbc=true`, `model.orb_cutoff=3.0`,
+`model.linear_basis=true` (the periodic-safe density decoder), and
+`initial_guess_pre_training_steps=0`. BOA's separate initial-guess pre-training
+opens a second datamodule on the same LMDB, which py-lmdb forbids in one
+process, so we disable it; the initial-guess module still trains jointly in the
+main loop.
+
+### Smoke run (CPU)
+
+`+trainer.devices=1` uses a leading `+` because the trainer config has no
+`devices` key by default. A short `hydra.run.dir` avoids a run-directory name
+built from the (long) override string.
+
+```bash
+python "$PROJECT_ROOT/boa/train.py" \
+    hydra.searchpath="[file://$(pwd)/boa_ft/configs]" \
+    hydra.run.dir=/tmp/boa_smoke \
+    experiment=lematrho \
+    data.basis_info.atomic_numbers="$ATOMIC_NUMBERS" \
+    trainer.accelerator=cpu +trainer.devices=1 \
+    trainer.max_steps=3 trainer.max_epochs=1 trainer.num_sanity_val_steps=0 \
+    data.datamodule.num_workers.train=0 data.datamodule.num_workers.val=0 \
+    data.datamodule.batch_size.train=2 data.datamodule.n_probe.train=512
+```
+
+## Adastra
+
+Use `submit_boa_adastra.sh` at the repo root (MI250, account `c1816212`, single
+GCD first). It sets the proxy, activates `venv311`, exports the three BOA path
+variables, reads `atomic_numbers.json`, and launches the Hydra command above.

boa_ft/__init__.py

@@ -0,0 +1,7 @@
+"""LeMat-Rho -> BOA (Basis Overlap Architecture) training arm.
+
+This package adapts the LeMat-Rho parquet charge-density dataset to the BOA
+codebase (github.com/sciai-lab/boa), which is editable-installed from a sibling
+clone at ``../boa`` together with its vendored ``scdp`` and ``sciai-dft``
+(``mldft``) packages. See ``boa_ft/README.md`` for setup and usage.
+"""

boa_ft/configs/data/datamodule/dataset/lematrho.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - transform: lematrho
+
+_target_: boa.data.dataset.LmdbDataset
+path: ${paths.data_dir}/${data.dataset_name}/data

boa_ft/configs/data/datamodule/dataset/transform/lematrho.yaml

@@ -0,0 +1,39 @@
+# Mirrors boa's default transform pipeline, but swaps the two open-boundary
+# AddRadiusEdgeIndex steps for boa_ft.transforms.PBCRadiusEdgeIndex so the graph
+# respects periodic boundary conditions (and drops the torch_cluster dependency).
+# The list order is kept identical so the ${..N.name} interpolations below still
+# point at the matching edge-index step.
+_target_: boa.data.transforms.MasterTransform
+transforms:
+  - _target_: boa.data.transforms.SampleProbe
+    n_probe: ${data.datamodule.n_probe.train}
+
+  - _target_: boa.data.transforms.ConvertToOFData
+    basis_info: ${data.basis_info}
+
+  - _target_: boa.data.transforms.ToTorch
+
+  - _target_: boa.data.transforms.AddMessagePassingMatrix
+    basis_info: ${data.basis_info}
+    type: overlap
+    remove_diagonal: False
+
+  - _target_: boa_ft.transforms.PBCRadiusEdgeIndex
+    radius: 6.0
+    name: message_edge_index
+
+  - _target_: boa.data.transforms.AddEdgeMatrices
+    basis_info: ${data.basis_info}
+    edge_name: ${..4.name}
+    name: message_edge_matrices
+
+  - _target_: boa_ft.transforms.PBCRadiusEdgeIndex
+    radius: 3.0
+    name: edge_index
+
+  - _target_: boa.data.transforms.AddEdgeMatrices
+    basis_info: ${data.basis_info}
+    edge_name: ${..6.name}
+    name: edge_matrices
+
+  - _target_: boa.data.transforms.ToTorch

boa_ft/configs/data/datamodule/lematrho.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - default
+  - dataset: lematrho
+
+_target_: boa.data.datamodule.ProbeDataModule
+
+batch_size:
+  train: 4
+  val: 4
+  test: 4
+
+# 15x15x15 grids hold 3375 points per material; 2000 probes subsamples each
+# sample. The datamodule default sets n_probe.val/test equal to n_probe.train.
+n_probe:
+  train: 2000

boa_ft/configs/data/lematrho.yaml

@@ -0,0 +1,18 @@
+defaults:
+  - datamodule: lematrho
+
+dataset_name: lematrho
+
+# Uncontracted def2-svp: a compact basis that keeps the per-sample pyscf
+# overlap (built in the dataloader) affordable for many-element solids.
+# atomic_numbers must list every element present in the LMDB. Preprocess writes
+# the exact list to <out-dir>/atomic_numbers.json; override this at launch, e.g.
+#   data.basis_info.atomic_numbers='[1, 8, 26]'
+# The default below is only a placeholder for smoke runs.
+basis_info:
+  _target_: boa.data.basis_info.BasisInfo.from_atomic_numbers_with_even_tempered_basis
+  basis: 'def2-svp'
+  atomic_numbers: [1, 8]
+  beta: None
+  even_tempered: false
+  uncontracted: True

boa_ft/configs/experiment/lematrho.yaml

@@ -0,0 +1,27 @@
+# @package _global_
+
+# LeMat-Rho periodic solids: enable the periodic density decoder and keep the
+# orbital cutoff tight (3.0 A) so the periodic image replication inside
+# GTOs.forward stays cheap for dense cells.
+defaults:
+  - override /data: lematrho
+
+name: lematrho
+
+# BOA runs its initial-guess pre-training as a separate Trainer that builds a
+# second datamodule on the same LMDB. py-lmdb refuses to open one environment
+# twice in a process, so the LMDB-backed path cannot pre-train in the same run.
+# We disable the separate phase; the initial-guess module still trains jointly
+# inside the main loop (it is part of the BOA net's forward).
+initial_guess_pre_training_steps: 0
+
+model:
+  pbc: true
+  orb_cutoff: 3.0
+  # Linear basis: sum independent atom-centered orbital contributions. BOA's
+  # non-linear (pairwise) orbital_inference matches every edge with its flip by
+  # unique-column bookkeeping that produces out-of-bounds indices under
+  # pbc=true (it was written for the non-periodic QM9 case). The linear path
+  # (orbital_inference_linear) has a clean periodic summation, so we use it for
+  # periodic solids.
+  linear_basis: true