kgcnn.molecule.dynamics package

Submodules

kgcnn.molecule.dynamics.ase_calc module

class kgcnn.molecule.dynamics.ase_calc.AtomsToGraphConverter(properties: Optional[dict] = None)

Bases: object

Convert ase.Atoms object to a GraphDict dictionary.

Simple tool to get named properties from ase.Atoms. Note that the actual graph indices and connections have to be generated with GraphPreProcessorBase instances.

Example usage:

import numpy as np
from ase import Atoms
from kgcnn.md.ase_calc import AtomsToGraphConverter
atoms = Atoms('N2', positions=[[0, 0, -1], [0, 0, 1]])
trans = AtomsToGraphConverter({
    "node_number": "get_atomic_numbers", "node_coordinates": "get_positions",
    "node_symbol": "get_chemical_symbols"})
print(trans(atoms))

__call__(atoms: Union[List[ase.atoms.Atoms], ase.atoms.Atoms]) → kgcnn.data.base.MemoryGraphList

Make GraphDict objects from ase.Atoms.

Parameters

atoms (list) – List of ase.Atoms objects or single ASE atoms object.

Returns

List of GraphDict objects.

Return type

MemoryGraphList

__init__(properties: Optional[dict] = None)

Set up AtomsToGraphConverter converter.

Parameters

properties (dict) – Dictionary of graph properties linked to ase.Atoms get attribute methods. Default is {“node_number”: “get_atomic_numbers”, “node_coordinates”: “get_positions”, “node_symbol”: “get_chemical_symbols”}.

get_config()

Get config for this class.

class kgcnn.molecule.dynamics.ase_calc.KgcnnSingleCalculator(model_predictor=None, atoms_converter: Optional[kgcnn.molecule.dynamics.ase_calc.AtomsToGraphConverter] = None, squeeze_energy: bool = True, **kwargs)

Bases: ase.calculators.calculator.Calculator

ASE calculator for machine learning models from kgcnn.

calculate(atoms=None, properties=None, system_changes=None)

Do the calculation.

properties: list of str

List of what needs to be calculated. Can be any combination of ‘energy’, ‘forces’, ‘stress’, ‘dipole’, ‘charges’, ‘magmom’ and ‘magmoms’.

system_changes: list of str

List of what has changed since last calculation. Can be any combination of these six: ‘positions’, ‘numbers’, ‘cell’, ‘pbc’, ‘initial_charges’ and ‘initial_magmoms’.

Subclasses need to implement this, but can ignore properties and system_changes if they want. Calculated properties should be inserted into results dictionary like shown in this dummy example:

self.results = {'energy': 0.0,
                'forces': np.zeros((len(atoms), 3)),
                'stress': np.zeros(6),
                'dipole': np.zeros(3),
                'charges': np.zeros(len(atoms)),
                'magmom': 0.0,
                'magmoms': np.zeros(len(atoms))}

The subclass implementation should first call this implementation to set the atoms attribute and create any missing directories.

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

kgcnn.molecule.dynamics.base module

class kgcnn.molecule.dynamics.base.MolDynamicsModelPredictor(model: Optional[keras.src.models.model.Model] = None, model_inputs: Optional[Union[list, dict]] = None, model_outputs: Optional[Union[list, dict]] = None, graph_preprocessors: Optional[List[Callable]] = None, graph_postprocessors: Optional[List[Callable]] = None, store_last_input: bool = False, store_last_output: bool = False, copy_graphs_in_store: bool = False, use_predict: bool = False, predict_verbose: Union[str, int] = 0, batch_size: int = 32, update_from_last_input: Optional[list] = None, update_from_last_input_skip: Optional[int] = None)

Bases: object

Model predictor class that adds pre- and postprocessors to the keras model to be able to add transformation steps to convert for example input and output representations to fit MD programs like ase . The MolDynamicsModelPredictor receives a MemoryGraphList in call and returns a MemoryGraphList .

__call__(graph_list: kgcnn.data.base.MemoryGraphList) → kgcnn.data.base.MemoryGraphList

Prediction of the model wrapper.

Parameters

graph_list (MemoryGraphList) – List of graphs to predict e.g. energies and forces.

Returns

List of general return graph dictionaries from model output.

Return type

MemoryGraphList

__init__(model: Optional[keras.src.models.model.Model] = None, model_inputs: Optional[Union[list, dict]] = None, model_outputs: Optional[Union[list, dict]] = None, graph_preprocessors: Optional[List[Callable]] = None, graph_postprocessors: Optional[List[Callable]] = None, store_last_input: bool = False, store_last_output: bool = False, copy_graphs_in_store: bool = False, use_predict: bool = False, predict_verbose: Union[str, int] = 0, batch_size: int = 32, update_from_last_input: Optional[list] = None, update_from_last_input_skip: Optional[int] = None)

Initialize MolDynamicsModelPredictor class.

Parameters

• model (tf.keras.Model) – Single trained keras model.

• model_inputs (list, dict) – List or single dictionary for model inputs.

• model_outputs (list, dict) – List of model output names or dictionary of output mappings from keras model output to the names in the return GraphDict .

• graph_preprocessors (list) – List of graph preprocessors, see kgcnn.graph.preprocessor .

• graph_postprocessors (list) – List of graph postprocessors, see kgcnn.graph.postprocessor .

• use_predict (bool) – Whether to use model.predict() or call method model() .

• batch_size (int) – Optional batch size for prediction.

• store_last_input (bool) – Whether to store last input graph list for model input. Default is False.

• store_last_output (bool) – Whether to store last output graph list from model output. Default is False.

• copy_graphs_in_store (bool) – Whether to make a copy of the graph lists or keep reference. Default is False.

• update_from_last_input (list) – List of graph properties to copy from last input into current input. This is placed before graph preprocessors. Default is None.

• update_from_last_input_skip (int) – If set to a value, this will skip the update from last input at given number of calls. Uses counter. Default is None.

_test_timing(graph_list: kgcnn.data.base.MemoryGraphList, repetitions: int = 100) → float

Evaluate timing for prediction.

Parameters

graph_list (MemoryGraphList) – List of graphs to predict e.g. energies and forces.

Returns

Time for one call.

Return type

float

static _translate_properties(properties, translation) → dict

Translate general model output.

Parameters

• properties (list, dict) – List of properties or dictionary of properties.

• translation (str, list, dict) – List of names or dictionary of name mappings like ‘{new_name: old_name}’.

Returns

Return dictionary with keys from translation.

Return type

dict

load(file_path: str)

save(file_path: str)

Module contents