refactor(ml): move inference behind a ModelPrediction seam

[junwen94]

Context

Core carries a hand-ported copy of the QRF95 feature pipeline. ml/features.py,
ml/kdistance_features.py, ml/cgcnn.py, ml/cgcnn_graph.py,
ml/atom_features.py, and ml/metallicity.py are about 700 lines that exist in
stfc/goldilocks-ml as well, because that is where they were trained. They have
already drifted: the same feature contract is called
comp_struct_soap_lattice_metal.v1 in goldilocks-ml, and
qrf_comp_struct_soap_lattice_metal — no version — in Core's
model_registry.toml and in both published PSDI records.

model_registry.toml looks configuration-driven but is not. soap_r_cut,
element_property_preset, metallicity_graph_radius and the rest only mean
anything because Core hardcodes one pipeline that consumes them. A model with
different descriptors needs new Core code, and Core has to know that predicting
a k-mesh requires loading a CGCNN metallicity checkpoint — a training-time
implementation detail of one model's feature vector.

This blocks the models on the roadmap. Metallicity, magnetism, smearing, and
exchange-correlation will each need a model, and on the current shape each one
lands as new featurisation code inside Core.

Relationship to #168. #168 owns how asset files reach the disk: the asset
store, domain registries, install/verify, and container profiles. This issue
owns what happens after that — who owns featurisation and inference, and what
crosses the Core/ml boundary. Nothing here proposes a new acquisition path.
The two reinforce each other: model.json now carries the feature
configuration and the digest-pinned supporting artifacts, so the ML registry
entry #168 is designing shrinks to a task, a source, and a location.

Relationship to #134. Deleting goldilocks_core.ml.* removes most of what
that issue defers.

Problem

1. Featurisation is duplicated between two repositories and has already drifted
   in the contract name.
2. Core must know model-internal details (a metallicity checkpoint is an input
   to the k-mesh feature vector).
3. Publishing a model requires a Core release, because the registry ships in
   Core's wheel.
4. advisors/kdistance_advisor.py and advisors/kindex_advisor.py implement
   the same step — a predicted quantity to a mesh — two different ways.
5. Nothing checks that an artifact and the installed code agree. A k-distance
   trained under one 2π convention and read under the other produces a
   plausible, wrong mesh silently.

Proposed approach

One prediction type, defined by Core

Core is the consumer, so Core owns the vocabulary. goldilocks-ml satisfies it
structurally and does not import Core, so there is no dependency cycle.

@dataclass(frozen=True, slots=True)
class ModelPrediction:
    parameter: str          # k_points | smearing | magnetism | xc | ...
    quantity: str           # k_distance | k_index | degauss | functional | ...
    value: float | int | bool | str | tuple[Any, ...]
    target_contract: str
    model_id: str
    confidence: float | None = None
    details: JsonDict | None = None      # recorded into Provenance.details
    warnings: tuple[str, ...] = ()       # recorded into Provenance.warnings

parameter matches the field in ParameterAdvice, so routing is a lookup, not
a branch per model. One type covers every advised parameter; a model for one
nothing covers yet adds a contract row and a resolver, not a new type.

The mesh ladder is Core's; a model picks a rung

build_kmesh_entries already enumerates every reachable mesh for a structure
and labels each with k_index, k_distance_interval, k_line_density_interval
and k_pra. That ladder is structure-determined and model-independent. One
resolver replaces both advisors:

def resolve_prediction(structure, prediction) -> KPointSelection

Two conventions must live in exactly one place there. k_distance rises as the
mesh coarsens while k_index, k_pra and k_line_density rise as it refines;
and k_distance_to_mesh uses the 2π reciprocal lattice while
mesh_to_k_line_density_interval uses the crystallographic one. Both are in
kmesh.py today and both are easy to get backwards per advisor.

The set of mesh-determining quantities is bounded by physics, unlike feature
pipelines, so Core owning it does not recreate the original problem.

A prediction is one value, never an interval

KPointSelection carries one grid, CalculationIntent has no conservatism
knob, and the study behind QRF95 reported its metrics against the median. The
interval currently survives only as prose inside provenance.reason. Choosing
which point to publish is a modelling decision; it stays with the model.

Uncertainty is demoted, not discarded. goldilocks-ml puts the interval in
details and, where a prediction is far wider than calibration led it to
expect, a sentence in warnings. Core copies both into Provenance verbatim
and branches on neither, so the judgement stays where the calibration is known.

Contracts are checked at load

load_model in goldilocks-ml already refuses an artifact whose target contract
has no defined parameter, whose trainer this build cannot serve, whose feature
schema it does not provide, whose pinned artifact digests do not match, or
whose estimator width disagrees with its own record. Each names what is
missing. An old install and a new artifact fail loudly instead of predicting
from the wrong vector.

Two upgrade paths, both leaving Core untouched

• A retrained model on the same feature contract is a data change: point the
  registry at the new record.
• A new feature contract requires a goldilocks-ml release, and the error
  names it. Core still does not change.

There is no third option. Shipping code inside the artifact means unpickling
arbitrary code, and a declarative feature DSL is what model_registry.toml
already attempts.

Dependencies

Core's base install already carries dscribe, matminer, pymatgen,
scikit-learn, sklearn-quantile, torch, and torch-geometric by the
decision in #32. Depending on goldilocks-ml[qrf95] replaces those direct pins
rather than adding to them, and removes the three-way manual synchronisation of
scikit-learn==1.7.2 across Core, goldilocks-ml, and the deposit manifests.
goldilocks_ml.inference imports on a base install, so ModelPrediction and
the contract table are readable without the scientific stack.

Phases

• Define the vocabulary. Add ModelPrediction and the StructureModel
  protocol to contracts.py, plus the parameter-to-resolver table. No
  behaviour change.
• Add the adapter path. Depend on goldilocks-ml[qrf95]; add a k-mesh
  advisor that calls load_model and resolve_prediction. Keep the
  existing path in place.
• Verify equivalence. Run both paths over the same structures and
  confirm identical meshes before anything is deleted.
• Unify the resolvers. Replace kdistance_advisor and kindex_advisor
  with the single ladder lookup, monotonicity handled once.
• Delete the duplicate. Remove goldilocks_core.ml.* featurisation and
  model loading with their tests, and drop the now-unused direct
  dependencies.
• Move the registry out of the wheel. Resolve in order: explicit path or
  GOLDILOCKS_MODEL_REGISTRY, user config, the registry shipped with
  goldilocks-ml, then Core's built-in fallback. Coordinate the file's shape
  with feat(assets): unify model and pseudopotential acquisition #168.

Acceptance criteria

• Core contains no matminer, SOAP, or CGCNN code.
• Core never names a metallicity checkpoint; the model's record declares it.
• An artifact whose contract Core cannot honour fails at load, naming what
  is missing, and a regression test covers each refusal.
• Publishing a retrained model on an existing feature contract changes no
  Python in either repository.
• import goldilocks_core loads no goldilocks_core.ml.* module (perf: lazy import goldilocks_core (defer ml.* / scipy) #134).
• Both paths produce identical meshes over the example structures before
  the old one is removed.
• uv run pytest, uv run ruff check src tests, and
  uv run pre-commit run --all-files are green.

Upstream state

The goldilocks-ml side is implemented on feat/2-qrf95-training
(stfc/goldilocks-ml#9): goldilocks_ml.inference with ModelPrediction,
load_model, the predictor registry, digest-verified artifact resolution, and
the contract checks. It needs a release before Core can depend on it.

stfc/goldilocks-ml#10 tracks a licence conflict in the ported code that must
be settled before either repository redistributes it.

---

Written by an agent on behalf of Junwen Yin.

[junwen94]

Two corrections after re-reading the tree.

The duplicated surface is larger than stated. goldilocks_core/ml/ is 1034
lines across ten modules, not ~700 across six. The list above misses
ml/inference.py, ml/models.py, and ml/model_registry.py, which are the
loading and dispatch layer rather than featurisation, but they go the same way.

features.py           236    kdistance_features.py  211    cgcnn.py     141
model_registry.py     118    models.py              120    metallicity.py 60
cgcnn_graph.py         53    atom_features.py        39    inference.py   46
__init__.py            10

advisors/kindex_advisor.py is the clearer example of the boundary problem
than the k-distance one. advise_kpoints() is three lines of pipeline that
belong on the other side:

features = extract_cslr_features(structure)
model = load_model(spec)
predicted_k_index = predict(model, features)

Everything after those three lines — generate_candidate_k_distances,
build_kmesh_entries, _select_kmesh_entry — is exactly what Core should keep.
The file is a clean cut line: delete the top, keep the bottom, and take the
predicted index from a ModelPrediction instead.

One naming note for whoever picks this up. The extra in goldilocks-ml is
being renamed from qrf95 to models, since it installs the stack for every
model rather than one of them. goldilocks-ml[qrf95] in the checklist above
will need to become goldilocks-ml[models].

[junwen94]

One more thing worth folding in, since this issue is what replaces the registry.

Core spells the k-points parameter two ways. ml/inference.py's
ContractSpec carries parameter="k_points", while model_registry.toml keys
its default off kpoints:

[defaults.kpoints]
[defaults.kpoints.features]
[defaults.kpoints.metallicity]

and ml/model_registry.py:74 reads data["defaults"]["kpoints"]. Nothing
breaks today because the two never meet, but the seam is exactly where they
start to: a prediction arrives naming its parameter, and Core routes on that
name.

goldilocks-ml has settled on k_points for the first segment of a release
name — k_points.k_distance.qrf.goldilocks_kdist_ultra.v1 — matching the
ContractSpec field rather than the registry key. Retiring
model_registry.toml in favour of a catalogue lookup is the moment to drop
kpoints and leave one spelling.

Worth stating in the acceptance criteria, since a silent mismatch here produces
"no advisor for this parameter" rather than an error anyone can read.

[junwen94]

Phase 1, and the metallicity slice of phase 2, are done on feat/174-model-prediction-seam — and this is no longer a local-only experiment, since goldilocks-ml 0.1.0 published today and the classifier's PSDI record (ba06w-n6a68) is public.

Phase 1: vocabulary. ModelPrediction, StructureModel, and PREDICTION_RESOLVERS landed in contracts.py, matching this issue's proposed shapes. No behaviour change.

Phase 2, metallicity only. runtime/models.py's MetallicityModel now calls goldilocks_ml.inference.load_model(...).predict(structure) and routes the result through PREDICTION_RESOLVERS["metallicity"], instead of hand-loading a CGCNN checkpoint and taking argmax. ml/qrf/metallicity.py's classify_metallicity/_electronic_character_from_probabilities are gone. Runtime() with no override now resolves the model from the asset store by default (new models/metallicity-is-metal asset, digest-pinned to the PSDI record) and raises AssetNotInstalled if it's missing — the same contract the QRF backend already has, no silent heuristic fallback. k-distance QRF (SOAP/matminer, the mesh ladder) is untouched, deliberately out of scope here.

Two bugs the old code had, structurally gone now that goldilocks-ml owns the decision:

1. The CGCNN checkpoint was a registered default asset and got downloaded, but its path was never threaded from Runtime into MetallicityModel at either the CLI or server entrypoint — production electronic_character was always the composition heuristic, never the model.
2. classify_metallicity took argmax — an implicit 0.5 threshold — with no path to read a decision threshold at all. The model's record documents threshold: 0.0478, deliberately low to catch ≥97% of true metals; argmax silently discarded that.

Ships clean. pyproject.toml depends on goldilocks-ml>=0.1.0 from PyPI, no path override. uv sync installs the real wheel; verified goldilocks_ml.__file__ resolves to site-packages, not a checkout. Full suite green (369 passed) against the published package, and a direct run against the real PSDI files confirms the same results as every earlier local check: Si → insulator (0.0094), Fe/Pt → metal (0.99+).

Test strategy: an autouse fixture stubs the model deterministically for the rest of the suite (offline, agrees with the composition heuristic on which elements are metallic); a new tests/integration/test_metallicity_asset.py installs the real PSDI asset over the network and skips rather than fails when unreachable — the mechanism to self-verify against a live model as more get added, without committing weights to the repo.

One unrelated fix along the way: assets/download.py's retry policy didn't cover 500, and PSDI's S3-compatible backend answers this specific file with a transient 500 often enough to notice — sometimes clearing on the very next request, sometimes surviving 3 retries. Added 500 to the retry list; it helps but doesn't fully cover the flakiness observed today, which looks backend-side rather than anything on Core's end.

Branch: feat/174-model-prediction-seam (pushed, not yet a PR).

[junwen94]

Working on the QRF95 record (q3bye-wep37) from the goldilocks-ml side. Four
things this issue does not yet record, one of which is a decision that has to be
made before the "delete the duplicate" phase.

The seam works today, against the published record

The "Upstream state" section is stale. goldilocks-ml 0.1.0 is tagged, on PyPI,
and feat/174-model-prediction-seam already declares goldilocks-ml>=0.1.0.

More usefully: the published record already carries model.json. The PSDI
files API lists six entries for q3bye-wep37, and model.json (10218 bytes)
is byte-identical to deposits/k_points/k_distance/qrf/model.json in
goldilocks-ml, as are README.md and manifest.json. Loading that directory
through the seam:

load_model(record_dir) -> QRF95Predictor
model_id: k_points.k_distance.qrf@comp_struct_soap_lattice_metal.v1

Si primitive (2 atoms)     value=0.2174 1/angstrom   details.interval=[0.1117, 0.5034]
Si conventional (8 atoms)  value=0.2191 1/angstrom   details.interval=[0.1117, 0.5034]
parameter=k_points  quantity=k_distance
contract=goldilocks.k_distance.mesh_lower_bound.2pi.v1

So the record is self-describing on the wire, not only in the repository. That
matters for the last phase: "move the registry out of the wheel" can read the
record instead of shipping a description of it.

Refreshed inventory

The body names ml/features.py, ml/kdistance_features.py and about 700
lines. The layout has changed and the duplicate has grown to 1184 lines on
upstream/main:

  207  ml/kindex/features.py         CSLR features, second copy
  200  ml/qrf/features.py            SOAP + Magpie, second copy
  154  ml/model_registry.py
  147  ml/qrf/inference.py           unpickle, quantiles, correction
  146  ml/qrf/cgcnn.py               CGCNN architecture, second copy
   96  ml/qrf/metallicity.py
   67  ml/registry.toml
   46  ml/qrf/cgcnn_graph.py
   32  ml/kindex/inference.py
   27  ml/qrf/atom_features.py

ml/kindex/ was not in the original inventory. The k-index model duplicated
the pattern rather than using the seam, so the count grows with each model
until this lands.

Two numbers that exist only in Core

This is the part that needs a decision. registry.toml holds:

[defaults.kpoints]
interval_confidence = 0.90

[defaults.kpoints.calibration]
correction = -0.0016

model.json for the same model records "calibration": null. These are fitted
quantities — a conformal correction and a coverage claim — living in the
consumer's config, and the published record does not know they exist. Deleting
goldilocks_core.ml.* deletes them.

The good news for the equivalence phase: correction is applied only to the
interval endpoints,

result = (float(median), float(lower - correction), float(upper + correction))

so it does not move the median, and k_distance_to_mesh(structure, median)
never sees it. Meshes are unaffected either way. (Worth noting the sign
while it is being decided: at -0.0016 this narrows the interval by 0.0016 on
each side rather than widening it, which is the opposite of what a conformal
correction usually does.)

interval_confidence = 0.90 is a different problem. It is the consumer
asserting the model's coverage. goldilocks-ml returns confidence=None for
this artifact and attaches a warning saying no conformal calibration is
recorded, so a prediction that today reports 0.90 confidence would start
reporting none.

Neither is a blocker for the mesh, both are provenance changes, and both want
an explicit decision rather than being dropped by the deletion: either they are
republished into the record as real calibration, or they are declared obsolete
with the interval. This is a modelling call, so it belongs upstream — happy to
take it in goldilocks-ml if that is the preference.

The name drift is inside a single record

The body says the contract is called qrf_comp_struct_soap_lattice_metal "in
Core's model_registry.toml and in both published PSDI records". It is
narrower and worse than that: q3bye-wep37 disagrees with itself.

  manifest.json   inference_requirements.feature_contract  qrf_comp_struct_soap_lattice_metal
  model.json      feature_schema                           comp_struct_soap_lattice_metal.v1

Only model.json is load-bearing — load_model compares it against the
installed contract and refuses a mismatch — so the drift is inert today. It
still means the record ships two names for one thing, and fixing it is a
goldilocks-ml republish, not a Core change.

For what it is worth, the values have not drifted at all. Every feature
parameter Core configures matches the goldilocks-ml default exactly:

  soap_r_cut                  10.0    10.0     qrf/features.py SOAP_DEFAULTS
  soap_n_max                     8       8
  soap_l_max                     6       6
  soap_sigma                   1.0     1.0
  lattice_symprec             0.01    0.01     qrf/features.py
  metallicity_graph_radius    10.0    10.0     qrf/embedding.py RADIUS
  metallicity_max_neighbors     12      12     qrf/embedding.py MAX_NEIGHBORS

Seven for seven, maintained by hand with nothing checking it. The equivalence
phase should pass on the first attempt; that it would is not currently
guaranteed by anything.

Artifacts still come from HuggingFace

Possibly #168's, recording it here because it sits in the same file. [defaults.kpoints]
and [defaults.kpoints.metallicity] resolve to huggingface.co/STFC-SCD/...
and huggingface.co/JunwenYin/..., pinned by revision. goldilocks-ml's
deposits/README.md states PSDI is the sole publication target and that
runtime consumers resolve PSDI record IDs with no secondary host. The SHA-256s
match the PSDI manifest, so the bytes are the same file; the source of truth in
config is not.

[defaults.electronic_character] on the seam branch is already in the target
shape — a target_contract, three PSDI URLs, one of them
role = "model_record". Applying that shape to [defaults.kpoints] is what
turns this from prose into a diff.

Written by an agent on behalf of Junwen Yin.