Target data Writers

The main entry point for writing target information is

metatrain.utils.data.writers.get_writer(filename: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = None, fileformat: str | None = None) → Writer[source]

Selects the appropriate writer based on the file extension.

For certain file suffixes, the systems will also be written (i.e xyz).

A path ending in a path separator (e.g. predictions/) is treated as a memory-mapped MemmapDataset directory rather than a single file. Since nothing exists on disk yet when a writer is selected, the trailing separator is the write-side equivalent of that check.

The capabilities of the model are used to infer the type (physical quantity) of the predictions. In this way, for example, position gradients of energies can be saved as forces.

For the moment, strain gradients of the energy are saved as stresses (and not as virials).

Parameters:

  • filename (str | Path) – name of the file to write, or a directory path ending in a path separator for a memmap dataset

  • capabilities (ModelCapabilities | None) – capabilities of the model

  • append (bool | None) – if True, the data will be appended to the file, if it exists. If False, the file will be overwritten. If None, the default behavior of the writer is used.

  • fileformat (str | None) – format of the target value file. If None the format is determined from the file extension.

Returns:

a Writer instance.

Raises:

ValueError – if filename both ends in a path separator and has a recognized file suffix (e.g. "predictions.zip/"), since it is then ambiguous whether a memmap directory or a file of that format was intended.

Return type:

Writer

Based on the provided filename the writer choses which child writer to use. The mapping which writer is used for which file type is stored in

metatrain.utils.data.writers.PREDICTIONS_WRITERS: Dict[str, WriterFactory] = {'.mts': <function _make_factory.<locals>.factory>, '.xyz': <function _make_factory.<locals>.factory>, '.zip': <function _make_factory.<locals>.factory>}

dict: dictionary mapping file suffixes to a prediction writer

Implemented Writers

Writer Abstract Class

class metatrain.utils.data.writers.Writer(filename: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = None)[source]

Bases: ABC

Parameters:
abstractmethod write(systems: List[System], predictions: Dict[str, TensorMap]) → None[source]

Write a single system and its predictions.

Parameters:

  • systems (List[System]) – List of systems to write.

  • predictions (Dict[str, TensorMap]) – Dictionary of TensorMaps with predictions for the systems.

Return type:

None

abstractmethod finish() → None[source]

Called after all writes. Optional to override.

Return type:

None

Available Implementations

The available implementations listed below represent concrete writers that inherit from the Writer abstract class.

class metatrain.utils.data.writers.ASEWriter(filename: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = False)[source]

Bases: Writer

Write systems and predictions to an ASE-compatible XYZ file.

Systems and predictions are converted and written to disk immediately in each write() call, avoiding unbounded memory growth for large datasets.

Parameters:

  • filename (str | Path) – Path to the output XYZ file.

  • capabilities (ModelCapabilities | None) – Model capabilities (unused, but matches base signature).

  • append (bool | None) – If True, append to the existing file, otherwise overwrite.

write(systems: List[System], predictions: Dict[str, TensorMap]) → None[source]

Convert systems and predictions to ASE Atoms and write to disk immediately.

Parameters:

  • systems (List[System]) – List of systems to write.

  • predictions (Dict[str, TensorMap]) – Dictionary of TensorMaps with predictions for the systems.

Return type:

None

finish() → None[source]

No-op: all data is written in write().

Return type:

None

class metatrain.utils.data.writers.DiskDatasetWriter(path: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = False)[source]

Bases: Writer

Write systems and predictions to a zip file, each system in a separate folder inside the zip.

Every finished zip carries a complete metadata/atom_counts.npy member (the atom count of each entry, in entry order), which is what makes the dataset usable with atom-count-based sampling. The metadata/ folder is reserved for such dataset-level informative files.

Parameters:

  • path (str | Path) – Path to the output zip file.

  • capabilities (ModelCapabilities | None) – Model capabilities.

  • append (bool | None) – If True, open the zip file in append mode.

write(systems: List[System], predictions: Dict[str, TensorMap]) → None[source]

Write a single (system, predictions) into the zip under a new folder “<index>/”.

The stored TensorMaps’ "system" sample labels always hold the entry number, whether the systems arrive batched or one at a time. Distinct labels are required to read the dataset back for training (per-sample maps are joined at collation).

Parameters:

  • systems (List[System]) – List of systems to write.

  • predictions (Dict[str, TensorMap]) – Dictionary of TensorMaps with predictions for the systems.

Return type:

None

finish() → None[source]

Write the complete per-structure atom-count file (metadata/atom_counts.npy) and close the zip file. Calling it again is a no-op.

Return type:

None

class metatrain.utils.data.writers.MetatensorWriter(filename: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = False)[source]

Bases: Writer

Write systems and predictions to Metatensor files (.mts).

Each write() call saves the batch predictions to temporary files on disk, avoiding unbounded memory growth. finish() loads the temp files, concatenates them with correct system label offsets, and writes the final output.

Parameters:

  • filename (str | Path) – Base filename for the output files. Each target will be saved in a separate file with the target name appended.

  • capabilities (ModelCapabilities | None) – Model capabilities.

  • append (bool | None) – Whether to append to existing files, unused here but kept for compatibility with the base class.

write(systems: List[System], predictions: Dict[str, TensorMap]) → None[source]

Save batch predictions to temporary files, freeing GPU/CPU memory immediately.

Parameters:

  • systems (List[System]) – List of systems to write.

  • predictions (Dict[str, TensorMap]) – Dictionary of TensorMaps with predictions for the systems.

Return type:

None

finish() → None[source]

Load temp files, shift system labels, join, and write final .mts files.

Return type:

None

class metatrain.utils.data.writers.MemmapWriter(path: str | Path, capabilities: ModelCapabilities | None = None, append: bool | None = False)[source]

Bases: Writer

Write systems and predictions to the on-disk layout consumed by metatrain.utils.data.dataset.MemmapDataset: a directory containing ns.npy, na.npy, x.bin, a.bin, c.bin, and one <target>.bin file per target (plus <target>_forces.bin/<target>_stress.bin for position/strain gradients). The resulting directory can be passed straight back to metatrain as a systems: read_from: dataset, which makes this format well suited for very large evaluation runs (e.g. >1M structures) that should not be re-read as a single in-memory XYZ file.

The final number of structures/atoms is only known once evaluation has finished, so every array is streamed to its .bin file (append-only, on every write() call) rather than buffered in memory or in a temporary directory. This keeps memory use and disk I/O to a single pass, at the cost of not knowing final shapes until finish(), where the small ns.npy/na.npy index files (which require the final structure/atom counts) are written and all files are closed.

Parameters:

  • path (str | Path) – Directory path to write into (e.g. predictions/), used as-is (unlike the other writers, no filename manipulation is performed). The directory is created if it does not exist yet. If it already exists and already contains any file, a FileExistsError is raised immediately rather than writing into (or overwriting) a possibly unrelated directory: this check happens up front, before any (potentially very long) evaluation runs, instead of lazily discovering a conflict only when the colliding file would have been written.

  • capabilities (ModelCapabilities | None) – Model capabilities (unused, but matches base signature).

  • append (bool | None) – Not supported for memmap datasets.

write(systems: List[System], predictions: Dict[str, TensorMap]) → None[source]

Append a batch of systems and predictions to the memmap dataset files.

Parameters:

  • systems (List[System]) – List of systems to write.

  • predictions (Dict[str, TensorMap]) – Dictionary of TensorMaps with predictions for the systems.

Return type:

None

finish() → None[source]

Write ns.npy/na.npy and close all open binary files.

Return type:

None