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.

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]

Accumulate systems and predictions to write them all at once in finish.

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 all accumulated systems and predictions to the XYZ file.

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

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]

Accumulate systems and predictions to write them all at once in finish.

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 all accumulated systems and predictions to Metatensor files.

Return type:

None