Feature nvalchemi d3 support

.github/workflows/CI.yaml

@@ -75,6 +75,11 @@ jobs:
 
           micromamba list
 
+      # Skip nvalchemiops if Python version is 3.10
+      - name: Install nvalchemi-toolkit-ops
+        if: matrix.python-version != '3.10'
+        run: pip install nvalchemi-toolkit-ops==0.3.0
+
       # this will figure out what datasets we will use in our tests
       # and save them to a file so we can hash this file to see if the cache needs updating
       # this will save "downloaded_datasets.txt" in the modelforge testing cache folder

devtools/conda-envs/test_env.yaml

@@ -43,4 +43,5 @@ dependencies:
       - pytorch2jax
       #- git+https://github.com/ArnNag/sake.git@nanometer
       - pytest-xdist
+      #- nvalchemi-toolkit-ops==0.3.0
 

docs/potentials.rst

@@ -281,7 +281,11 @@ Rather than having all 3 of these energy contributions returned by the core netw
 
 
     [potential.postprocessing_parameter.per_system_vdw_energy]
-    maximum_interaction_radius = "10.0 angstrom"
+    maximum_interaction_radius = "10.0 angstrom"  # this will only take effect when d3_engine is set to "nvalchemiops"
+                                                  # otherwise the "tad-dftd3" will handle neighbor list internally
+    parameter_set: str = "wB97M-D3(BJ)"
+    d3_engine = "nvalchemiops"
+    d3_parameters_path = "None"
 
     # this will be added to the per_system_energy that results from the `per_atom_energy` reduction operation
     [potential.postprocessing_parameter.sum_per_system_energy]

modelforge/data/dftd3_parameters/dftd3_parameters.yaml

@@ -0,0 +1,4 @@
+dftd3_parameters_0:
+  file_name: "dftd3_parameters_0.pt"
+  generation_date: "2026-02-23"
+  original_url: "https://www.chemie.uni-bonn.de/grimme/de/software/dft-d3/dftd3.tgz"

modelforge/potential/parameters.py

@@ -192,6 +192,9 @@ class DispersionPotential(ParametersBase):
     # note maximum interaction radius should be passed as a string with units or unit.Quantity;
     # it will be converted to float in appropriate unit system
     maximum_interaction_radius: float
+    parameter_set: str = "wB97M-D3(BJ)"
+    d3_engine: str = "nvalchemiops"
+    d3_parameters_path: str = "None"
 
     converted_units = field_validator(
         "maximum_interaction_radius",

modelforge/potential/potential.py

@@ -125,7 +125,7 @@ class PostProcessing(torch.nn.Module):
 
     def __init__(
         self,
-        postprocessing_parameter: Dict[str, Dict[str, bool]],
+        postprocessing_parameter: Dict[str, Dict[str, Any]],
         dataset_statistic: Dict[str, Dict[str, float]],
     ):
         """
@@ -215,6 +215,15 @@ def __init__(
                     ],
                     length_conversion_factor=length_conversion_factor,
                     energy_conversion_factor=energy_conversion_factor,
+                    parameter_set=postprocessing_parameter["per_system_vdw_energy"][
+                        "parameter_set"
+                    ],
+                    d3_engine=postprocessing_parameter["per_system_vdw_energy"][
+                        "d3_engine"
+                    ],
+                    d3_parameters_path=postprocessing_parameter[
+                        "per_system_vdw_energy"
+                    ]["d3_parameters_path"],
                 )
             )
             self._registered_properties.append("per_system_vdw_energy")
@@ -315,9 +324,6 @@ def __init__(
             Forward pass for the potential model, computing energy and forces that accepts individual tensors rather than NNPInput class, necessary for JIT compiled model.
         forward(input_data: NNPInput) -> Dict[str, torch.Tensor]
             Forward pass for the potential model, computing energy and forces.
-
-
-
         """
 
         super().__init__()
@@ -328,13 +334,14 @@ def __init__(
         )
         # note cannot jit compile the dispersion interactions as tad-dftd3 is not compatible with torchscript
         if "per_system_vdw_energy" in postprocessing._registered_properties:
+            # double check if nvalchemiops works with JIT
             log.warning(
-                "JIT compiling the postprocessing module with vdw interactions will not work."
+                "JIT compiling the postprocessing module with vdw interactions will not work if using tad-dftd3."
             )
+            log.warning("tad-dftd3 packaged is not compatible with torchscript.")
             log.warning(
-                "The tad-DFTD3 packaged used is not compatible with torchscript."
+                "Use nvalchemiops as the DFTD3 engine or disable JIT compilation by setting jit=False."
             )
-            log.warning("Disabling JIT compilation by setting jit=False.")
 
         self.postprocessing = (
             torch.jit.script(postprocessing) if jit else postprocessing
@@ -431,13 +438,12 @@ def _add_pairlist(
         core_output["local_d_ij"] = pairlist_output.local_cutoff.d_ij
         core_output["local_r_ij"] = pairlist_output.local_cutoff.r_ij
 
-        # this is commented out for now, as the vdw energy is calculated via
-        # DFTD3 which handles pair calculations internally
-        # If we have other vdw implementations, we can uncomment this
-        # if "per_system_vdw_energy" in self.postprocessing._registered_properties:
-        #     core_output["vdw_pair_indices"] = pairlist_output.vdw_cutoff.pair_indices
-        #     core_output["vdw_d_ij"] = pairlist_output.vdw_cutoff.d_ij
-        #     core_output["vdw_r_ij"] = pairlist_output.vdw_cutoff.r_ij
+        if (
+            "per_system_vdw_energy" in self.postprocessing._registered_properties
+        ):  # FIXME: no need when running tad-dftd3
+            core_output["vdw_pair_indices"] = pairlist_output.vdw_cutoff.pair_indices
+            core_output["vdw_d_ij"] = pairlist_output.vdw_cutoff.d_ij
+            core_output["vdw_r_ij"] = pairlist_output.vdw_cutoff.r_ij
 
         if (
             "per_system_electrostatic_energy"
@@ -725,9 +731,9 @@ def setup_potential(
     # - local_cutoff is the maximum interaction radius for the NNP core network (i.e., the local interaction radius)
     # this is always required.
     # - vdw_cutoff is the cutoff for the van der Waals interactions, which is only required in the vdw interactions are
-    # included as a post processing step.
+    # included as a post-processing step.
     # - electrostatic_cutoff is the cutoff for the electrostatic interactions, which is only required in the electrostatic
-    # interactions are included as a post processing step.
+    # interactions are included as a post-processing step.
     #
     # note zbl potential does not require a unique cutoff definition; it will use local cutoff and then calculates
     # the zbl potential based on the radii of the two atoms in the pair.
@@ -744,7 +750,6 @@ def setup_potential(
         )
         use_electrostatic_cutoff = True
 
-    # note vdw isn't implemented yet, but we can add it later
     if "per_system_vdw_energy" in postprocessing._registered_properties:
         vdw_cutoff = (
             potential_parameter.postprocessing_parameter.per_system_vdw_energy.maximum_interaction_radius
@@ -756,13 +761,12 @@ def setup_potential(
     if use_training_mode_neighborlist:
         from modelforge.potential.neighbors import NeighborListForTraining
 
-        # note vdw_cutoff is not being used as this is handled internally by the DFTD3 implementation
         neighborlist = NeighborListForTraining(
             local_cutoff=local_cutoff,
-            # vdw_cutoff=vdw_cutoff,
+            vdw_cutoff=vdw_cutoff,
             electrostatic_cutoff=electrostatic_cutoff,
             local_only_unique_pairs=only_unique_pairs,
-            # use_vdw_cutoff=use_vdw_cutoff,
+            use_vdw_cutoff=use_vdw_cutoff,
             use_electrostatic_cutoff=use_electrostatic_cutoff,
         )
     else:
@@ -772,14 +776,13 @@ def setup_potential(
 
         from modelforge.potential.neighbors import NeighborlistForInference
 
-        # note vdw_cutoff is not being used as this is handled internally by the DFTD3 implementation
         neighborlist = NeighborlistForInference(
             local_cutoff=local_cutoff,
-            # vdw_cutoff=vdw_cutoff,
+            vdw_cutoff=vdw_cutoff,
             electrostatic_cutoff=electrostatic_cutoff,
             displacement_function=displacement_function,
             local_only_unique_pairs=only_unique_pairs,
-            # use_vdw_cutoff=use_vdw_cutoff,
+            use_vdw_cutoff=use_vdw_cutoff,
             use_electrostatic_cutoff=use_electrostatic_cutoff,
         )
         # we can set the strategy here before passing this to the Potential