Training new emulators¶

This notebook shows how to fit and register a new artifact. Unlike predicting, training is COSMA-bound: it reads per-realisation measurements through the halocat loaders, so it needs that package and the simulation data. Heavy fits belong on slurm (account dp004, partition cosma8-serial) — never the login node.

The cells here are illustrative and not executed (they would require the data). Run them in the cosemu env on COSMA.

The training pipeline¶

Every trained property is a plugin with the same four-step flow, which the CLI train verb runs end to end:

assemble(imodels, iboxes)  ->  build(bundle)  ->  validate(artifact, bundle)  ->  register(...)

• assemble — pull the measured training set from halocat into a TrainingBundle (X = θ design, Y = the measured quantity, Yerr = across-box SEM, the coordinate grid, and meta).
• build — fit the artifact (a weighted-PCA + per-component GP, or a boost).
• validate — compute the accuracy block (interior LOO, fiducial OOS).
• register — pickle the artifact, write meta.json, pin dependency hashes, and update manifest.json.

The easy way: the CLI¶

For a registered property at a (gravity, redshift), one command does everything. Submit it via slurm:

micromamba run -n cosemu python3 -m haloemu.cli train hmf \
    --gravity LCDM --redshift 0.25 --imodels 1-64 --iboxes 1-5

--imodels / --iboxes select the design models and boxes. The haloemu/properties/*.slurm scripts are worked examples that chain measure → train → validate → plot with --dependency=afterok.

The programmatic way¶

The CLI is a thin wrapper over the plugin. Doing it by hand gives you the bundle, the artifact, and the accuracy dict to inspect before registering.

In [ ]:

import os
os.environ.setdefault("JAX_PLATFORMS", "cpu")

from haloemu.properties import get_plugin_class
from haloemu.core.util import parse_int_list
from haloemu.registry import Registry

prop, gravity, redshift = "hmf", "LCDM", 0.25
cls = get_plugin_class(prop)                 # resolve the plugin
plugin = cls(gravity, redshift)

# 1. assemble the training set from halocat (needs the data on COSMA)
bundle = plugin.assemble(parse_int_list("1-64"), parse_int_list("1-5"))
print("X:", bundle.X.shape, "Y:", bundle.Y.shape, "coord:", bundle.coord.shape)
print("theta_keys:", bundle.theta_keys)

# 2. fit
artifact = plugin.build(bundle)

# 3. validate -> the accuracy block stored in the manifest
accuracy = plugin.validate(artifact, bundle)
print("accuracy:", accuracy)

import os os.environ.setdefault("JAX_PLATFORMS", "cpu") from haloemu.properties import get_plugin_class from haloemu.core.util import parse_int_list from haloemu.registry import Registry prop, gravity, redshift = "hmf", "LCDM", 0.25 cls = get_plugin_class(prop) # resolve the plugin plugin = cls(gravity, redshift) # 1. assemble the training set from halocat (needs the data on COSMA) bundle = plugin.assemble(parse_int_list("1-64"), parse_int_list("1-5")) print("X:", bundle.X.shape, "Y:", bundle.Y.shape, "coord:", bundle.coord.shape) print("theta_keys:", bundle.theta_keys) # 2. fit artifact = plugin.build(bundle) # 3. validate -> the accuracy block stored in the manifest accuracy = plugin.validate(artifact, bundle) print("accuracy:", accuracy)

In [ ]:

# 4. register: persists the pickle + meta.json + manifest entry.
#    Use a scratch root so you do not overwrite the shipped artifacts while
#    experimenting; drop `root=` to register into the package.
reg = Registry(root="/tmp/haloemu_artifacts_demo")
entry = reg.register(prop, gravity, redshift,
                     artifact=artifact, kind="trained",
                     accuracy=accuracy)
print("registered", entry["key"], "->", entry["sha256"][:12])

# predict from the freshly trained artifact
import numpy as np
print(np.asarray(reg.predict(prop, gravity, redshift,
                             [0.31, 0.677, 0.967, 0.83])).shape)

# 4. register: persists the pickle + meta.json + manifest entry. # Use a scratch root so you do not overwrite the shipped artifacts while # experimenting; drop `root=` to register into the package. reg = Registry(root="/tmp/haloemu_artifacts_demo") entry = reg.register(prop, gravity, redshift, artifact=artifact, kind="trained", accuracy=accuracy) print("registered", entry["key"], "->", entry["sha256"][:12]) # predict from the freshly trained artifact import numpy as np print(np.asarray(reg.predict(prop, gravity, redshift, [0.31, 0.677, 0.967, 0.83])).shape)

Tuning the fit¶

The GP hyper-parameters live in each property module (e.g. velocity.py DEFAULT_HYPER). The suite-wide defaults — expsquared kernel, constant mean, and n_restarts = 10 — were chosen by the ablation gates and are optimal across the suite; matern and a linear mean are rejected. Raising n_components or n_restarts is the usual lever (the f(R) hmf screened tail was an under-convergence bug fixed by n_restarts 3 → 10).

Evidence harnesses (each with a matching *.slurm) live in haloemu/properties/: ablation_sweep.py, restarts_sweep.py, pk_xi_validation.py, fRn1_fiducial_oos.py, velocity_ablation.py, …

Training f(R) artifacts (seed-paired boosts)¶

Most fRn1 artifacts are not fit directly — they emulate the boost over the matched-seed ΛCDM twin and compose onto the pinned ΛCDM base. The plugin's assemble/build switch to the boost representation automatically when gravity="fRn1" for the boost-winning properties; b_cum and the velocity m10 moment stay direct 5-parameter GPs. Train the ΛCDM base first, then the f(R) boost (its registration pins the base's sha256).

micromamba run -n cosemu python3 -m haloemu.cli train pk_mm --gravity LCDM --redshift 0.25
micromamba run -n cosemu python3 -m haloemu.cli train pk_mm --gravity fRn1 --redshift 0.25

To improve the f(R) sampling, generate_frn1_infill.py produces an ordered maximin infill of new 5D cosmologies (biased to the screening transition); see the f(R) gravity page.

Derived properties¶

b_diff and xi_hh are derived — recipes over upstream artifacts, with no weights to train. haloemu train b_diff exits with a note; instead they are "built" by re-validating against their (re-pinned) dependencies. Retrain their upstreams (b_cum, hmf, xi_mm, r_ab, xi_hh_smallr) and the derived predictions follow.

Where it lands¶

haloemu/artifacts/<property>/<gravity>/z<z:.2f>/{emulator.pkl, meta.json}
haloemu/manifest.json     # one entry per artifact: accuracy, deps, code_version

Inspect with haloemu list or reg.entry(prop, gravity, z). A new artifact is only "shipped" once its gate passes — see Accuracy & gates.