RESP Charges

Dockerfile

@@ -78,19 +78,13 @@ COPY . /code/timemachine/
 WORKDIR /code/timemachine/
 RUN pip install --no-cache-dir -e . && rm -rf ./build
 
-# Container with only cuda base, half the size of the timemachine_cuda_dev container
-# Need to copy curand/cudart as these are dependencies of the Timemachine GPU code
-FROM docker.io/nvidia/cuda:12.4.1-base-ubuntu20.04 AS timemachine
+FROM docker.io/nvidia/cuda:12.4.1-devel-ubuntu20.04 AS timemachine
 ARG LIBXRENDER_VERSION
 ARG LIBXEXT_VERSION
 RUN (apt-get update || true) && apt-get install --no-install-recommends -y libxrender1=${LIBXRENDER_VERSION} libxext-dev=${LIBXEXT_VERSION} \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy curand libraries from image, only require cudart and curand
-COPY --from=timemachine_cuda_dev /usr/local/cuda/targets/x86_64-linux/lib/libcurand* /usr/local/cuda/targets/x86_64-linux/lib/
-COPY --from=timemachine_cuda_dev /usr/local/cuda/lib64/libcurand* /usr/local/cuda/lib64/
-
 COPY --from=timemachine_cuda_dev /opt/conda/ /opt/conda/
 COPY --from=timemachine_cuda_dev /code/ /code/
 COPY --from=timemachine_cuda_dev /root/.bashrc /root/.bashrc

environment.yml

@@ -8,3 +8,5 @@ dependencies:
   # Include numpy/scipy to reduce the size of the Docker container
   - numpy=2.2.2
   - scipy=1.15.1
+  - openff-toolkit=0.16.9
+  - openff-recharge=0.5.3

requirements.txt

@@ -7,3 +7,7 @@ rdkit==2024.9.4
 --extra-index-url https://pypi.anaconda.org/OpenEye/simple
 openeye-toolkits==2020.2.0
 openmm==8.2.0
+gpu4pyscf-cuda12x==1.4.3
+cutensor-cu12==2.2.0
+Auto3D==2.3.1
+torch==2.7.0

tests/test_handlers.py

@@ -23,8 +23,6 @@
 from timemachine.testsystems import fetch_freesolv
 from timemachine.utils import path_to_internal_file
 
-pytestmark = [pytest.mark.nocuda]
-
 TEST_FEATURE_SIZE = 16
 
 
@@ -123,6 +121,7 @@ def env_nn_args():
     return enc_unflatten_str, params, props
 
 
+@pytest.mark.nocuda
 def test_harmonic_bond():
     patterns = [
         ["[#6X4:1]-[#6X4:2]", 0.1, 0.2],
@@ -250,6 +249,7 @@ def test_harmonic_bond():
     assert bond_params.shape == (0, 2)
 
 
+@pytest.mark.nocuda
 def test_harmonic_angle():
     patterns = [
         ["[*:1]-[#8:2]-[*:3]", 0.1, 0.2],
@@ -273,6 +273,7 @@ def test_harmonic_angle():
     assert angle_params.shape == (0, 3)
 
 
+@pytest.mark.nocuda
 def test_proper_torsion():
     # proper torsions have a variadic number of terms
 
@@ -323,6 +324,7 @@ def test_proper_torsion():
     assert proper_idxs.shape == (0, 4)
 
 
+@pytest.mark.nocuda
 def test_improper_torsion():
     patterns = [
         ["[*:1]~[#6X3:2](~[*:3])~[*:4]", 1.5341333333333333, 3.141592653589793, 2.0],
@@ -378,6 +380,7 @@ def test_improper_torsion():
         assert idxs[0] < idxs[-1]
 
 
+@pytest.mark.nocuda
 def test_exclusions():
     mol = Chem.MolFromSmiles("FC(F)=C(F)F")
     exc_idxs, scales = nonbonded.generate_exclusion_idxs(mol, scale12=0.0, scale13=0.2, scale14_q=0.25, scale14_lj=0.75)
@@ -428,6 +431,7 @@ def test_exclusions():
     np.testing.assert_equal(scales, expected_scales)
 
 
+@pytest.mark.nocuda
 def test_am1bcc_parameterization():
     # currently takes no parameters
     smirks = []
@@ -456,6 +460,30 @@ def test_am1bcc_parameterization():
     # assert vjp_fn(charges_adjoints) == None
 
 
+def test_resp_parameterization():
+    # currently takes no parameters
+    smirks = []
+    params = []
+    props = None
+
+    cache_key = nonbonded.RESP_CHARGE_CACHE
+    resp = nonbonded.RESPHandler(smirks, params, props)
+    mol = Chem.AddHs(Chem.MolFromSmiles("C1CNCOC1F"))
+
+    assert not mol.HasProp(cache_key)
+
+    charges = resp.parameterize(mol)
+
+    assert len(charges) == mol.GetNumAtoms()
+    assert mol.HasProp(cache_key)
+
+    cached_charges = resp.parameterize(mol)
+    np.testing.assert_equal(charges, cached_charges)
+
+    new_charges, vjp_fn = jax.vjp(functools.partial(resp.partial_parameterize, None, mol))
+
+
+@pytest.mark.nocuda
 def test_am1_ccc():
     patterns = [
         ["[#6X4:1]-[#1:2]", 0.46323257920556493],
@@ -564,6 +592,7 @@ def test_am1_ccc():
     np.testing.assert_array_equal(es_params_from_cache, es_params)
 
 
+@pytest.mark.nocuda
 def test_simple_charge_handler():
     patterns = [
         ["[#1:1]", 99.0],
@@ -679,6 +708,7 @@ def test_gbsa_handler():
     assert np.all(adjoints[mask] == 0.0)
 
 
+@pytest.mark.nocuda
 def test_am1ccc_throws_error_on_phosphorus():
     """Temporary, until phosphorus patterns are added to AM1CCC port"""
     ff = Forcefield.load_default()
@@ -692,6 +722,7 @@ def test_am1ccc_throws_error_on_phosphorus():
     assert "unsupported element" in str(e)
 
 
+@pytest.mark.nocuda
 @pytest.mark.parametrize(
     "am1bcc_ff", ["smirnoff_1_1_0_am1bcc.py", "smirnoff_2_0_0_am1bcc.py", "smirnoff_2_2_0_am1bcc.py"]
 )
@@ -717,6 +748,7 @@ def test_am1bcc_handles_phosphorus(am1bcc_ff):
     )
 
 
+@pytest.mark.nocuda
 def test_am1_differences():
     ff = Forcefield.load_default()
 
@@ -754,6 +786,7 @@ def test_am1_differences():
             assert 0
 
 
+@pytest.mark.nocuda
 def test_am1elf10_conformer_independence():
     with path_to_internal_file("timemachine.testsystems.data", "ligands_40.sdf") as path_to_ligand:
         mols = utils.read_sdf(path_to_ligand)
@@ -780,6 +813,7 @@ def test_am1elf10_conformer_independence():
         assert np.sum(delta_charges) == pytest.approx(0.0)
 
 
+@pytest.mark.nocuda
 def test_trans_carboxlic_acid():
     # Test fallback to turn off hydrogen sampling if charge generation failed
     # due to trans-COOH
@@ -795,6 +829,7 @@ def test_trans_carboxlic_acid():
     assert np.sum(delta_charges) == pytest.approx(0.0)
 
 
+@pytest.mark.nocuda
 def test_freesolv_failures():
     # Test failures for 3 cases in freesolv in openeye toolkits 2022.2.2
     with path_to_internal_file("timemachine.testsystems.data", "freesolv_omega_failures.sdf") as path_to_ligand:
@@ -805,6 +840,7 @@ def test_freesolv_failures():
         assert am1elf10_charges is not None
 
 
+@pytest.mark.nocuda
 def test_compute_or_load_oe_charges():
     """Loop over test ligands, asserting that charges are stored in expected property and that the same charges are
     returned on repeated calls"""
@@ -844,6 +880,7 @@ def assert_permutation_equivariance(mol, fxn, perm):
     np.testing.assert_allclose(qs_perm, qs[perm])
 
 
+@pytest.mark.nocuda
 @pytest.mark.parametrize("mol_idx", [0, 1, 2, 3, 4, 5])
 def test_partial_charge_equivariance_on_freesolv(mol_idx):
     ff = Forcefield.load_default()
@@ -857,6 +894,7 @@ def test_partial_charge_equivariance_on_freesolv(mol_idx):
     assert_permutation_equivariance(mol, ff.q_handle.parameterize, perm)
 
 
+@pytest.mark.nocuda
 def test_charging_compounds_with_non_zero_charge():
     patterns = [
         ["[#6a:1]:[#6a:2]", 0.0],
@@ -884,6 +922,7 @@ def test_charging_compounds_with_non_zero_charge():
     np.testing.assert_almost_equal(np.sum(es_params) / np.sqrt(ONE_4PI_EPS0), -1.0, decimal=5)
 
 
+@pytest.mark.nocuda
 def test_precomputed_charge_handler():
     with path_to_internal_file("timemachine.testsystems.water_exchange", "bb_centered_espaloma.sdf") as path_to_ligand:
         mol = utils.read_sdf(path_to_ligand)[0]
@@ -995,6 +1034,7 @@ def test_precomputed_charge_handler():
     )
 
 
+@pytest.mark.nocuda
 def test_compute_or_load_bond_smirks_matches():
     """Loop over test ligands, asserting that
     * verify no cache key
@@ -1038,6 +1078,7 @@ def test_compute_or_load_bond_smirks_matches():
         np.testing.assert_array_equal(fresh_types, cached_types)
 
 
+@pytest.mark.nocuda
 def test_apply_bond_charge_corrections():
     """Assert that applying random bond charge corrections does not change net charge"""
 
@@ -1067,6 +1108,7 @@ def test_apply_bond_charge_corrections():
         np.testing.assert_almost_equal(final_net_charge, initial_net_charge)
 
 
+@pytest.mark.nocuda
 def test_lennard_jones_handler():
     patterns = [
         ["[#1:1]", 99.0, 999.0],
@@ -1144,6 +1186,7 @@ def test_lennard_jones_handler():
     assert np.all(adjoints[mask] == 0.0)
 
 
+@pytest.mark.nocuda
 def test_symmetric_am1ccc():
     """Assert that (symmetric_bond_smarts, +1.0) has same behavior as (symmetric_bond_smarts, 0.0) on one test mol"""
 
@@ -1164,6 +1207,7 @@ def test_symmetric_am1ccc():
     np.testing.assert_array_equal(test_charges, ref_charges)
 
 
+@pytest.mark.nocuda
 def test_nn_handler():
     with path_to_internal_file("timemachine.testsystems.data", "ligands_40.sdf") as path_to_ligand:
         all_mols = utils.read_sdf(path_to_ligand)
@@ -1196,6 +1240,7 @@ def loss_fn(params):
     print("jit grad", jax.jit(grad_fn)(params))  # also a few seconds
 
 
+@pytest.mark.nocuda
 def test_harmonic_bonds_complete():
     """On a test molecule containing [oxygen] ~ [halogen] bonds,
     assert that a ValueError is raised."""
@@ -1210,6 +1255,7 @@ def test_harmonic_bonds_complete():
     assert "missing bonds" in str(e)
 
 
+@pytest.mark.nocuda
 @pytest.mark.parametrize("is_nn", [True, False])
 @pytest.mark.parametrize(
     "protein_path_and_symmetries",
@@ -1312,6 +1358,7 @@ def test_env_bcc_peptide_symmetries(protein_path_and_symmetries, is_nn, env_nn_a
             np.testing.assert_almost_equal(raw_charges[other], raw_charges[first])
 
 
+@pytest.mark.nocuda
 @pytest.mark.nightly(reason="Slow")
 @pytest.mark.parametrize("is_nn", [True, False])
 @pytest.mark.parametrize("protein_path", ["5dfr_solv_equil.pdb", "hif2a_nowater_min.pdb"])

tests/test_serialization.py

@@ -281,6 +281,21 @@ def test_am1bcc():
     assert am1.props == am1.props
 
 
+def test_resp():
+    smirks = []
+    params = []
+    props = None
+
+    resp = nonbonded.RESPHandler(smirks, params, props)
+    obj = resp.serialize()
+    all_handlers, _, _ = deserialize_handlers(bin_to_str(obj))
+
+    resp = all_handlers[0]
+    np.testing.assert_equal(resp.smirks, resp.smirks)
+    np.testing.assert_equal(resp.params, resp.params)
+    assert resp.props == resp.props
+
+
 def test_am1ccc():
     patterns = [
         ["[#6X4:1]-[#1:2]", 0.46323257920556493],

timemachine/ff/__init__.py

@@ -58,6 +58,7 @@ class Forcefield:
             nonbonded.AM1BCCCCCHandler,
             nonbonded.PrecomputedChargeHandler,
             nonbonded.NNHandler,
+            nonbonded.RESPHandler,
         ]
     ]
     q_handle_intra: Optional[
@@ -68,6 +69,7 @@ class Forcefield:
             nonbonded.AM1BCCCCCIntraHandler,
             nonbonded.PrecomputedChargeHandler,
             nonbonded.NNHandler,
+            nonbonded.RESPIntraHandler,
         ]
     ]
 
@@ -201,6 +203,7 @@ def from_handlers(
                     nonbonded.AM1CCCIntraHandler,
                     nonbonded.AM1BCCIntraHandler,
                     nonbonded.SimpleChargeIntraHandler,
+                    nonbonded.RESPIntraHandler,
                     nonbonded.PrecomputedChargeIntraHandler,
                 ),
             ):
@@ -220,6 +223,7 @@ def from_handlers(
                     nonbonded.AM1BCCCCCHandler,
                     nonbonded.AM1CCCHandler,
                     nonbonded.AM1BCCHandler,
+                    nonbonded.RESPHandler,
                     nonbonded.SimpleChargeHandler,
                     nonbonded.PrecomputedChargeHandler,
                     nonbonded.NNHandler,
@@ -258,6 +262,8 @@ def from_handlers(
                 q_handle_intra = nonbonded.PrecomputedChargeIntraHandler(
                     q_handle.smirks, q_handle.params, q_handle.props
                 )
+            elif isinstance(q_handle, nonbonded.RESPHandler):
+                q_handle_intra = nonbonded.RESPIntraHandler(q_handle.smirks, q_handle.params, q_handle.props)
             else:
                 raise ValueError(f"Unsupported charge handler {q_handle}")