Merge branch 'master' into nonbonded-interaction-groups

.git-blame-ignore-revs

@@ -1,2 +1,30 @@
 # Migrate code style to black
 4545c577a8dae5827a64594e8bfb100f7e0423c8
+
+# Move fe/ to timemachine/fe/
+277cc7463d22af28e5d80478ed3628d0070ffd10
+
+# Move ff/ to timemachine/ff/
+f661c86f108acceeefc067c918dc007cd73219e9
+
+# Move optimize/ to timemachine/optimize/
+e21cbb8f048840dc24723b60ca3351dbeac8d207
+
+# Move md/ to timemachine/md/
+2508c3ff76b9b7c781909626ef4fbb2f9e514f80
+
+# Move training/ to attic/training
+772c37329dc4f16b238c425a37045b059747c946
+
+# Move scripts/ to attic/scripts
+1a721dd3f05d6011cf028b0588e066682d38ba59
+
+# Move testsystems/ to timemachine/testsystems/
+#   and datasets/ to timeamchine/datasets/
+1e414a0e778948441cd3d01e51c62380230c96ea
+
+# Move parallel/ to timemachine/parallel
+4b814f2ee02be99a4ee4a4e8a4b55fc03bc28adf
+
+# Apply isort fixes
+61e068297ce26083b1168535ba908f9987696605

requirements.txt

@@ -1,7 +1,7 @@
 bootstrapped==0.0.2
 grpcio==1.30.0
 grpcio-tools==1.30.0
-protobuf==3.12.2
+protobuf==3.15.0
 numpy==1.21.4
 scipy==1.6.0
 pymbar==3.0.5

setup.py

@@ -97,9 +97,9 @@ def build_extension(self, ext):
         "dev": ["black==21.10b0", "isort==5.10.1", "flake8==4.0.1", "pre-commit", "grpcio==1.43.0"],
         "test": ["pytest", "pytest-cov"],
     },
-    # package_data={
-    #     "sample": ["package_data.dat"],
-    # },
+    package_data={
+        "datasets": ["timemachine/datasets"],
+    },
     # entry_points={
     #     "console_scripts": [
     #         "sample=sample:main",

tests/common.py

@@ -432,6 +432,7 @@ def compare_forces(
         ref_du_dx, ref_du_dp, ref_du_dl = grad_fn(x, params, box, lamb)
 
         for combo in itertools.product([False, True], repeat=4):
+
             (compute_du_dx, compute_du_dp, compute_du_dl, compute_u) = combo
 
             # do each computation twice to check determinism
@@ -445,7 +446,7 @@ def compare_forces(
             if compute_du_dl:
                 np.testing.assert_allclose(ref_du_dl, test_du_dl, rtol=rtol)
             if compute_du_dp:
-                np.testing.assert_allclose(ref_du_dp, test_du_dp, rtol=rtol)
+                np.testing.assert_allclose(ref_du_dp, test_du_dp, rtol=rtol, atol=atol)
 
             test_du_dx_2, test_du_dp_2, test_du_dl_2, test_u_2 = test_impl.execute_selective(
                 x, params, box, lamb, compute_du_dx, compute_du_dp, compute_du_dl, compute_u

tests/test_bonded.py

@@ -152,7 +152,7 @@ def test_harmonic_bond(self, n_particles=64, n_bonds=35, dim=3):
         box = np.eye(3) * 100
 
         # specific to harmonic bond force
-        relative_tolerance_at_precision = {np.float32: 2e-5, np.float64: 1e-9}
+        relative_tolerance_at_precision = {np.float64: 1e-7, np.float32: 2e-5}
 
         for precision, rtol in relative_tolerance_at_precision.items():
             test_potential = potentials.HarmonicBond(bond_idxs)

tests/test_datasets.py

@@ -0,0 +1,18 @@
+from timemachine.datasets import fetch_freesolv
+
+
+def test_fetch_freesolv():
+    """assert expected number of molecules loaded -- with unique names and expected property annotations"""
+    mols = fetch_freesolv()
+
+    # expected number of mols loaded
+    assert len(mols) == 642
+
+    # expected mol properties present, interpretable as floats
+    for mol in mols:
+        props = mol.GetPropsAsDict()
+        _, _ = float(props["dG"]), float(props["dG_err"])
+
+    # unique names
+    names = [mol.GetProp("_Name") for mol in mols]
+    assert len(set(names)) == len(names)

tests/test_enhanced.py

@@ -0,0 +1,10 @@
+from timemachine.datasets import fetch_freesolv
+from timemachine.md.enhanced import identify_rotatable_bonds
+
+
+def test_identify_rotatable_bonds_runs_on_freesolv():
+    """pass if no runtime errors are encountered"""
+    mols = fetch_freesolv()
+
+    for mol in mols:
+        _ = identify_rotatable_bonds(mol)

tests/test_handlers.py

@@ -175,6 +175,17 @@ def test_proper_torsion():
     mask = np.argwhere(torsion_params > 90)
     assert np.all(ff_adjoints[mask] == 0.0)
 
+    # assert expected shape when no matches are found
+    null_smirks = ["[#1:1]~[#1:2]~[#1:3]~[#1:4]"]  # should never match anything
+    null_params = np.zeros((len(null_smirks), 3))
+    null_counts = np.ones(len(null_smirks), dtype=int)
+
+    assigned_params, proper_idxs = bonded.ProperTorsionHandler.static_parameterize(
+        null_params, null_smirks, null_counts, mol
+    )
+    assert assigned_params.shape == (0, 3)
+    assert proper_idxs.shape == (0, 4)
+
 
 def test_improper_torsion():
 
@@ -220,6 +231,14 @@ def test_improper_torsion():
     mask = np.argwhere(torsion_params > 90)
     assert np.all(ff_adjoints[mask] == 0.0)
 
+    # assert expected shape when no matches are found
+    null_smirks = ["[#1:1]~[#1:2](~[#1:3])~[#1:4]"]  # should never match anything
+    null_params = np.zeros((len(null_smirks), 3))
+
+    assigned_params, improper_idxs = bonded.ImproperTorsionHandler.static_parameterize(null_params, null_smirks, mol)
+    assert assigned_params.shape == (0, 3)
+    assert improper_idxs.shape == (0, 4)
+
 
 def test_exclusions():
 

tests/test_nonbonded.py

@@ -184,10 +184,11 @@ def test_correctness(self):
         """
         Test against the reference jax platform for correctness.
         """
-
         # we can't go bigger than this due to memory limitations in the the reference platform.
         for N in [33, 65, 231, 1050, 4080]:
 
+            np.random.seed(2022)
+
             test_conf = self.host_conf[:N]
 
             # strip out parts of the system
@@ -218,7 +219,7 @@ def test_correctness(self):
                 self.cutoff,
             )
 
-            for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
+            for precision, rtol, atol in [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]:
 
                 self.compare_forces(
                     test_conf,
@@ -227,7 +228,8 @@ def test_correctness(self):
                     self.lamb,
                     ref_nonbonded_fn,
                     test_nonbonded_fn,
-                    rtol,
+                    rtol=rtol,
+                    atol=atol,
                     precision=precision,
                 )
 
@@ -306,7 +308,7 @@ def test_nblist_box_resize(self):
         # the rebuild is triggered as long as the box *changes*.
         for test_box in [big_box, box]:
 
-            for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
+            for precision, rtol, atol in [(np.float64, 1e-8, 1e-10), (np.float32, 1e-4, 3e-5)]:
 
                 self.compare_forces(
                     host_conf,
@@ -315,7 +317,8 @@ def test_nblist_box_resize(self):
                     lamb,
                     ref_nonbonded_fn,
                     test_nonbonded_fn,
-                    rtol,
+                    rtol=rtol,
+                    atol=atol,
                     precision=precision,
                 )
 
@@ -457,7 +460,7 @@ def test_nonbonded(self):
             lambda_plane_idxs = np.random.randint(low=-2, high=2, size=N, dtype=np.int32)
             lambda_offset_idxs = np.random.randint(low=-2, high=2, size=N, dtype=np.int32)
 
-            for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
+            for precision, rtol, atol in [(np.float64, 1e-8, 3e-11), (np.float32, 1e-4, 3e-5)]:
 
                 for cutoff in [1.0]:
                     # E = 0 # DEBUG!
@@ -470,7 +473,15 @@ def test_nonbonded(self):
                         print("lambda", lamb, "cutoff", cutoff, "precision", precision, "xshape", coords.shape)
 
                         self.compare_forces(
-                            coords, charge_params, box, lamb, ref_potential, test_potential, rtol, precision=precision
+                            coords,
+                            charge_params,
+                            box,
+                            lamb,
+                            ref_potential,
+                            test_potential,
+                            rtol=rtol,
+                            atol=atol,
+                            precision=precision,
                         )
 
     def test_nonbonded_with_box_smaller_than_cutoff(self):

tests/test_parameter_interpolation.py

@@ -41,7 +41,7 @@ def test_nonbonded(self):
         lambda_offset_idxs = np.random.randint(low=0, high=2, size=N, dtype=np.int32)
 
         for lamb in [0.0, 0.2, 1.0]:
-            for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
+            for precision, rtol, atol in [(np.float64, 1e-8, 3e-11), (np.float32, 1e-4, 3e-6)]:
 
                 # E = 0 # DEBUG!
                 qlj_src, ref_potential, test_potential = prepare_water_system(
@@ -67,7 +67,8 @@ def test_nonbonded(self):
                     lamb,
                     ref_interpolated_potential,
                     test_interpolated_potential,
-                    rtol,
+                    rtol=rtol,
+                    atol=atol,
                     precision=precision,
                 )
 
@@ -127,7 +128,7 @@ def u_reference(x, params, box, lamb):
             qlj = interpolate_params(lamb, qlj_src, qlj_dst)
             return ref_potential(x, qlj, box, lamb)
 
-        for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
+        for precision, rtol, atol in [(np.float64, 1e-8, 1e-11), (np.float32, 1e-4, 1e-6)]:
 
             for lamb in [0.0, 0.2, 0.6, 0.7, 0.8, 1.0]:
 
@@ -147,5 +148,13 @@ def u_reference(x, params, box, lamb):
                 )
 
                 self.compare_forces(
-                    coords, qlj, box, lamb, u_reference, test_interpolated_potential, rtol, precision=precision
+                    coords,
+                    qlj,
+                    box,
+                    lamb,
+                    u_reference,
+                    test_interpolated_potential,
+                    rtol=rtol,
+                    atol=atol,
+                    precision=precision,
                 )

timemachine/cpp/src/integrator.cu

@@ -37,9 +37,9 @@ __global__ void update_forward(
     const int N,
     const int D,
     const RealType ca,
-    const RealType *cbs, // N x 3, not P x N x 3, but we could just pass in the first index
-    const RealType *ccs, // N
-    const RealType *noise,
+    const RealType *cbs,   // N
+    const RealType *ccs,   // N
+    const RealType *noise, // N x 3
     RealType *x_t,
     RealType *v_t,
     const unsigned long long *du_dx,
@@ -74,6 +74,7 @@ void LangevinIntegrator::step_fwd(
     size_t n_blocks = (N_ * D + tpb - 1) / tpb;
     dim3 dimGrid_dx(n_blocks, D);
 
+    curandErrchk(curandSetStream(cr_rng_, stream));
     curandErrchk(templateCurandNormal(cr_rng_, d_noise_, round_up_even(N_ * D), 0.0, 1.0));
 
     update_forward<double>

timemachine/cpp/src/integrator.hpp

@@ -16,9 +16,9 @@ class Integrator {
 class LangevinIntegrator : public Integrator {
 
 private:
-    double dt_;
-    double N_;
-    double ca_;
+    const int N_;
+    const double dt_;
+    const double ca_;
     double *d_cbs_;
     double *d_ccs_;
     double *d_noise_;

timemachine/datasets/__init__.py

@@ -0,0 +1,3 @@
+from timemachine.datasets.utils import fetch_freesolv
+
+__all__ = ["fetch_freesolv"]

timemachine/datasets/utils.py

@@ -0,0 +1,11 @@
+from importlib import resources
+from typing import List
+
+from rdkit import Chem
+
+
+def fetch_freesolv() -> List[Chem.Mol]:
+    with resources.path("timemachine", "datasets") as datasets_path:
+        freesolv_path = str(datasets_path / "freesolv" / "freesolv.sdf")
+        supplier = Chem.SDMolSupplier(freesolv_path, removeHs=False)
+        return [mol for mol in supplier]

timemachine/ff/handlers/bonded.py

@@ -26,10 +26,6 @@ def generate_vd_idxs(mol, smirks):
     return bond_idxs, param_idxs
 
 
-def parameterize_ligand(params, param_idxs):
-    return params[param_idxs]
-
-
 # its trivial to re-use this for everything except the ImproperTorsions
 class ReversibleBondHandler(SerializableMixIn):
     def __init__(self, smirks, params, props):
@@ -139,7 +135,15 @@ def static_parameterize(params, smirks, counts, mol):
 
         scatter_idxs = np.array(scatter_idxs)
 
-        return params[scatter_idxs], np.repeat(torsion_idxs, repeats, axis=0).astype(np.int32)
+        # if no matches found, return arrays that can still be concatenated as expected
+        if len(param_idxs) > 0:
+            assigned_params = params[scatter_idxs]
+            proper_idxs = np.repeat(torsion_idxs, repeats, axis=0).astype(np.int32)
+        else:
+            assigned_params = params[:0]  # empty slice with same dtype, other dimensions
+            proper_idxs = np.zeros((0, 4), dtype=np.int32)
+
+        return assigned_params, proper_idxs
 
     def serialize(self):
         list_params = []
@@ -213,4 +217,12 @@ def make_key(idxs):
 
         param_idxs = np.array(param_idxs)
 
-        return params[param_idxs], np.array(improper_idxs, dtype=np.int32)
+        # if no matches found, return arrays that can still be concatenated as expected
+        if len(param_idxs) > 0:
+            assigned_params = params[param_idxs]
+            improper_idxs = np.array(improper_idxs, dtype=np.int32)
+        else:
+            assigned_params = params[:0]  # empty slice with same dtype, other dimensions
+            improper_idxs = np.zeros((0, 4), dtype=np.int32)
+
+        return assigned_params, improper_idxs

timemachine/ff/handlers/nonbonded.py

@@ -155,10 +155,6 @@ def generate_nonbonded_idxs(mol, smirks):
     return param_idxs
 
 
-def parameterize_ligand(params, param_idxs):
-    return params[param_idxs]
-
-
 def compute_or_load_am1_charges(mol):
     """Unless already cached in mol's "AM1Cache" property, use OpenEye to compute AM1 partial charges."""
 

timemachine/md/enhanced.py

@@ -42,10 +42,10 @@ def identify_rotatable_bonds(mol):
 
     """
     pattern = Chem.MolFromSmarts("[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]")
-    matches = mol.GetSubstructMatches(pattern)
+    matches = mol.GetSubstructMatches(pattern, uniquify=1)
 
     # sanity check
-    assert len(matches) == rdMolDescriptors.CalcNumRotatableBonds(mol)
+    assert len(matches) >= rdMolDescriptors.CalcNumRotatableBonds(mol)
 
     sorted_matches = set()