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 _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 _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)