mdtraj mol.xyz unit (#17)

openmm · Nov 6, 2020 · f9716fd · f9716fd
1 parent 667a282
commit f9716fd
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 5 deletions.
diff --git a/ani/BenchmarkCudaANISymmetryFunctions.cu b/ani/BenchmarkCudaANISymmetryFunctions.cu
@@ -6,6 +6,8 @@
 #include <stdexcept>
 #include <string>
 #include <vector>
+#include <stdio.h>
+#include <time.h>
 
 using namespace std;
 
@@ -149,7 +151,14 @@ int main(int argc, char* argv[]) {
             {12.5, 3.1625, 14.1, 2.74889}
         };
         CudaANISymmetryFunctions ani(species.size(), 7, 5.1, 3.5, periodicBoxVectors.size() > 0, species, radialFunctions, angularFunctions, true);
+        clock_t start, finish;
+        double duration;
+        start = clock();
         runBenchmark(ani, stoi(argv[2]), positions, species, periodicBoxVectors);
+        finish = clock();
+        duration = (double)(finish - start) / CLOCKS_PER_SEC;
+        printf("  %f s\n", duration);
+        printf("  %f ms/it\n", duration/stoi(argv[2])*1000);
     }
     catch (const exception& e) {
         cout << e.what() << endl;

diff --git a/pytorch/BenchmarkTorchANISymmetryFunctions.py b/pytorch/BenchmarkTorchANISymmetryFunctions.py
@@ -9,7 +9,7 @@
 
 mol = mdtraj.load('molecules/2iuz_ligand.mol2')
 species = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-positions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
+positions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
 
 nnp = torchani.models.ANI2x(periodic_table_index=True, model_index=None).to(device)
 speciesPositions = nnp.species_converter((species, positions))

diff --git a/pytorch/README.md b/pytorch/README.md
@@ -20,7 +20,7 @@ device = torch.device('cuda')
 # 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, dtype=torch.float32, requires_grad=True, device=device)
+positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
 
 # Construct ANI-2x and replace its native featurizer with NNPOps implementation
 nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)

diff --git a/pytorch/SymmetryFunctions.py b/pytorch/SymmetryFunctions.py
@@ -48,7 +48,7 @@ class TorchANISymmetryFunctions(torch.nn.Module):
         # 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, dtype=torch.float32, requires_grad=True, device=device)
+        >>> positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
 
         # Construct ANI-2x and replace its native featurizer with NNPOps implementation
         >>> nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)

diff --git a/pytorch/TestSymmetryFunctions.py b/pytorch/TestSymmetryFunctions.py
@@ -37,7 +37,7 @@ def test_compare_with_native(deviceString, molFile):
 
     mol = mdtraj.load(f'molecules/{molFile}_ligand.mol2')
     atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
+    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
 
     nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
     energy_ref = nnp((atomicNumbers, atomicPositions)).energies
@@ -64,7 +64,7 @@ def test_model_serialization(deviceString, molFile):
 
     mol = mdtraj.load(f'molecules/{molFile}_ligand.mol2')
     atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
+    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
 
     nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
     nnp_ref.aev_computer = TorchANISymmetryFunctions(nnp_ref.aev_computer)