Reorganize ml/ by model family + feature-set registry as more models land

[junwen94]

Cleanup / design: reorganize ml/ as more models land, and introduce a feature-set registry keyed on ModelSpec.feature_set.

ml/ positioning (for reference)

ml/ is core's inference machinery — everything needed to turn a Structure into a model prediction, for models trained elsewhere (goldilocks-models). Core loads and runs models; it does not train them. It holds four concerns:

• feature extraction (structure → numbers): features.py (CSLR), kdistance_features.py (QRF comp/struct/soap/lattice + assembly), atom_features.py (atom embeddings)
• model architectures: cgcnn.py, cgcnn_graph.py
• model loading (artifact → object): models.py (joblib/HF), metallicity.py (torch ckpt)
• inference: inference.py (predict)

Dependency direction is one-way: advisors/ (stage backends) → ml/ (reusable blocks). ml/ never imports pipeline stages. This keeps "run a model" decoupled from "plug a model into the pipeline."

Problem

The QRF k-points port (#31) grew ml/ to a flat directory that mixes generic infra (models.py, inference.py) with model-specific code (CGCNN architecture, QRF vs CSLR feature pipelines). As more models arrive (ALIGNN k-points #30, magnetic #34), the flat layout will sprawl and it will be unclear which files belong to which model.

This is the "different models need different feature sets" theme (#33/#34) showing up in the directory structure.

Proposal (do when the 3rd model family lands; not urgent)

1. Group by model family (or by concern), e.g.

   ml/
     models.py inference.py         # generic loading + inference
     kdistance/  # QRF: features + cgcnn/graph/atom_features it uses
     cslr/       # k-index features
     metallicity/  # CGCNN metallicity (also serves #34)
   

   (CGCNN architecture is shared by kdistance metal-features and metallicity — factor it so it is not duplicated.)
2. feature-set registry: a map ModelSpec.feature_set -> extractor so an advisor resolves the right feature pipeline from the spec, and adding a model = registering a feature set (no pipeline change). This is the mechanism referenced in Architecture: separate k-point sampling mode (task-driven) from density (structure/ML), and where ML k-point prediction belongs #33.

Constraints

• Keep import-stable APIs for goldilocks-models (kmesh.build_kmesh_entries, the feature-extraction entry point). A reorg must preserve or re-export these.
• Upstream-facing (shared module layout) — coordinate.

Low priority cleanup. Related: #26, #30, #31, #33, #34.

---

Written by an agent on behalf of Junwen Yin.