Python API¶
TrainCraft exposes a curated public API via import traincraft. Every stage
that the CLI runs is available as a standalone Python function.
Top-level functions¶
import traincraft as tc
# Config
cfg = tc.load_config("my_run.toml") # → TrainCraftConfig
cfg2 = tc.loads_config(toml_string) # → TrainCraftConfig
# Pipeline stages (pure functions — same ones the CLI calls)
seeds = tc.build_geometries(cfg.geometry) # N seed structures
structure = tc.build_geometry(cfg.geometry) # exactly-one convenience
calc = tc.make_calculator(cfg.calculator)
frames = tc.run_sampling(structure, calc, job, cfg.sampling)
selected = tc.run_funnel(frames, cfg.selection)
labeled = tc.label_frames(selected, cfg.labeling.calculator, out_dir=out)
result = tc.run_training(labeled, cfg.training, job)
quality = tc.run_validation(labeled, calc_cfg, cfg.validation, out_dir=out)
summary = tc.run_pipeline(cfg) # the whole pipeline
# Dataset IO
tc.write_frames("out.extxyz", frames)
frames = tc.read_frames("dataset.extxyz")
Structure¶
Structure
dataclass
¶
Source code in src/traincraft/core/structure.py
n_fragments
property
¶
Number of distinct mobile fragments (excludes framework atoms).
from_ase
classmethod
¶
to_ase
¶
Return a copy of the atoms with properties/provenance in info.
Source code in src/traincraft/core/structure.py
copy
¶
set_fragments
¶
to_pymatgen
¶
from_pymatgen
classmethod
¶
Build a :class:Structure from a pymatgen Structure/Molecule.
to_rdkit
¶
Return an RDKit Mol with bonds perceived (non-periodic only).
from_rdkit
classmethod
¶
Build a :class:Structure from one conformer of an RDKit Mol.
Source code in src/traincraft/core/structure.py
Provenance¶
Provenance
dataclass
¶
Source code in src/traincraft/core/provenance.py
Workspace and Job¶
Workspace
¶
Owns an absolute run directory and hands out sub-directories/jobs.
Source code in src/traincraft/core/workspace.py
subdir
¶
Job
dataclass
¶
Source code in src/traincraft/core/workspace.py
Geometry¶
build_geometries
¶
Resolve a :class:GeometryConfig into its list of seed structures.
Source code in src/traincraft/geometry/__init__.py
build_geometry
¶
Resolve a config declaring exactly one structure (raise otherwise).
Convenience for callers that need a single seed; the pipeline itself uses
:func:build_geometries so nothing is silently dropped.
Source code in src/traincraft/geometry/__init__.py
Converter¶
ase_to_pymatgen
¶
Convert ASE Atoms to a pymatgen Structure (periodic) or Molecule.
The choice is driven by periodicity: an Atoms periodic in all three
directions becomes a Structure; otherwise a Molecule (the cell is
dropped, since a partially periodic slab/wire has no pymatgen analogue).
Source code in src/traincraft/core/converter.py
pymatgen_to_ase
¶
Convert a pymatgen Structure or Molecule to ASE Atoms.
ase_to_rdkit
¶
Convert non-periodic ASE Atoms to an RDKit Mol with bonds perceived.
Bonds are inferred from the 3D geometry by RDKit's DetermineBonds (the
xyz2mol algorithm). Raises if the structure is periodic in any direction.
Source code in src/traincraft/core/converter.py
rdkit_to_ase
¶
Convert one conformer of an RDKit Mol to ASE Atoms.
conf_id selects which embedded conformer to read (default: the first).
Source code in src/traincraft/core/converter.py
Fragment helpers¶
get_fragments
¶
Return the per-atom fragment array, or None if unset.
set_fragments
¶
Attach/overwrite the per-atom fragment array (length must equal len(atoms)).
Source code in src/traincraft/core/fragments.py
infer_fragments
¶
infer_fragments(atoms: Atoms, scale: float = 1.2, framework_mask: ndarray | None = None) -> np.ndarray
Assign fragment ids by connected components of a covalent-radius graph.
Two atoms bond if distance < scale * (r_cov[i] + r_cov[j]).
framework_mask (optional bool array, length == len(atoms)): atoms marked
True are forced to FRAMEWORK (-1) and excluded from the connectivity graph.
Returns the array; does NOT mutate atoms.
Source code in src/traincraft/core/fragments.py
fragment_ids
¶
Sorted list of mobile fragment ids (excludes FRAMEWORK == -1).
fragment_mask
¶
Boolean mask selecting atoms of fragment fid.
Registry¶
register
¶
Decorator: register obj under (kind, name).
Source code in src/traincraft/core/registry.py
get
¶
available
¶
Dataset¶
Dataset
¶
Source code in src/traincraft/datasets/dataset.py
write_frames
¶
Source code in src/traincraft/datasets/io.py
read_frames
¶
Source code in src/traincraft/datasets/io.py
Validation¶
run_validation
¶
Predict every reference property on frames with calc_cfg and
compare; write quality_report.json under out_dir when given.
A frame whose core (energy/forces) prediction fails is skipped and recorded, consistent with the labeling stage; optional properties (stress, dipole, polarizability) that the calculator cannot produce are dropped per property.
Source code in src/traincraft/validation/validate.py
141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 | |
ValidationResult
dataclass
¶
Outcome of a validation run (mirrors quality_report.json).
Source code in src/traincraft/validation/validate.py
load_reference_frames
¶
Read labeled frames from either TrainCraft (tc_*) or training
(REF_*) extended-XYZ files.
Source code in src/traincraft/validation/validate.py
resolve_calculator
¶
The model to validate: explicit [validation.calculator], else the
model the train stage recorded in model/manifest.json.
Source code in src/traincraft/validation/validate.py
Run state¶
Event logs + run index (see the CLI's runs/status
for the command-line view). One events.jsonl per run is the source of truth;
the SQLite index is a disposable cache.
run_status
¶
Fold a run's event log into its current status.
Returns {"name", "status", "engine", "planned", "stages": {stage: {...}},
"started", "updated"} where status is one of empty | running |
failed | completed | partial (partial: some stages done, plan unknown or
unfinished, nothing currently running).
Source code in src/traincraft/state.py
RunIndex
¶
SQLite index over every run below outdir — a disposable cache.
refresh() re-derives rows only for runs whose events.jsonl changed
since the last refresh, and drops rows whose run directories vanished.
Only ever open this on the submit/login side (single writer); compute jobs
write event logs, never the index.
Source code in src/traincraft/state.py
refresh
¶
Sync the index with the event logs on disk; return #rows updated.
Source code in src/traincraft/state.py
runs
¶
All indexed runs, most recently updated first.
Source code in src/traincraft/state.py
log_event
¶
Append one event record to the run's log (best-effort: never raises OSError).
Source code in src/traincraft/state.py
read_events
¶
All events of a run, oldest first (empty if the run has no log).
Source code in src/traincraft/state.py
Service layer¶
The HTTP app and MCP server are thin wrappers over traincraft.service.core
(see Serve Runs to Agents & UIs).
create_app
¶
Source code in src/traincraft/service/app.py
create_mcp
¶
Source code in src/traincraft/service/mcp_server.py
list_runs
¶
Every run under outdir (index refreshed first), newest first.
Source code in src/traincraft/service/core.py
get_run
¶
get_report
¶
The run's validation quality_report.json (metrics, checks, passed).
Source code in src/traincraft/service/core.py
get_manifest
¶
A stage's manifest.json (stages that write one: label, train).
Source code in src/traincraft/service/core.py
validate_config_text
¶
Validate a TOML config string; never raises.
Returns {"valid": bool, "name", "stages", "errors"} — the same check as
traincraft validate, but on config text so callers need no file.
Source code in src/traincraft/service/core.py
submit_config
¶
Start the pipeline for a config file; return immediately.
Slurm configs are submitted as dependency-chained jobs; anything else runs
as a detached local process (output to <run>/local.log). Either way
progress is observable through the run's event log (get_run).