Merge pull request #33 from Benzoin96485/simulation

NVT Simulation

README.md

@@ -56,7 +56,9 @@ Enerzyme reads the `<model directory>` for the model configuration, load the mod
 
 ## Simulation
 
-Scanning on the distance between two atoms is supported.
+Supported simulation types:
+- Flexible scanning on the distance between two atoms.
+- Constrained Langevin MD
 
 ```bash
 enerzyme simulate -c <configuration yaml file> -o <output directory> -m <model directory>

enerzyme/config/nvt_md.yaml

@@ -0,0 +1,19 @@
+Simulation:
+  environment: "ase"                                # simulation environment: ase
+  dtype: float64                                    # pytorch data type throughout the computation: float32 (single) / float64 (double), default float64, for better numerical behavior during optimization
+  cuda: true                                        # if cuda is searched and used if possible
+  task: md                                          # molecular dynamics
+  idx_start_from: 1                                 # for your confusion of indices starting from 0 or 1
+  neighbor_list: full                               # you want this to be the same as how the model is trained
+  fs_in_t: 1                                        # numerical value of one femtosecond in the time step, friction, etc.
+  Hartree_in_E: 1.0                                 # numerical value of one Hartree in the model
+  integrate:                                        
+    integrator: Langevin                            # md integrator: Langevin
+    time_step: 0.5                                  # md time step
+    temperature_in_K: 300                           # temperature in NVT ensemble
+    friction: 0.01                                  # friction coefficient for Langevin dynamics
+    n_step: 100000                                  # number of maximum step
+System:
+  structure_file: "/your/initial/configuration.xyz" # initial configuration of simulation
+  charge: -1                                        # total charge of the structure (only useful if the model is charge-aware)
+  multiplicity: 1                                   # multiplicity of the structure (only useful if the model is spin-aware)

enerzyme/data/transform.py

@@ -39,7 +39,7 @@ def load_atomic_energy(atomic_energy_path: str) -> pd.DataFrame:
 
 
 class AtomicEnergyTransform:
-    def __init__(self, atomic_energy_path: str, *args, **kwargs) -> None:
+    def __init__(self, atomic_energy_path: str, simulation_mode=False, *args, **kwargs) -> None:
         self.atomic_energies = load_atomic_energy(atomic_energy_path)
         self.transform_type = "shift"
 
@@ -54,8 +54,8 @@ def transform(self, new_input: Dict[str, Iterable]) -> None:
 
     def inverse_transform(self, new_output: Dict[str, Iterable]) -> None:
         if len(new_output["Za"]) == 1:
-            for i in tqdm(range(len(new_output["E"]))):
-                new_output["E"][i] -= sum(self.atomic_energies.loc[new_output["Za"][0]]["atomic_energy"])
+            for i in range(len(new_output["E"])):
+                new_output["E"][i] += sum(self.atomic_energies.loc[new_output["Za"][0]]["atomic_energy"])
         else:
             for i in range(len(new_output["E"])):
                 new_output["E"][i] += sum(self.atomic_energies.loc[new_output["Za"][i]]["atomic_energy"])

enerzyme/tasks/simulation.py

@@ -1,15 +1,16 @@
 import os.path as osp
+from addict import Dict
 import numpy as np
 import ase
 import torch
 from ase.constraints import FixAtoms, FixBondLengths
 from ase.calculators.calculator import Calculator, all_changes
 from ase.optimize import BFGS
+from ase.units import Hartree, fs
 from ..data import full_neighbor_list
 from ..utils import logger
 from .trainer import DTYPE_MAPPING, _decorate_batch_output, _decorate_batch_input
 
-
 class ASECalculator(Calculator):
     implemented_properties = ["energy", "forces", "dipole", "charges"]
 
@@ -19,6 +20,11 @@ def __init__(
         restart=None,
         label=None,
         atoms=None,
+        device=None,
+        dtype=None,
+        transform=None,
+        neighbor_list_type=None,
+        Hartree_in_E=1,
         **params
     ):
         Calculator.__init__(
@@ -27,15 +33,16 @@ def __init__(
         self.model = model
         self.N = 0
         self.positions = None
-        self.device = self.parameters.device
+        self.device = device
         self.pbc = np.array([False])
         self.cell = None
         self.cell_offsets = None
-        self.dtype = self.parameters.dtype
-        self.neighbor_list_type = self.parameters.neighbor_list_type
-        self.transform = self.parameters.transform
+        self.dtype = dtype
+        self.neighbor_list_type = neighbor_list_type
+        self.transform = transform
+        self.Hartree_in_E = Hartree_in_E
 
-    def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
+    def calculate(self, atoms=None, properties=["energy", "forces", "dipole", "charges"], system_changes=all_changes):
         Calculator.calculate(self, atoms, properties, system_changes)
         features = {
             "Q": self.parameters.charge,
@@ -60,9 +67,9 @@ def calculate(self, atoms=None, properties=["energy"], system_changes=all_change
             targets=None
         )
         self.transform.inverse_transform(output)
-        self.results["energy"] = output["E"][0]
+        self.results["energy"] = output["E"][0] / self.Hartree_in_E * Hartree
         self.results["dipole"] = output["M2"][0]
-        self.results["forces"] = output["Fa"][0]
+        self.results["forces"] = output["Fa"][0] / self.Hartree_in_E * Hartree
         self.results["charges"] = output["Qa"][0]
 
 
@@ -75,8 +82,11 @@ def __init__(self, config, model, out_dir, transform):
         self.multiplicity = config.System.multiplicity
         self.transform = transform
         self.idx_start_from = config.Simulation.get("idx_start_from", 1)
+        self.integrate_config = config.Simulation.get("integrate", None)
         self.constraint_config = config.Simulation.get("constraint", None)
         self.sampling_config = config.Simulation.get("sampling", None)
+        self.fs_in_t = config.Simulation.get("fs_in_t", 1)
+        self.log_interval = config.Simulation.get("log_interval", 20)
         self.neighbor_list_type = config.Simulation.get("neighbor_list", "full")
         self.cuda = config.Simulation.get('cuda', False)
         self.dtype = DTYPE_MAPPING[config.Simulation.get("dtype", "float64")]
@@ -85,7 +95,7 @@ def __init__(self, config, model, out_dir, transform):
         self.model = model.to(self.device).type(self.dtype)
         self.model.eval()
         self.out_dir = out_dir
-        self.simulation_config = {k: v for k, v in config.Simulation.items() if not hasattr(self, k)}
+        # self.simulation_config = {k: (v.to_dict() if hasattr(v, "to_dict") else v) for k, v in config.Simulation.items() if not hasattr(self, k)}
         # initialize
         getattr(self, f"_init_{self.environment}_env")()
 
@@ -99,7 +109,7 @@ def _init_ase_env(self):
             dtype=self.dtype,
             neighbor_list_type=self.neighbor_list_type,
             transform=self.transform,
-            **self.simulation_config
+            #**self.simulation_config
         )
         self.system.calc = calculator
         if self.constraint_config is not None:
@@ -134,5 +144,21 @@ def write_xyz(atoms=None):
                 logger.info(f"Final energy: {self.system.get_potential_energy()}")
                 logger.info(f"Final distance: {self.system.get_distance(i0, i1)}")
 
+    def _run_md(self) -> None:
+        if self.integrate_config.integrator.lower() == "langevin":
+            from ase.md.langevin import Langevin
+            dyn = Langevin(self.system, 
+                timestep=self.integrate_config.time_step * fs / self.fs_in_t,
+                temperature_K=self.integrate_config.temperature_in_K,
+                friction=self.integrate_config.friction,
+                logfile="-",
+                loginterval=self.log_interval
+            )
+            def write_xyz(atoms=None):
+                ase.io.write(osp.join(self.out_dir, f"md.traj.xyz"), atoms, append=True)
+            dyn.attach(write_xyz, interval=self.log_interval, atoms=self.system) 
+            dyn.run(self.integrate_config.n_step)
+
+
     def run(self):
         getattr(self, f"_run_{self.task}")()