Coverage for openxps/couplings/inner_product_coupling.py: 100%
70 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-13 22:08 +0000
1"""
2Inner product coupling.
4.. module:: openxps.couplings.inner_product_coupling
5 :platform: Linux, MacOS, Windows
6 :synopsis: The inner product of two separable vector-valued functions
8.. classauthor:: Charlles Abreu <craabreu@gmail.com>
10"""
12import typing as t
14import cvpack
15import openmm as mm
16from openmm import _openmm as mmswig
17from openmm import unit as mmunit
19from openxps.utils import Function, preprocess_args
21from ..dynamical_variable import DynamicalVariable
22from .base import Coupling
25class InnerProductCoupling(Coupling):
26 r"""Coupling defined by the inner product of two separable vector-valued functions.
28 Use this class if the coupling energy is the inner product of a vector-valued
29 function :math:`{\boldsymbol \lambda}` of the extended dynamical variables
30 :math:`{\bf s}` and a vector-valued function :math:`{\bf u}` of the physical
31 coordinates :math:`{\bf r}`, i.e.,
33 .. math::
35 U = {\boldsymbol \lambda}({\bf s}) \cdot {\bf u}({\bf r})
36 = \sum_{i=1}^n \lambda_i({\bf s}) u_i({\bf r})
38 where :math:`n` is the dimensionality of the vector-valued functions.
40 Parameters
41 ----------
42 forces
43 A sequence of :OpenMM:`Force` objects whose sum depends linearly on a vector of
44 :math:`n` global context parameters.
45 dynamical_variables
46 The dynamical variables of the extended phase-space system.
47 functions
48 A dictionary defining global context parameters (keys) as mathematical functions
49 (values) of the dynamical variables. It is not necessary to include identity
50 functions (i.e., ``key=value``).
51 **parameters
52 Named parameters that appear in the mathematical expressions.
54 Examples
55 --------
56 >>> from copy import copy
57 >>> from collections import namedtuple
58 >>> import cvpack
59 >>> import openxps as xps
60 >>> from openmm import unit
61 >>> from math import pi
62 >>> import openmm as mm
63 >>> from openmmtools import testsystems
64 >>> model = testsystems.AlanineDipeptideVacuum()
66 A function to remove Coulomb interactions:
68 >>> Parameters = namedtuple(
69 ... "Parameters",
70 ... ["index", "charge", "sigma", "epsilon"],
71 ... )
72 >>> def remove_coulomb_interaction(
73 ... force: mm.NonbondedForce, p1: Parameters, p2: Parameters
74 ... ) -> None:
75 ... force.addException(
76 ... p1.index,
77 ... p2.index,
78 ... 0.0,
79 ... (p1.sigma + p2.sigma)/2,
80 ... (p1.epsilon * p2.epsilon).sqrt(),
81 ... )
83 Remove carbonyl oxygen <=> amide hydrogen interactions:
85 >>> nbforce = next(
86 ... f for f in model.system.getForces()
87 ... if isinstance(f, mm.NonbondedForce)
88 ... )
89 >>> O1 = Parameters(5, *nbforce.getParticleParameters(5))
90 >>> H1 = Parameters(7, *nbforce.getParticleParameters(7))
91 >>> O2 = Parameters(15, *nbforce.getParticleParameters(15))
92 >>> H2 = Parameters(17, *nbforce.getParticleParameters(17))
93 >>> remove_coulomb_interaction(nbforce, O1, H2)
94 >>> remove_coulomb_interaction(nbforce, O2, H1)
96 Add scaled Coulomb interactions:
98 >>> force = mm.CustomBondForce(f"scaling*chargeProd/r")
99 >>> _ = force.addGlobalParameter("scaling", 1.0)
100 >>> _ = force.addEnergyParameterDerivative("scaling")
101 >>> _ = force.addPerBondParameter("chargeProd")
102 >>> _ = force.addBond(O1.index, H2.index, [O1.charge*H2.charge])
103 >>> _ = force.addBond(O2.index, H1.index, [O2.charge*H1.charge])
105 Create a coupling between the dynamical variable and the nonbonded force:
107 >>> lambda_dv = xps.DynamicalVariable(
108 ... name="lambda",
109 ... unit=unit.dimensionless,
110 ... mass=1.0 * unit.dalton * unit.nanometer**2,
111 ... bounds=xps.ReflectiveBounds(0.0, 1.0, unit.dimensionless),
112 ... )
113 >>> coupling = xps.InnerProductCoupling(
114 ... [force],
115 ... [lambda_dv],
116 ... functions={"scaling": "(1-cos(alpha*lambda))/2"},
117 ... alpha=pi*unit.radian,
118 ... )
119 >>> context = xps.ExtendedSpaceContext(
120 ... xps.ExtendedSpaceSystem(model.system, coupling),
121 ... xps.LockstepIntegrator(mm.VerletIntegrator(1.0 * mmunit.femtosecond)),
122 ... mm.Platform.getPlatformByName("Reference"),
123 ... )
124 >>> context.setPositions(model.positions)
125 >>> context.setDynamicalVariableValues([1.0])
126 """
128 @preprocess_args
129 def __init__(
130 self,
131 forces: t.Iterable[mm.Force],
132 dynamical_variables: t.Iterable[DynamicalVariable],
133 functions: t.Optional[dict[str, str]] = None,
134 **parameters: mmunit.Quantity,
135 ) -> None:
136 forces = [
137 cvpack.OpenMMForceWrapper(
138 force,
139 mmunit.kilojoule_per_mole,
140 name=force.getName(),
141 )
142 if not isinstance(force, cvpack.CollectiveVariable)
143 else force
144 for force in forces
145 ]
146 self._functions = [
147 Function(name, expression, **parameters)
148 for name, expression in (functions or {}).items()
149 ]
150 self._parameters = parameters
151 super().__init__(forces, dynamical_variables)
152 self._dynamic_parameters = self._getDynamicParameters()
154 def __getstate__(self) -> dict[str, t.Any]:
155 state = super().__getstate__()
156 state["functions"] = self._functions
157 state["parameters"] = self._parameters
158 return state
160 def __setstate__(self, state: dict[str, t.Any]) -> None:
161 self._functions = state["functions"]
162 self._parameters = state["parameters"]
163 super().__setstate__(state)
164 self._dynamic_parameters = self._getDynamicParameters()
166 def _getDynamicParameters(self) -> set[str]:
167 all_force_parameters = {
168 force.getGlobalParameterName(index)
169 for force in self._forces
170 for index in range(force.getNumGlobalParameters())
171 }
172 function_names = {fn.getName() for fn in self._functions}
173 if self._functions:
174 function_variables = set.union(
175 *(fn.getVariables() for fn in self._functions),
176 )
177 else:
178 function_variables = set()
180 if missing := function_names - all_force_parameters:
181 raise ValueError(
182 "These functions are not global parameters in the provided forces: "
183 + ", ".join(missing)
184 )
186 all_dvs = {dv.name for dv in self._dynamical_variables}
188 if functions_missing_dvs := [
189 fn.getName() for fn in self._functions if not (fn.getVariables() & all_dvs)
190 ]:
191 raise ValueError(
192 "These functions do not depend on any dynamical variables: "
193 + ", ".join(functions_missing_dvs)
194 )
196 dvs_in_force_parameters = all_dvs & all_force_parameters
197 dvs_in_function_variables = all_dvs & function_variables
199 if dvs_in_both := dvs_in_function_variables & dvs_in_force_parameters:
200 raise ValueError(
201 "These dynamical variables are both function variables and global "
202 f"context parameters: {', '.join(dvs_in_both)}"
203 )
204 if dvs_in_neither := all_dvs - (
205 dvs_in_function_variables | dvs_in_force_parameters
206 ):
207 raise ValueError(
208 "These dynamical variables are neither function variables nor global "
209 "context parameters: " + ", ".join(dvs_in_neither)
210 )
212 dynamic_parameters = dvs_in_force_parameters | function_names
213 for force in self._forces:
214 force_parameters = {
215 force.getGlobalParameterName(index)
216 for index in range(force.getNumGlobalParameters())
217 }
218 has_derivatives = {
219 force.getEnergyParameterDerivativeName(index)
220 for index in range(force.getNumEnergyParameterDerivatives())
221 }
222 if missing := force_parameters & dynamic_parameters - has_derivatives:
223 raise ValueError(
224 "The following parameters require a derivative and are present in "
225 f"force {force.getName()}, but no derivative was requested: "
226 + ", ".join(missing)
227 )
228 return dynamic_parameters
230 @staticmethod
231 def _derivativeName(parameter: str) -> str:
232 return "derivative_with_respect_to_" + parameter
234 def _createFlippedForce(self) -> mm.CustomCVForce:
235 flipped_force = mm.CustomCVForce(
236 "+".join(f"{p}*{self._derivativeName(p)}" for p in self._dynamic_parameters)
237 )
238 all_dvs = [dv.name for dv in self._dynamical_variables]
239 for parameter in self._dynamic_parameters:
240 flipped_force.addGlobalParameter(self._derivativeName(parameter), 0.0)
241 for fn in self._functions:
242 flipped_force.addCollectiveVariable(
243 fn.getName(), fn.createCollectiveVariable(all_dvs)
244 )
245 for dv in self._dynamical_variables:
246 if dv.name in self._dynamic_parameters:
247 flipped_force.addCollectiveVariable(
248 dv.name, dv.createCollectiveVariable(self._dv_indices[dv.name])
249 )
250 return flipped_force
252 def updateExtensionContext(
253 self,
254 physical_context: mm.Context,
255 extension_context: mm.Context,
256 ) -> None:
257 state = mmswig.Context_getState(physical_context, mm.State.ParameterDerivatives)
258 for name, value in mmswig.State_getEnergyParameterDerivatives(state).items():
259 if name in self._dynamic_parameters:
260 mmswig.Context_setParameter(
261 extension_context, self._derivativeName(name), value
262 )
265InnerProductCoupling.registerTag("!openxps.InnerProductCoupling")