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

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.

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)

Bases: ABC

Parameters:

• filename (str | Path)

• capabilities (ModelCapabilities | None)

• append (bool | None)

abstractmethod write(systems: List[System], predictions: Dict[str, TensorMap])

Write a single system and its predictions.

Parameters:

• systems (List[System])

• predictions (Dict[str, TensorMap])

abstractmethod finish()

Called after all writes. Optional to override.

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)

Bases: Writer

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

Parameters:

• filename (str | Path)

• capabilities (ModelCapabilities | None)

• append (bool | None)

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

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

Parameters:

• systems (List[System])

• predictions (Dict[str, TensorMap])

finish()

Write all accumulated systems and predictions to the XYZ file.

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

Bases: Writer

Parameters:

• path (str | Path)

• capabilities (ModelCapabilities | None)

• append (bool | None)

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

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

Parameters:

• systems (List[System])

• predictions (Dict[str, TensorMap])

finish()

Called after all writes. Optional to override.

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

Bases: Writer

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

Parameters:

• filename (str | Path)

• capabilities (ModelCapabilities | None)

• append (bool | None)

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

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

Parameters:

• systems (List[System])

• predictions (Dict[str, TensorMap])

finish()

Write all accumulated systems and predictions to Metatensor files.