def test_nested_input():
dict_protocol = DummyInputOutputProtocol("dict_protocol")
dict_protocol.input_value = {"a": ThermodynamicState(1.0 * unit.kelvin)}
quantity_protocol = DummyInputOutputProtocol("quantity_protocol")
quantity_protocol.input_value = ProtocolPath("output_value[a].temperature",
dict_protocol.id)
schema = WorkflowSchema()
schema.protocol_schemas = [dict_protocol.schema, quantity_protocol.schema]
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as temporary_directory:
with DaskLocalCluster() as calculation_backend:
results_futures = workflow_graph.execute(temporary_directory,
calculation_backend)
assert len(results_futures) == 1
result = results_futures[0].result()
assert isinstance(result, WorkflowResult)
def test_from_schema():
protocol_a = DummyProtocol("protocol_a")
protocol_a.input_value = 1 * unit.kelvin
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema]
workflow = Workflow.from_schema(schema, {}, unique_id="")
assert workflow is not None
rebuilt_schema = workflow.schema
rebuilt_schema.outputs_to_store = UNDEFINED
assert rebuilt_schema.json(format=True) == schema.json(format=True)
def test_unique_ids():
protocol_a = DummyProtocol("protocol-a")
protocol_a.input_value = 1
group_a = ProtocolGroup("group-a")
group_a.add_protocols(protocol_a)
group_b = ProtocolGroup("group-b")
group_b.add_protocols(protocol_a)
schema = WorkflowSchema()
schema.protocol_schemas = [group_a.schema, group_b.schema]
with pytest.raises(ValueError) as error_info:
schema.validate()
assert "Several protocols in the schema have the same id" in str(
error_info.value)
assert "protocol-a" in str(error_info.value)
def test_index_replicated_protocol():
replicator = ProtocolReplicator("replicator")
replicator.template_values = ["a", "b", "c", "d"]
replicated_protocol = DummyInputOutputProtocol(
f"protocol_{replicator.placeholder_id}")
replicated_protocol.input_value = ReplicatorValue(replicator.id)
schema = WorkflowSchema()
schema.protocol_replicators = [replicator]
schema.protocol_schemas = [replicated_protocol.schema]
for index in range(len(replicator.template_values)):
indexing_protocol = DummyInputOutputProtocol(
f"indexing_protocol_{index}")
indexing_protocol.input_value = ProtocolPath("output_value",
f"protocol_{index}")
schema.protocol_schemas.append(indexing_protocol.schema)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
def test_simple_workflow_graph(calculation_backend, compute_resources,
exception):
expected_value = (1 * unit.kelvin).plus_minus(0.1 * unit.kelvin)
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = expected_value
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema, protocol_b.schema]
schema.final_value_source = ProtocolPath("output_value", protocol_b.id)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as directory:
if calculation_backend is not None:
with DaskLocalCluster() as calculation_backend:
if exception:
with pytest.raises(AssertionError):
workflow_graph.execute(directory, calculation_backend,
compute_resources)
return
else:
results_futures = workflow_graph.execute(
directory, calculation_backend, compute_resources)
assert len(results_futures) == 1
result = results_futures[0].result()
else:
result = workflow_graph.execute(directory, calculation_backend,
compute_resources)[0]
if exception:
with pytest.raises(AssertionError):
workflow_graph.execute(directory, calculation_backend,
compute_resources)
return
assert isinstance(result, WorkflowResult)
assert result.value.value == expected_value.value
def test_schema_serialization(calculation_layer, property_type):
"""Tests serialisation and deserialization of a calculation schema."""
schema = registered_calculation_schemas[calculation_layer][property_type]
if callable(schema):
schema = schema()
json_schema = schema.json()
schema_from_json = WorkflowSchema.parse_json(json_schema)
property_recreated_json = schema_from_json.json()
assert json_schema == property_recreated_json
def test_nested_replicators():
dummy_schema = WorkflowSchema()
dummy_protocol = DummyReplicableProtocol("dummy_$(rep_a)_$(rep_b)")
dummy_protocol.replicated_value_a = ReplicatorValue("rep_a")
dummy_protocol.replicated_value_b = ReplicatorValue("rep_b")
dummy_schema.protocol_schemas = [dummy_protocol.schema]
replicator_a = ProtocolReplicator(replicator_id="rep_a")
replicator_a.template_values = ["a", "b"]
replicator_b = ProtocolReplicator(replicator_id="rep_b")
replicator_b.template_values = [1, 2]
dummy_schema.protocol_replicators = [replicator_a, replicator_b]
dummy_schema.validate()
dummy_property = create_dummy_property(Density)
dummy_metadata = Workflow.generate_default_metadata(
dummy_property, "smirnoff99Frosst-1.1.0.offxml", [])
dummy_workflow = Workflow(dummy_metadata, "")
dummy_workflow.schema = dummy_schema
assert len(dummy_workflow.protocols) == 4
assert dummy_workflow.protocols["dummy_0_0"].replicated_value_a == "a"
assert dummy_workflow.protocols["dummy_0_1"].replicated_value_a == "a"
assert dummy_workflow.protocols["dummy_1_0"].replicated_value_a == "b"
assert dummy_workflow.protocols["dummy_1_1"].replicated_value_a == "b"
assert dummy_workflow.protocols["dummy_0_0"].replicated_value_b == 1
assert dummy_workflow.protocols["dummy_0_1"].replicated_value_b == 2
assert dummy_workflow.protocols["dummy_1_0"].replicated_value_b == 1
assert dummy_workflow.protocols["dummy_1_1"].replicated_value_b == 2
print(dummy_workflow.schema)
def test_workflow_with_groups():
expected_value = (1 * unit.kelvin).plus_minus(0.1 * unit.kelvin)
protocol_a = DummyInputOutputProtocol("protocol_a")
protocol_a.input_value = expected_value
protocol_b = DummyInputOutputProtocol("protocol_b")
protocol_b.input_value = ProtocolPath("output_value", protocol_a.id)
conditional_group = ConditionalGroup("conditional_group")
conditional_group.add_protocols(protocol_a, protocol_b)
condition = ConditionalGroup.Condition()
condition.right_hand_value = 2 * unit.kelvin
condition.type = ConditionalGroup.Condition.Type.LessThan
condition.left_hand_value = ProtocolPath("output_value.value",
conditional_group.id,
protocol_b.id)
conditional_group.add_condition(condition)
schema = WorkflowSchema()
schema.protocol_schemas = [conditional_group.schema]
schema.final_value_source = ProtocolPath("output_value",
conditional_group.id,
protocol_b.id)
schema.validate()
workflow = Workflow({})
workflow.schema = schema
workflow_graph = workflow.to_graph()
with tempfile.TemporaryDirectory() as directory:
with DaskLocalCluster() as calculation_backend:
results_futures = workflow_graph.execute(directory,
calculation_backend)
assert len(results_futures) == 1
result = results_futures[0].result()
assert isinstance(result, WorkflowResult)
assert result.value.value == expected_value.value
def test_replicated_ids():
replicator = ProtocolReplicator("replicator-a")
protocol_a = DummyProtocol("protocol-a")
protocol_a.input_value = 1
group_a = ProtocolGroup(f"group-a-{replicator.placeholder_id}")
group_a.add_protocols(protocol_a)
schema = WorkflowSchema()
schema.protocol_schemas = [group_a.schema]
schema.protocol_replicators = [replicator]
with pytest.raises(ValueError) as error_info:
schema.validate()
assert (
f"The children of replicated protocol {group_a.id} must also contain the "
"replicators placeholder" in str(error_info.value))
def test_advanced_nested_replicators():
dummy_schema = WorkflowSchema()
replicator_a = ProtocolReplicator(replicator_id="replicator_a")
replicator_a.template_values = ["a", "b"]
replicator_b = ProtocolReplicator(
replicator_id=f"replicator_b_{replicator_a.placeholder_id}")
replicator_b.template_values = ProtocolPath(
f"dummy_list[{replicator_a.placeholder_id}]", "global")
dummy_protocol = DummyReplicableProtocol(f"dummy_"
f"{replicator_a.placeholder_id}_"
f"{replicator_b.placeholder_id}")
dummy_protocol.replicated_value_a = ReplicatorValue(replicator_a.id)
dummy_protocol.replicated_value_b = ReplicatorValue(replicator_b.id)
dummy_schema.protocol_schemas = [dummy_protocol.schema]
dummy_schema.protocol_replicators = [replicator_a, replicator_b]
dummy_schema.validate()
dummy_property = create_dummy_property(Density)
dummy_metadata = Workflow.generate_default_metadata(
dummy_property, "smirnoff99Frosst-1.1.0.offxml", [])
dummy_metadata["dummy_list"] = [[1], [2]]
dummy_workflow = Workflow(dummy_metadata, "")
dummy_workflow.schema = dummy_schema
assert len(dummy_workflow.protocols) == 2
assert dummy_workflow.protocols["dummy_0_0"].replicated_value_a == "a"
assert dummy_workflow.protocols["dummy_0_0"].replicated_value_b == 1
assert dummy_workflow.protocols["dummy_1_0"].replicated_value_a == "b"
assert dummy_workflow.protocols["dummy_1_0"].replicated_value_b == 2
def default_simulation_schema(absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_molecules=1000):
"""Returns the default calculation schema to use when estimating
this class of property from direct simulations.
Parameters
----------
absolute_tolerance: pint.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_molecules: int
The number of molecules to use in the simulation.
Returns
-------
SimulationSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = SimulationSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
# Define the protocol which will extract the average dielectric constant
# from the results of a simulation.
extract_dielectric = ExtractAverageDielectric("extract_dielectric")
extract_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
# Define the protocols which will run the simulation itself.
use_target_uncertainty = (absolute_tolerance != UNDEFINED
or relative_tolerance != UNDEFINED)
protocols, value_source, output_to_store = generate_base_simulation_protocols(
extract_dielectric,
use_target_uncertainty,
n_molecules=n_molecules,
)
# Make sure the input of the analysis protcol is properly hooked up.
extract_dielectric.system_path = ProtocolPath(
"system_path", protocols.assign_parameters.id)
# Dielectric constants typically take longer to converge, so we need to
# reflect this in the maximum number of convergence iterations.
protocols.converge_uncertainty.max_iterations = 400
# Set up the gradient calculations. For dielectric constants, we need to use
# a slightly specialised reweighting protocol which we set up here.
coordinate_source = ProtocolPath("output_coordinate_file",
protocols.equilibration_simulation.id)
trajectory_source = ProtocolPath(
"trajectory_file_path",
protocols.converge_uncertainty.id,
protocols.production_simulation.id,
)
statistics_source = ProtocolPath(
"statistics_file_path",
protocols.converge_uncertainty.id,
protocols.production_simulation.id,
)
gradient_mbar_protocol = ReweightDielectricConstant("gradient_mbar")
gradient_mbar_protocol.reference_dipole_moments = [
ProtocolPath(
"dipole_moments",
protocols.converge_uncertainty.id,
extract_dielectric.id,
)
]
gradient_mbar_protocol.reference_volumes = [
ProtocolPath("volumes", protocols.converge_uncertainty.id,
extract_dielectric.id)
]
gradient_mbar_protocol.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
gradient_mbar_protocol.reference_reduced_potentials = statistics_source
(
gradient_group,
gradient_replicator,
gradient_source,
) = generate_gradient_protocol_group(
gradient_mbar_protocol,
ProtocolPath("force_field_path", "global"),
coordinate_source,
trajectory_source,
statistics_source,
)
# Build the workflow schema.
schema = WorkflowSchema()
schema.protocol_schemas = [
protocols.build_coordinates.schema,
protocols.assign_parameters.schema,
protocols.energy_minimisation.schema,
protocols.equilibration_simulation.schema,
protocols.converge_uncertainty.schema,
protocols.extract_uncorrelated_trajectory.schema,
protocols.extract_uncorrelated_statistics.schema,
gradient_group.schema,
]
schema.protocol_replicators = [gradient_replicator]
schema.outputs_to_store = {"full_system": output_to_store}
schema.gradients_sources = [gradient_source]
schema.final_value_source = value_source
calculation_schema.workflow_schema = schema
return calculation_schema
def default_simulation_schema(absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_molecules=2000):
"""Returns the default calculation schema to use when estimating
this class of property from direct simulations.
Parameters
----------
absolute_tolerance: pint.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_molecules: int
The number of molecules to use in the simulation.
Returns
-------
SimulationSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = SimulationSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
use_target_uncertainty = (absolute_tolerance != UNDEFINED
or relative_tolerance != UNDEFINED)
# Setup the fully solvated systems.
build_full_coordinates = coordinates.BuildCoordinatesPackmol(
"build_solvated_coordinates")
build_full_coordinates.substance = ProtocolPath("substance", "global")
build_full_coordinates.max_molecules = n_molecules
assign_full_parameters = forcefield.BaseBuildSystem(
"assign_solvated_parameters")
assign_full_parameters.force_field_path = ProtocolPath(
"force_field_path", "global")
assign_full_parameters.substance = ProtocolPath("substance", "global")
assign_full_parameters.coordinate_file_path = ProtocolPath(
"coordinate_file_path", build_full_coordinates.id)
# Perform a quick minimisation of the full system to give
# YANK a better starting point for its minimisation.
energy_minimisation = openmm.OpenMMEnergyMinimisation(
"energy_minimisation")
energy_minimisation.system_path = ProtocolPath(
"system_path", assign_full_parameters.id)
energy_minimisation.input_coordinate_file = ProtocolPath(
"coordinate_file_path", build_full_coordinates.id)
equilibration_simulation = openmm.OpenMMSimulation(
"equilibration_simulation")
equilibration_simulation.ensemble = Ensemble.NPT
equilibration_simulation.steps_per_iteration = 100000
equilibration_simulation.output_frequency = 10000
equilibration_simulation.timestep = 2.0 * unit.femtosecond
equilibration_simulation.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
equilibration_simulation.system_path = ProtocolPath(
"system_path", assign_full_parameters.id)
equilibration_simulation.input_coordinate_file = ProtocolPath(
"output_coordinate_file", energy_minimisation.id)
# Create a substance which only contains the solute (e.g. for the
# vacuum phase simulations).
filter_solvent = miscellaneous.FilterSubstanceByRole("filter_solvent")
filter_solvent.input_substance = ProtocolPath("substance", "global")
filter_solvent.component_roles = [Component.Role.Solvent]
filter_solute = miscellaneous.FilterSubstanceByRole("filter_solute")
filter_solute.input_substance = ProtocolPath("substance", "global")
filter_solute.component_roles = [Component.Role.Solute]
# Setup the solute in vacuum system.
build_vacuum_coordinates = coordinates.BuildCoordinatesPackmol(
"build_vacuum_coordinates")
build_vacuum_coordinates.substance = ProtocolPath(
"filtered_substance", filter_solute.id)
build_vacuum_coordinates.max_molecules = 1
assign_vacuum_parameters = forcefield.BaseBuildSystem(
"assign_parameters")
assign_vacuum_parameters.force_field_path = ProtocolPath(
"force_field_path", "global")
assign_vacuum_parameters.substance = ProtocolPath(
"filtered_substance", filter_solute.id)
assign_vacuum_parameters.coordinate_file_path = ProtocolPath(
"coordinate_file_path", build_vacuum_coordinates.id)
# Set up the protocol to run yank.
run_yank = yank.SolvationYankProtocol("run_solvation_yank")
run_yank.solute = ProtocolPath("filtered_substance", filter_solute.id)
run_yank.solvent_1 = ProtocolPath("filtered_substance",
filter_solvent.id)
run_yank.solvent_2 = Substance()
run_yank.thermodynamic_state = ProtocolPath("thermodynamic_state",
"global")
run_yank.steps_per_iteration = 500
run_yank.checkpoint_interval = 50
run_yank.solvent_1_coordinates = ProtocolPath(
"output_coordinate_file", equilibration_simulation.id)
run_yank.solvent_1_system = ProtocolPath("system_path",
assign_full_parameters.id)
run_yank.solvent_2_coordinates = ProtocolPath(
"coordinate_file_path", build_vacuum_coordinates.id)
run_yank.solvent_2_system = ProtocolPath("system_path",
assign_vacuum_parameters.id)
# Set up the group which will run yank until the free energy has been determined to within
# a given uncertainty
conditional_group = groups.ConditionalGroup("conditional_group")
conditional_group.max_iterations = 20
if use_target_uncertainty:
condition = groups.ConditionalGroup.Condition()
condition.type = groups.ConditionalGroup.Condition.Type.LessThan
condition.right_hand_value = ProtocolPath("target_uncertainty",
"global")
condition.left_hand_value = ProtocolPath(
"estimated_free_energy.error", conditional_group.id,
run_yank.id)
conditional_group.add_condition(condition)
# Define the total number of iterations that yank should run for.
total_iterations = miscellaneous.MultiplyValue("total_iterations")
total_iterations.value = 2000
total_iterations.multiplier = ProtocolPath("current_iteration",
conditional_group.id)
# Make sure the simulations gets extended after each iteration.
run_yank.number_of_iterations = ProtocolPath("result",
total_iterations.id)
conditional_group.add_protocols(total_iterations, run_yank)
# Define the full workflow schema.
schema = WorkflowSchema()
schema.protocol_schemas = [
build_full_coordinates.schema,
assign_full_parameters.schema,
energy_minimisation.schema,
equilibration_simulation.schema,
filter_solvent.schema,
filter_solute.schema,
build_vacuum_coordinates.schema,
assign_vacuum_parameters.schema,
conditional_group.schema,
]
schema.final_value_source = ProtocolPath("estimated_free_energy",
conditional_group.id,
run_yank.id)
calculation_schema.workflow_schema = schema
return calculation_schema
def default_reweighting_schema(
absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_effective_samples=50,
):
"""Returns the default calculation schema to use when estimating
this property by reweighting existing data.
Parameters
----------
absolute_tolerance: openff.evaluator.unit.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_effective_samples: int
The minimum number of effective samples to require when
reweighting the cached simulation data.
Returns
-------
ReweightingSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = ReweightingSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
protocols, data_replicator = generate_base_reweighting_protocols(
statistical_inefficiency=AverageDielectricConstant(
"average_dielectric_$(data_replicator)"
),
reweight_observable=ReweightDielectricConstant("reweight_dielectric"),
)
protocols.zero_gradients.input_observables = ProtocolPath(
"output_observables[Volume]",
protocols.join_observables.id,
)
protocols.statistical_inefficiency.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
protocols.reweight_observable.required_effective_samples = n_effective_samples
# We don't need to perform bootstrapping as this protocol is only used to
# calculate the statistical inefficiency and equilibration time. The
# re-weighting protocol will instead compute the bootstrapped uncertainties.
protocols.statistical_inefficiency.bootstrap_iterations = 1
# Set up a protocol to re-evaluate the dipole moments at the target state
# and concatenate the into a single array.
compute_dipoles = ComputeDipoleMoments("compute_dipoles_$(data_replicator)")
compute_dipoles.parameterized_system = ProtocolPath(
"parameterized_system", protocols.build_target_system.id
)
compute_dipoles.trajectory_path = ProtocolPath(
"trajectory_file_path", protocols.unpack_stored_data.id
)
compute_dipoles.gradient_parameters = ProtocolPath(
"parameter_gradient_keys", "global"
)
join_dipoles = ConcatenateObservables("join_dipoles")
join_dipoles.input_observables = ProtocolPath(
"dipole_moments",
compute_dipoles.id,
)
# Point the dielectric protocols to the volumes and dipole moments.
protocols.statistical_inefficiency.volumes = ProtocolPath(
"observables[Volume]", protocols.unpack_stored_data.id
)
protocols.statistical_inefficiency.dipole_moments = ProtocolPath(
"dipole_moments", compute_dipoles.id
)
# Make sure to decorrelate the dipole moments.
decorrelate_dipoles = DecorrelateObservables("decorrelate_dipoles")
decorrelate_dipoles.time_series_statistics = ProtocolPath(
"time_series_statistics", protocols.statistical_inefficiency.id
)
decorrelate_dipoles.input_observables = ProtocolPath(
"output_observables", join_dipoles.id
)
protocols.reweight_observable.dipole_moments = ProtocolPath(
"output_observables", decorrelate_dipoles.id
)
protocols.reweight_observable.volumes = ProtocolPath(
"output_observables", protocols.decorrelate_observable.id
)
protocols.reweight_observable.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
schema = WorkflowSchema()
schema.protocol_schemas = [
*(x.schema for x in protocols),
compute_dipoles.schema,
join_dipoles.schema,
decorrelate_dipoles.schema,
]
schema.protocol_replicators = [data_replicator]
schema.final_value_source = ProtocolPath(
"value", protocols.reweight_observable.id
)
calculation_schema.workflow_schema = schema
return calculation_schema
def default_simulation_schema(
absolute_tolerance=UNDEFINED, relative_tolerance=UNDEFINED, n_molecules=1000
):
"""Returns the default calculation schema to use when estimating
this class of property from direct simulations.
Parameters
----------
absolute_tolerance: openff.evaluator.unit.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_molecules: int
The number of molecules to use in the simulation.
Returns
-------
SimulationSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = SimulationSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
use_target_uncertainty = (
absolute_tolerance != UNDEFINED or relative_tolerance != UNDEFINED
)
# Define the protocols which will run the simulation itself.
protocols, value_source, output_to_store = generate_simulation_protocols(
AverageDielectricConstant("average_dielectric"),
use_target_uncertainty,
n_molecules=n_molecules,
)
# Add a protocol to compute the dipole moments and pass these to
# the analysis protocol.
compute_dipoles = ComputeDipoleMoments("compute_dipoles")
compute_dipoles.parameterized_system = ProtocolPath(
"parameterized_system", protocols.assign_parameters.id
)
compute_dipoles.trajectory_path = ProtocolPath(
"trajectory_file_path", protocols.production_simulation.id
)
compute_dipoles.gradient_parameters = ProtocolPath(
"parameter_gradient_keys", "global"
)
protocols.converge_uncertainty.add_protocols(compute_dipoles)
protocols.analysis_protocol.volumes = ProtocolPath(
f"observables[{ObservableType.Volume.value}]",
protocols.production_simulation.id,
)
protocols.analysis_protocol.dipole_moments = ProtocolPath(
"dipole_moments",
compute_dipoles.id,
)
# Build the workflow schema.
schema = WorkflowSchema()
schema.protocol_schemas = [
protocols.build_coordinates.schema,
protocols.assign_parameters.schema,
protocols.energy_minimisation.schema,
protocols.equilibration_simulation.schema,
protocols.converge_uncertainty.schema,
protocols.decorrelate_trajectory.schema,
protocols.decorrelate_observables.schema,
]
schema.outputs_to_store = {"full_system": output_to_store}
schema.final_value_source = value_source
calculation_schema.workflow_schema = schema
return calculation_schema
def test_workflow_layer():
"""Test the `WorkflowLayer` calculation layer. As the `SimulationLayer`
is the simplest implementation of the abstract layer, we settle for
testing this."""
properties_to_estimate = [
create_dummy_property(Density),
create_dummy_property(Density),
]
# Create a very simple workflow which just returns some placeholder
# value.
estimated_value = Observable(
(1 * unit.kelvin).plus_minus(0.1 * unit.kelvin))
protocol_a = DummyProtocol("protocol_a")
protocol_a.input_value = estimated_value
schema = WorkflowSchema()
schema.protocol_schemas = [protocol_a.schema]
schema.final_value_source = ProtocolPath("output_value", protocol_a.id)
layer_schema = SimulationSchema()
layer_schema.workflow_schema = schema
options = RequestOptions()
options.add_schema("SimulationLayer", "Density", layer_schema)
batch = server.Batch()
batch.queued_properties = properties_to_estimate
batch.options = options
with tempfile.TemporaryDirectory() as directory:
with temporarily_change_directory(directory):
# Create a directory for the layer.
layer_directory = "simulation_layer"
os.makedirs(layer_directory)
# Set-up a simple storage backend and add a force field to it.
force_field = SmirnoffForceFieldSource.from_path(
"smirnoff99Frosst-1.1.0.offxml")
storage_backend = LocalFileStorage()
batch.force_field_id = storage_backend.store_force_field(
force_field)
# Create a simple calculation backend to test with.
with DaskLocalCluster() as calculation_backend:
def dummy_callback(returned_request):
assert len(returned_request.estimated_properties) == 2
assert len(returned_request.exceptions) == 0
simulation_layer = SimulationLayer()
simulation_layer.schedule_calculation(
calculation_backend,
storage_backend,
layer_directory,
batch,
dummy_callback,
True,
)
def test_group_replicators():
dummy_schema = WorkflowSchema()
replicator_id = "replicator"
dummy_replicated_protocol = DummyInputOutputProtocol(
f"dummy_$({replicator_id})")
dummy_replicated_protocol.input_value = ReplicatorValue(replicator_id)
dummy_group = ProtocolGroup("dummy_group")
dummy_group.add_protocols(dummy_replicated_protocol)
dummy_protocol_single_value = DummyInputOutputProtocol(
f"dummy_single_$({replicator_id})")
dummy_protocol_single_value.input_value = ProtocolPath(
"output_value", dummy_group.id, dummy_replicated_protocol.id)
dummy_protocol_list_value = AddValues("dummy_list")
dummy_protocol_list_value.values = ProtocolPath(
"output_value", dummy_group.id, dummy_replicated_protocol.id)
dummy_schema.protocol_schemas = [
dummy_group.schema,
dummy_protocol_single_value.schema,
dummy_protocol_list_value.schema,
]
replicator = ProtocolReplicator(replicator_id)
replicator.template_values = [
(1.0 * unit.kelvin).plus_minus(1.0 * unit.kelvin),
(2.0 * unit.kelvin).plus_minus(2.0 * unit.kelvin),
]
dummy_schema.protocol_replicators = [replicator]
dummy_schema.validate()
dummy_property = create_dummy_property(Density)
dummy_metadata = Workflow.generate_default_metadata(
dummy_property, "smirnoff99Frosst-1.1.0.offxml", [])
dummy_workflow = Workflow(dummy_metadata, "")
dummy_workflow.schema = dummy_schema
assert len(dummy_workflow.protocols) == 4
assert (dummy_workflow.protocols[dummy_group.id].protocols["dummy_0"].
input_value.value == replicator.template_values[0].value)
assert (dummy_workflow.protocols[dummy_group.id].protocols["dummy_1"].
input_value.value == replicator.template_values[1].value)
assert dummy_workflow.protocols[
"dummy_single_0"].input_value == ProtocolPath("output_value",
dummy_group.id,
"dummy_0")
assert dummy_workflow.protocols[
"dummy_single_1"].input_value == ProtocolPath("output_value",
dummy_group.id,
"dummy_1")
assert len(dummy_workflow.protocols["dummy_list"].values) == 2
assert dummy_workflow.protocols["dummy_list"].values[0] == ProtocolPath(
"output_value", dummy_group.id, "dummy_0")
assert dummy_workflow.protocols["dummy_list"].values[1] == ProtocolPath(
"output_value", dummy_group.id, "dummy_1")
def default_reweighting_schema(
absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_effective_samples=50,
):
"""Returns the default calculation schema to use when estimating
this property by reweighting existing data.
Parameters
----------
absolute_tolerance: pint.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_effective_samples: int
The minimum number of effective samples to require when
reweighting the cached simulation data.
Returns
-------
ReweightingSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = ReweightingSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
data_replicator_id = "data_replicator"
# Set up a protocol to extract the dielectric constant from the stored data.
extract_dielectric = ExtractAverageDielectric(
f"calc_dielectric_$({data_replicator_id})")
# For the dielectric constant, we employ a slightly more advanced reweighting
# protocol set up for calculating fluctuation properties.
reweight_dielectric = ReweightDielectricConstant("reweight_dielectric")
reweight_dielectric.reference_dipole_moments = ProtocolPath(
"uncorrelated_values", extract_dielectric.id)
reweight_dielectric.reference_volumes = ProtocolPath(
"uncorrelated_volumes", extract_dielectric.id)
reweight_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
reweight_dielectric.bootstrap_uncertainties = True
reweight_dielectric.bootstrap_iterations = 200
reweight_dielectric.required_effective_samples = n_effective_samples
protocols, data_replicator = generate_base_reweighting_protocols(
extract_dielectric, reweight_dielectric, data_replicator_id)
# Make sure input is taken from the correct protocol outputs.
extract_dielectric.system_path = ProtocolPath(
"system_path", protocols.build_reference_system.id)
extract_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", protocols.unpack_stored_data.id)
# Set up the gradient calculations
coordinate_path = ProtocolPath("output_coordinate_path",
protocols.concatenate_trajectories.id)
trajectory_path = ProtocolPath("output_trajectory_path",
protocols.concatenate_trajectories.id)
statistics_path = ProtocolPath("statistics_file_path",
protocols.reduced_target_potential.id)
reweight_dielectric_template = copy.deepcopy(reweight_dielectric)
(
gradient_group,
gradient_replicator,
gradient_source,
) = generate_gradient_protocol_group(
reweight_dielectric_template,
ProtocolPath("force_field_path", "global"),
coordinate_path,
trajectory_path,
statistics_path,
replicator_id="grad",
effective_sample_indices=ProtocolPath("effective_sample_indices",
reweight_dielectric.id),
)
schema = WorkflowSchema()
schema.protocol_schemas = [
*(x.schema for x in protocols),
gradient_group.schema,
]
schema.protocol_replicators = [data_replicator, gradient_replicator]
schema.gradients_sources = [gradient_source]
schema.final_value_source = ProtocolPath("value",
protocols.mbar_protocol.id)
calculation_schema.workflow_schema = schema
return calculation_schema
def default_simulation_schema(existing_schema=None):
"""Returns the default calculation schema to use when estimating
this class of property from direct simulations.
Parameters
----------
existing_schema: SimulationSchema, optional
An existing schema whose settings to use. If set,
the schema's `workflow_schema` will be overwritten
by this method.
Returns
-------
SimulationSchema
The schema to follow when estimating this property.
"""
calculation_schema = SimulationSchema()
if existing_schema is not None:
assert isinstance(existing_schema, SimulationSchema)
calculation_schema = copy.deepcopy(existing_schema)
schema = WorkflowSchema(
property_type=HostGuestBindingAffinity.__name__)
schema.id = "{}{}".format(HostGuestBindingAffinity.__name__, "Schema")
# Initial coordinate and topology setup.
filter_ligand = miscellaneous.FilterSubstanceByRole("filter_ligand")
filter_ligand.input_substance = ProtocolPath("substance", "global")
filter_ligand.component_roles = [Component.Role.Ligand]
# We only support substances with a single guest ligand.
filter_ligand.expected_components = 1
schema.protocols[filter_ligand.id] = filter_ligand.schema
# Construct the protocols which will (for now) take as input a set of host coordinates,
# and generate a set of charges for them.
filter_receptor = miscellaneous.FilterSubstanceByRole(
"filter_receptor")
filter_receptor.input_substance = ProtocolPath("substance", "global")
filter_receptor.component_roles = [Component.Role.Receptor]
# We only support substances with a single host receptor.
filter_receptor.expected_components = 1
schema.protocols[filter_receptor.id] = filter_receptor.schema
# Perform docking to position the guest within the host.
perform_docking = coordinates.BuildDockedCoordinates("perform_docking")
perform_docking.ligand_substance = ProtocolPath(
"filtered_substance", filter_ligand.id)
perform_docking.receptor_coordinate_file = ProtocolPath(
"receptor_mol2", "global")
schema.protocols[perform_docking.id] = perform_docking.schema
# Solvate the docked structure using packmol
filter_solvent = miscellaneous.FilterSubstanceByRole("filter_solvent")
filter_solvent.input_substance = ProtocolPath("substance", "global")
filter_solvent.component_roles = [Component.Role.Solvent]
schema.protocols[filter_solvent.id] = filter_solvent.schema
solvate_complex = coordinates.SolvateExistingStructure(
"solvate_complex")
solvate_complex.max_molecules = 1000
solvate_complex.substance = ProtocolPath("filtered_substance",
filter_solvent.id)
solvate_complex.solute_coordinate_file = ProtocolPath(
"docked_complex_coordinate_path", perform_docking.id)
schema.protocols[solvate_complex.id] = solvate_complex.schema
# Assign force field parameters to the solvated complex system.
build_solvated_complex_system = forcefield.BaseBuildSystem(
"build_solvated_complex_system")
build_solvated_complex_system.force_field_path = ProtocolPath(
"force_field_path", "global")
build_solvated_complex_system.coordinate_file_path = ProtocolPath(
"coordinate_file_path", solvate_complex.id)
build_solvated_complex_system.substance = ProtocolPath(
"substance", "global")
build_solvated_complex_system.charged_molecule_paths = [
ProtocolPath("receptor_mol2", "global")
]
schema.protocols[build_solvated_complex_system.
id] = build_solvated_complex_system.schema
# Solvate the ligand using packmol
solvate_ligand = coordinates.SolvateExistingStructure("solvate_ligand")
solvate_ligand.max_molecules = 1000
solvate_ligand.substance = ProtocolPath("filtered_substance",
filter_solvent.id)
solvate_ligand.solute_coordinate_file = ProtocolPath(
"docked_ligand_coordinate_path", perform_docking.id)
schema.protocols[solvate_ligand.id] = solvate_ligand.schema
# Assign force field parameters to the solvated ligand system.
build_solvated_ligand_system = forcefield.BaseBuildSystem(
"build_solvated_ligand_system")
build_solvated_ligand_system.force_field_path = ProtocolPath(
"force_field_path", "global")
build_solvated_ligand_system.coordinate_file_path = ProtocolPath(
"coordinate_file_path", solvate_ligand.id)
build_solvated_ligand_system.substance = ProtocolPath(
"substance", "global")
schema.protocols[build_solvated_ligand_system.
id] = build_solvated_ligand_system.schema
# Employ YANK to estimate the binding free energy.
yank_protocol = yank.LigandReceptorYankProtocol("yank_protocol")
yank_protocol.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
yank_protocol.number_of_iterations = 2000
yank_protocol.steps_per_iteration = 500
yank_protocol.checkpoint_interval = 10
yank_protocol.verbose = True
yank_protocol.force_field_path = ProtocolPath("force_field_path",
"global")
yank_protocol.ligand_residue_name = ProtocolPath(
"ligand_residue_name", perform_docking.id)
yank_protocol.receptor_residue_name = ProtocolPath(
"receptor_residue_name", perform_docking.id)
yank_protocol.solvated_ligand_coordinates = ProtocolPath(
"coordinate_file_path", solvate_ligand.id)
yank_protocol.solvated_ligand_system = ProtocolPath(
"system_path", build_solvated_ligand_system.id)
yank_protocol.solvated_complex_coordinates = ProtocolPath(
"coordinate_file_path", solvate_complex.id)
yank_protocol.solvated_complex_system = ProtocolPath(
"system_path", build_solvated_complex_system.id)
schema.protocols[yank_protocol.id] = yank_protocol.schema
# Define where the final values come from.
schema.final_value_source = ProtocolPath("estimated_free_energy",
yank_protocol.id)
calculation_schema.workflow_schema = schema
return calculation_schema