Detach parameter graph before OpenMM conversion (#113)

SimonBoothroyd · Jul 27, 2024 · fe5d679 · fe5d679
1 parent 5220607
commit fe5d679
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 9 deletions.
diff --git a/smee/converters/openmm/_openmm.py b/smee/converters/openmm/_openmm.py
@@ -104,7 +104,7 @@ def create_openmm_system(
                 parent_idxs = [i + start_idx for i in key.orientation_atom_indices]
 
                 local_frame_coords = smee.geometry.polar_to_cartesian_coords(
-                    v_sites.parameters[[parameter_idx], :]
+                    v_sites.parameters[[parameter_idx], :].detach()
                 )
                 origin, x_dir, y_dir = v_sites.weights[parameter_idx]
 

diff --git a/smee/converters/openmm/nonbonded.py b/smee/converters/openmm/nonbonded.py
@@ -117,12 +117,15 @@ def _build_vdw_lookup(
     for i, j in itertools.product(range(n_params), range(n_params)):
         if (i, j) in exceptions:
             parameters = {
-                col: potential.parameters[exceptions[i, j], col_idx]
+                col: potential.parameters[exceptions[i, j], col_idx].detach()
                 for col, col_idx in parameter_col_to_idx.items()
             }
         else:
             parameters = _eval_mixing_fn(
-                potential, mixing_fn, potential.parameters[i], potential.parameters[j]
+                potential,
+                mixing_fn,
+                potential.parameters[i].detach(),
+                potential.parameters[j].detach(),
             )
 
         unit_conversion = {
@@ -259,7 +262,7 @@ def _add_parameters_to_vdw_without_lookup(
 
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameter_map = topology.parameters[potential.type]
-        parameters = parameter_map.assignment_matrix @ potential.parameters
+        parameters = parameter_map.assignment_matrix @ potential.parameters.detach()
 
         for _ in range(n_copies):
             for parameter in parameters:
@@ -465,7 +468,7 @@ def convert_lj_potential(
 
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameter_map = topology.parameters[potential.type]
-        parameters = parameter_map.assignment_matrix @ potential.parameters
+        parameters = parameter_map.assignment_matrix @ potential.parameters.detach()
 
         for _ in range(n_copies):
             for epsilon, sigma in parameters:
@@ -536,7 +539,7 @@ def convert_coulomb_potential(
 
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameter_map = topology.parameters[potential.type]
-        parameters = parameter_map.assignment_matrix @ potential.parameters
+        parameters = parameter_map.assignment_matrix @ potential.parameters.detach()
 
         for _ in range(n_copies):
             for charge in parameters:

diff --git a/smee/converters/openmm/valence.py b/smee/converters/openmm/valence.py
@@ -24,7 +24,7 @@ def convert_bond_potential(
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameters = (
             topology.parameters[potential.type].assignment_matrix @ potential.parameters
-        )
+        ).detach()
 
         for _ in range(n_copies):
             atom_idxs = topology.parameters[potential.type].particle_idxs + idx_offset
@@ -56,7 +56,7 @@ def _convert_angle_potential(
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameters = (
             topology.parameters[potential.type].assignment_matrix @ potential.parameters
-        )
+        ).detach()
 
         for _ in range(n_copies):
             atom_idxs = topology.parameters[potential.type].particle_idxs + idx_offset
@@ -92,7 +92,7 @@ def convert_torsion_potential(
     for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
         parameters = (
             topology.parameters[potential.type].assignment_matrix @ potential.parameters
-        )
+        ).detach()
 
         for _ in range(n_copies):
             atom_idxs = topology.parameters[potential.type].particle_idxs + idx_offset