Remove TorchANI replacement modules and dependency

.github/workflows/ci.yml

@@ -25,23 +25,32 @@ jobs:
           # Oldest supported versions
           - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.0)
             enable_cuda: true
+            set_torch_dir: true
             cuda: "11.8.0"
             gcc: "10.3.*"
             nvcc: "11.8"
             python: "3.10.*"
-            torchani: "2.2.*"
             pytorch: "2.0.*"
 
-          # Latest supported versions (with CUDA)
-          - name: Linux (CUDA 12, Python 3.13, PyTorch 2.5)
+          - name: Linux (CUDA 12.6, Python 3.13, PyTorch 2.5)
             enable_cuda: true
+            set_torch_dir: true
             cuda: "12.6.0"
             gcc: "10.3.*"
             nvcc: "12.*"
             python: "3.13.*"
-            torchani: "2.2.*"
             pytorch: "2.5.*"
 
+          # Newest supported versions
+          - name: Linux (CUDA 12.9, Python 3.13, PyTorch 2.9)
+            enable_cuda: true
+            set_torch_dir: false
+            cuda: "12.9.0"
+            gcc: "14.*"
+            nvcc: "12.9"
+            python: "3.13.*"
+            pytorch: "2.9.*"
+
     steps:
     - name: Check out
       uses: actions/checkout@v2
@@ -50,7 +59,7 @@ jobs:
       uses: jlumbroso/free-disk-space@main
 
     - name: Install CUDA Toolkit
-      uses: Jimver/cuda-toolkit@v0.2.21
+      uses: Jimver/cuda-toolkit@v0.2.30
       with:
         cuda: ${{ matrix.cuda }}
         linux-local-args: '["--toolkit", "--override"]'
@@ -68,7 +77,6 @@ jobs:
     - name: Prepare dependencies
       run: |
         sed -i -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
-               -e "/torchani/c\  - torchani ${{ matrix.torchani }}" \
                -e "/python/c\  - python ${{ matrix.python }}" \
                -e "/pytorch-gpu/c\  - pytorch-gpu ${{ matrix.pytorch }}" \
                environment.yml
@@ -91,9 +99,20 @@ jobs:
       run: |
         conda activate nnpops
         mkdir build && cd build
+        SET_TORCH_DIR=""
+        if [ "${{ matrix.set_torch_dir }}" = "true" ]; then
+            # For older PyTorch, this is necessary; for newer PyTorch, CMake can
+            # find it automatically, and including this breaks the build.
+            SET_TORCH_DIR="-DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch"
+        fi
+        # For newer PyTorch, it is necessary to point CMake to the nvcc that we
+        # installed directly onto the system; otherwise, it will pick up the
+        # one pulled in by conda after installing pytorch-gpu, and break since
+        # the CUDA headers will be missing.
         cmake .. \
               -DENABLE_CUDA=${{ matrix.enable_cuda }} \
-              -DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch \
+              -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
+              $SET_TORCH_DIR \
               -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
         make install
 

CMakeLists.txt

@@ -16,7 +16,6 @@ enable_testing()
 # Source files of the library
 set(SRC_FILES src/ani/CpuANISymmetryFunctions.cpp
               src/ani/CudaANISymmetryFunctions.cu
-              src/pytorch/BatchedNN.cpp
               src/pytorch/CFConv.cpp
               src/pytorch/CFConvNeighbors.cpp
               src/pytorch/SymmetryFunctions.cpp
@@ -75,8 +74,8 @@ add_custom_target(copy_test ALL
 add_custom_command(
   TARGET copy_test POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E copy_directory
-  ${CMAKE_SOURCE_DIR}/src/pytorch/molecules
-  ${CMAKE_BINARY_DIR}/test/molecules
+  ${CMAKE_SOURCE_DIR}/src/pytorch/test_data
+  ${CMAKE_BINARY_DIR}/test/test_data
 )
 
 # Copy each test to the binary dir and add it
@@ -103,12 +102,8 @@ add_test(TestGetNeighborPairs  pytest -v --doctest-modules ${CMAKE_BINARY_DIR}/t
 # Installation
 install(TARGETS ${LIBRARY} DESTINATION ${Python3_SITEARCH}/${NAME})
 install(FILES src/pytorch/__init__.py
-              src/pytorch/BatchedNN.py
               src/pytorch/CFConv.py
               src/pytorch/CFConvNeighbors.py
-              src/pytorch/EnergyShifter.py
-              src/pytorch/OptimizedTorchANI.py
-              src/pytorch/SpeciesConverter.py
               src/pytorch/SymmetryFunctions.py
         DESTINATION ${Python3_SITEARCH}/${NAME})
 install(FILES src/pytorch/neighbors/__init__.py

README.md

@@ -39,7 +39,7 @@ If you don't have `conda`, we recommend installing [Miniconda](https://docs.cond
 #### Prerequisites
 
 - *CUDA Toolkit* (https://developer.nvidia.com/cuda-downloads)
-- *Miniconda* (https://docs.conda.io/en/latest/miniconda.html#linux-installers)
+- *Miniconda* (https://docs.conda.io/en/latest/miniconda.html)
 
 #### Build & install
 
@@ -48,24 +48,20 @@ If you don't have `conda`, we recommend installing [Miniconda](https://docs.cond
 $ git clone https://github.com/openmm/NNPOps.git
 ```
 
-- Set `CUDA_HOME`
-```bash
-$ export CUDA_HOME=/usr/local/cuda-11.2
-```
-
-- Crate and activate a *Conda* environment
+- Create and activate a Conda environment
 ```bash
 $ cd NNPOps
 $ conda env create -n nnpops -f environment.yml
 $ conda activate nnpops
 ```
 
-- Configure, build, and install
+- Configure, build, and install (note: for some older PyTorch versions, you may
+  need to add `-DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch`
+  as an argument to `cmake`; for newer PyTorch versions, this is unnecessary,
+  and may actually cause CMake configuration to fail)
 ```bash
 $ mkdir build && cd build
-$ cmake .. \
-        -DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch \
-        -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
+$ cmake .. -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
 $ make install
 ```
 
@@ -74,56 +70,12 @@ $ make install
 $ ctest --verbose
 ```
 
-## Usage
-
-Accelerated [*TorchANI*](https://aiqm.github.io/torchani/) operations:
-- [`torchani.AEVComputer`](https://aiqm.github.io/torchani/api.html?highlight=speciesaev#torchani.AEVComputer)
-- [`torchani.neurochem.NeuralNetwork`](https://aiqm.github.io/torchani/api.html#module-torchani.neurochem)
-
-### Example
-
-```python
-import mdtraj
-import torch
-import torchani
-
-from NNPOps.SpeciesConverter import TorchANISpeciesConverter
-from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-from NNPOps.BatchedNN import TorchANIBatchedNN
-from NNPOps.EnergyShifter import TorchANIEnergyShifter
-
-from NNPOps import OptimizedTorchANI
+## Operations
 
-device = torch.device('cuda')
+The following optimized operations are present in NNPOps and accessible from
+Python using the listed classes or functions:
 
-# Load a molecule
-molecule = mdtraj.load('molecule.mol2')
-species = torch.tensor([[atom.element.atomic_number for atom in molecule.top.atoms]], device=device)
-positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-# Construct ANI-2x and replace its operations with the optimized ones
-nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-nnp.species_converter = TorchANISpeciesConverter(nnp.species_converter, species).to(device)
-nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species).to(device)
-nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
-nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, species).to(device)
-
-# Compute energy and forces
-energy = nnp((species, positions)).energies
-energy.backward()
-forces = -positions.grad.clone()
-
-print(energy, forces)
-
-# Alternatively, all the optimizations can be applied with OptimizedTorchANI
-nnp2 = torchani.models.ANI2x(periodic_table_index=True).to(device)
-nnp2 = OptimizedTorchANI(nnp2, species).to(device)
-
-# Compute energy and forces again
-energy = nnp2((species, positions)).energies
-positions.grad.zero_()
-energy.backward()
-forces = -positions.grad.clone()
-
-print(energy, forces)
-```
+- ANI symmetry functions: `NNPOps.SymmetryFunctions.ANISymmetryFunctions`
+- Continuous filter convolution (CFConv): `NNPOps.CFConv.CFConv`, `NNPOps.CFConv.CFConvNeighbors`
+- Neighbor pair enumeration: `NNPOps.neighbors.getNeighborPairs()`
+- Particle mesh Ewald (PME): `NNPOps.pme.PME`

devtools/test_data/generate.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+
+"""
+Uses TorchANI to generate test cases in `src/pytorch/test_data/*.pt` for the
+ANI symmetry functions.  Requires TorchANI and mdtraj to be installed.
+"""
+
+import mdtraj
+import torch
+import torchani
+
+def main():
+    for name in ("1hvj", "1hvk", "2iuz", "3hkw", "3hky", "3lka", "3o99"):
+        generate_molecule_test_case(f"{name}_ligand.mol2", f"{name}.pt")
+    generate_molecule_test_case("water.pdb", "water.pt", True)
+
+def generate_molecule_test_case(in_path, out_path, pbc=False):
+    molecule = mdtraj.load(in_path)
+    atomic_numbers = torch.tensor([atom.element.atomic_number for atom in molecule.top.atoms])
+    atomic_positions = torch.tensor(molecule.xyz[0] * 10, requires_grad=True)
+    cell = torch.tensor(molecule.unitcell_vectors[0] * 10) if pbc else None
+
+    nnp = torchani.models.ANI2x()
+
+    parameters = dict(
+        numSpecies=nnp.aev_computer.num_species,
+        Rcr=nnp.aev_computer.radial.cutoff,
+        Rca=nnp.aev_computer.angular.cutoff,
+        EtaR=nnp.aev_computer.radial.eta.tolist(),
+        ShfR=nnp.aev_computer.radial.shifts.tolist(),
+        EtaA=nnp.aev_computer.angular.eta.tolist(),
+        Zeta=nnp.aev_computer.angular.zeta.tolist(),
+        ShfA=nnp.aev_computer.angular.shifts.tolist(),
+        ShfZ=nnp.aev_computer.angular.sections.tolist(),
+        atomSpecies=nnp.species_converter(atomic_numbers).tolist(),
+    )
+    output = nnp.aev_computer.forward(
+        torch.tensor(parameters["atomSpecies"]).unsqueeze(0),
+        atomic_positions.unsqueeze(0),
+        cell,
+        None if cell is None else torch.tensor([True, True, True])
+    )[0]
+    total = torch.sum(output)
+    total.backward()
+    testcase = dict(
+        parameters=parameters,
+        positions=atomic_positions.detach(),
+        cell=cell,
+        output=output,
+        grad=atomic_positions.grad,
+    )
+    torch.save(testcase, out_path)
+
+if __name__ == "__main__":
+    main()

environment.yml

@@ -2,12 +2,9 @@ channels:
   - conda-forge
 dependencies:
   - cmake >=3.20
-  - gxx_linux-64 10.3.*
+  - gxx_linux-64
   - make
-  - mdtraj
-  - torchani 2.2.*
   - pytest
-  - python 3.10.*
-  - pytorch-gpu 2.0.*
-  - sysroot_linux-64 2.17
-
+  - python
+  - pytorch-gpu
+  - sysroot_linux-64