Merge pull request #348 from deepmodeling/damp

feat: implement dft damping for ase calc

Author: anyangml

lambench/metrics/direct_task_weights.yml

@@ -35,8 +35,8 @@ MD22:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
-  energy_std: 0.007941836149915322
-  force_std: 1.1391327961625524
+  energy_std: 0.008959619353114803
+  force_std: 1.1964522496892305
   virial_std: null
 REANN_CO2_Ni100:
   domain: Catalysis

lambench/metrics/results/metadata.json

@@ -53,7 +53,7 @@
     },
     "MD22": {
       "DISPLAY_NAME": "MD22",
-      "DESCRIPTION": "Dataset containing MD trajectories of the 42-atom tetrapeptide Ac-Ala3-NHMe from the MD22 benchmark set. Calculations were performed using FHI-aims and i-Pi software at the DFT-PBE+MBD level of theory. Trajectories were sampled at temperatures between 400-500 K at 1 fs resolution. [https://www.science.org/doi/10.1126/sciadv.adf0873]",
+      "DESCRIPTION": "Dataset containing MD trajectories of the 42-atom tetrapeptide Ac-Ala3-NHMe from the MD22 benchmark set. Calculations were performed using FHI-aims and i-Pi software at the DFT-PBE+MBD level of theory. The dataset was relabeled using Gaussian with PBE/6-31G(d). Trajectories were sampled at temperatures between 400-500 K at 1 fs resolution. [https://www.science.org/doi/10.1126/sciadv.adf0873]",
       "domain": "Molecules",
       "energy_rmse": {
         "DISPLAY_NAME": "E RMSE (meV)",

lambench/models/ase_models.py

@@ -2,7 +2,7 @@
 import logging
 from functools import cached_property
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Callable, Literal, Optional
 
 import dpdata
 import numpy as np
@@ -12,10 +12,12 @@
 )
 from ase import Atoms
 from ase.calculators.calculator import Calculator
+from ase.calculators.mixing import SumCalculator
 from ase.constraints import FixSymmetry
 from ase.filters import FrechetCellFilter
 from ase.io import write
 from ase.optimize import FIRE
+from dftd3.ase import DFTD3
 from tqdm import tqdm
 
 from lambench.models.basemodel import BaseLargeAtomModel
@@ -179,7 +181,7 @@ def evaluate(
             import torch
 
             torch.set_default_dtype(torch.float32)
-            return self.run_ase_dptest(self, task.test_data)
+            return self.run_ase_dptest(self, task.test_data, task.dispersion_correction)
         elif isinstance(task, CalculatorTask):
             if task.task_name == "nve_md":
                 from lambench.tasks.calculator.nve_md.nve_md import (
@@ -265,7 +267,12 @@ def evaluate(
             )
 
     @staticmethod
-    def run_ase_dptest(model: ASEModel, test_data: Path) -> dict:
+    def run_ase_dptest(
+        model: ASEModel,
+        test_data: Path,
+        dispersion_correction: Literal["d3bj", "d3zero"] | None = None,
+        # check all supported levels at dftd3.qcschema._available_levels
+    ) -> dict:
         # Add fparam for charge and spin multiplicity if needed
         datatype = DataType(
             "fparam",
@@ -277,6 +284,10 @@ def run_ase_dptest(model: ASEModel, test_data: Path) -> dict:
         dpdata.LabeledSystem.register_data_type(datatype)
 
         calc = model.calc
+        if dispersion_correction:
+            calc = SumCalculator(
+                [calc, DFTD3(method="PBE", dispersion_correction=dispersion_correction)]
+            )
 
         energy_err = []
         energy_pre = []

lambench/tasks/direct/direct_tasks.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import ClassVar
+from typing import ClassVar, Literal
 from lambench.tasks.base_task import BaseTask
 from lambench.databases.direct_predict_table import DirectPredictRecord
 
@@ -12,6 +12,8 @@ class DirectPredictTask(BaseTask):
 
     record_type: ClassVar = DirectPredictRecord
     task_config: ClassVar = Path(__file__).parent / "direct_tasks.yml"
+    dispersion_correction: Literal["d3bj", "d3zero"] | None = None
 
     def __init__(self, task_name: str, **kwargs):
         super().__init__(task_name=task_name, test_data=kwargs["test_data"])
+        self.dispersion_correction = kwargs.get("dispersion_correction")

lambench/tasks/direct/direct_tasks.yml

@@ -1,27 +1,29 @@
 ANI:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/ANI"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/ANI"
 HEA25_S:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/HEA25S"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/HEA25S"
 HEA25_bulk:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/HEA25"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/HEA25"
 MoS2:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/MoS2"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/MoS2"
 MD22:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/MD22"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/MD22"
 REANN_CO2_Ni100:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/REANN_CO2_Ni100"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/REANN_CO2_Ni100"
 NequIP_NC_2022:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/NequIP_NC_2022"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/NequIP_NC_2022"
 AIMD-Chig:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/AIMD_chig"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/AIMD_chig"
 Cu_MgO_catalysts:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/Cu_MgO_CO2"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/Cu_MgO_CO2"
+  dispersion_correction: d3zero
 Si_ZEO22:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/Si_ZEO22"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/Si_ZEO22"
+  dispersion_correction: d3bj
 HPt_NC_2022:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/HPt_NC2022"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/HPt_NC2022"
 Ca_batteries_CM2021:
-  test_data: "/bohr/lambench-ood-zwtr/v3/LAMBench-TestData-v3/Ca_batteries"
+  test_data: "/bohr/lambench-ood-zwtr/v4/LAMBench-TestData-v3/Ca_batteries"
 ## DEPRECATED
 # Collision:
 #   test_data: "/bohr/lambench-ood-zwtr/v2/LAMBench-TestData-v2/Collision"

lambench/workflow/dflow.py

@@ -12,7 +12,11 @@
 from dflow.plugins.bohrium import BohriumDatasetsArtifact, create_job_group
 from dflow.plugins.dispatcher import DispatcherExecutor
 from dflow.python import OP, Artifact, PythonOPTemplate
+
 import dpdata
+import dftd3
+import cffi
+import pycparser
 
 import lambench
 from lambench.models.basemodel import BaseLargeAtomModel
@@ -60,7 +64,8 @@ def submit_tasks_dflow(
                 image=model.virtualenv,
                 envs={k: v for k, v in os.environ.items() if k.startswith("MYSQL")},
                 python_packages=[
-                    Path(package.__path__[0]) for package in [lambench, dpdata]
+                    Path(package.__path__[0])
+                    for package in [lambench, dpdata, dftd3, cffi, pycparser]
                 ],
             ),
             parameters={

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
 	"tqdm",
 	"dpdata @ git+https://github.com/deepmodeling/dpdata.git@devel#egg=dpdata",
 	"pandas",
+	"dftd3"
 ]
 
 authors = [

tests/metrics/test_visualization.py

@@ -14,7 +14,7 @@ def test_aggregate_ood_results_for_one_model(
     model.show_calculator_task = False
     aggregator = ResultsFetcher()
     result = aggregator.aggregate_ood_results_for_one_model(model=model)
-    np.testing.assert_almost_equal(result["Molecules"], 0.234724350, decimal=5)
+    np.testing.assert_almost_equal(result["Molecules"], desired=0.22748765, decimal=5)
     np.testing.assert_almost_equal(result["Inorganic Materials"], 0.2972349, decimal=5)
     assert result["Catalysis"] is None
     with caplog.at_level(logging.WARNING):