From 500681a701c2cd2b3342d4d57e69ad2444bc6fe8 Mon Sep 17 00:00:00 2001 From: steinmig Date: Tue, 22 Jul 2025 11:04:23 -0400 Subject: [PATCH 1/5] more information in batch --- nff/io/ase.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/nff/io/ase.py b/nff/io/ase.py index f38c84fc..b805d351 100644 --- a/nff/io/ase.py +++ b/nff/io/ase.py @@ -1,5 +1,6 @@ """ASE wrapper for the Neural Force Field.""" +from typing import List import copy import numpy as np @@ -85,7 +86,6 @@ def convert_props_units(self, target_unit): self.props = const.convert_units(self.props, conversion_factor) self.props.update({"units": target_unit}) - return def get_mol_nbrs(self, r_cut=95): """Dense directed neighbor list for each molecule, in case that's needed @@ -248,6 +248,7 @@ def get_batch(self): if self.pbc.any(): self.props["cell"] = torch.Tensor(np.array(self.cell)) self.props["lattice"] = self.cell.tolist() + self.props["pbc"] = self.pbc.tolist() self.props["nxyz"] = torch.Tensor(self.get_nxyz()) if self.props.get("num_atoms") is None: @@ -259,9 +260,11 @@ def get_batch(self): if self.mol_idx is not None: self.props["mol_idx"] = self.mol_idx + self.props["device"] = self.device + return self.props - def get_list_atoms(self): + def get_list_atoms(self) -> List[Atoms]: """Returns a list of ASE Atoms objects, each representing a molecule in the system. Returns: @@ -285,7 +288,7 @@ def get_list_atoms(self): cells = torch.split(torch.Tensor(self.props["lattice"]), 3) else: cells = torch.unsqueeze(torch.Tensor(np.array(self.cell)), 0).repeat(len(mol_split_idx), 1, 1) - Atoms_list = [] + atoms_list = [] for i, molecule_xyz in enumerate(positions): atoms = Atoms( @@ -299,9 +302,9 @@ def get_list_atoms(self): # of any of the atoms atoms.set_masses(masses[i]) - Atoms_list.append(atoms) + atoms_list.append(atoms) - return Atoms_list + return atoms_list def update_num_atoms(self): """Update the number of atoms in the system. From 7457f3f53c03f1a3c345c14ae4edcd6e22a918ed Mon Sep 17 00:00:00 2001 From: steinmig Date: Tue, 22 Jul 2025 11:05:13 -0400 Subject: [PATCH 2/5] type and force inclusion fixes --- nff/io/bias_calculators.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/nff/io/bias_calculators.py b/nff/io/bias_calculators.py index f12be966..ad9084cf 100644 --- a/nff/io/bias_calculators.py +++ b/nff/io/bias_calculators.py @@ -47,7 +47,7 @@ class BiasBase(NeuralFF): def __init__( self, model, - cv_defs: list[dict], + cv_defs: List[dict], equil_temp: float = 300.0, device="cpu", en_key="energy", @@ -290,6 +290,8 @@ def calculate( requires_stress = "stress" in self.properties if requires_stress: kwargs["requires_stress"] = True + if "forces" in self.properties: + kwargs["requires_forces"] = True if getattr(self, "model_kwargs", None) is not None: kwargs.update(self.model_kwargs) @@ -298,8 +300,12 @@ def calculate( # change energy and force to numpy array and eV model_energy = prediction[self.en_key].detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL) - if grad_key in prediction: - model_grad = prediction[grad_key].detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL) + gradient = prediction.get(grad_key) + forces = prediction.get("forces") + if gradient is not None: + model_grad = gradient.detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL) + elif forces is not None: + model_grad = - forces.detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL) else: raise KeyError(grad_key) @@ -389,7 +395,7 @@ class with neural force field def __init__( self, model, - cv_defs: list[dict], + cv_defs: List[dict], dt: float, friction_per_ps: float, equil_temp: float = 300.0, @@ -570,7 +576,7 @@ class with neural force field def __init__( self, model, - cv_defs: list[dict], + cv_defs: List[dict], dt: float, friction_per_ps: float, amd_parameter: float, @@ -801,7 +807,7 @@ class WTMeABF(eABF): def __init__( self, model, - cv_defs: list[dict], + cv_defs: List[dict], dt: float, friction_per_ps: float, equil_temp: float = 300.0, @@ -982,7 +988,7 @@ class AttractiveBias(NeuralFF): def __init__( self, model, - cv_defs: list[dict], + cv_defs: List[dict], gamma=1.0, device="cpu", en_key="energy", From 67e852ae86c85334f65190dd6c9295a6af641d0a Mon Sep 17 00:00:00 2001 From: steinmig Date: Tue, 22 Jul 2025 11:05:45 -0400 Subject: [PATCH 3/5] add new classmethod to load foundation models based on path --- nff/nn/models/mace.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/nff/nn/models/mace.py b/nff/nn/models/mace.py index 2a332c1d..f9bc8472 100644 --- a/nff/nn/models/mace.py +++ b/nff/nn/models/mace.py @@ -265,6 +265,15 @@ def load_foundations( NffScaleMACE: NffScaleMACE foundational model. """ mace_model_path = get_mace_mp_model_path(model) + return cls.load_foundations_path(mace_model_path, map_location, default_dtype) + + @classmethod + def load_foundations_path( + cls, + mace_model_path: str, + map_location: str = "cpu", + default_dtype: Literal["", "float32", "float64"] = "float32", + ) -> NffScaleMACE: mace_model = torch.load(mace_model_path, map_location=map_location) init_params = get_init_kwargs_from_model(mace_model) model_dtype = get_model_dtype(mace_model) From 348055ebbef76bd3d6009c91c2d4bcc36cb06689 Mon Sep 17 00:00:00 2001 From: steinmig Date: Tue, 22 Jul 2025 11:06:03 -0400 Subject: [PATCH 4/5] cleanup --- nff/utils/constants.py | 28 ++++++++++------------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/nff/utils/constants.py b/nff/utils/constants.py index 4238a5d9..a566c85d 100644 --- a/nff/utils/constants.py +++ b/nff/utils/constants.py @@ -191,26 +191,18 @@ ELEC_CONFIG = { - "1": [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "6": [2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "7": [2, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "8": [2, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "9": [2, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "11": [2, 2, 6, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "14": [2, 2, 6, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "16": [2, 2, 6, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "17": [2, 2, 5, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - "86": [2, 2, 6, 2, 6, 2, 10, 6, 2, 10, 6, 2, 14, 10, 6, 2, 6, 10, 4], + 1: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 6: [2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 7: [2, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 8: [2, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 9: [2, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 11: [2, 2, 6, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 14: [2, 2, 6, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 16: [2, 2, 6, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 17: [2, 2, 5, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 86: [2, 2, 6, 2, 6, 2, 10, 6, 2, 10, 6, 2, 14, 10, 6, 2, 6, 10, 4], } -ELEC_CONFIG = {int(key): val for key, val in ELEC_CONFIG.items()} - - -# with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), -# "elec_config.json"), "r") as f: -# ELEC_CONFIG = json.load(f) -# ELEC_CONFIG = {int(key): val for key, val in ELEC_CONFIG.items()} - def convert_units(props, conversion_dict): """Converts dictionary of properties to the desired units. From d65510f191806146942c7f8a4c9b804010b0b8d1 Mon Sep 17 00:00:00 2001 From: steinmig Date: Tue, 22 Jul 2025 11:06:34 -0400 Subject: [PATCH 5/5] introduce general potential --- nff/io/ase_calcs.py | 15 ++++------ nff/io/potential.py | 71 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 10 deletions(-) create mode 100644 nff/io/potential.py diff --git a/nff/io/ase_calcs.py b/nff/io/ase_calcs.py index b19a0918..1844ddde 100644 --- a/nff/io/ase_calcs.py +++ b/nff/io/ase_calcs.py @@ -7,11 +7,12 @@ given geometry. Both calculators are used to perform molecular dynamics simulations and geometry optimizations using NFF AtomsBatch objects. """ +from __future__ import annotations import os import sys from collections import Counter -from typing import List, Union +from typing import List, Union, TYPE_CHECKING import numpy as np import torch @@ -33,6 +34,8 @@ from nff.utils.cuda import batch_detach, batch_to from nff.utils.geom import batch_compute_distance, compute_distances from nff.utils.scatter import compute_grad +if TYPE_CHECKING: + from nff.io.potential import Potential HARTREE_TO_EV = HARTREE_TO_KCAL_MOL / EV_TO_KCAL_MOL @@ -53,7 +56,7 @@ class NeuralFF(Calculator): def __init__( self, - model, + model: Potential, device="cpu", jobdir=None, en_key="energy", @@ -134,15 +137,10 @@ def calculate( Calculator.calculate(self, atoms, self.properties, system_changes) # run model - # atomsbatch = AtomsBatch(atoms) - # batch_to(atomsbatch.get_batch(), self.device) batch = batch_to(atoms.get_batch(), self.device) # add keys so that the readout function can calculate these properties - # print("Properties: ", self.properties, "\n\n") - # print("en_key:", self.en_key) grad_key = self.en_key + "_grad" - # print("grad_key:", grad_key) batch[self.en_key] = [] batch[grad_key] = [] @@ -158,7 +156,6 @@ def calculate( kwargs.update(self.model_kwargs) prediction = self.model(batch, **kwargs) - # print(prediction.keys()) # change energy and force to numpy array conversion_factor: dict = const.conversion_factors.get((self.model_units, self.prediction_units), const.DEFAULT) @@ -1159,9 +1156,7 @@ def calculate( system_changes (default from ase) """ - # print("calculating ...") self.step += 1 - # print("step ", self.step, self.step*0.0005) if not any(isinstance(self.model, i) for i in UNDIRECTED): check_directed(self.model, atoms) diff --git a/nff/io/potential.py b/nff/io/potential.py new file mode 100644 index 00000000..2960838f --- /dev/null +++ b/nff/io/potential.py @@ -0,0 +1,71 @@ +from __future__ import annotations +from copy import deepcopy +from typing import Optional, Callable + +from ase.calculators.calculator import Calculator, all_changes +from ase.symbols import Symbols +import numpy as np +import torch + +from nff.io.ase_calcs import AtomsBatch + + + +class Potential(torch.nn.Module): + pass + + +class AsePotential(Potential): + + def __init__(self, calculator: Calculator, embedding_fun: Optional[Callable[[AtomsBatch], torch.Tensor]] = None) \ + -> None: + super().__init__() + self.calculator = calculator + self.embedding_fun = embedding_fun + + def __call__(self, batch: dict, **kwargs): + properties = ["energy"] + if kwargs.get("requires_stress", False): + properties.append("stress") + if kwargs.get("requires_forces", False): + properties.append("forces") + if kwargs.get("requires_dipole", False): + properties.append("dipole") + if kwargs.get("requires_charges", False): + properties.append("charges") + if kwargs.get("requires_embedding", False): + if self.embedding_fun is None: + raise RuntimeError("Required embedding but no embedding function provided.") + embedding = self.embedding_fun(batch) + else: + embedding = None + + nxyz = batch.get("nxyz") + if nxyz is None: + raise RuntimeError("Batch is missing 'nxyz' key.") + pbc = batch.get("pbc") + if pbc is not None: + pbc = np.array(pbc, dtype=bool) + cell = np.array(batch.get("cell")).reshape(3, 3) + else: + cell = None + atoms_batch = AtomsBatch( + symbols=Symbols(nxyz[:,0].detach().cpu().numpy()), + positions=nxyz[:, 1:4].detach().cpu().numpy(), + pbc=pbc, + cell=cell, + device=batch.get("device", "cpu") + ) + + self.calculator.calculate(atoms_batch, properties=properties, system_changes=all_changes) + results = deepcopy(self.calculator.results) + for key, value in results.items(): + if isinstance(value, str): + continue + if not hasattr(value, "__iter__"): + results[key] = torch.tensor([value], device=atoms_batch.device) + else: + results[key] = torch.tensor(value, device=atoms_batch.device) + if embedding is not None: + results["embedding"] = embedding + return results \ No newline at end of file