def test_frame_join_fail(observable_frames, expected_raises, expected_message):
with expected_raises as error_info:
ObservableFrame.join(*observable_frames)
assert (expected_message is None and error_info is None
or expected_message in str(error_info.value))
def test_frame_join():
gradient_unit = unit.mole / unit.kilojoule
observable_frames = [
ObservableFrame({
"Temperature":
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 = ObservableFrame.join(*observable_frames)
assert len(joined) == 4
assert numpy.allclose(
joined["Temperature"].value,
numpy.arange(4).reshape(-1, 1) * unit.kelvin,
)
assert numpy.allclose(
joined["Temperature"].gradients[0].value,
numpy.arange(4).reshape(-1, 1) * unit.kelvin * gradient_unit,
)
def process_successful_property(physical_property, layer_directory, **_):
"""Return a result as if the property had been successfully estimated."""
dummy_data_directory = path.join(layer_directory, "good_dummy_data")
makedirs(dummy_data_directory, exist_ok=True)
dummy_stored_object = StoredSimulationData()
dummy_stored_object.substance = physical_property.substance
dummy_stored_object.thermodynamic_state = physical_property.thermodynamic_state
dummy_stored_object.property_phase = physical_property.phase
dummy_stored_object.force_field_id = ""
dummy_stored_object.coordinate_file_name = ""
dummy_stored_object.trajectory_file_name = ""
dummy_stored_object.observables = ObservableFrame()
dummy_stored_object.statistical_inefficiency = 1.0
dummy_stored_object.number_of_molecules = 10
dummy_stored_object.source_calculation_id = ""
dummy_stored_object_path = path.join(layer_directory,
"good_dummy_data.json")
with open(dummy_stored_object_path, "w") as file:
json.dump(dummy_stored_object, file, cls=TypedJSONEncoder)
return_object = CalculationLayerResult()
return_object.physical_property = physical_property
return_object.data_to_store = [(dummy_stored_object_path,
dummy_data_directory)]
return return_object
def test_frame_constructor(observables):
observable_frame = ObservableFrame(observables)
assert all(observable_type in observable_frame
for observable_type in observables)
assert all(
observable_frame[observable_type] == observables[observable_type]
for observable_type in observables)
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 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_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_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_from_openmm(pressure):
observable_frame = ObservableFrame.from_openmm(
get_data_filename("test/statistics/openmm_statistics.csv"), pressure)
expected_types = {*ObservableType} - {ObservableType.ReducedPotential}
if pressure is None:
expected_types -= {ObservableType.Enthalpy}
assert {*observable_frame} == expected_types
assert len(observable_frame) == 10
expected_values = {
ObservableType.PotentialEnergy:
7934.831868494968 * unit.kilojoule / unit.mole,
ObservableType.KineticEnergy:
5939.683117957521 * unit.kilojoule / unit.mole,
ObservableType.TotalEnergy:
13874.51498645249 * unit.kilojoule / unit.mole,
ObservableType.Temperature: 286.38157154881503 * unit.kelvin,
ObservableType.Volume: 26.342326662784938 * unit.nanometer**3,
ObservableType.Density:
0.6139877476363793 * unit.gram / unit.milliliter,
}
for observable_type, expected_value in expected_values.items():
assert numpy.isclose(observable_frame[observable_type].value[0],
expected_value)
if pressure is not None:
expected_enthalpy = (13874.51498645249 * unit.kilojoule / unit.mole +
pressure * 26.342326662784938 *
unit.nanometer**3 * unit.avogadro_constant)
assert numpy.isclose(observable_frame["Enthalpy"].value[0],
expected_enthalpy)
def test_frame_validate_key(key, expected):
assert ObservableFrame._validate_key(key) == expected
def test_divide_observables(value_a, value_b, expected_value):
_compare_observables(value_a / value_b, expected_value)
@pytest.mark.parametrize(
"observables",
[
{
"Temperature": ObservableArray(value=numpy.ones(2) * unit.kelvin)
},
{
ObservableType.Temperature:
ObservableArray(value=numpy.ones(2) * unit.kelvin)
},
ObservableFrame({
ObservableType.Temperature:
ObservableArray(value=numpy.ones(2) * unit.kelvin)
}),
],
)
def test_frame_constructor(observables):
observable_frame = ObservableFrame(observables)
assert all(observable_type in observable_frame
for observable_type in observables)
assert all(
observable_frame[observable_type] == observables[observable_type]
for observable_type in observables)
def test_frame_round_trip():
def _evaluate_energies(
thermodynamic_state: ThermodynamicState,
system: "openmm.System",
trajectory: "Trajectory",
compute_resources: ComputeResources,
enable_pbc: bool = True,
high_precision: bool = True,
) -> ObservableFrame:
"""Evaluates the reduced and potential energies of each frame in a trajectory
using the specified system and at a particular state.
Parameters
----------
thermodynamic_state
The thermodynamic state to evaluate the reduced potentials at.
system
The system to evaluate the energies using.
trajectory
A trajectory of configurations to evaluate.
compute_resources: ComputeResources
The compute resources available to execute on.
enable_pbc
Whether PBC are enabled. This controls whether box vectors
are set or not.
high_precision
Whether to compute the energies using double precision.
Returns
-------
The array containing the evaluated potentials.
"""
from simtk import openmm
from simtk import unit as simtk_unit
integrator = openmm.VerletIntegrator(0.1 * simtk_unit.femtoseconds)
platform = setup_platform_with_resources(compute_resources, high_precision)
openmm_context = openmm.Context(system, integrator, platform)
potentials = np.zeros(trajectory.n_frames, dtype=np.float64)
reduced_potentials = np.zeros(trajectory.n_frames, dtype=np.float64)
temperature = pint_quantity_to_openmm(thermodynamic_state.temperature)
beta = 1.0 / (simtk_unit.BOLTZMANN_CONSTANT_kB * temperature)
pressure = pint_quantity_to_openmm(thermodynamic_state.pressure)
for frame_index in range(trajectory.n_frames):
positions = trajectory.xyz[frame_index]
if enable_pbc:
box_vectors = trajectory.openmm_boxes(frame_index)
openmm_context.setPeriodicBoxVectors(*box_vectors)
openmm_context.setPositions(positions)
state = openmm_context.getState(getEnergy=True)
potential_energy = state.getPotentialEnergy()
unreduced_potential = potential_energy / simtk_unit.AVOGADRO_CONSTANT_NA
if pressure is not None and enable_pbc:
unreduced_potential += pressure * state.getPeriodicBoxVolume()
potentials[frame_index] = potential_energy.value_in_unit(
simtk_unit.kilojoule_per_mole)
reduced_potentials[frame_index] = unreduced_potential * beta
potentials *= unit.kilojoule / unit.mole
reduced_potentials *= unit.dimensionless
observables_frame = ObservableFrame({
ObservableType.PotentialEnergy:
ObservableArray(potentials),
ObservableType.ReducedPotential:
ObservableArray(reduced_potentials),
})
return observables_frame
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 create_dummy_simulation_data(
directory_path,
substance,
force_field_id="dummy_ff_id",
coordinate_file_name="output.pdb",
trajectory_file_name="trajectory.dcd",
statistical_inefficiency=1.0,
phase=PropertyPhase.Liquid,
number_of_molecules=1,
calculation_id=None,
):
"""Creates a dummy `StoredSimulationData` object and
the corresponding data directory.
Parameters
----------
directory_path: str
The path to the dummy data directory to create.
substance: Substance
force_field_id
coordinate_file_name
trajectory_file_name
statistics_file_name
statistical_inefficiency
phase
number_of_molecules
calculation_id
Returns
-------
StoredSimulationData
The dummy stored data object.
"""
os.makedirs(directory_path, exist_ok=True)
data = StoredSimulationData()
data.substance = substance
data.force_field_id = force_field_id
data.thermodynamic_state = ThermodynamicState(1.0 * unit.kelvin)
data.property_phase = phase
data.coordinate_file_name = coordinate_file_name
data.trajectory_file_name = trajectory_file_name
data.observables = ObservableFrame()
with open(os.path.join(directory_path, coordinate_file_name), "w") as file:
file.write("")
with open(os.path.join(directory_path, trajectory_file_name), "w") as file:
file.write("")
data.statistical_inefficiency = statistical_inefficiency
data.number_of_molecules = number_of_molecules
if calculation_id is None:
calculation_id = str(uuid.uuid4())
data.source_calculation_id = calculation_id
return data
]
concatenate_protocol.execute(temporary_directory, ComputeResources())
final_trajectory = mdtraj.load(
concatenate_protocol.output_trajectory_path, top=coordinate_path)
assert len(final_trajectory) == len(original_trajectory) * 2
@pytest.mark.parametrize(
"observables",
[
[ObservableArray(value=np.zeros((2, 3)) * unit.kelvin)],
[ObservableArray(value=np.zeros((2, 3)) * unit.kelvin)] * 2,
[
ObservableFrame({
"Temperature":
ObservableArray(value=np.zeros((2, 3)) * unit.kelvin)
})
],
[
ObservableFrame({
"Temperature":
ObservableArray(value=np.zeros((2, 3)) * unit.kelvin)
})
] * 2,
],
)
def test_concatenate_observables(observables):
concatenate_protocol = ConcatenateObservables("")
concatenate_protocol.input_observables = observables
concatenate_protocol.execute()