Skip to content

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
@dataclass
class Structure:
    atoms: Atoms
    properties: dict[str, Any] = field(default_factory=dict)
    provenance: Provenance = field(default_factory=Provenance)

    @property
    def hash(self) -> str:
        """Content hash from composition + geometry (rounded for stability)."""
        a = self.atoms
        payload = {
            "numbers": a.get_atomic_numbers().tolist(),
            "positions": np.round(a.get_positions(), 4).tolist(),
            "cell": np.round(np.asarray(a.get_cell()), 4).tolist(),
            "pbc": a.get_pbc().tolist(),
        }
        blob = json.dumps(payload, sort_keys=True).encode()
        return hashlib.sha1(blob).hexdigest()[:16]

    @classmethod
    def from_ase(cls, atoms: Atoms, **kwargs: Any) -> Structure:
        return cls(atoms=atoms.copy(), **kwargs)

    def to_ase(self, with_properties: bool = True) -> Atoms:
        """Return a copy of the atoms with properties/provenance in ``info``."""
        atoms = self.atoms.copy()
        atoms.info["tc_provenance"] = self.provenance.to_dict()
        atoms.info["tc_hash"] = self.hash
        if with_properties:
            for key, value in self.properties.items():
                if key == "forces" and value is not None:
                    atoms.arrays["tc_forces"] = np.asarray(value)
                else:
                    atoms.info[f"tc_{key}"] = value
        return atoms

    def copy(self) -> Structure:
        return Structure(
            atoms=self.atoms.copy(),
            properties=dict(self.properties),
            provenance=Provenance.from_dict(self.provenance.to_dict()),
        )

    # --- fragment identity helpers ----------------------------------------
    @property
    def fragments(self):
        """Per-atom fragment array, or None if unset."""
        from .fragments import get_fragments
        return get_fragments(self.atoms)

    def set_fragments(self, frag) -> None:
        """Attach/overwrite the per-atom fragment array."""
        from .fragments import set_fragments
        set_fragments(self.atoms, frag)

    @property
    def n_fragments(self) -> int:
        """Number of distinct mobile fragments (excludes framework atoms)."""
        from .fragments import fragment_ids
        return len(fragment_ids(self.atoms))

    # --- interop (see core.converter) -------------------------------------
    def to_pymatgen(self):
        """Return a pymatgen ``Structure`` (periodic) or ``Molecule``."""
        from .converter import ase_to_pymatgen
        return ase_to_pymatgen(self.atoms)

    def to_rdkit(self, charge: int = 0):
        """Return an RDKit ``Mol`` with bonds perceived (non-periodic only)."""
        from .converter import ase_to_rdkit
        return ase_to_rdkit(self.atoms, charge=charge)

    @classmethod
    def from_pymatgen(cls, obj, **kwargs: Any) -> Structure:
        """Build a :class:`Structure` from a pymatgen ``Structure``/``Molecule``."""
        from .converter import pymatgen_to_ase
        return cls.from_ase(pymatgen_to_ase(obj), **kwargs)

    @classmethod
    def from_rdkit(cls, mol, conf_id: int = 0, **kwargs: Any) -> Structure:
        """Build a :class:`Structure` from one conformer of an RDKit ``Mol``."""
        from .converter import rdkit_to_ase
        return cls.from_ase(rdkit_to_ase(mol, conf_id=conf_id), **kwargs)

hash property

hash: str

Content hash from composition + geometry (rounded for stability).

fragments property

fragments

Per-atom fragment array, or None if unset.

n_fragments property

n_fragments: int

Number of distinct mobile fragments (excludes framework atoms).

from_ase classmethod

from_ase(atoms: Atoms, **kwargs: Any) -> Structure
Source code in src/traincraft/core/structure.py
@classmethod
def from_ase(cls, atoms: Atoms, **kwargs: Any) -> Structure:
    return cls(atoms=atoms.copy(), **kwargs)

to_ase

to_ase(with_properties: bool = True) -> Atoms

Return a copy of the atoms with properties/provenance in info.

Source code in src/traincraft/core/structure.py
def to_ase(self, with_properties: bool = True) -> Atoms:
    """Return a copy of the atoms with properties/provenance in ``info``."""
    atoms = self.atoms.copy()
    atoms.info["tc_provenance"] = self.provenance.to_dict()
    atoms.info["tc_hash"] = self.hash
    if with_properties:
        for key, value in self.properties.items():
            if key == "forces" and value is not None:
                atoms.arrays["tc_forces"] = np.asarray(value)
            else:
                atoms.info[f"tc_{key}"] = value
    return atoms

copy

copy() -> Structure
Source code in src/traincraft/core/structure.py
def copy(self) -> Structure:
    return Structure(
        atoms=self.atoms.copy(),
        properties=dict(self.properties),
        provenance=Provenance.from_dict(self.provenance.to_dict()),
    )

set_fragments

set_fragments(frag) -> None

Attach/overwrite the per-atom fragment array.

Source code in src/traincraft/core/structure.py
def set_fragments(self, frag) -> None:
    """Attach/overwrite the per-atom fragment array."""
    from .fragments import set_fragments
    set_fragments(self.atoms, frag)

to_pymatgen

to_pymatgen()

Return a pymatgen Structure (periodic) or Molecule.

Source code in src/traincraft/core/structure.py
def to_pymatgen(self):
    """Return a pymatgen ``Structure`` (periodic) or ``Molecule``."""
    from .converter import ase_to_pymatgen
    return ase_to_pymatgen(self.atoms)

from_pymatgen classmethod

from_pymatgen(obj, **kwargs: Any) -> Structure

Build a :class:Structure from a pymatgen Structure/Molecule.

Source code in src/traincraft/core/structure.py
@classmethod
def from_pymatgen(cls, obj, **kwargs: Any) -> Structure:
    """Build a :class:`Structure` from a pymatgen ``Structure``/``Molecule``."""
    from .converter import pymatgen_to_ase
    return cls.from_ase(pymatgen_to_ase(obj), **kwargs)

to_rdkit

to_rdkit(charge: int = 0)

Return an RDKit Mol with bonds perceived (non-periodic only).

Source code in src/traincraft/core/structure.py
def to_rdkit(self, charge: int = 0):
    """Return an RDKit ``Mol`` with bonds perceived (non-periodic only)."""
    from .converter import ase_to_rdkit
    return ase_to_rdkit(self.atoms, charge=charge)

from_rdkit classmethod

from_rdkit(mol, conf_id: int = 0, **kwargs: Any) -> Structure

Build a :class:Structure from one conformer of an RDKit Mol.

Source code in src/traincraft/core/structure.py
@classmethod
def from_rdkit(cls, mol, conf_id: int = 0, **kwargs: Any) -> Structure:
    """Build a :class:`Structure` from one conformer of an RDKit ``Mol``."""
    from .converter import rdkit_to_ase
    return cls.from_ase(rdkit_to_ase(mol, conf_id=conf_id), **kwargs)

Provenance

Provenance dataclass

Source code in src/traincraft/core/provenance.py
@dataclass
class Provenance:
    origin: str = "generated"
    source: str | None = None  # e.g. "builder:nanotube", "source:file"
    transforms: list[str] = field(default_factory=list)
    calculator: str | None = None  # method that produced ``properties``
    level_of_theory: dict[str, Any] = field(default_factory=dict)
    seed: int | None = None
    parents: list[str] = field(default_factory=list)  # parent structure hashes
    extra: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.origin not in ORIGINS:
            raise ValueError(f"origin must be one of {ORIGINS}, got {self.origin!r}")

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Provenance:
        known = {f for f in cls.__dataclass_fields__}  # noqa: C416
        return cls(**{k: v for k, v in data.items() if k in known})

to_dict

to_dict() -> dict[str, Any]
Source code in src/traincraft/core/provenance.py
def to_dict(self) -> dict[str, Any]:
    return asdict(self)

from_dict classmethod

from_dict(data: dict[str, Any]) -> Provenance
Source code in src/traincraft/core/provenance.py
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Provenance:
    known = {f for f in cls.__dataclass_fields__}  # noqa: C416
    return cls(**{k: v for k, v in data.items() if k in known})

Workspace and Job

Workspace

Owns an absolute run directory and hands out sub-directories/jobs.

Source code in src/traincraft/core/workspace.py
class Workspace:
    """Owns an absolute run directory and hands out sub-directories/jobs."""

    def __init__(self, root: str | Path):
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)

    def subdir(self, *parts: str) -> Path:
        p = self.root.joinpath(*parts)
        p.mkdir(parents=True, exist_ok=True)
        return p

    def job(self, *parts: str) -> Job:
        return Job(dir=self.subdir(*parts))

subdir

subdir(*parts: str) -> Path
Source code in src/traincraft/core/workspace.py
def subdir(self, *parts: str) -> Path:
    p = self.root.joinpath(*parts)
    p.mkdir(parents=True, exist_ok=True)
    return p

job

job(*parts: str) -> Job
Source code in src/traincraft/core/workspace.py
def job(self, *parts: str) -> Job:
    return Job(dir=self.subdir(*parts))

Job dataclass

Source code in src/traincraft/core/workspace.py
@dataclass
class Job:
    dir: Path

    @property
    def marker(self) -> Path:
        return self.dir / ".tc_done"

    def done(self) -> bool:
        return self.marker.exists()

    def mark_done(self) -> None:
        self.marker.write_text("ok\n")

    def path(self, *parts: str) -> Path:
        return self.dir.joinpath(*parts)

path

path(*parts: str) -> Path
Source code in src/traincraft/core/workspace.py
def path(self, *parts: str) -> Path:
    return self.dir.joinpath(*parts)

Geometry

build_geometries

build_geometries(geom_cfg) -> list[Structure]

Resolve a :class:GeometryConfig into its list of seed structures.

Source code in src/traincraft/geometry/__init__.py
def build_geometries(geom_cfg) -> list[Structure]:
    """Resolve a :class:`GeometryConfig` into its list of seed structures."""
    if geom_cfg.replicate == 1:
        return _build_once(geom_cfg, None)
    structures: list[Structure] = []
    for i in range(geom_cfg.replicate):
        structures.extend(_build_once(geom_cfg, i))
    return structures

build_geometry

build_geometry(geom_cfg) -> Structure

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
def build_geometry(geom_cfg) -> Structure:
    """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.
    """
    structures = build_geometries(geom_cfg)
    if len(structures) != 1:
        raise ValueError(
            f"geometry expands to {len(structures)} structures "
            "(replicate/conformers/multi-frame source); use build_geometries()"
        )
    return structures[0]

Converter

ase_to_pymatgen

ase_to_pymatgen(atoms: Atoms) -> Structure | Molecule

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
def ase_to_pymatgen(atoms: Atoms) -> Structure | Molecule:
    """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).
    """
    adaptor = _require_pymatgen()
    if bool(np.all(atoms.get_pbc())):
        return adaptor.get_structure(atoms)
    return adaptor.get_molecule(atoms)

pymatgen_to_ase

pymatgen_to_ase(obj: Structure | Molecule) -> Atoms

Convert a pymatgen Structure or Molecule to ASE Atoms.

Source code in src/traincraft/core/converter.py
def pymatgen_to_ase(obj: Structure | Molecule) -> Atoms:
    """Convert a pymatgen ``Structure`` or ``Molecule`` to ASE ``Atoms``."""
    adaptor = _require_pymatgen()
    return adaptor.get_atoms(obj)

ase_to_rdkit

ase_to_rdkit(atoms: Atoms, charge: int = 0) -> Mol

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
def ase_to_rdkit(atoms: Atoms, charge: int = 0) -> Mol:
    """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.
    """
    if bool(np.any(atoms.get_pbc())):
        raise ValueError(
            "ase_to_rdkit needs a non-periodic structure; got pbc="
            f"{atoms.get_pbc().tolist()}. RDKit molecules are not periodic."
        )
    Chem, rdDetermineBonds = _require_rdkit()

    buf = io.StringIO()
    write(buf, atoms, format="xyz")
    mol = Chem.MolFromXYZBlock(buf.getvalue())
    if mol is None:
        raise ValueError("RDKit could not parse the structure as an XYZ molecule")
    try:
        rdDetermineBonds.DetermineBonds(mol, charge=charge)
    except ValueError as e:
        raise ValueError(
            f"RDKit could not perceive bonds (charge={charge}): {e}. "
            "Try passing the correct total charge."
        ) from e
    return mol

rdkit_to_ase

rdkit_to_ase(mol: Mol, conf_id: int = 0) -> Atoms

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
def rdkit_to_ase(mol: Mol, conf_id: int = 0) -> Atoms:
    """Convert one conformer of an RDKit ``Mol`` to ASE ``Atoms``.

    ``conf_id`` selects which embedded conformer to read (default: the first).
    """
    if mol.GetNumConformers() == 0:
        raise ValueError(
            "RDKit molecule has no conformers; embed one first "
            "(e.g. AllChem.EmbedMolecule)."
        )
    conf = mol.GetConformer(conf_id)
    positions = np.asarray(conf.GetPositions())
    symbols = [atom.GetSymbol() for atom in mol.GetAtoms()]
    return Atoms(symbols=symbols, positions=positions)

Fragment helpers

get_fragments

get_fragments(atoms: Atoms) -> np.ndarray | None

Return the per-atom fragment array, or None if unset.

Source code in src/traincraft/core/fragments.py
def get_fragments(atoms: Atoms) -> np.ndarray | None:
    """Return the per-atom fragment array, or None if unset."""
    if FRAGMENT_KEY in atoms.arrays:
        return atoms.arrays[FRAGMENT_KEY].astype(int)
    return None

set_fragments

set_fragments(atoms: Atoms, frag: ndarray | list[int]) -> None

Attach/overwrite the per-atom fragment array (length must equal len(atoms)).

Source code in src/traincraft/core/fragments.py
def set_fragments(atoms: Atoms, frag: np.ndarray | list[int]) -> None:
    """Attach/overwrite the per-atom fragment array (length must equal len(atoms))."""
    frag = np.asarray(frag, dtype=int)
    if frag.shape != (len(atoms),):
        raise ValueError(f"fragment array must have shape ({len(atoms)},), got {frag.shape}")
    atoms.set_array(FRAGMENT_KEY, frag)

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
def infer_fragments(
    atoms: Atoms,
    scale: float = 1.2,
    framework_mask: np.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`.
    """
    from ase.neighborlist import NeighborList, natural_cutoffs
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import connected_components

    n = len(atoms)
    result = np.full(n, FRAMEWORK, dtype=int)

    # Determine which atoms are mobile (not in the framework mask).
    mobile = np.ones(n, dtype=bool)
    if framework_mask is not None:
        framework_mask = np.asarray(framework_mask, dtype=bool)
        if framework_mask.shape != (n,):
            raise ValueError(
                f"framework_mask must have shape ({n},), got {framework_mask.shape}"
            )
        mobile[framework_mask] = False

    mobile_idx = np.where(mobile)[0]
    if len(mobile_idx) == 0:
        return result

    cutoffs = natural_cutoffs(atoms, mult=scale)
    nl = NeighborList(cutoffs, self_interaction=False, bothways=True)
    nl.update(atoms)

    # Build adjacency only among mobile atoms.
    idx_map = {int(i): j for j, i in enumerate(mobile_idx)}
    m = len(mobile_idx)
    rows, cols = [], []
    for local, global_i in enumerate(mobile_idx):
        neighbours, _ = nl.get_neighbors(global_i)
        for global_j in neighbours:
            if global_j in idx_map:
                rows.append(local)
                cols.append(idx_map[global_j])

    if rows:
        data = np.ones(len(rows), dtype=np.int8)
        adj = csr_matrix((data, (rows, cols)), shape=(m, m))
    else:
        adj = csr_matrix((m, m), dtype=np.int8)

    n_components, labels = connected_components(adj, directed=False)
    for local, global_i in enumerate(mobile_idx):
        result[global_i] = int(labels[local])

    return result

fragment_ids

fragment_ids(atoms: Atoms) -> list[int]

Sorted list of mobile fragment ids (excludes FRAMEWORK == -1).

Source code in src/traincraft/core/fragments.py
def fragment_ids(atoms: Atoms) -> list[int]:
    """Sorted list of mobile fragment ids (excludes FRAMEWORK == -1)."""
    frag = get_fragments(atoms)
    if frag is None:
        return []
    return sorted(int(i) for i in np.unique(frag) if i != FRAMEWORK)

fragment_mask

fragment_mask(atoms: Atoms, fid: int) -> np.ndarray

Boolean mask selecting atoms of fragment fid.

Source code in src/traincraft/core/fragments.py
def fragment_mask(atoms: Atoms, fid: int) -> np.ndarray:
    """Boolean mask selecting atoms of fragment `fid`."""
    frag = get_fragments(atoms)
    if frag is None:
        raise ValueError("no fragment array set on these atoms")
    return frag == fid

Registry

register

register(kind: str, name: str, *, capabilities: Iterable[str] | None = None)

Decorator: register obj under (kind, name).

Source code in src/traincraft/core/registry.py
def register(kind: str, name: str, *, capabilities: Iterable[str] | None = None):
    """Decorator: register ``obj`` under ``(kind, name)``."""

    def decorator(obj):
        bucket = _REGISTRY.setdefault(kind, {})
        if name in bucket:
            raise RegistryError(f"{kind} {name!r} is already registered")
        bucket[name] = {"obj": obj, "capabilities": set(capabilities or ())}
        return obj

    return decorator

get

get(kind: str, name: str)
Source code in src/traincraft/core/registry.py
def get(kind: str, name: str):
    try:
        return _REGISTRY[kind][name]["obj"]
    except KeyError:
        raise RegistryError(
            f"unknown {kind} {name!r}; available: {available(kind)}"
        ) from None

available

available(kind: str) -> list[str]
Source code in src/traincraft/core/registry.py
def available(kind: str) -> list[str]:
    return sorted(_REGISTRY.get(kind, {}))

Dataset

Dataset

Source code in src/traincraft/datasets/dataset.py
class Dataset:
    def __init__(self, path: str | Path):
        path = Path(path)
        if path.suffix != ".extxyz":
            path = path.with_suffix(".extxyz")
        self.path = path
        self._frames: list[Structure] = []
        self._hashes: set[str] = set()

    def append(self, structures: list[Structure]) -> int:
        """Add new frames, skipping exact duplicates. Returns count added."""
        added = 0
        for s in structures:
            h = s.hash
            if h in self._hashes:
                continue
            self._hashes.add(h)
            self._frames.append(s)
            added += 1
        return added

    def filter(self, origin: str | None = None) -> list[Structure]:
        if origin is None:
            return list(self._frames)
        return [s for s in self._frames if s.provenance.origin == origin]

    @property
    def frames(self) -> list[Structure]:
        return list(self._frames)

    def __len__(self) -> int:
        return len(self._frames)

    def write(self) -> Path:
        return write_frames(self.path, self._frames)

    @classmethod
    def load(cls, path: str | Path) -> Dataset:
        ds = cls(path)
        ds.append(read_frames(ds.path))
        return ds

append

append(structures: list[Structure]) -> int

Add new frames, skipping exact duplicates. Returns count added.

Source code in src/traincraft/datasets/dataset.py
def append(self, structures: list[Structure]) -> int:
    """Add new frames, skipping exact duplicates. Returns count added."""
    added = 0
    for s in structures:
        h = s.hash
        if h in self._hashes:
            continue
        self._hashes.add(h)
        self._frames.append(s)
        added += 1
    return added

filter

filter(origin: str | None = None) -> list[Structure]
Source code in src/traincraft/datasets/dataset.py
def filter(self, origin: str | None = None) -> list[Structure]:
    if origin is None:
        return list(self._frames)
    return [s for s in self._frames if s.provenance.origin == origin]

write

write() -> Path
Source code in src/traincraft/datasets/dataset.py
def write(self) -> Path:
    return write_frames(self.path, self._frames)

write_frames

write_frames(path: str | Path, structures: list[Structure], append: bool = False) -> Path
Source code in src/traincraft/datasets/io.py
def write_frames(path: str | Path, structures: list[Structure], append: bool = False) -> Path:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    images = [_to_atoms(s) for s in structures]
    write(str(path), images, format="extxyz", append=append)
    return path

read_frames

read_frames(path: str | Path) -> list[Structure]
Source code in src/traincraft/datasets/io.py
def read_frames(path: str | Path) -> list[Structure]:
    images = read(str(path), index=":", format="extxyz")
    if not isinstance(images, list):
        images = [images]
    out: list[Structure] = []
    for atoms in images:
        props: dict = {}
        if "tc_forces" in atoms.arrays:
            props["forces"] = np.asarray(atoms.arrays["tc_forces"])
        for key in list(atoms.info):
            if key.startswith("tc_") and key not in _META_KEYS:
                props[key[3:]] = _decode(atoms.info[key])
        raw = atoms.info.get("tc_provenance")
        prov = Provenance.from_dict(json.loads(raw)) if raw else Provenance()
        out.append(Structure(atoms=atoms, properties=props, provenance=prov))
    return out

Validation

run_validation

run_validation(frames: list[Structure], calc_cfg, cfg, *, out_dir=None) -> ValidationResult

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
def run_validation(frames: list[Structure], calc_cfg, cfg, *, out_dir=None) -> ValidationResult:
    """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.
    """
    if cfg.max_frames is not None:
        frames = frames[: cfg.max_frames]
    if not frames:
        raise ValueError("validation has no reference frames to compare against")

    calc = make_calculator(calc_cfg)
    requested = set(cfg.properties) if cfg.properties else None

    def want(prop: str) -> bool:
        return requested is None or prop in requested

    e_ref: list[float] = []
    e_pred: list[float] = []
    n_atoms: list[int] = []
    f_ref: list[np.ndarray] = []
    f_pred: list[np.ndarray] = []
    f_rmse_per_frame: list[float] = []
    symbols: list[list[str]] = []
    tensor_pairs: dict[str, list[tuple[np.ndarray, np.ndarray]]] = {}
    failures: list[dict] = []

    for i, s in enumerate(frames):
        ref = s.properties
        atoms = s.to_ase(with_properties=False)
        atoms.calc = calc
        try:
            ref_e = ref.get("energy")
            if want("energy") and ref_e is not None:
                e_pred.append(float(atoms.get_potential_energy()))
                e_ref.append(float(ref_e))
                n_atoms.append(len(atoms))
            ref_f = _as_array(ref.get("forces"))
            if want("forces") and ref_f is not None:
                pred_f = np.asarray(atoms.get_forces())
                f_pred.append(pred_f)
                f_ref.append(ref_f)
                symbols.append(atoms.get_chemical_symbols())
                f_rmse_per_frame.append(float(np.sqrt(np.mean((pred_f - ref_f) ** 2))))
        except Exception as exc:
            failures.append({"frame": i, "error": f"{type(exc).__name__}: {exc}"})
            logger.warning("validation frame %d failed, skipping: %s", i, exc)
            continue

        # Optional tensor properties: drop per property if the calculator
        # cannot produce them, without voiding the frame's E/F comparison.
        for prop in ("stress", "dipole", "polarizability"):
            ref_value = _as_array(ref.get(prop))
            if not want(prop) or ref_value is None:
                continue
            try:
                if prop == "stress":
                    if not atoms.get_pbc().any():
                        continue
                    pair = (_full_stress(ref_value), _full_stress(atoms.get_stress()))
                elif prop == "dipole":
                    pair = (ref_value.ravel(), np.asarray(atoms.get_dipole_moment()).ravel())
                else:
                    pol = getattr(calc, "results", {}).get("polarizability")
                    if pol is None:
                        continue
                    pair = (ref_value.ravel(), np.asarray(pol).ravel())
            except Exception as exc:
                logger.debug("frame %d: no %s from calculator (%s)", i, prop, exc)
                continue
            tensor_pairs.setdefault(prop, []).append(pair)

    metrics: dict = {}
    if e_ref:
        metrics.update(energy_metrics(e_ref, e_pred, n_atoms))
    if f_ref:
        metrics.update(force_metrics(f_ref, f_pred, symbols))
    for prop, pairs in tensor_pairs.items():
        metrics.update(array_metrics(prop, [r for r, _ in pairs], [p for _, p in pairs]))

    if not metrics:
        first = failures[0]["error"] if failures else "no reference properties found"
        raise RuntimeError(
            f"validation produced no comparable predictions over {len(frames)} frames "
            f"({first})"
        )

    checks, passed = _apply_thresholds(metrics, cfg.thresholds)

    report = {
        "calculator": calc_cfg.model_dump(),
        "n_frames": len(frames),
        "n_compared": len(frames) - len(failures),
        "n_failed": len(failures),
        "failures": failures,
        "properties": sorted(
            {m.split("_")[0] for m in metrics if isinstance(metrics[m], (int, float))}
        ),
        "metrics": metrics,
        "thresholds": dict(cfg.thresholds),
        "checks": checks,
        "passed": passed,
        "parity": {
            "energy": [[r, p] for r, p in zip(e_ref, e_pred, strict=True)],
            "n_atoms": n_atoms,
            "forces_rmse_per_frame": f_rmse_per_frame,
        },
    }

    report_path = None
    if out_dir is not None:
        out_dir = Path(out_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        report_path = out_dir / REPORT_NAME
        report_path.write_text(json.dumps(report, indent=2))

    return ValidationResult(
        metrics=metrics,
        checks=checks,
        passed=passed,
        n_frames=len(frames),
        n_failed=len(failures),
        report=report,
        report_path=report_path,
    )

ValidationResult dataclass

Outcome of a validation run (mirrors quality_report.json).

Source code in src/traincraft/validation/validate.py
@dataclass
class ValidationResult:
    """Outcome of a validation run (mirrors ``quality_report.json``)."""

    metrics: dict
    checks: list[dict]
    passed: bool | None  # None when no thresholds were configured
    n_frames: int
    n_failed: int
    report: dict = field(default_factory=dict)
    report_path: Path | None = None

load_reference_frames

load_reference_frames(path: str | Path) -> list[Structure]

Read labeled frames from either TrainCraft (tc_*) or training (REF_*) extended-XYZ files.

Source code in src/traincraft/validation/validate.py
def load_reference_frames(path: str | Path) -> list[Structure]:
    """Read labeled frames from either TrainCraft (``tc_*``) or training
    (``REF_*``) extended-XYZ files."""
    images = ase_read(str(path), index=":", format="extxyz")
    if not isinstance(images, list):
        images = [images]
    if images and any(
        k.startswith("REF_") for k in (*images[0].info, *images[0].arrays)
    ):
        return [Structure(atoms=a, properties=_ref_properties(a)) for a in images]
    return read_frames(path)

resolve_calculator

resolve_calculator(cfg, model_manifest: Path)

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
def resolve_calculator(cfg, model_manifest: Path):
    """The model to validate: explicit ``[validation.calculator]``, else the
    model the ``train`` stage recorded in ``model/manifest.json``."""
    if cfg.calculator is not None:
        return cfg.calculator
    if model_manifest.exists():
        manifest = json.loads(model_manifest.read_text())
        model_path = manifest.get("model_path")
        if model_path:
            from ..config.models import MaceCalc

            return MaceCalc(model_path=model_path)
    raise ValueError(
        "validation needs a model to test: set [validation.calculator] explicitly, "
        "or run a [training] stage first so the trained model is picked up from "
        f"{model_manifest}"
    )

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

run_status(run_dir: str | Path) -> dict

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
def run_status(run_dir: str | Path) -> dict:
    """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).
    """
    run_dir = Path(run_dir)
    events = read_events(run_dir)
    stages: dict[str, dict] = {}
    engine = None
    planned: list[str] | None = None
    started = updated = None

    for e in events:
        ts = e.get("ts")
        started = started or ts
        updated = ts or updated
        name, ev = e.get("stage"), e.get("event")
        data = {k: v for k, v in e.items() if k not in ("stage", "event", "ts")}
        if name == "run":
            engine = data.get("engine", engine)
            planned = data.get("stages", planned)
            continue
        if ev == "start":
            stages[name] = {"state": "running", "started": ts, **data}
        elif ev in _STAGE_STATE:
            st = stages.setdefault(name, {})
            st.update(data)
            st["state"] = _STAGE_STATE[ev]

    label = stages.get("label")
    if label is not None and label["state"] == "running":
        progress = _label_progress(run_dir, label.get("started"))
        if progress is not None:
            label["progress"] = progress

    states = [st["state"] for st in stages.values()]
    if not stages:
        status = "empty"
    elif "failed" in states:
        status = "failed"
    elif "running" in states:
        status = "running"
    elif planned and all(
        stages.get(p, {}).get("state") in ("done", "cached") for p in planned
    ):
        status = "completed"
    elif planned:
        status = "partial"
    else:  # no plan recorded (standalone stage runs): all seen stages finished
        status = "partial"

    last_stage = next(reversed(stages), None)
    return {
        "name": run_dir.name,
        "dir": str(run_dir),
        "status": status,
        "engine": engine,
        "planned": planned,
        "stages": stages,
        "last_stage": last_stage,
        "started": started,
        "updated": updated,
    }

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
class 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.
    """

    def __init__(self, outdir: str | Path):
        self.outdir = Path(outdir)
        self.path = self.outdir / INDEX_NAME

    def _connect(self) -> sqlite3.Connection:
        con = sqlite3.connect(self.path)
        con.execute(
            """CREATE TABLE IF NOT EXISTS runs (
                   dir TEXT PRIMARY KEY, name TEXT, status TEXT, last_stage TEXT,
                   engine TEXT, started TEXT, updated TEXT,
                   events_mtime REAL, detail TEXT)"""
        )
        return con

    def refresh(self) -> int:
        """Sync the index with the event logs on disk; return #rows updated."""
        con = self._connect()
        try:
            known = dict(con.execute("SELECT dir, events_mtime FROM runs"))
            seen, updated = set(), 0
            for events in sorted(self.outdir.glob(f"*/{EVENTS_NAME}")):
                run_dir = events.parent
                seen.add(str(run_dir))
                mtime = events.stat().st_mtime
                if known.get(str(run_dir)) == mtime:
                    continue
                st = run_status(run_dir)
                con.execute(
                    "REPLACE INTO runs VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                    (
                        str(run_dir), st["name"], st["status"], st["last_stage"],
                        st["engine"], st["started"], st["updated"],
                        mtime, json.dumps(st),
                    ),
                )
                updated += 1
            gone = set(known) - seen
            if gone:
                con.executemany(
                    "DELETE FROM runs WHERE dir = ?", [(d,) for d in gone]
                )
            con.commit()
            return updated
        finally:
            con.close()

    def runs(self) -> list[dict]:
        """All indexed runs, most recently updated first."""
        con = self._connect()
        try:
            rows = con.execute(
                "SELECT detail FROM runs ORDER BY updated DESC, name"
            ).fetchall()
            return [json.loads(detail) for (detail,) in rows]
        finally:
            con.close()

refresh

refresh() -> int

Sync the index with the event logs on disk; return #rows updated.

Source code in src/traincraft/state.py
def refresh(self) -> int:
    """Sync the index with the event logs on disk; return #rows updated."""
    con = self._connect()
    try:
        known = dict(con.execute("SELECT dir, events_mtime FROM runs"))
        seen, updated = set(), 0
        for events in sorted(self.outdir.glob(f"*/{EVENTS_NAME}")):
            run_dir = events.parent
            seen.add(str(run_dir))
            mtime = events.stat().st_mtime
            if known.get(str(run_dir)) == mtime:
                continue
            st = run_status(run_dir)
            con.execute(
                "REPLACE INTO runs VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    str(run_dir), st["name"], st["status"], st["last_stage"],
                    st["engine"], st["started"], st["updated"],
                    mtime, json.dumps(st),
                ),
            )
            updated += 1
        gone = set(known) - seen
        if gone:
            con.executemany(
                "DELETE FROM runs WHERE dir = ?", [(d,) for d in gone]
            )
        con.commit()
        return updated
    finally:
        con.close()

runs

runs() -> list[dict]

All indexed runs, most recently updated first.

Source code in src/traincraft/state.py
def runs(self) -> list[dict]:
    """All indexed runs, most recently updated first."""
    con = self._connect()
    try:
        rows = con.execute(
            "SELECT detail FROM runs ORDER BY updated DESC, name"
        ).fetchall()
        return [json.loads(detail) for (detail,) in rows]
    finally:
        con.close()

log_event

log_event(run_dir: str | Path, stage: str, event: str, **data) -> None

Append one event record to the run's log (best-effort: never raises OSError).

Source code in src/traincraft/state.py
def log_event(run_dir: str | Path, stage: str, event: str, **data) -> None:
    """Append one event record to the run's log (best-effort: never raises OSError)."""
    record = {
        "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "stage": stage,
        "event": event,
        **data,
    }
    try:
        run_dir = Path(run_dir)
        run_dir.mkdir(parents=True, exist_ok=True)
        with (run_dir / EVENTS_NAME).open("a") as fh:
            fh.write(json.dumps(record, default=str) + "\n")
    except OSError:  # state logging must never kill the science
        logger.warning("could not append to %s", run_dir / EVENTS_NAME, exc_info=True)

read_events

read_events(run_dir: str | Path) -> list[dict]

All events of a run, oldest first (empty if the run has no log).

Source code in src/traincraft/state.py
def read_events(run_dir: str | Path) -> list[dict]:
    """All events of a run, oldest first (empty if the run has no log)."""
    path = Path(run_dir) / EVENTS_NAME
    if not path.exists():
        return []
    events = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            events.append(json.loads(line))
        except json.JSONDecodeError:  # torn write at a crash: skip the line
            logger.debug("skipping malformed event line in %s", path)
    return events

Service layer

The HTTP app and MCP server are thin wrappers over traincraft.service.core (see Serve Runs to Agents & UIs).

create_app

create_app(outdir: str | Path = 'runs')
Source code in src/traincraft/service/app.py
def create_app(outdir: "str | Path" = "runs"):
    try:
        from fastapi import FastAPI, HTTPException
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "the HTTP service needs fastapi — install with: pixi install -e service"
        ) from exc

    outdir = Path(outdir)
    app = FastAPI(
        title="TrainCraft",
        version=__version__,
        description=f"Run state, quality reports and pipeline submission over `{outdir}`.",
    )

    def _found(fn, *args):
        try:
            return fn(*args)
        except FileNotFoundError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @app.get("/runs")
    def runs() -> list[dict]:
        """All runs with their current status (index refreshed lazily)."""
        return core.list_runs(outdir)

    @app.get("/runs/{name}")
    def run_detail(name: str) -> dict:
        """Per-stage status of one run, from its event log."""
        return _found(core.get_run, outdir, name)

    @app.get("/runs/{name}/events")
    def run_events(name: str) -> list[dict]:
        """The raw event log (oldest first)."""
        return _found(core.get_events, outdir, name)

    @app.get("/runs/{name}/report")
    def quality_report(name: str) -> dict:
        """The validation quality report (metrics, threshold checks, passed)."""
        return _found(core.get_report, outdir, name)

    @app.get("/runs/{name}/manifest/{stage}")
    def stage_manifest(name: str, stage: str) -> dict:
        """A stage's manifest.json (label, train)."""
        return _found(core.get_manifest, outdir, name, stage)

    @app.post("/config/validate")
    def validate_config(body: ConfigText) -> dict:
        """Validate TOML config text (same check as `traincraft validate`)."""
        return core.validate_config_text(body.toml)

    @app.post("/runs")
    def submit(body: Submission) -> dict:
        """Start the pipeline for a server-side config file; returns immediately."""
        try:
            return core.submit_config(body.config_path, force=body.force)
        except FileNotFoundError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        except Exception as exc:  # bad config, unloadable sections, …
            raise HTTPException(status_code=422, detail=str(exc)) from exc

    return app

create_mcp

create_mcp(outdir: str | Path = 'runs')
Source code in src/traincraft/service/mcp_server.py
def create_mcp(outdir: str | Path = "runs"):
    try:
        from mcp.server.fastmcp import FastMCP
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "the MCP server needs the mcp package — install with: pixi install -e service"
        ) from exc

    outdir = Path(outdir)
    server = FastMCP(
        "traincraft",
        instructions=(
            "TrainCraft generates MLIP training datasets and fine-tunes MACE "
            f"models. These tools operate on the runs directory `{outdir}`. "
            "Judge a trained model ONLY by its quality report "
            "(quality_report tool) — never from training logs alone."
        ),
    )

    @server.tool()
    def list_runs() -> list[dict]:
        """List every TrainCraft run with its status (running/completed/failed/…),
        last active stage, engine and timestamps. Start here to find a run."""
        return core.list_runs(outdir)

    @server.tool()
    def run_status(name: str) -> dict:
        """Per-stage status of one run: state (pending/running/done/cached/failed),
        frame counts, wall times, errors, and live labeling progress. Use this to
        answer 'where is my run?' — do not parse logs."""
        return core.get_run(outdir, name)

    @server.tool()
    def quality_report(name: str) -> dict:
        """The run's validation quality report: RMSE/MAE per property,
        per-element force errors, threshold checks and the overall pass/fail.
        Read this before making ANY claim about a trained model's quality, and
        report the metrics + verdict to the user."""
        return core.get_report(outdir, name)

    @server.tool()
    def stage_manifest(name: str, stage: str) -> dict:
        """A stage's manifest.json. stage='label' → level of theory + counts
        (n_attempted/n_failed/n_reused) + per-frame errors; stage='train' → the
        exact rendered training command, heads, and split sizes."""
        return core.get_manifest(outdir, name, stage)

    @server.tool()
    def validate_config(toml_text: str) -> dict:
        """Validate TrainCraft TOML config text without running anything.
        Returns {valid, name, stages, errors}. ALWAYS validate before submitting;
        fix every error rather than guessing at field names."""
        return core.validate_config_text(toml_text)

    @server.tool()
    def submit_run(config_path: str, force: bool = False) -> dict:
        """Start the pipeline for a config FILE PATH (Slurm if configured, else a
        detached local process) and return immediately. Confirm with the user
        before submitting anything expensive (DFT labeling, MACE training).
        Track progress afterwards with run_status."""
        return core.submit_config(config_path, force=force)

    return server

list_runs

list_runs(outdir: str | Path) -> list[dict]

Every run under outdir (index refreshed first), newest first.

Source code in src/traincraft/service/core.py
def list_runs(outdir: str | Path) -> list[dict]:
    """Every run under ``outdir`` (index refreshed first), newest first."""
    index = RunIndex(outdir)
    index.refresh()
    return [
        {k: st.get(k) for k in ("name", "status", "last_stage", "engine", "started", "updated")}
        for st in index.runs()
    ]

get_run

get_run(outdir: str | Path, name: str) -> dict

Full per-stage status of one run, straight from its event log.

Source code in src/traincraft/service/core.py
def get_run(outdir: str | Path, name: str) -> dict:
    """Full per-stage status of one run, straight from its event log."""
    return run_status(_run_dir(outdir, name))

get_report

get_report(outdir: str | Path, name: str) -> dict

The run's validation quality_report.json (metrics, checks, passed).

Source code in src/traincraft/service/core.py
def get_report(outdir: str | Path, name: str) -> dict:
    """The run's validation ``quality_report.json`` (metrics, checks, passed)."""
    path = _run_dir(outdir, name) / "validation" / REPORT_NAME
    if not path.exists():
        raise FileNotFoundError(f"run {name!r} has no quality report (no validate stage yet)")
    return json.loads(path.read_text())

get_manifest

get_manifest(outdir: str | Path, name: str, stage: str) -> dict

A stage's manifest.json (stages that write one: label, train).

Source code in src/traincraft/service/core.py
def get_manifest(outdir: str | Path, name: str, stage: str) -> dict:
    """A stage's ``manifest.json`` (stages that write one: label, train)."""
    if stage not in _MANIFESTS:
        raise FileNotFoundError(
            f"stage {stage!r} writes no manifest; choose from {sorted(_MANIFESTS)}"
        )
    sub, fname = _MANIFESTS[stage]
    path = _run_dir(outdir, name) / sub / fname
    if not path.exists():
        raise FileNotFoundError(f"run {name!r} has no {stage} manifest yet")
    return json.loads(path.read_text())

validate_config_text

validate_config_text(text: str) -> dict

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
def validate_config_text(text: str) -> dict:
    """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.
    """
    try:
        cfg = loads_config(text)
    except Exception as exc:
        return {"valid": False, "name": None, "stages": [], "errors": str(exc)}
    return {
        "valid": True,
        "name": cfg.run.name,
        "stages": enabled_stages(cfg),
        "errors": None,
    }

submit_config

submit_config(config_path: str | Path, *, force: bool = False) -> dict

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).

Source code in src/traincraft/service/core.py
def submit_config(config_path: str | Path, *, force: bool = False) -> dict:
    """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``).
    """
    config_path = Path(config_path).resolve()
    cfg = load_config(config_path)
    ws = workspace_for(cfg)

    if cfg.orchestration is not None and cfg.orchestration.engine == "slurm":
        jobs = submit_slurm(cfg, str(config_path))
        return {
            "run": cfg.run.name,
            "dir": str(ws.root),
            "engine": "slurm",
            "jobs": {j.stage: j.job_id for j in jobs},
        }

    cmd = [sys.executable, "-m", "traincraft.cli", "run", str(config_path)]
    if force:
        cmd.append("--force")
    log = ws.root / "local.log"
    with log.open("ab") as fh:
        proc = subprocess.Popen(
            cmd, stdout=fh, stderr=subprocess.STDOUT, start_new_session=True
        )
    return {
        "run": cfg.run.name,
        "dir": str(ws.root),
        "engine": "local",
        "pid": proc.pid,
        "log": str(log),
    }