marathon-train 0.2.0__py3-none-any.whl

marathon/README.md

@@ -0,0 +1,209 @@
+# Technical readme
+
+Welcome to the technical `marathon` readme! This is where we keep high-level conventions and so on. While the idea is that the subpackages are pretty independent, there are some implicit contracts that bind them together which we explain here. Individual modules make more targeted assumptions where needed; those are documented in their own `README.md` files.
+
+## Data model
+
+The most basic internal data object is a `Sample`, which is a `namedtuple` of `(structure, labels)`. The `structure` dict describes the atomic geometry and its neighbor graph:
+
+```python
+structure = {
+    "positions": ...,      # (num_atoms, 3)
+    "atomic_numbers": ..., # (num_atoms,) -- Z values
+    "cell": ...,           # (3, 3)
+    "charges": ...,        # (num_atoms,)
+    "centers": ...,        # (num_pairs,) -- i indices of neighbor pairs
+    "others": ...,         # (num_pairs,) -- j indices of neighbor pairs
+    "displacements": ...,  # (num_pairs, 3) -- R_ij vectors
+    "cell_shifts": ...,    # (num_pairs, 3) -- periodic image shifts
+    "pbc": ...,            # (3,) -- periodic boundary conditions
+    "max_neighbors": ...,  # int -- max neighbors per atom
+}
+```
+
+The `labels` dict contains target properties:
+
+```python
+labels = {
+    "energy": ...,    # scalar -- total energy (eV)
+    "forces": ...,    # (num_atoms, 3) -- eV/Å
+    "stress": ...,    # (3, 3) -- dU/dε (eV, see below)
+    "num_atoms": ..., # int
+}
+```
+
+The neighbor list is computed by [`vesin`](https://github.com/Luthaf/vesin) from positions and a cutoff radius. The `displacements` array contains the actual $R_{ij} = R_j - R_i + S \cdot \text{cell}$ vectors. In the typical case, these are the quantities that the model operates on and that we differentiate through to obtain forces (see `evaluate/README.md`).
+
+The next step is a `Batch`, which is another `namedtuple` that collates one or more `Sample`s, padded and index-offset so that an MPNN can process them as one big disconnected graph:
+
+```python
+Batch(
+    atomic_numbers,      # (num_atoms,) -- Z_i
+    displacements,       # (num_pairs, 3) -- R_ij
+    centers,             # (num_pairs,) -- i indices, offset per structure
+    others,              # (num_pairs,) -- j indices, offset per structure
+    atom_to_structure,   # (num_atoms,) -- which structure each atom belongs to
+    pair_to_structure,   # (num_pairs,) -- which structure each pair belongs to
+    structure_mask,      # (num_structures,) -- False for padding
+    atom_mask,           # (num_atoms,) -- False for padding
+    pair_mask,           # (num_pairs,) -- False for padding
+    labels,              # dict of batched labels + masks
+)
+```
+
+It is expected that models that require something less off-the-shelf implement their own `Batch` class and related infrastructure. We try, as much as possible, to be agnostic to the internals of the batch. Only parts of the code that *must* explicitly interact with it care about internals, for example some parts of `marathon.evaluate`.
+
+## Padding and masking
+
+Padding adds a single extra "padding structure" at the end of the batch. Its atoms fill the remaining allocated slots; its pairs point both endpoints to a single padding atom (creating self-loops that don't participate in message passing). Displacements for padding pairs are zero.
+
+**Downstream code must handle zero displacements** ($R_{ij} = 0$). Division by $|R_{ij}|$ will blow up if you don't mask.
+
+Three boolean masks distinguish real from padding: `structure_mask`, `atom_mask`, `pair_mask`. Additionally, each label `key` in `batch.labels` has a corresponding `key + "_mask"` (e.g. `energy_mask`, `forces_mask`) that is `True` where the label is valid and `False` where it's missing or padding. This handles the case where some structures have, say, energy labels but not stress labels. NaN values in input labels are treated as missing data and get mask `False`.
+
+Depending on the model architecture, it may be better to always have `Batch`es with the exact same shape (so they can be stacked and processed with `vmap`/`scan`), or it may be better to batch dynamically to varying shapes. The former avoids recompilation; the latter reduces wasted computation. `marathon.grain` provides both strategies (`ToFixedShapeBatch` and `ToFixedLengthBatch`).
+
+
+## Label conventions
+
+All *per-structure* properties are stored as genuinely per-structure quantities -- not per atom and not per volume:
+
+- **Energy** is total energy (eV), not energy/atom.
+- **Stress** is the ASE stress multiplied by volume, i.e. $dU/d\varepsilon$ in energy units (eV), not pressure. If you need stress in force/area units, divide by volume.
+- **Forces** are per-atom by nature (eV/Å).
+
+The properties/normalisation system (see below) can be used to change display units and normalise by atom count for loss computation, but it implicitly relies on inputs following this convention. The pipelines process `ase.Atoms` as inputs, and expect the energy there to be total energy and the stress the standard `ase` stress definition; we handle the conversion.
+
+Everything up to generating a `Batch` is done in pure `numpy` in configurable precision (default: float64/int64). It is up to downstream code to convert to `jax` arrays and potentially cast into other `dtypes`. Some of the infrastructure is therefore also compatible with `pytorch`.
+
+
+## Properties System
+
+Marathon supports custom properties beyond the standard `energy`, `forces`, `stress`. The key idea is that users define **one** `PROPERTIES` dict and **one** `NORMALIZATION` dict, then pass them to all relevant infrastructure--each module extracts what it needs.
+
+### Property definition
+
+A property has these fields (not all are required everywhere):
+
+```python
+PROPERTIES = {
+    "energy": {
+        "shape": (1,),                    # for data loading/batching
+        "storage": "atoms.calc",          # for DataSource: where in ase.Atoms
+        "report_unit": (1000, "meV"),     # for emit: (scale, unit)
+        "symbol": "E",                    # for emit: short symbol
+    },
+    "forces": {
+        "shape": ("atom", 3),
+        "storage": "atoms.calc",
+        "report_unit": (1000, "meV/Å"),
+        "symbol": "F",
+    },
+    # ...
+}
+```
+
+- `shape`: tuple where `"atom"` placeholder gets replaced with num_atoms. **Only leading `"atom"` dimension is supported** (e.g., `("atom", 3)` works, but `(3, "atom")` does not).
+- `storage`: where to read/write from ase.Atoms: `"atoms.calc"`, `"atoms.arrays"`, `"atoms.info"`
+- `report_unit`: `(scale_factor, base_unit)` for converting internal units (eV) to display units (meV)
+- `symbol`: short symbol for console/log output
+
+**Note on storage:** `atoms.calc` is reserved for `energy`, `forces`, and `stress` (these are what ASE calculators provide). Custom properties should use `atoms.arrays` (per-atom data) or `atoms.info` (per-structure data). We don't guarantee that other properties will work correctly through calculators.
+
+### Normalization
+
+The normalization dict specifies which properties are divided by atom count:
+
+```python
+NORMALIZATION = {
+    "energy": "atom",
+    "stress": "atom",
+}
+```
+
+- `"atom"` means divide by num_atoms during loss/metrics and append `/atom` to units
+- Properties not listed (like forces) are not modified. Setting normalization to a quantity that's of dimension `[atom,...]` will yield undefined behaviour or crash.
+
+### Where configs are used
+
+| Module | Uses from PROPERTIES | Uses NORMALIZATION |
+|--------|---------------------|-------------------|
+| `data/` batching | `shape` | - |
+| `grain/DataSource` | `shape`, `storage` | - |
+| `evaluate/` loss | - | yes |
+| `evaluate/` metrics | - | yes |
+| `emit/` logging | `report_unit`, `symbol` | yes (for unit suffix) |
+| `emit/` plotting | `report_unit` | yes (for unit suffix) |
+
+### Custom properties example
+
+To add dipole moment as a custom property:
+
+```python
+# One unified PROPERTIES dict
+PROPERTIES = {
+    "energy": {
+        "shape": (1,), "storage": "atoms.calc",
+        "report_unit": (1000, "meV"), "symbol": "E",
+    },
+    "forces": {
+        "shape": ("atom", 3), "storage": "atoms.calc",
+        "report_unit": (1000, "meV/Å"), "symbol": "F",
+    },
+    "dipole": {
+        "shape": (3,), "storage": "atoms.info",  # per-structure, not from calc
+        "report_unit": (1, "Debye"), "symbol": "μ",
+    },
+}
+
+# One unified NORMALIZATION dict
+NORMALIZATION = {
+    "energy": "atom",
+    # dipole not listed -> no /atom suffix, no per-atom normalization
+}
+
+# Pass to all relevant places:
+from marathon.emit import Txt
+from marathon.emit.pretty import format_metrics
+from marathon.evaluate import get_loss_fn
+
+logger = Txt(keys=["energy", "dipole"], properties=PROPERTIES, normalization=NORMALIZATION)
+loss_fn = get_loss_fn(predict_fn, weights={"energy": 1.0, "dipole": 1.0}, normalization=NORMALIZATION)
+msg = format_metrics(metrics, properties=PROPERTIES, normalization=NORMALIZATION)
+```
+
+The defaults are split across modules for convenience (`data/properties.py`, `evaluate/properties.py`, `emit/properties.py`), but users can combine them into a single dict.
+
+
+## Serialization conventions
+
+Model architectures are stored as YAML "spec dicts" -- a single-entry dict mapping a fully-qualified class path to its `__init__` kwargs:
+
+```yaml
+my_package.models.MACE:
+  num_features: 128
+  num_layers: 3
+```
+
+`marathon.io.to_dict(module)` serialises any dataclass (including `flax.nn.Module`); `marathon.io.from_dict(spec)` reconstructs it by dynamically importing the class. `parent` and `name` fields are excluded (these are Flax-internal).
+
+Numerical data (weights, optimizer state) uses `msgpack` via Flax's serialization. Checkpoints combine both: the architecture lives in `model.yaml`, the weights in `model.msgpack`, training state in `state.msgpack`, and metrics in `metrics.yaml`. See `emit/README.md` for details.
+
+
+## Code style
+
+Formatting is handled by `ruff` with a line length of 92. The line length is enforced by the formatter only, not the linter — so lines that can't be shortened automatically are left alone. We suppress rules that get in the way of research code: short variable names (`E741`), lambdas (`E731`), non-top-level imports (`E402`), and line length warnings (`E501`).
+
+Import ordering groups `numpy` and `jax` before other third-party packages (configured via `isort` sections in `pyproject.toml`). Lazy imports are used for optional dependencies (`grain`, `matplotlib`, `wandb`, `numba`) and occasionally to avoid circular imports or heavy startup costs.
+
+The code is deliberately kept **copy-and-pasteable** — individual files should be easy to pull into a script or notebook for quick hacking. This is part of the reason for inline testing, avoiding deep relative import chains, and keeping modules self-contained where possible.
+
+More specifically:
+
+- **Descriptive names, minimal docstrings.** Docstrings only where behaviour isn't obvious from the signature and context. Comments explain *why*, not *what*.
+- **Lambdas and comprehensions** are used freely. Short closures are preferred over named helper functions when the logic is local and throwaway.
+- **`namedtuple`** for data objects (`Sample`, `Batch`). Plain dicts for labels and properties configs. Dataclasses (via `flax.nn.Module` or the `frozen` helper in `utils.py`) for model components.
+- **Factory functions** follow the `get_*_fn` pattern (e.g. `get_loss_fn`, `get_predict_fn`, `get_metrics_fn`). They return closures that capture configuration, keeping the hot path clean.
+- **Inline tests** at the bottom of many modules, separated from production code by a `# -- test --` comment. Always use exactly this marker. These tests run on import and serve as both documentation and a basic smoke test. They're not a substitute for `pytest` but catch regressions immediately.
+- **`__all__`** in `__init__.py` files defines the public API of each subpackage.
+- **Private functions** are prefixed with `_` (e.g. `_huber`). No `__double_underscore` mangling.

marathon/__init__.py

@@ -0,0 +1,3 @@
+from opsis import Comms
+
+comms = Comms("marathon")

marathon/data/README.md

@@ -0,0 +1,27 @@
+# `marathon.data`: Samples, batching, and data splits
+
+Converts `ase.Atoms` objects into the internal `Sample` and `Batch` representations described in the top-level README. This module is pure numpy and does not depend on JAX (except `splits.py`, which uses `jax.random` for reproducible shuffling).
+
+## `sample`: `ase.Atoms` → `Sample`
+
+`to_structure(atoms, cutoff)` computes the neighbor list (via `vesin`) and packs geometry into a structure dict. `to_labels(atoms, ...)` extracts target properties using the `storage` field from the properties config. `to_sample` combines both.
+
+For standard properties (`energy`, `forces`, `stress`), extraction from the ASE calculator handles the stress convention (multiply by volume to get $dU/d\varepsilon$). Custom properties are read from `atoms.info` or `atoms.arrays` as configured.
+
+## `batching`: `[Sample]` → `Batch`
+
+`batch_samples` collates a list of `Sample`s into a single `Batch` namedtuple with padding. Index arrays (`centers`, `others`) are offset per structure so the batch looks like one disconnected graph. Labels are batched separately by `batch_labels`, which also creates per-key masks (NaN → `False`).
+
+## `sizes`: Buffer allocation
+
+`determine_max_sizes` scans samples to find worst-case atom and pair counts for a given batch size, then rounds up via `next_size` to a JIT-friendly number.
+
+Typically not needed and superseded by more complex `grain` machinery -- but this is nice for very small all-GPU training runs.
+
+## `splits`: Train/valid/test splits
+
+`get_splits` draws non-overlapping index sets using `jax.random` for reproducibility.
+
+## `properties`: Shape and storage specs
+
+`DEFAULT_PROPERTIES` defines `shape` and `storage` for the standard properties. `deduce_shape` resolves the `"atom"` placeholder at batch time.

marathon/data/__init__.py

@@ -0,0 +1,27 @@
+import os
+from pathlib import Path
+
+from .batching import Batch, batch_labels, batch_samples
+from .sample import Sample, to_sample, to_structure
+from .sizes import determine_max_sizes
+from .splits import get_splits
+
+dataset_folder = os.environ.get("DATASETS")
+
+if dataset_folder is not None:
+    datasets = Path(dataset_folder)
+else:
+    datasets = Path.cwd()
+
+
+__all__ = [
+    "Sample",
+    "Batch",
+    "batch_samples",
+    "batch_labels",
+    "determine_max_sizes",
+    "to_sample",
+    "to_structure",
+    "get_splits",
+    "datasets",
+]

marathon/data/batching.py

@@ -0,0 +1,236 @@
+import numpy as np
+
+from collections import namedtuple
+
+from .properties import DEFAULT_PROPERTIES, deduce_shape, is_per_atom
+from .sample import Sample
+
+# Padded, index-offset collation of multiple Samples into one disconnected graph.
+Batch = namedtuple(
+    "Batch",
+    (
+        "atomic_numbers",  # Z_i
+        "displacements",  # R_ij
+        "centers",  # i
+        "others",  # j
+        "atom_to_structure",  # map atomic_numbers to original structures
+        "pair_to_structure",  # map displacements to original structures
+        "structure_mask",  # False for padding
+        "atom_mask",  # False for padding
+        "pair_mask",  # False for padding
+        "labels",
+    ),
+)
+
+
+def batch_samples(
+    samples,
+    num_atoms,
+    num_pairs,
+    keys,
+    num_structures=None,
+    float_dtype=None,
+    int_dtype=None,
+    properties=DEFAULT_PROPERTIES,
+):
+    """Collate samples into a Batch, padding to fixed num_atoms/num_pairs with masks.
+
+    num_atoms/num_pairs must exceed the real totals (padding needs at least one extra slot).
+    """
+    if float_dtype is None:
+        float_dtype = samples[0].structure["displacements"].dtype
+    if int_dtype is None:
+        int_dtype = samples[0].structure["centers"].dtype
+
+    if num_structures is None:
+        num_structures = len(samples) + 1
+    else:
+        num_input_structures = len(samples)
+        assert num_input_structures + 1 <= num_structures
+
+    atomic_numbers = np.zeros(num_atoms, dtype=int_dtype)
+    displacements = np.zeros((num_pairs, 3), dtype=float_dtype)
+    centers = np.zeros(num_pairs, dtype=int_dtype)
+    others = np.zeros(num_pairs, dtype=int_dtype)
+    atom_to_structure = np.zeros(num_atoms, dtype=int_dtype)
+    pair_to_structure = np.zeros(num_pairs, dtype=int_dtype)
+    structure_mask = np.zeros(num_structures, dtype=bool)
+    atom_mask = np.zeros(num_atoms, dtype=bool)
+    pair_mask = np.zeros(num_pairs, dtype=bool)
+
+    labels = batch_labels(
+        [sample.labels for sample in samples],
+        num_structures,
+        num_atoms,
+        keys,
+        float_dtype=float_dtype,
+        int_dtype=int_dtype,
+        properties=properties,
+    )
+
+    atom_offset = 0
+    pair_offset = 0
+    for i, sample in enumerate(samples):
+        s = sample.structure
+
+        _n_atoms = s["atomic_numbers"].shape[0]
+        _n_pairs = s["displacements"].shape[0]
+
+        atom_slice = slice(atom_offset, atom_offset + _n_atoms)
+        pair_slice = slice(pair_offset, pair_offset + _n_pairs)
+
+        atomic_numbers[atom_slice] = s["atomic_numbers"]
+        displacements[pair_slice] = s["displacements"]
+        centers[pair_slice] = s["centers"] + atom_offset
+        others[pair_slice] = s["others"] + atom_offset
+
+        atom_to_structure[atom_slice] = i
+        pair_to_structure[pair_slice] = i
+
+        structure_mask[i] = True
+        atom_mask[atom_slice] = True
+        pair_mask[pair_slice] = True
+
+        atom_offset += _n_atoms
+        pair_offset += _n_pairs
+
+    assert atom_offset < len(atom_mask), "no room for padding!"
+
+    # now we add the padding
+
+    # skip atomic_numbers -- there is no element 0
+    # skip displacements -- already zero
+    centers[pair_offset:] = atom_offset
+    others[pair_offset:] = atom_offset
+
+    atom_to_structure[atom_offset:] = num_structures - 1
+    pair_to_structure[pair_offset:] = num_structures - 1
+
+    # skip masks -- already False
+
+    # skip labels -- already zero
+
+    return Batch(
+        atomic_numbers,
+        displacements,
+        centers,
+        others,
+        atom_to_structure,
+        pair_to_structure,
+        structure_mask,
+        atom_mask,
+        pair_mask,
+        labels,
+    )
+
+
+def batch_labels(
+    list_of_labels,
+    num_structures,
+    num_atoms,
+    keys,
+    float_dtype=np.float64,
+    int_dtype=np.int64,
+    properties=DEFAULT_PROPERTIES,
+):
+    """Stack label dicts into padded arrays with per-key NaN-aware masks."""
+    labels = {}
+
+    for key in keys:
+        if key not in properties:
+            raise KeyError(f"unknown key: {key}")
+
+        shape = deduce_shape(num_structures, num_atoms, properties[key]["shape"])
+        labels[key] = np.zeros(shape, dtype=float_dtype)
+        labels[key + "_mask"] = labels[key].astype(bool)
+
+    labels["num_atoms"] = np.ones(num_structures, dtype=int_dtype)
+
+    atom_offset = 0
+    for i, l in enumerate(list_of_labels):
+        num_atoms = l["num_atoms"]
+        atom_slice = slice(atom_offset, atom_offset + num_atoms)
+
+        labels["num_atoms"][i] = num_atoms
+
+        for key in keys:
+            per_atom = is_per_atom(properties[key]["shape"])
+            values = l[key]
+            if not np.isnan(values).any():
+                if per_atom:
+                    labels[key][atom_slice] = values
+                    labels[key + "_mask"][atom_slice] = True
+                else:
+                    # scalar per-structure: squeeze to avoid deprecation warning
+                    labels[key][i] = np.squeeze(values)
+                    labels[key + "_mask"][i] = True
+            # else: stays zero, mask False
+
+        atom_offset += num_atoms
+
+    return labels
+
+
+# -- test --
+
+test_samples = [
+    Sample(
+        dict(
+            displacements=np.ones((6, 3)),
+            atomic_numbers=np.array([0, 0, 0]),
+            centers=np.array([0, 0, 1, 1, 2, 2]),
+            others=np.array([1, 2, 0, 2, 0, 1]),
+        ),
+        {
+            "energy": 0.1,
+            "forces": np.ones((3, 3)),
+            "num_atoms": 3,
+        },
+    ),
+    Sample(
+        dict(
+            displacements=np.ones((2, 3)),
+            atomic_numbers=np.array([0, 0]),
+            centers=np.array([0, 1]),
+            others=np.array([1, 0]),
+        ),
+        {
+            "energy": 0.2,
+            "forces": np.ones((2, 3)),
+            "num_atoms": 2,
+        },
+    ),
+]
+
+num_pairs = 10
+num_atoms = 8
+
+test_batch = batch_samples(test_samples, num_atoms, num_pairs, ["energy", "forces"])
+
+
+assert test_batch.displacements.shape[0] == num_pairs
+assert test_batch.centers.shape[0] == num_pairs
+assert test_batch.others.shape[0] == num_pairs
+assert test_batch.atomic_numbers.shape[0] == num_atoms
+np.testing.assert_equal(
+    test_batch.displacements, np.concatenate((np.ones((8, 3)), np.zeros((2, 3))))
+)
+
+np.testing.assert_equal(test_batch.centers, np.array([0, 0, 1, 1, 2, 2, 3, 4, 5, 5]))
+np.testing.assert_equal(test_batch.others, np.array([1, 2, 0, 2, 0, 1, 4, 3, 5, 5]))
+np.testing.assert_equal(test_batch.atomic_numbers, np.array([0, 0, 0, 0, 0, 0, 0, 0]))
+np.testing.assert_equal(test_batch.atom_to_structure, np.array([0, 0, 0, 1, 1, 2, 2, 2]))
+np.testing.assert_equal(
+    test_batch.pair_to_structure, np.array([0, 0, 0, 0, 0, 0, 1, 1, 2, 2])
+)
+
+np.testing.assert_equal(
+    test_batch.atom_mask, np.array([True, True, True, True, True, False, False, False])
+)
+np.testing.assert_equal(test_batch.structure_mask, np.array([True, True, False]))
+
+np.testing.assert_equal(test_batch.labels["energy"], np.array([0.1, 0.2, 0.0]))
+np.testing.assert_equal(test_batch.labels["energy_mask"], np.array([True, True, False]))
+assert test_batch.labels["forces"].shape == (8, 3)
+
+np.testing.assert_array_equal(test_batch.labels["num_atoms"], np.array([3, 2, 1]))

marathon/data/properties.py

@@ -0,0 +1,54 @@
+DEFAULT_PROPERTIES = {
+    "energy": {"shape": (1,), "storage": "atoms.calc"},
+    "forces": {"shape": ("atom", 3), "storage": "atoms.calc"},
+    "stress": {"shape": (3, 3), "storage": "atoms.calc"},
+}
+
+
+def deduce_shape(num_structures, num_atoms, shape):
+    # batch context: resolves shape spec, prepends num_structures
+    # for non-atom properties.
+    # see also: marathon.grain.data_source.properties.deduce_shape
+
+    # scalar per structure: no dummy dimension
+    if shape == (1,):
+        shape = ()
+
+    # replace "atom" -> num_atoms, keep (optional) trailing shape
+    if (len(shape) > 0) and (shape[0] == "atom"):
+        shape = (num_atoms,) + shape[1:]
+    else:
+        shape = (num_structures,) + shape
+
+    return shape
+
+
+def is_per_atom(shape):
+    return (len(shape) > 0) and (shape[0] == "atom")
+
+
+# -- test --
+
+
+def test_properties():
+    # deduce_shape: scalar per structure -> shape (num_structures,)
+    # (1,) removes dummy dim, then prepends num_structures
+    assert deduce_shape(5, 10, (1,)) == (5,)
+
+    # deduce_shape: per-atom property
+    assert deduce_shape(5, 10, ("atom", 3)) == (10, 3)
+    assert deduce_shape(5, 10, ("atom",)) == (10,)
+
+    # deduce_shape: per-structure tensor property
+    assert deduce_shape(5, 10, (3, 3)) == (5, 3, 3)
+    assert deduce_shape(5, 10, (6,)) == (5, 6)
+
+    # is_per_atom
+    assert is_per_atom(("atom", 3)) is True
+    assert is_per_atom(("atom",)) is True
+    assert is_per_atom((3, 3)) is False
+    assert is_per_atom((1,)) is False
+    assert is_per_atom(()) is False
+
+
+test_properties()