def test_average_free_energies_protocol():
"""Tests adding together two free energies."""
delta_g_one = Observable(
value=(-10.0 * unit.kilocalorie / unit.mole).plus_minus(
1.0 * unit.kilocalorie / unit.mole),
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "sigma"),
value=0.1 * unit.kilocalorie / unit.mole / unit.angstrom,
)
],
)
delta_g_two = Observable(
value=(-20.0 * unit.kilocalorie / unit.mole).plus_minus(
2.0 * unit.kilocalorie / unit.mole),
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "sigma"),
value=0.2 * unit.kilocalorie / unit.mole / unit.angstrom,
)
],
)
thermodynamic_state = ThermodynamicState(298 * unit.kelvin,
1 * unit.atmosphere)
sum_protocol = AverageFreeEnergies("")
sum_protocol.values = [delta_g_one, delta_g_two]
sum_protocol.thermodynamic_state = thermodynamic_state
sum_protocol.execute()
result_value = sum_protocol.result.value.to(unit.kilocalorie / unit.mole)
result_uncertainty = sum_protocol.result.error.to(unit.kilocalorie /
unit.mole)
assert isinstance(sum_protocol.result, Observable)
assert result_value.magnitude == pytest.approx(-20.0, abs=0.2)
assert result_uncertainty.magnitude == pytest.approx(2.0, abs=0.2)
assert (sum_protocol.confidence_intervals[0] > result_value >
sum_protocol.confidence_intervals[1])
gradient_value = sum_protocol.result.gradients[0].value.to(
unit.kilocalorie / unit.mole / unit.angstrom)
beta = 1.0 / (298.0 * unit.kelvin * unit.molar_gas_constant).to(
unit.kilocalorie / unit.mole)
assert np.isclose(
gradient_value.magnitude,
(0.1 * np.exp(-beta.magnitude * -10.0) +
0.2 * np.exp(-beta.magnitude * -20.0)) /
(np.exp(-beta.magnitude * -10.0) + np.exp(-beta.magnitude * -20.0)),
)
def test_gradient_division():
gradient_a = ParameterGradient(
ParameterGradientKey("vdW", "[#1:1]", "epsilon"), 2.0 * unit.kelvin)
result = gradient_a / 2.0
assert np.isclose(result.value.to(unit.kelvin).magnitude, 1.0)
gradient_c = ParameterGradient(
ParameterGradientKey("vdW", "[#1:1]", "epsilon"), 1.0 * unit.kelvin)
with pytest.raises(ValueError):
gradient_a / gradient_c
def _execute(self, directory, available_resources):
force_field_source = ForceFieldSource.from_json(self.force_field_path)
if not isinstance(force_field_source, SmirnoffForceFieldSource):
raise ValueError("Only SMIRNOFF force fields are supported.")
force_field = force_field_source.to_force_field()
parameter_units = {
gradient_key: openmm_quantity_to_pint(
getattr(
force_field.get_parameter_handler(
gradient_key.tag).parameters[gradient_key.smirks],
gradient_key.attribute,
)).units
for gradient_key in self.gradient_parameters
}
self.input_observables.clear_gradients()
if isinstance(self.input_observables, Observable):
self.output_observables = Observable(
value=self.input_observables.value,
gradients=[
ParameterGradient(
key=gradient_key,
value=(0.0 * self.input_observables.value.units /
parameter_units[gradient_key]),
) for gradient_key in self.gradient_parameters
],
)
elif isinstance(self.input_observables, ObservableArray):
self.output_observables = ObservableArray(
value=self.input_observables.value,
gradients=[
ParameterGradient(
key=gradient_key,
value=(
numpy.zeros(self.input_observables.value.shape) *
self.input_observables.value.units /
parameter_units[gradient_key]),
) for gradient_key in self.gradient_parameters
],
)
else:
raise NotImplementedError()
def test_bootstrap(data_values, expected_error, sub_counts):
def bootstrap_function(values: ObservableArray) -> Observable:
return Observable(
value=values.value.mean().plus_minus(0.0 * values.value.units),
gradients=[
ParameterGradient(gradient.key, numpy.mean(gradient.value))
for gradient in values.gradients
],
)
data = ObservableArray(
value=data_values,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=data_values,
)
],
)
average = bootstrap(bootstrap_function, 1000, 1.0, sub_counts, values=data)
assert numpy.isclose(average.value, data.value.mean())
assert numpy.isclose(average.gradients[0].value, data.value.mean())
if expected_error is not None:
assert numpy.isclose(average.error, expected_error, rtol=0.1)
def test_observable_array_join():
gradient_unit = unit.mole / unit.kilojoule
observables = [
ObservableArray(
value=(numpy.arange(2) + i * 2) * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=(numpy.arange(2) + i * 2) * unit.kelvin *
gradient_unit,
)
],
) for i in range(2)
]
joined = ObservableArray.join(*observables)
assert len(joined) == 4
assert numpy.allclose(joined.value,
numpy.arange(4).reshape(-1, 1) * unit.kelvin)
assert numpy.allclose(
joined.gradients[0].value,
numpy.arange(4).reshape(-1, 1) * unit.kelvin * gradient_unit,
)
def test_observable_array_valid_initializer(
value: unit.Quantity,
gradient_values: List[unit.Quantity],
expected_value: unit.Quantity,
expected_gradient_values: List[unit.Quantity],
):
observable = ObservableArray(
value,
[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=gradient_value,
) for gradient_value in gradient_values
],
)
# noinspection PyUnresolvedReferences
assert observable.value.shape == expected_value.shape
assert numpy.allclose(observable.value, expected_value)
assert all(observable.gradients[i].value.shape ==
expected_gradient_values[i].shape
for i in range(len(expected_gradient_values)))
assert all(
numpy.allclose(observable.gradients[i].value,
expected_gradient_values[i])
for i in range(len(expected_gradient_values)))
def _execute(self, directory, available_resources):
import mdtraj
charges = self._extract_charges(self.parameterized_system.system)
charge_derivatives = self._compute_charge_derivatives(len(charges))
dipole_moments = []
dipole_gradients = {key: [] for key in self.gradient_parameters}
for chunk in mdtraj.iterload(
self.trajectory_path, top=self.parameterized_system.topology_path, chunk=50
):
xyz = chunk.xyz.transpose(0, 2, 1) * unit.nanometers
dipole_moments.extend(xyz.dot(charges))
for key in self.gradient_parameters:
dipole_gradients[key].extend(xyz.dot(charge_derivatives[key]))
self.dipole_moments = ObservableArray(
value=np.vstack(dipole_moments),
gradients=[
ParameterGradient(key=key, value=np.vstack(dipole_gradients[key]))
for key in self.gradient_parameters
],
)
def bootstrap_function(values: ObservableArray) -> Observable:
return Observable(
value=values.value.mean().plus_minus(0.0 * values.value.units),
gradients=[
ParameterGradient(gradient.key, numpy.mean(gradient.value))
for gradient in values.gradients
],
)
def _compute_weights(
mbar: pymbar.MBAR,
target_reduced_potentials: ObservableArray) -> ObservableArray:
"""Return the values that each sample at the target state should be weighted
by.
Parameters
----------
mbar
A pre-computed MBAR object encoded information from the reference states.
target_reduced_potentials
The reduced potentials at the target state.
Returns
-------
The values to weight each sample by.
"""
from scipy.special import logsumexp
u_kn = target_reduced_potentials.value.to(
unit.dimensionless).magnitude.T
log_denominator_n = logsumexp(mbar.f_k - mbar.u_kn.T,
b=mbar.N_k,
axis=1)
f_hat = -logsumexp(-u_kn - log_denominator_n, axis=1)
# Calculate the weights
weights = np.exp(f_hat - u_kn - log_denominator_n) * unit.dimensionless
# Compute the gradients of the weights.
weight_gradients = []
for gradient in target_reduced_potentials.gradients:
gradient_value = gradient.value.magnitude.flatten()
# Compute the numerator of the gradient. We need to specifically ask for the
# sign of the exp sum as the numerator may be negative.
d_f_hat_numerator, d_f_hat_numerator_sign = logsumexp(
-u_kn - log_denominator_n,
b=gradient_value,
axis=1,
return_sign=True)
d_f_hat_d_theta = d_f_hat_numerator_sign * np.exp(
d_f_hat_numerator + f_hat)
d_weights_d_theta = ((d_f_hat_d_theta - gradient_value) * weights *
gradient.value.units)
weight_gradients.append(
ParameterGradient(key=gradient.key, value=d_weights_d_theta.T))
return ObservableArray(value=weights.T, gradients=weight_gradients)
def test_gradient_subtraction():
gradient_a = ParameterGradient(
ParameterGradientKey("vdW", "[#1:1]", "epsilon"), 1.0 * unit.kelvin)
gradient_b = ParameterGradient(
ParameterGradientKey("vdW", "[#1:1]", "epsilon"), 2.0 * unit.kelvin)
result = gradient_a - gradient_b
assert np.isclose(result.value.to(unit.kelvin).magnitude, -1.0)
result = gradient_b - gradient_a
assert np.isclose(result.value.to(unit.kelvin).magnitude, 1.0)
gradient_c = ParameterGradient(
ParameterGradientKey("vdW", "[#6:1]", "epsilon"), 1.0 * unit.kelvin)
with pytest.raises(ValueError):
gradient_a - gradient_c
with pytest.raises(ValueError):
gradient_c - gradient_a
with pytest.raises(ValueError):
gradient_a - 1.0
def test_observable_array_join_single():
gradient_unit = unit.mole / unit.kilojoule
joined = ObservableArray.join(
ObservableArray(
value=(numpy.arange(2)) * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=(numpy.arange(2)) * unit.kelvin * gradient_unit,
)
],
))
assert len(joined) == 2
def _mock_observable(
value: ValueType,
gradient_values: List[Tuple[str, str, str, ValueType]],
object_type: Union[Type[Observable], Type[ObservableArray]],
):
return object_type(
value=value,
gradients=[
ParameterGradient(
key=ParameterGradientKey(tag, smirks, attribute),
value=value * unit.kelvin,
) for tag, smirks, attribute, value in gradient_values
],
)
def test_observable_array_subset():
observable = ObservableArray(
value=numpy.arange(4) * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=numpy.arange(4) * unit.kelvin,
)
],
)
subset = observable.subset([1, 3])
assert len(subset) == 2
assert numpy.allclose(subset.value,
numpy.array([[1.0], [3.0]]) * unit.kelvin)
assert numpy.allclose(subset.gradients[0].value,
numpy.array([[1.0], [3.0]]) * unit.kelvin)
def subset(self, indices: Iterable[int]) -> "ObservableArray":
"""Extracts the subset of the values stored for this observable at the
specified indices.
Parameters
----------
indices
The indices of the entries to extract.
Returns
-------
The subset of the observable values.
"""
return self.__class__(
value=self._value[indices],
gradients=[
ParameterGradient(key=gradient.key, value=gradient.value[indices])
for gradient in self._gradients
],
)
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)
def test_observable_array_round_trip(value):
observable = ObservableArray(
value=value * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=value * 2.0 * unit.kelvin,
)
],
)
round_tripped: ObservableArray = json.loads(json.dumps(
observable, cls=TypedJSONEncoder),
cls=TypedJSONDecoder)
assert isinstance(round_tripped, ObservableArray)
assert numpy.isclose(observable.value, round_tripped.value)
assert len(observable.gradients) == len(round_tripped.gradients)
assert observable.gradients[0] == round_tripped.gradients[0]
def test_observable_round_trip():
observable = Observable(
value=(0.1 * unit.kelvin).plus_minus(0.2 * unit.kelvin),
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=0.2 * unit.kelvin,
)
],
)
round_tripped: Observable = json.loads(json.dumps(observable,
cls=TypedJSONEncoder),
cls=TypedJSONDecoder)
assert isinstance(round_tripped, Observable)
assert numpy.isclose(observable.value, round_tripped.value)
assert numpy.isclose(observable.error, round_tripped.error)
assert len(observable.gradients) == len(round_tripped.gradients)
assert observable.gradients[0] == round_tripped.gradients[0]
def test_frame_subset():
observable_frame = ObservableFrame({
"Temperature":
ObservableArray(
value=numpy.arange(4) * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=numpy.arange(4) * unit.kelvin,
)
],
)
})
subset = observable_frame.subset([1, 3])
assert len(subset) == 2
assert numpy.allclose(subset["Temperature"].value,
numpy.array([[1.0], [3.0]]) * unit.kelvin)
assert numpy.allclose(
subset["Temperature"].gradients[0].value,
numpy.array([[1.0], [3.0]]) * unit.kelvin,
)
def join(cls, *observables: "ObservableArray") -> "ObservableArray":
"""Concatenates multiple observables together in the order that they appear in
the args list.
Parameters
----------
observables
The observables to join.
Returns
-------
The concatenated observable object.
"""
if len(observables) < 1:
raise ValueError("At least one observable must be provided.")
if len(observables) == 1:
return observables[0]
expected_gradients = {gradient.key for gradient in observables[0].gradients}
expected_gradient_units = {
gradient.key: gradient.value.units for gradient in observables[0].gradients
}
# Ensure the arrays contain the same observables.
if not all(
observable.value.dimensionality == observables[0].value.dimensionality
for observable in observables
):
raise ValueError("The observables must all have compatible units.")
# Ensure the arrays contain gradients for the same FF parameters.
if not all(
{gradient.key for gradient in observable.gradients} == expected_gradients
for observable in observables
):
raise ValueError(
"The observables must contain gradient information for the same "
"parameters."
)
# Ensure the gradients are all in the same units.
if not all(
{gradient.key: gradient.value.units for gradient in observable.gradients}
== expected_gradient_units
for observable in observables
):
raise ValueError(
"The gradients of each of the observables must have the same units."
)
return ObservableArray(
value=numpy.concatenate(
[
observable.value.to(observables[0].value.units).magnitude
for observable in observables
]
)
* observables[0].value.units,
gradients=[
ParameterGradient(
key=gradient_key,
value=numpy.concatenate(
[
next(
x for x in observable.gradients if x.key == gradient_key
)
.value.to(expected_gradient_units[gradient_key])
.magnitude
for observable in observables
]
)
* expected_gradient_units[gradient_key],
)
for gradient_key in expected_gradients
],
)
"str" * unit.kelvin,
[],
pytest.raises(TypeError),
"The value must be a unit-wrapped integer, float or numpy array.",
),
(
numpy.ones((2, 2, 2)) * unit.kelvin,
[],
pytest.raises(ValueError),
"The wrapped array must not contain more than two dimensions.",
),
(
None,
[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=numpy.ones((2, 2)) * unit.kelvin,
),
],
pytest.raises(ValueError),
"A valid value must be provided.",
),
(
1.0 * unit.kelvin,
[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#6:1]", "epsilon"),
value=numpy.ones(1),
),
],
pytest.raises(TypeError),
"The gradient values must be unit-wrapped integers, floats or numpy arrays.",
def _compute_gradients(
gradient_parameters: List[ParameterGradientKey],
observables: ObservableFrame,
force_field: "ForceField",
thermodynamic_state: ThermodynamicState,
topology: "Topology",
trajectory: "Trajectory",
compute_resources: ComputeResources,
enable_pbc: bool = True,
perturbation_amount: float = 0.0001,
):
"""Computes the gradients of the provided observables with respect to
the set of specified force field parameters using the central difference
finite difference method.
Notes
-----
The ``observables`` object will be modified in-place.
Parameters
----------
gradient_parameters
The parameters to differentiate with respect to.
observables
The observables to differentiate.
force_field
The full set force field parameters which contain the parameters to
differentiate.
thermodynamic_state
The state at which the trajectory was sampled
topology
The topology of the system the observables were collected for.
trajectory
The trajectory over which the observables were collected.
compute_resources
The compute resources available for the computations.
enable_pbc
Whether PBC should be enabled when re-evaluating system energies.
perturbation_amount
The amount to perturb for the force field parameter by.
"""
from simtk import openmm
gradients = defaultdict(list)
observables.clear_gradients()
if enable_pbc:
# Make sure the PBC are set on the topology otherwise the cut-off will be
# set incorrectly.
topology.box_vectors = trajectory.openmm_boxes(0)
for parameter_key in gradient_parameters:
# Build the slightly perturbed systems.
reverse_system, reverse_parameter_value = system_subset(
parameter_key, force_field, topology, -perturbation_amount)
forward_system, forward_parameter_value = system_subset(
parameter_key, force_field, topology, perturbation_amount)
# Perform a cheap check to try and catch most cases where the systems energy
# does not depend on this parameter.
reverse_xml = openmm.XmlSerializer.serialize(reverse_system)
forward_xml = openmm.XmlSerializer.serialize(forward_system)
if not enable_pbc:
disable_pbc(reverse_system)
disable_pbc(forward_system)
reverse_parameter_value = openmm_quantity_to_pint(
reverse_parameter_value)
forward_parameter_value = openmm_quantity_to_pint(
forward_parameter_value)
# Evaluate the energies using the reverse and forward sub-systems.
if reverse_xml != forward_xml:
reverse_energies = _evaluate_energies(
thermodynamic_state,
reverse_system,
trajectory,
compute_resources,
enable_pbc,
)
forward_energies = _evaluate_energies(
thermodynamic_state,
forward_system,
trajectory,
compute_resources,
enable_pbc,
)
else:
zeros = np.zeros(len(trajectory))
reverse_energies = forward_energies = ObservableFrame({
ObservableType.PotentialEnergy:
ObservableArray(
zeros * unit.kilojoule / unit.mole,
[
ParameterGradient(
key=parameter_key,
value=(zeros * unit.kilojoule / unit.mole /
reverse_parameter_value.units),
)
],
),
ObservableType.ReducedPotential:
ObservableArray(
zeros * unit.dimensionless,
[
ParameterGradient(
key=parameter_key,
value=(zeros * unit.dimensionless /
reverse_parameter_value.units),
)
],
),
})
potential_gradient = ParameterGradient(
key=parameter_key,
value=(forward_energies[ObservableType.PotentialEnergy].value -
reverse_energies[ObservableType.PotentialEnergy].value) /
(forward_parameter_value - reverse_parameter_value),
)
reduced_potential_gradient = ParameterGradient(
key=parameter_key,
value=(forward_energies[ObservableType.ReducedPotential].value -
reverse_energies[ObservableType.ReducedPotential].value) /
(forward_parameter_value - reverse_parameter_value),
)
gradients[ObservableType.PotentialEnergy].append(potential_gradient)
gradients[ObservableType.TotalEnergy].append(potential_gradient)
gradients[ObservableType.Enthalpy].append(potential_gradient)
gradients[ObservableType.ReducedPotential].append(
reduced_potential_gradient)
if ObservableType.KineticEnergy in observables:
gradients[ObservableType.KineticEnergy].append(
ParameterGradient(
key=parameter_key,
value=(
np.zeros(potential_gradient.value.shape) *
observables[ObservableType.KineticEnergy].value.units /
reverse_parameter_value.units),
))
if ObservableType.Density in observables:
gradients[ObservableType.Density].append(
ParameterGradient(
key=parameter_key,
value=(np.zeros(potential_gradient.value.shape) *
observables[ObservableType.Density].value.units /
reverse_parameter_value.units),
))
if ObservableType.Volume in observables:
gradients[ObservableType.Volume].append(
ParameterGradient(
key=parameter_key,
value=(np.zeros(potential_gradient.value.shape) *
observables[ObservableType.Volume].value.units /
reverse_parameter_value.units),
))
for observable_type in observables:
observables[observable_type] = ObservableArray(
value=observables[observable_type].value,
gradients=gradients[observable_type],
)
def _bootstrap_function(self, **kwargs: ObservableArray) -> Observable:
"""The function to perform on the data set being sampled by
bootstrapping.
Parameters
----------
observables
The bootstrap sample values.
Returns
-------
The result of evaluating the data.
"""
# The simple base function only supports a single observable.
assert len(kwargs) == 1
# Compute the mean observable.
sample_observable = next(iter(kwargs.values()))
mean_observable = np.mean(sample_observable.value, axis=0)
if sample_observable.value.shape[1] > 1:
mean_observable = mean_observable.reshape(1, -1)
else:
mean_observable = mean_observable.item()
# Retrieve the potential gradients for easy access
potential_gradients = {
gradient.key: gradient.value
for gradient in (
[]
if self.potential_energies == UNDEFINED
else self.potential_energies.gradients
)
}
observable_gradients = {
gradient.key: gradient
for gradient in (
[]
if self.potential_energies == UNDEFINED
else sample_observable.gradients
)
}
# Compute the mean gradients.
gradients = []
for gradient_key in observable_gradients:
gradient = observable_gradients[gradient_key]
value = np.mean(gradient.value, axis=0) - self.thermodynamic_state.beta * (
np.mean(
sample_observable.value * potential_gradients[gradient.key],
axis=0,
)
- (
np.mean(sample_observable.value, axis=0)
* np.mean(potential_gradients[gradient.key], axis=0)
)
)
if sample_observable.value.shape[1] > 1:
value = value.reshape(1, -1)
else:
value = value.item()
gradients.append(ParameterGradient(key=gradient.key, value=value))
return_type = (
Observable if sample_observable.value.shape[1] == 1 else ObservableArray
)
return return_type(value=mean_observable, gradients=gradients)
def compute_dielectric_constant(
dipole_moments: ObservableArray,
volumes: ObservableArray,
temperature: unit.Quantity,
average_function,
) -> Observable:
"""A function to compute the average dielectric constant from an array of
dipole moments and an array of volumes, whereby the average values of the
observables are computed using a custom function.
Parameters
----------
dipole_moments
The dipole moments array.
volumes
The volume array.
temperature
The temperature at which the dipole_moments and volumes were sampled.
average_function
The function to use when evaluating the average of an observable.
Returns
-------
The average value of the dielectric constant.
"""
dipole_moments_sqr = dipole_moments * dipole_moments
dipole_moments_sqr = ObservableArray(
value=dipole_moments_sqr.value.sum(axis=1),
gradients=[
ParameterGradient(gradient.key, gradient.value.sum(axis=1))
for gradient in dipole_moments_sqr.gradients
],
)
avg_sqr_dipole_moments = average_function(observable=dipole_moments_sqr)
avg_sqr_dipole_moments = ObservableArray(
avg_sqr_dipole_moments.value, avg_sqr_dipole_moments.gradients
)
avg_dipole_moment = average_function(observable=dipole_moments)
avg_dipole_moment_sqr = avg_dipole_moment * avg_dipole_moment
avg_dipole_moment_sqr = ObservableArray(
value=avg_dipole_moment_sqr.value.sum(axis=1),
gradients=[
ParameterGradient(gradient.key, gradient.value.sum(axis=1))
for gradient in avg_dipole_moment_sqr.gradients
],
)
avg_volume = average_function(observable=volumes)
avg_volume = ObservableArray(avg_volume.value, avg_volume.gradients)
dipole_variance = avg_sqr_dipole_moments - avg_dipole_moment_sqr
prefactor = 1.0 / (3.0 * E0 * unit.boltzmann_constant * temperature)
dielectric_constant = 1.0 * unit.dimensionless + prefactor * (
dipole_variance / avg_volume
)
return Observable(
value=dielectric_constant.value.item().to(unit.dimensionless),
gradients=[
ParameterGradient(
gradient.key,
gradient.value.item().to(
unit.dimensionless
* gradient.value.units
/ dielectric_constant.value.units
),
)
for gradient in dielectric_constant.gradients
],
)
def _execute(self, directory, available_resources):
from scipy.special import logsumexp
default_unit = unit.kilocalorie / unit.mole
boltzmann_factor = (
self.thermodynamic_state.temperature * unit.molar_gas_constant
)
boltzmann_factor.ito(default_unit)
beta = 1.0 / boltzmann_factor
values = [
(-beta * value.value.to(default_unit)).to(unit.dimensionless).magnitude
for value in self.values
]
# Compute the mean.
mean = logsumexp(values)
# Compute the gradients of the mean.
value_gradients = [
{gradient.key: -beta * gradient.value for gradient in value.gradients}
for value in self.values
]
value_gradients_by_key = {
gradient_key: [
gradients_by_key[gradient_key] for gradients_by_key in value_gradients
]
for gradient_key in value_gradients[0]
}
mean_gradients = []
for gradient_key, gradient_values in value_gradients_by_key.items():
expected_unit = value_gradients[0][gradient_key].units
d_log_mean_numerator, d_mean_numerator_sign = logsumexp(
values,
b=[x.to(expected_unit).magnitude for x in gradient_values],
return_sign=True,
)
d_mean_numerator = d_mean_numerator_sign * np.exp(d_log_mean_numerator)
d_mean_d_theta = d_mean_numerator / np.exp(mean)
mean_gradients.append(
ParameterGradient(
key=gradient_key,
value=-boltzmann_factor * d_mean_d_theta * expected_unit,
)
)
# Compute the standard error and 95% CI
cycle_result = np.empty(self.bootstrap_cycles)
for cycle_index, cycle in enumerate(range(self.bootstrap_cycles)):
cycle_values = np.empty(len(self.values))
for value_index, value in enumerate(self.values):
cycle_mean = value.value.to(default_unit).magnitude
cycle_sem = value.error.to(default_unit).magnitude
sampled_value = np.random.normal(cycle_mean, cycle_sem) * default_unit
cycle_values[value_index] = (
(-beta * sampled_value).to(unit.dimensionless).magnitude
)
# ΔG° = -RT × Log[ Σ_{n} exp(-βΔG°_{n}) ]
cycle_result[cycle_index] = logsumexp(cycle_values)
mean = -boltzmann_factor * mean
sem = np.std(-boltzmann_factor * cycle_result)
confidence_intervals = np.empty(2)
sorted_statistics = np.sort(cycle_result)
confidence_intervals[0] = sorted_statistics[int(0.025 * self.bootstrap_cycles)]
confidence_intervals[1] = sorted_statistics[int(0.975 * self.bootstrap_cycles)]
confidence_intervals = -boltzmann_factor * confidence_intervals
self.result = Observable(value=mean.plus_minus(sem), gradients=mean_gradients)
self.confidence_intervals = confidence_intervals
def _initialize(self, value: unit.Quantity, gradients: List[ParameterGradient]):
expected_types = (int, float, numpy.ndarray)
if value is not None:
if not isinstance(value, unit.Quantity) or not isinstance(
value.magnitude, expected_types
):
raise TypeError(
"The value must be a unit-wrapped integer, float or numpy array."
)
if not issubclass(type(value.magnitude), numpy.ndarray):
value = numpy.array([value.magnitude]) * value.units
# Ensure the inner array has a uniform shape.
if value.magnitude.ndim > 2:
raise ValueError(
"The wrapped array must not contain more than two dimensions."
)
if value.magnitude.ndim < 2:
value = value.reshape(-1, 1)
reshaped_gradients = []
if gradients is not None:
if value is None:
raise ValueError("A valid value must be provided.")
# Make sure the value and gradients have the same wrapped type.
if not all(
isinstance(gradient.value, unit.Quantity)
and isinstance(gradient.value.magnitude, expected_types)
for gradient in gradients
):
raise TypeError(
"The gradient values must be unit-wrapped integers, floats "
"or numpy arrays."
)
# Make sure the gradient values are all numpy arrays and make sure the each
# have the same shape as the value.
for gradient in gradients:
gradient_value = gradient.value.magnitude
if not isinstance(gradient.value.magnitude, numpy.ndarray):
gradient_value = numpy.array([gradient_value])
if gradient_value.ndim < 2:
gradient_value = gradient_value.reshape(-1, 1)
if gradient_value.ndim > 2:
raise ValueError(
"Gradient values must not contain more than two dimensions."
)
if value.magnitude.shape[1] != gradient_value.shape[1]:
raise ValueError(
f"Gradient values should be {value.magnitude.shape[1]}-"
f"dimensional to match the dimensionality of the value."
)
if gradient_value.shape[0] != value.magnitude.shape[0]:
raise ValueError(
f"Gradient values should have a length of "
f"{value.magnitude.shape[0]} to match the length of the value."
)
reshaped_gradients.append(
ParameterGradient(
key=gradient.key,
value=unit.Quantity(gradient_value, gradient.value.units),
)
)
super(ObservableArray, self)._initialize(value, reshaped_gradients)
def _reweight_observables(
self,
weights: ObservableArray,
mbar: pymbar.MBAR,
target_reduced_potentials: ObservableArray,
**observables: ObservableArray,
) -> typing.Union[ObservableArray, Observable]:
"""A function which computes the average value of an observable using
weights computed from MBAR and from a set of component observables.
Parameters
----------
weights
The MBAR weights
observables
The component observables which may be combined to yield the final
average observable of interest.
mbar
A pre-computed MBAR object encoded information from the reference states.
This will be used to compute the std error when not bootstrapping.
target_reduced_potentials
The reduced potentials at the target state. This will be used to compute
the std error when not bootstrapping.
Returns
-------
The re-weighted average observable.
"""
observable = observables.pop("observable")
assert len(observables) == 0
return_type = ObservableArray if observable.value.shape[
1] > 1 else Observable
weighted_observable = weights * observable
average_value = weighted_observable.value.sum(axis=0)
average_gradients = [
ParameterGradient(key=gradient.key,
value=gradient.value.sum(axis=0))
for gradient in weighted_observable.gradients
]
if return_type == Observable:
average_value = average_value.item()
average_gradients = [
ParameterGradient(key=gradient.key,
value=gradient.value.item())
for gradient in average_gradients
]
else:
average_value = average_value.reshape(1, -1)
average_gradients = [
ParameterGradient(key=gradient.key,
value=gradient.value.reshape(1, -1))
for gradient in average_gradients
]
if self.bootstrap_uncertainties is False:
# Unfortunately we need to re-compute the average observable for now
# as pymbar does not expose an easier way to compute the average
# uncertainty.
observable_dimensions = observable.value.shape[1]
assert observable_dimensions == 1
results = mbar.computeExpectations(
observable.value.T.magnitude,
target_reduced_potentials.value.T.magnitude,
state_dependent=True,
)
uncertainty = results[1][-1] * observable.value.units
average_value = average_value.plus_minus(uncertainty)
return return_type(value=average_value, gradients=average_gradients)
from openff.evaluator.substances import Component, ExactAmount, MoleFraction, Substance
@pytest.mark.parametrize(
"values",
[
[random.randint(1, 10) for _ in range(10)],
[random.random() for _ in range(10)],
[random.random() * unit.kelvin for _ in range(10)],
[
(random.random() * unit.kelvin).plus_minus(random.random() * unit.kelvin)
for x in range(10)
],
[
ParameterGradient(
ParameterGradientKey("a", "b", "c"), random.random() * unit.kelvin
)
for _ in range(10)
],
],
)
def test_add_values_protocol(values):
with tempfile.TemporaryDirectory() as temporary_directory:
add_quantities = AddValues("add")
add_quantities.values = values
add_quantities.execute(temporary_directory, ComputeResources())
assert add_quantities.result == reduce(operator.add, values)