core.datasets.ase_datasets

core.datasets.ase_datasets#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

AseAtomsDataset

This is an abstract Dataset that includes helpful utilities for turning

AseReadDataset

This Dataset uses ase.io.read to load data from a directory on disk.

AseReadMultiStructureDataset

This Dataset can read multiple structures from each file using ase.io.read.

AseDBDataset

This Dataset connects to an ASE Database, allowing the storage of atoms objects

Functions#

apply_one_tags(atoms[, skip_if_nonzero, skip_always])

This function will apply tags of 1 to an ASE atoms object.

Module Contents#

core.datasets.ase_datasets.apply_one_tags(atoms: ase.Atoms, skip_if_nonzero: bool = True, skip_always: bool = False)#

This function will apply tags of 1 to an ASE atoms object. It is used as an atoms_transform in the datasets contained in this file.

Certain models will treat atoms differently depending on their tags. For example, GemNet-OC by default will only compute triplet and quadruplet interactions for atoms with non-zero tags. This model throws an error if there are no tagged atoms. For this reason, the default behavior is to tag atoms in structures with no tags.

Parameters:

  • skip_if_nonzero (bool) – If at least one atom has a nonzero tag, do not tag any atoms

  • skip_always (bool) – Do not apply any tags. This arg exists so that this function can be disabled without needing to pass a callable (which is currently difficult to do with main.py)

class core.datasets.ase_datasets.AseAtomsDataset(config: dict, atoms_transform: Callable[[ase.Atoms, Any, Ellipsis], ase.Atoms] = apply_one_tags)#

Bases: fairchem.core.datasets.base_dataset.BaseDataset, abc.ABC

This is an abstract Dataset that includes helpful utilities for turning ASE atoms objects into OCP-usable data objects. This should not be instantiated directly as get_atoms_object and load_dataset_get_ids are not implemented in this base class.

Derived classes must add at least two things:

self.get_atoms_object(id): a function that takes an identifier and returns a corresponding atoms object

self.load_dataset_get_ids(config: dict): This function is responsible for any initialization/loads

of the dataset and importantly must return a list of all possible identifiers that can be passed into self.get_atoms_object(id)

Identifiers need not be any particular type.

a2g_args#
a2g#
key_mapping#
transforms#
atoms_transform#
ids#
num_samples#
__getitem__(idx)#
abstract get_atoms(idx: str | int) → ase.Atoms#
abstract _load_dataset_get_ids(config)#
abstract get_relaxed_energy(identifier)#
sample_property_metadata(num_samples: int = 100) → dict#
get_metadata(attr, idx)#
class core.datasets.ase_datasets.AseReadDataset(config: dict, atoms_transform: Callable[[ase.Atoms, Any, Ellipsis], ase.Atoms] = apply_one_tags)#

Bases: AseAtomsDataset

This Dataset uses ase.io.read to load data from a directory on disk. This is intended for small-scale testing and demonstrations of OCP. Larger datasets are better served by the efficiency of other dataset types such as LMDB.

For a full list of ASE-readable filetypes, see https://wiki.fysik.dtu.dk/ase/ase/io/io.html

Parameters:

  • config (dict) –

    src (str): The source folder that contains your ASE-readable files

    pattern (str): Filepath matching each file you want to read

    ex. “/POSCAR”, “.cif”, “.xyz” search recursively with two wildcards: “*/POSCAR” or “**/*.cif”

    a2g_args (dict): Keyword arguments for fairchem.core.preprocessing.AtomsToGraphs()

    default options will work for most users

    If you are using this for a training dataset, set “r_energy”:True, “r_forces”:True, and/or “r_stress”:True as appropriate In that case, energy/forces must be in the files you read (ex. OUTCAR)

    ase_read_args (dict): Keyword arguments for ase.io.read()

    keep_in_memory (bool): Store data in memory. This helps avoid random reads if you need

    to iterate over a dataset many times (e.g. training for many epochs). Not recommended for large datasets.

    include_relaxed_energy (bool): Include the relaxed energy in the resulting data object.

    The relaxed structure is assumed to be the final structure in the file (e.g. the last frame of a .traj).

    atoms_transform_args (dict): Additional keyword arguments for the atoms_transform callable

    transform_args (dict): Additional keyword arguments for the transform callable

    key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used

    in the model with the corresponding property as it was named in the dataset. Only need to use if the name is different.

  • atoms_transform (callable, optional) – Additional preprocessing function applied to the Atoms object. Useful for applying tags, for example.

_load_dataset_get_ids(config) → list[pathlib.Path]#
get_atoms(idx: str | int) → ase.Atoms#
get_relaxed_energy(identifier) → float#
class core.datasets.ase_datasets.AseReadMultiStructureDataset(config: dict, atoms_transform: Callable[[ase.Atoms, Any, Ellipsis], ase.Atoms] = apply_one_tags)#

Bases: AseAtomsDataset

This Dataset can read multiple structures from each file using ase.io.read. The disadvantage is that all files must be read at startup. This is a significant cost for large datasets.

This is intended for small-scale testing and demonstrations of OCP. Larger datasets are better served by the efficiency of other dataset types such as LMDB.

For a full list of ASE-readable filetypes, see https://wiki.fysik.dtu.dk/ase/ase/io/io.html

Parameters:

  • config (dict) –

    src (str): The source folder that contains your ASE-readable files

    pattern (str): Filepath matching each file you want to read

    ex. “.traj”, “.xyz” search recursively with two wildcards: “/POSCAR” or “/*.cif”

    index_file (str): Filepath to an indexing file, which contains each filename

    and the number of structures contained in each file. For instance:

    /path/to/relaxation1.traj 200 /path/to/relaxation2.traj 150

    This will overrule the src and pattern that you specify!

    a2g_args (dict): Keyword arguments for fairchem.core.preprocessing.AtomsToGraphs()

    default options will work for most users

    If you are using this for a training dataset, set “r_energy”:True, “r_forces”:True, and/or “r_stress”:True as appropriate In that case, energy/forces must be in the files you read (ex. OUTCAR)

    ase_read_args (dict): Keyword arguments for ase.io.read()

    keep_in_memory (bool): Store data in memory. This helps avoid random reads if you need

    to iterate over a dataset many times (e.g. training for many epochs). Not recommended for large datasets.

    include_relaxed_energy (bool): Include the relaxed energy in the resulting data object.

    The relaxed structure is assumed to be the final structure in the file (e.g. the last frame of a .traj).

    use_tqdm (bool): Use TQDM progress bar when initializing dataset

    atoms_transform_args (dict): Additional keyword arguments for the atoms_transform callable

    transform_args (dict): Additional keyword arguments for the transform callable

    key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used

    in the model with the corresponding property as it was named in the dataset. Only need to use if the name is different.

  • atoms_transform (callable, optional) – Additional preprocessing function applied to the Atoms object. Useful for applying tags, for example.

  • transform (callable, optional) – Additional preprocessing function for the Data object

_load_dataset_get_ids(config) → list[str]#
get_atoms(idx: str) → ase.Atoms#
sample_property_metadata(num_samples: int = 100) → dict#
get_relaxed_energy(identifier) → float#
class core.datasets.ase_datasets.AseDBDataset(config: dict, atoms_transform: Callable[[ase.Atoms, Any, Ellipsis], ase.Atoms] = apply_one_tags)#

Bases: AseAtomsDataset

This Dataset connects to an ASE Database, allowing the storage of atoms objects with a variety of backends including JSON, SQLite, and database server options.

For more information, see: https://databases.fysik.dtu.dk/ase/ase/db/db.html

Parameters:

  • config (dict) –

    src (str): Either

    • the path an ASE DB,

    • the connection address of an ASE DB,

    • a folder with multiple ASE DBs,

    • a list of folders with ASE DBs

    • a glob string to use to find ASE DBs, or

    • a list of ASE db paths/addresses.

    If a folder, every file will be attempted as an ASE DB, and warnings are raised for any files that can’t connect cleanly

    Note that for large datasets, ID loading can be slow and there can be many ids, so it’s advised to make loading the id list as easy as possible. There is not an obvious way to get a full list of ids from most ASE dbs besides simply looping through the entire dataset. See the AseLMDBDataset which was written with this usecase in mind.

    connect_args (dict): Keyword arguments for ase.db.connect()

    select_args (dict): Keyword arguments for ase.db.select()

    You can use this to query/filter your database

    a2g_args (dict): Keyword arguments for fairchem.core.preprocessing.AtomsToGraphs()

    default options will work for most users

    If you are using this for a training dataset, set “r_energy”:True, “r_forces”:True, and/or “r_stress”:True as appropriate In that case, energy/forces must be in the database

    keep_in_memory (bool): Store data in memory. This helps avoid random reads if you need

    to iterate over a dataset many times (e.g. training for many epochs). Not recommended for large datasets.

    atoms_transform_args (dict): Additional keyword arguments for the atoms_transform callable

    transforms (dict[str, dict]): Dictionary specifying data transforms as {transform_function: config}

    where config is a dictionary specifying arguments to the transform_function

    key_mapping (dict[str, str]): Dictionary specifying a mapping between the name of a property used

    in the model with the corresponding property as it was named in the dataset. Only need to use if the name is different.

  • atoms_transform (callable, optional) – Additional preprocessing function applied to the Atoms object. Useful for applying tags, for example.

  • transform (callable, optional) – deprecated?

_load_dataset_get_ids(config: dict) → list[int]#
get_atoms(idx: int) → ase.Atoms#

Get atoms object corresponding to datapoint idx. Useful to read other properties not in data object. :param idx: index in dataset :type idx: int

Returns:

ASE atoms corresponding to datapoint idx

Return type:

atoms

static connect_db(address: str | pathlib.Path, connect_args: dict | None = None) → ase.db.core.Database#
__del__()#
sample_property_metadata(num_samples: int = 100) → dict#
abstract get_relaxed_energy(identifier)#
Contents