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_observables_join_fail(observables, expected_raises, expected_message):
with expected_raises as error_info:
ObservableArray.join(*observables)
assert (expected_message is None and error_info is None
or expected_message in str(error_info.value))
def _bootstrap_function(self,
**observables: ObservableArray) -> Observable:
"""Re-weights a set of reference observables to the target state.
Parameters
-------
observables
The observables to reweight, in addition to the reference and target
reduced potentials.
"""
reference_reduced_potentials = observables.pop(
"reference_reduced_potentials")
target_reduced_potentials = observables.pop(
"target_reduced_potentials")
# Construct the mbar object using the specified reference reduced potentials.
# These may be the input values or values which have been sampled during
# bootstrapping, hence why it is not precomputed once.
mbar = pymbar.MBAR(
reference_reduced_potentials.value.to(
unit.dimensionless).magnitude.T,
self.frame_counts,
verbose=False,
relative_tolerance=1e-12,
)
# Compute the MBAR weights.
weights = self._compute_weights(mbar, target_reduced_potentials)
return self._reweight_observables(weights, mbar,
target_reduced_potentials,
**observables)
def _bootstrap_function(self, **kwargs: ObservableArray):
return compute_dielectric_constant(
kwargs.pop("dipole_moments"),
kwargs.pop("volumes"),
self.thermodynamic_state.temperature,
super(AverageDielectricConstant, self)._bootstrap_function,
)
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 _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 test_average_dielectric_constant():
with tempfile.TemporaryDirectory() as temporary_directory:
average_observable = AverageDielectricConstant("")
average_observable.dipole_moments = ObservableArray(
np.zeros((1, 3)) * unit.elementary_charge * unit.nanometer)
average_observable.volumes = ObservableArray(
np.ones((1, 1)) * unit.nanometer**3)
average_observable.thermodynamic_state = ThermodynamicState(
298.15 * unit.kelvin, 1.0 * unit.atmosphere)
average_observable.bootstrap_iterations = 1
average_observable.execute(temporary_directory)
assert np.isclose(average_observable.value.value,
1.0 * unit.dimensionless)
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_reweight_observables():
with tempfile.TemporaryDirectory() as directory:
reweight_protocol = ReweightObservable("")
reweight_protocol.observable = ObservableArray(value=np.zeros(10) *
unit.kelvin)
reweight_protocol.reference_reduced_potentials = [
ObservableArray(value=np.zeros(10) * unit.dimensionless)
]
reweight_protocol.frame_counts = [10]
reweight_protocol.target_reduced_potentials = ObservableArray(
value=np.zeros(10) * unit.dimensionless)
reweight_protocol.bootstrap_uncertainties = True
reweight_protocol.required_effective_samples = 0
reweight_protocol.execute(directory, ComputeResources())
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 test_frame_set_invalid_item(observable_frame, key, value, expected_raises,
expected_message):
with expected_raises as error_info:
observable_frame[key] = ObservableArray(value=value)
assert (expected_message is None and error_info is None
or expected_message in str(error_info.value))
def test_observable_array_invalid_initializer(value, gradients,
expected_raises,
expected_message):
with expected_raises as error_info:
ObservableArray(value, gradients)
assert expected_message in str(error_info.value)
def test_average_observable():
with tempfile.TemporaryDirectory() as temporary_directory:
average_observable = AverageObservable("")
average_observable.observable = ObservableArray(1.0 * unit.kelvin)
average_observable.bootstrap_iterations = 1
average_observable.execute(temporary_directory)
assert np.isclose(average_observable.value.value, 1.0 * unit.kelvin)
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 _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_decorrelate_observables():
with tempfile.TemporaryDirectory() as temporary_directory:
protocol = DecorrelateObservables("")
protocol.input_observables = ObservableArray(
np.ones((10, 1)) * unit.nanometer**3)
protocol.time_series_statistics = TimeSeriesStatistics(10, 4, 2.0, 2)
protocol.execute(temporary_directory)
assert len(protocol.output_observables) == 4
def test_reweight_dielectric_constant():
with tempfile.TemporaryDirectory() as directory:
reweight_protocol = ReweightDielectricConstant("")
reweight_protocol.dipole_moments = ObservableArray(
value=np.zeros((10, 3)) * unit.elementary_charge * unit.nanometers)
reweight_protocol.volumes = ObservableArray(value=np.ones((10, 1)) *
unit.nanometer**3)
reweight_protocol.reference_reduced_potentials = [
ObservableArray(value=np.zeros(10) * unit.dimensionless)
]
reweight_protocol.target_reduced_potentials = ObservableArray(
value=np.zeros(10) * unit.dimensionless)
reweight_protocol.thermodynamic_state = ThermodynamicState(
298.15 * unit.kelvin, 1.0 * unit.atmosphere)
reweight_protocol.frame_counts = [10]
reweight_protocol.bootstrap_uncertainties = True
reweight_protocol.required_effective_samples = 0
reweight_protocol.execute(directory, ComputeResources())
def test_frame_round_trip():
observable_frame = ObservableFrame(
{"Temperature": ObservableArray(value=numpy.ones(2) * unit.kelvin)})
round_tripped: ObservableFrame = json.loads(json.dumps(
observable_frame, cls=TypedJSONEncoder),
cls=TypedJSONDecoder)
assert isinstance(round_tripped, ObservableFrame)
assert {*observable_frame} == {*round_tripped}
assert len(observable_frame) == len(round_tripped)
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 _reweight_observables(
self,
weights: ObservableArray,
mbar: pymbar.MBAR,
target_reduced_potentials: ObservableArray,
**observables: ObservableArray,
) -> Observable:
volumes = observables.pop("volumes")
dipole_moments = observables.pop("dipole_moments")
dielectric_constant = compute_dielectric_constant(
dipole_moments,
volumes,
self.thermodynamic_state.temperature,
functools.partial(
super(ReweightDielectricConstant, self)._reweight_observables,
weights=weights,
mbar=mbar,
target_reduced_potentials=target_reduced_potentials,
),
)
return dielectric_constant
def test_frame_magic_functions(key):
observable_frame = ObservableFrame()
assert len(observable_frame) == 0
observable_frame[key] = ObservableArray(value=numpy.ones(1) * unit.kelvin)
assert len(observable_frame) == 1
assert key in observable_frame
assert {*observable_frame} == {ObservableType.Temperature}
del observable_frame[key]
assert len(observable_frame) == 0
assert key not in observable_frame
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_compute_gradients(tmpdir, smirks, all_zeros):
# Load a short trajectory.
coordinate_path = get_data_filename("test/trajectories/water.pdb")
trajectory_path = get_data_filename("test/trajectories/water.dcd")
trajectory = mdtraj.load_dcd(trajectory_path, coordinate_path)
observables = ObservableFrame({
"PotentialEnergy":
ObservableArray(
np.zeros(len(trajectory)) * unit.kilojoule / unit.mole)
})
_compute_gradients(
[ParameterGradientKey("vdW", smirks, "epsilon")],
observables,
ForceField("openff-1.2.0.offxml"),
ThermodynamicState(298.15 * unit.kelvin, 1.0 * unit.atmosphere),
Topology.from_mdtraj(trajectory.topology, [Molecule.from_smiles("O")]),
trajectory,
ComputeResources(),
True,
)
assert len(
observables["PotentialEnergy"].gradients[0].value) == len(trajectory)
if all_zeros:
assert np.allclose(
observables["PotentialEnergy"].gradients[0].value,
0.0 * unit.kilojoule / unit.kilocalorie,
)
else:
assert not np.allclose(
observables["PotentialEnergy"].gradients[0].value,
0.0 * unit.kilojoule / unit.kilocalorie,
)
def _compute_final_observables(self, temperature,
pressure) -> ObservableFrame:
"""Converts the openmm statistic csv file into an openff-evaluator
``ObservableFrame`` and computes additional missing data, such as reduced
potentials and derivatives of the energies with respect to any requested
force field parameters.
Parameters
----------
temperature: openff.evaluator.unit.Quantity
The temperature that the simulation is being run at.
pressure: openff.evaluator.unit.Quantity
The pressure that the simulation is being run at.
"""
observables = ObservableFrame.from_openmm(self._local_statistics_path,
pressure)
reduced_potentials = (
observables[ObservableType.PotentialEnergy].value /
unit.avogadro_constant)
if pressure is not None:
pv_terms = pressure * observables[ObservableType.Volume].value
reduced_potentials += pv_terms
beta = 1.0 / (unit.boltzmann_constant * temperature)
observables[ObservableType.ReducedPotential] = ObservableArray(
value=(beta * reduced_potentials).to(unit.dimensionless))
if pressure is not None:
observables[ObservableType.Enthalpy] = observables[
ObservableType.TotalEnergy] + observables[
ObservableType.Volume] * pressure * (
1.0 * unit.avogadro_constant)
return observables
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 test_zero_gradient():
with tempfile.TemporaryDirectory() as directory:
force_field_path = os.path.join(directory, "ff.json")
with open(force_field_path, "w") as file:
file.write(build_tip3p_smirnoff_force_field().json())
gradient_key = ParameterGradientKey("vdW", "[#1]-[#8X2H2+0:1]-[#1]",
"epsilon")
zero_gradients = ZeroGradients("")
zero_gradients.input_observables = ObservableArray(value=0.0 *
unit.kelvin)
zero_gradients.gradient_parameters = [gradient_key]
zero_gradients.force_field_path = force_field_path
zero_gradients.execute()
assert len(zero_gradients.output_observables.gradients) == 1
assert zero_gradients.output_observables.gradients[
0].key == gradient_key
assert np.allclose(
zero_gradients.output_observables.gradients[0].value, 0.0)
def _execute(self, directory, available_resources):
# Retrieve the observables to reweight.
observables = self._observables()
if len(observables) == 0:
raise ValueError("There were no observables to reweight.")
if len(self.frame_counts) != len(self.reference_reduced_potentials):
raise ValueError(
"A frame count must be provided for each reference state.")
expected_frames = sum(self.frame_counts)
if any(
len(input_array) != expected_frames for input_array in [
self.target_reduced_potentials,
*self.reference_reduced_potentials,
*observables.values(),
]):
raise ValueError(
f"The length of the input arrays do not match the expected length "
f"specified by the frame counts ({expected_frames}).")
# Concatenate the reduced reference potentials into a single array.
# We ignore the gradients of the reference state potential as these
# should be all zero.
reference_reduced_potentials = ObservableArray(value=np.hstack([
reduced_potentials.value
for reduced_potentials in self.reference_reduced_potentials
]))
# Ensure that there is enough effective samples to re-weight.
self.effective_samples = self._compute_effective_samples(
reference_reduced_potentials)
if self.effective_samples < self.required_effective_samples:
raise ValueError(
f"There was not enough effective samples to reweight - "
f"{self.effective_samples} < {self.required_effective_samples}"
)
if self.bootstrap_uncertainties:
self.value = bootstrap(
self._bootstrap_function,
self.bootstrap_iterations,
1.0,
self.frame_counts,
reference_reduced_potentials=reference_reduced_potentials,
target_reduced_potentials=self.target_reduced_potentials,
**observables,
)
else:
self.value = self._bootstrap_function(
reference_reduced_potentials=reference_reduced_potentials,
target_reduced_potentials=self.target_reduced_potentials,
**observables,
)
def test_observable_array_len():
assert len(ObservableArray(value=numpy.arange(5) * unit.kelvin)) == 5
@pytest.mark.parametrize(
"observables, expected_raises, expected_message",
[
(
[],
pytest.raises(ValueError),
"At least one observable must be provided.",
),
(
[
ObservableArray(value=numpy.ones(1) * unit.kelvin),
ObservableArray(value=numpy.ones(1) * unit.pascal),
],
pytest.raises(ValueError),
"The observables must all have compatible units.",
),
(
[
ObservableArray(
value=numpy.ones(2) * unit.kelvin,
gradients=[
ParameterGradient(
key=ParameterGradientKey("vdW", "[#1:1]", "sigma"),
value=numpy.ones(2) * unit.kelvin / unit.angstrom,
)
],
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)
def test_observable_array_len():
assert len(ObservableArray(value=numpy.arange(5) * unit.kelvin)) == 5
