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