Batch D3 and SevenNetD3Model for Torch-Sim interface

.pre-commit-config.yaml

@@ -65,5 +65,6 @@ repos:
               ^example_inputs/data/|
               ^sevenn/pair_e3gnn/comm_brick.cpp|
               ^sevenn/pair_e3gnn/pair_d3.cu|
-              ^sevenn/pair_e3gnn/pair_d3_for_ase.cu
+              ^sevenn/pair_e3gnn/pair_d3_for_ase.cu|
+              ^sevenn/pair_e3gnn/pair_d3_batch.cu
           )

CHANGELOG.md

@@ -6,6 +6,7 @@ All notable changes to this project will be documented in this file.
 - Support OpenEquivariance
 - Per-atom stress (atomic virial) support in LAMMPS pair_e3gnn and ASE calculator
 - `compute_atomic_virial` option in `SevenNetCalculator`
+- Add Batch D3 and SevenNetD3Model in Torch-Sim interface
 
 ### Changed
 - **[Breaking]** Rename optional dependency group `mliap` into `mliap12` (reflecting its CUDA 12.x dependency).

docs/source/user_guide/ase_calculator.md

@@ -1,19 +1,65 @@
 (ase_calculator)=
 # ASE calculator
 
+### `SevenNetCalculator`
+
 SevenNet provides an ASE interface via the ASE calculator. Models can be loaded using the following Python code:
 ```python
 from sevenn.calculator import SevenNetCalculator
 # The 'modal' argument is required if the model is trained with multi-fidelity learning enabled.
 calc_omni = SevenNetCalculator(model='7net-omni', modal='mpa')
 ```
-SevenNet also supports CUDA-accelerated D3 calculations.
-For more information about D3, follow [here](../user_guide/d3.md))
+
+## D3 dispersion correction
+
+SevenNet also supports a CUDA-accelerated Grimme D3 van der Waals corrections through ASE calculators.
+For the method, units, supported functionals/damping types, and limitations, see {doc}`./d3`.
+
+### `SevenNetD3Calculator` (SevenNet + D3)
+
+It sums a `SevenNetCalculator` and a `D3Calculator`, returning combined energy, forces, and stress.
+
 ```python
+# To obtain 7net-0 + D3(BJ) results
 from sevenn.calculator import SevenNetD3Calculator
-calc = SevenNetD3Calculator(model='7net-0', device='cuda')
+calc = SevenNetD3Calculator(
+           model='7net-0',
+           device='cuda',
+           functional_name='pbe',  # 7net-0 is trained with PBE results
+           damping_type='damp_bj',
+       )
 ```
 
+It accepts every `SevenNetCalculator` argument (`model`, `modal`, `enable_cueq`/`enable_flash`/`enable_oeq`, ...) plus the D3 parameters listed below.
+
+### `D3Calculator` (D3 only)
+
+`D3Calculator` computes only the D3 dispersion term. Use it directly when you want the dispersion contribution on its own, or combine it with another calculator via ASE's `SumCalculator` (this is exactly what `SevenNetD3Calculator` does internally):
+
+```python
+from ase.calculators.mixing import SumCalculator
+from sevenn.calculator import SevenNetCalculator, D3Calculator
+
+calc = SumCalculator([SevenNetCalculator(model='7net-0'), D3Calculator()])
+```
+
+### D3 parameters
+
+Both calculators share the same D3 knobs:
+
+| Parameter | Default | Description |
+| --- | --- | --- |
+| `damping_type` | `'damp_bj'` | Becke-Johnson damping |
+| `functional_name` | `'pbe'` | XC functional where the D3 parameters are fitted to |
+| `vdw_cutoff` | `9000` | Squared energy/force cutoff in Bohr² (≈ 50.20 Å) |
+| `cn_cutoff` | `1600` | Squared coordination-number cutoff in Bohr² (≈ 21.17 Å) |
+
+See {doc}`./d3` for the full explanation of each option and the list of supported functionals.
+
+:::{caution}
+The D3 calculators require a CUDA GPU (compute capability ≥ 6.0); CPU + D3 is not supported. There is also a hard limit of 46,340 atoms per call. See the Cautions section of {doc}`./d3`.
+:::
+
 Use enable_cueq, enable_flash, or enable_oeq to use cuEquivariance, flashTP, or OpenEquivariance for faster inference.
 For more information about accelerators, follow [here](./accelerator.md)
 ```python

docs/source/user_guide/torchsim.md

@@ -30,7 +30,7 @@ model = SevenNetModel(model="7net-omni", modal="mpa", enable_oeq=True)
 # or enable_cueq=True or enable_flash=True
 ```
 
-The `device` parameter defaults to `'auto'` (CUDA if available, otherwise CPU).
+The `device` parameter defaults to `auto` (CUDA if available, otherwise CPU).
 
 ### Batched MD
 
@@ -59,3 +59,44 @@ relaxed_state = ts.optimize(
     optimizer=ts.Optimizer.fire,
 )
 ```
+
+## SevenNet + D3 dispersion
+
+`SevenNetD3Model` adds the CUDA-accelerated Grimme D3 dispersion correction on top of `SevenNetModel`. It is a drop-in replacement: it accepts the same model specifiers and accelerator flags, and plugs into `ts.integrate` / `ts.optimize` the same way.
+
+```python
+from sevenn.torchsim import SevenNetD3Model
+
+model = SevenNetD3Model(model="7net-omni", modal="mpa", **d3_kwargs)
+
+final_state = ts.integrate(
+    system=[atoms] * 10,
+    model=model,
+    n_steps=100,
+    timestep=0.002,
+    temperature=300,
+    integrator=ts.Integrator.nvt_langevin,
+)
+```
+
+It requires a CUDA GPU (the D3 backends are GPU-only) and accepts the same D3 parameters as the ASE calculators (`damping_type`, `functional_name`, `vdw_cutoff`, `cn_cutoff`); see {doc}`./d3` and the [`D3Calculator` documentation](./ase_calculator.md#d3-dispersion-correction).
+
+### Choosing the D3 backend
+
+D3 can be evaluated two ways, selected by `d3_mode`:
+
+- `serial`: a per-system loop over the ASE-style `D3Calculator`.
+- `batch`: a single batched CUDA kernel launch that computes D3 for all systems at once.
+- `auto` (default): use `batch` when the number of systems in the batch exceeds `d3_batch_threshold` (default `4`), otherwise `serial`.
+
+The `auto` heuristic exists because the two paths have different trade-offs: the batched kernel amortizes its per-call overhead across many systems, so it wins for large batches, while the serial loop is cheaper for the few-system case. Tune the cutoff with `d3_batch_threshold`, or force a single backend with `d3_mode`:
+
+```python
+# Always use the batched CUDA kernel
+model = SevenNetD3Model(model="7net-omni", modal="mpa", d3_mode="batch")
+
+# Auto, but switch to batched D3 only above 8 systems
+model = SevenNetD3Model(
+    model="7net-omni", modal="mpa", d3_mode="auto", d3_batch_threshold=8,
+)
+```

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
 ]
 [project.optional-dependencies]
 test = ["pytest", "pytest-cov>=5", "ipython"]
-torchsim = ["torch-sim-atomistic>=0.5.2; python_version >= '3.12'"]
+torchsim = ["torch-sim-atomistic>=0.6.0; python_version >= '3.12'"]
 cueq12 = ["cuequivariance>=0.6.0; python_version >= '3.10'", "cuequivariance-torch>=0.6.0; python_version >= '3.10'", "cuequivariance-ops-torch-cu12; python_version >= '3.10'"]
 cueq13 = ["cuequivariance>=0.7.0; python_version >= '3.10'", "cuequivariance-torch>=0.7.0; python_version >= '3.10'", "cuequivariance-ops-torch-cu13; python_version >= '3.10'"]
 oeq = ["openequivariance"]