m3gnet.models package

Module contents

Graph pretrained

class m3gnet.models.BasePotential(*args, **kwargs)

Bases: keras.engine.training.Model, abc.ABC

Potential abstract class

call(graph: Union[m3gnet.graph._types.MaterialGraph, pymatgen.core.structure.Structure, pymatgen.core.structure.Molecule, List], include_forces: bool = True, include_stresses: bool = True)

Apply the potential to a graph Args:

graph (Union[MaterialGraph, Structure, Molecule, List]): a

structure, molecule, graph or a list repr of a graph

include_forces (bool): whether to include forces as outputs include_stresses (bool: whether to include stresses as outputs

Returns: energy [forces, stress]

get_ef(obj: Union[pymatgen.core.structure.Structure, pymatgen.core.structure.Molecule, ase.atoms.Atoms, m3gnet.graph._types.MaterialGraph, List]) → tuple

get energy and force from a Structure, a graph or a list repr of a graph Args:

obj (Union[StructureOrMolecule, MaterialGraph, List]): a structure,

material graph or list repr of a graph

Returns:

get_ef_tensor(graph: List[tensorflow.python.framework.ops.Tensor]) → tuple

get energy and force from a list repr of graph Args:

graph (List[tf.Tensor]: a list repr of a graph

Returns:

get_efs(obj: Union[pymatgen.core.structure.Structure, pymatgen.core.structure.Molecule, ase.atoms.Atoms, m3gnet.graph._types.MaterialGraph, List], include_stresses: bool = True)

get energy and force from a Structure, a graph or a list repr of a graph Args:

obj (Union[StructureOrMolecule, MaterialGraph, List]): a structure,

material graph or list repr of a graph

include_stresses (bool): whether to include stress

Returns:

get_efs_tensor(graph: List[tensorflow.python.framework.ops.Tensor], include_stresses: bool = True) → tuple

get energy and force from a list repr of a graph Args:

graph (List[tf.Tensor]): a list repr of a graph include_stresses (bool): whether to include stress

Returns:

get_energies(graph: List)

Compute the energy of a MaterialGraph Args:

graph: List, a graph from structure

Returns: energy values, size [Ns]

get_forces(graph: List)

Compute forces of a graph given the atom positions Args:

graph: List, a graph from structure

Returns: forces [Na, 3]

get_stresses(graph: List)

Compute stress of a graph given the atom positions Args:

graph: List, a graph from structure

Returns: stresses [Ns, 3, 3]

class m3gnet.models.GraphModelMixin(*args, **kwargs)

Bases: keras.engine.base_layer.Layer

GraphModelMixin adds the following functionality to a graph model

• predict_structure

• predict_structures

• predict_graph

• predict_graphs

predict_graph(graph: Union[m3gnet.graph._types.MaterialGraph, List]) → tensorflow.python.framework.ops.Tensor

predict properties from a graph Args:

graph (Union[MaterialGraph, List]): a material graph, either in

object or list repr

Returns: predicted property

predict_graphs(graph_list: List[Union[m3gnet.graph._types.MaterialGraph, List]], batch_size: int = 128) → tensorflow.python.framework.ops.Tensor

predict properties from graphs Args:

graph_list (List[Union[MaterialGraph, List]]): a list of material

graph, either in object or list repr

batch_size (int): batch size

Returns: predicted properties

predict_structure(structure: Union[pymatgen.core.structure.Structure, pymatgen.core.structure.Molecule, ase.atoms.Atoms]) → tensorflow.python.framework.ops.Tensor

predict properties from structure Args:

structure (StructureOrMolecule): a pymatgen Structure/Molecule,

or ase.Atoms

Returns: predicted values

predict_structures(structures: List[Union[pymatgen.core.structure.Structure, pymatgen.core.structure.Molecule, ase.atoms.Atoms]], batch_size: int = 128) → tensorflow.python.framework.ops.Tensor

predict properties from structures Args:

structures (list[StructureOrMolecule]): a list of structures batch_size (int): batch size for the prediction

Returns: predicted values

class m3gnet.models.M3GNet(*args, **kwargs)

Bases: m3gnet.models._base.GraphModelMixin, keras.engine.training.Model

The main M3GNet model

call(graph: List, **kwargs) → tensorflow.python.framework.ops.Tensor

Args:

graph (list): list repr of a MaterialGraph **kwargs:

Returns:

classmethod from_config(config: dict) → m3gnet.models._m3gnet.M3GNet

Construct the model from a config dict Args:

config (dict): config dict from get_config method

Returns: new M3GNet instance

classmethod from_dir(dirname: str, custom_objects: Optional[dict] = None) → m3gnet.models._m3gnet.M3GNet

Load the model from a directory

Args:

dirname (str): directory to save the model custom_objects (dict): dictionary for custom object

Returns: M3GNet model

get_config()

Get config dict for serialization Returns:

classmethod load(model_name: str = 'MP-2021.2.8-EFS') → m3gnet.models._m3gnet.M3GNet

Load the model weights from pre-trained model Args:

model_name (str): model name or the path for saved model. Defaults to “MP-2021.2.8-EFS”.

Returns: M3GNet object.

save(dirname: str)

Saves the model to a directory.

Args:

dirname (str): directory to save the model

set_element_refs(element_refs: numpy.ndarray)

Set element reference for the property Args:

element_refs (np.ndarray): element reference value for the

extensive property

class m3gnet.models.M3GNetCalculator(potential: m3gnet.models._base.Potential, compute_stress: bool = True, stress_weight: float = 1.0, **kwargs)

Bases: ase.calculators.calculator.Calculator

M3GNet calculator based on ase Calculator

calculate(atoms: Optional[ase.atoms.Atoms] = None, properties: Optional[list] = None, system_changes: Optional[list] = None)

Args:

atoms (ase.Atoms): ase Atoms object properties (list): list of properties to calculate system_changes (list): monitor which properties of atoms were

changed for new calculation. If not, the previous calculation results will be loaded.

Returns:

implemented_properties: List[str] = ['energy', 'free_energy', 'forces', 'stress']

Properties calculator can handle (energy, forces, …)

class m3gnet.models.MolecularDynamics(atoms: ase.atoms.Atoms, potential: Union[m3gnet.models._base.Potential, str] = 'MP-2021.2.8-EFS', ensemble: str = 'nvt', temperature: int = 300, timestep: float = 1.0, pressure: float = 6.324209121801212e-07, taut: Optional[float] = None, taup: Optional[float] = None, compressibility_au: Optional[float] = None, trajectory: Optional[Union[str, ase.io.trajectory.Trajectory]] = None, logfile: Optional[str] = None, loginterval: int = 1, append_trajectory: bool = False)

Bases: object

Molecular dynamics class

run(steps: int)

Thin wrapper of ase MD run Args:

steps (int): number of MD steps

Returns:

set_atoms(atoms: ase.atoms.Atoms)

Set new atoms to run MD Args:

atoms (Atoms): new atoms for running MD

Returns:

class m3gnet.models.Potential(*args, **kwargs)

Bases: m3gnet.models._base.BasePotential

Defines the Potential class. New potential should subclass from this class and define the “get_energies” method

get_energies(graph: List) → tensorflow.python.framework.ops.Tensor

get energies from a list repr of a graph Args:

graph (List): list repr of a graph

Returns:

class m3gnet.models.Relaxer(potential: Optional[Union[m3gnet.models._base.Potential, str]] = None, optimizer: Union[ase.optimize.optimize.Optimizer, str] = 'FIRE', relax_cell: bool = True, stress_weight: float = 0.01)

Bases: object

Relaxer is a class for structural relaxation

relax(atoms: ase.atoms.Atoms, fmax: float = 0.1, steps: int = 500, traj_file: Optional[str] = None, interval=1, verbose=False, **kwargs)

Args:

atoms (Atoms): the atoms for relaxation fmax (float): total force tolerance for relaxation convergence.

Here fmax is a sum of force and stress forces

steps (int): max number of steps for relaxation traj_file (str): the trajectory file for saving interval (int): the step interval for saving the trajectories **kwargs:

Returns: