"""
A collection of protocols for calculating the gradients of observables
with respect to force field parameters.
"""
import abc
import typing
import pint
from openff.evaluator.attributes import UNDEFINED
from openff.evaluator.forcefield import ParameterGradient, ParameterGradientKey
from openff.evaluator.substances import Substance
from openff.evaluator.thermodynamics import ThermodynamicState
from openff.evaluator.workflow import Protocol, workflow_protocol
from openff.evaluator.workflow.attributes import InputAttribute, OutputAttribute
[docs]@workflow_protocol()
class BaseGradientPotentials(Protocol, abc.ABC):
"""A base class for protocols which will evaluate the reduced potentials of a
series of configurations using a set of force field parameters which have been
slightly increased and slightly decreased. These are mainly useful when
estimating gradients with respect to force field parameters using the central
difference method.
"""
force_field_path = InputAttribute(
docstring="The path to the force field which contains the parameters to "
"differentiate the observable with respect to. When reweighting "
"observables, this should be the `target` force field.",
type_hint=str,
default_value=UNDEFINED,
)
statistics_path = InputAttribute(
docstring="The path to a statistics array containing potentials "
"evaluated at each frame of the trajectory using the input "
"`force_field_path` and at the input `thermodynamic_state`.",
type_hint=str,
default_value=UNDEFINED,
)
thermodynamic_state = InputAttribute(
docstring="The thermodynamic state to estimate the gradients at. When "
"reweighting observables, this should be the `target` state.",
type_hint=ThermodynamicState,
default_value=UNDEFINED,
)
substance = InputAttribute(
docstring="The substance which describes the composition of the system.",
type_hint=Substance,
default_value=UNDEFINED,
)
coordinate_file_path = InputAttribute(
docstring="A path to a PDB coordinate file which describes the topology of "
"the system.",
type_hint=str,
default_value=UNDEFINED,
)
trajectory_file_path = InputAttribute(
docstring="A path to the trajectory of configurations",
type_hint=str,
default_value=UNDEFINED,
)
enable_pbc = InputAttribute(
docstring="If true, periodic boundary conditions will be enabled when "
"re-evaluating the reduced potentials.",
type_hint=bool,
default_value=True,
)
parameter_key = InputAttribute(
docstring="The key of the parameter to differentiate with respect to.",
type_hint=ParameterGradientKey,
default_value=UNDEFINED,
)
perturbation_scale = InputAttribute(
docstring="The amount to perturb the parameter by, such that "
"p_new = p_old * (1 +/- `perturbation_scale`)",
type_hint=float,
default_value=1.0e-4,
)
use_subset_of_force_field = InputAttribute(
docstring="If true, the reduced potentials will be estimated using "
"a system which only contains the parameters of interest, e.g. if the "
"gradient of interest is with respect to the VdW epsilon parameter, then "
"all valence / electrostatic terms will be ignored.",
type_hint=bool,
default_value=True,
)
effective_sample_indices = InputAttribute(
docstring="This a placeholder input which is not currently implemented.",
type_hint=list,
default_value=UNDEFINED,
optional=True,
)
reverse_potentials_path = OutputAttribute(
docstring="A file path to the energies evaluated using the parameters"
"perturbed in the reverse direction.",
type_hint=str,
)
forward_potentials_path = OutputAttribute(
docstring="A file path to the energies evaluated using the parameters"
"perturbed in the forward direction.",
type_hint=str,
)
reverse_parameter_value = OutputAttribute(
docstring="The value of the parameter perturbed in the reverse direction.",
type_hint=pint.Quantity,
)
forward_parameter_value = OutputAttribute(
docstring="The value of the parameter perturbed in the forward direction.",
type_hint=pint.Quantity,
)
[docs]@workflow_protocol()
class CentralDifferenceGradient(Protocol):
"""A protocol which employs the central diference method
to estimate the gradient of an observable A, such that
grad = (A(x-h) - A(x+h)) / (2h)
Notes
-----
The `values` input must either be a list of pint.Quantity, a ProtocolPath to a list
of pint.Quantity, or a list of ProtocolPath which each point to a pint.Quantity.
"""
parameter_key = InputAttribute(
docstring="The key of the parameter to differentiate with respect to.",
type_hint=ParameterGradientKey,
default_value=UNDEFINED,
)
reverse_observable_value = InputAttribute(
docstring="The value of the observable evaluated using the parameters"
"perturbed in the reverse direction.",
type_hint=typing.Union[pint.Quantity, pint.Measurement],
default_value=UNDEFINED,
)
forward_observable_value = InputAttribute(
docstring="The value of the observable evaluated using the parameters"
"perturbed in the forward direction.",
type_hint=typing.Union[pint.Quantity, pint.Measurement],
default_value=UNDEFINED,
)
reverse_parameter_value = InputAttribute(
docstring="The value of the parameter perturbed in the reverse direction.",
type_hint=pint.Quantity,
default_value=UNDEFINED,
)
forward_parameter_value = InputAttribute(
docstring="The value of the parameter perturbed in the forward direction.",
type_hint=pint.Quantity,
default_value=UNDEFINED,
)
gradient = OutputAttribute(
docstring="The estimated gradient", type_hint=ParameterGradient
)
def _execute(self, directory, available_resources):
if self.forward_parameter_value < self.reverse_parameter_value:
raise ValueError(
f"The forward parameter value ({self.forward_parameter_value}) must "
f"be larger than the reverse value ({self.reverse_parameter_value})."
)
reverse_value = self.reverse_observable_value
forward_value = self.forward_observable_value
if isinstance(reverse_value, pint.Measurement):
reverse_value = reverse_value.value
if isinstance(forward_value, pint.Measurement):
forward_value = forward_value.value
gradient = (forward_value - reverse_value) / (
self.forward_parameter_value - self.reverse_parameter_value
)
self.gradient = ParameterGradient(self.parameter_key, gradient)