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#
This is an abstract Dataset that includes helpful utilities for turning |
|
This Dataset uses ase.io.read to load data from a directory on disk. |
|
This Dataset can read multiple structures from each file using ase.io.read. |
|
This Dataset connects to an ASE Database, allowing the storage of atoms objects |
Functions#
|
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)#