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

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

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

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.

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>/”.

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]

Close the zip file.

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