def test_workflow_schema_merging(
calculation_layer, property_type, workflow_merge_function
):
"""Tests that two of the exact the same calculations get merged into one
by the `WorkflowGraph`."""
schema = registered_calculation_schemas[calculation_layer][property_type]
if callable(schema):
schema = schema()
if not isinstance(schema, WorkflowCalculationSchema):
pytest.skip("Not a `WorkflowCalculationSchema`.")
property_class = getattr(openff.evaluator.properties, property_type)
dummy_property = create_dummy_property(property_class)
global_metadata = create_dummy_metadata(dummy_property, calculation_layer)
workflow_a = Workflow(global_metadata, "workflow_a")
workflow_a.schema = schema.workflow_schema
workflow_b = Workflow(global_metadata, "workflow_b")
workflow_b.schema = schema.workflow_schema
workflow_graph = workflow_merge_function(workflow_a, workflow_b)
workflow_graph_a = workflow_a.to_graph()
workflow_graph_b = workflow_b.to_graph()
dependants_graph_a = workflow_graph_a._protocol_graph._build_dependants_graph(
workflow_graph_a.protocols, False, apply_reduction=True
)
dependants_graph_b = workflow_graph_b._protocol_graph._build_dependants_graph(
workflow_graph_b.protocols, False, apply_reduction=True
)
ordered_dict_a = OrderedDict(sorted(dependants_graph_a.items()))
ordered_dict_a = {key: sorted(value) for key, value in ordered_dict_a.items()}
ordered_dict_b = OrderedDict(sorted(dependants_graph_b.items()))
ordered_dict_b = {key: sorted(value) for key, value in ordered_dict_b.items()}
merge_order_a = graph.topological_sort(ordered_dict_a)
merge_order_b = graph.topological_sort(ordered_dict_b)
assert len(workflow_graph.protocols) == len(workflow_a.protocols)
for protocol_id in workflow_a.protocols:
assert protocol_id in workflow_graph.protocols
for protocol_id_A, protocol_id_B in zip(merge_order_a, merge_order_b):
assert protocol_id_A == protocol_id_B
assert (
workflow_a.protocols[protocol_id_A].schema.json()
== workflow_b.protocols[protocol_id_B].schema.json()
)
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_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_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_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_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_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 _build_workflow_graph(
cls,
working_directory,
storage_backend,
properties,
force_field_path,
parameter_gradient_keys,
options,
):
"""Construct a graph of the protocols needed to calculate a set of
properties.
Parameters
----------
working_directory: str
The local directory in which to store all local,
temporary calculation data from this graph.
storage_backend: StorageBackend
The backend used to store / retrieve data from previous calculations.
properties : list of PhysicalProperty
The properties to attempt to compute.
force_field_path : str
The path to the force field parameters to use in the workflow.
parameter_gradient_keys: list of ParameterGradientKey
A list of references to all of the parameters which all observables
should be differentiated with respect to.
options: RequestOptions
The options to run the workflows with.
"""
provenance = {}
workflows = []
for index, physical_property in enumerate(properties):
logger.info(f"Building workflow {index} of {len(properties)}")
property_type = type(physical_property).__name__
# Make sure a schema has been defined for this class of property
# and this layer.
if (property_type not in options.calculation_schemas
or cls.__name__
not in options.calculation_schemas[property_type]):
continue
schema = options.calculation_schemas[property_type][cls.__name__]
# Make sure the calculation schema is the correct type for this layer.
assert isinstance(schema, WorkflowCalculationSchema)
assert isinstance(schema, cls.required_schema_type())
global_metadata = cls._get_workflow_metadata(
working_directory,
physical_property,
force_field_path,
parameter_gradient_keys,
storage_backend,
schema,
)
if global_metadata is None:
# Make sure we have metadata returned for this
# property, e.g. we have data to reweight if
# required.
continue
workflow = Workflow(global_metadata, physical_property.id)
workflow.schema = schema.workflow_schema
workflows.append(workflow)
provenance[physical_property.id] = CalculationSource(
fidelity=cls.__name__, provenance=workflow.schema)
workflow_graph = WorkflowGraph()
workflow_graph.add_workflows(*workflows)
return workflow_graph, provenance
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 test_density_dielectric_merging(workflow_merge_function):
substance = Substance.from_components("C")
density = openff.evaluator.properties.Density(
thermodynamic_state=ThermodynamicState(
temperature=298 * unit.kelvin, pressure=1 * unit.atmosphere
),
phase=PropertyPhase.Liquid,
substance=substance,
value=10 * unit.gram / unit.mole,
uncertainty=1 * unit.gram / unit.mole,
)
dielectric = openff.evaluator.properties.DielectricConstant(
thermodynamic_state=ThermodynamicState(
temperature=298 * unit.kelvin, pressure=1 * unit.atmosphere
),
phase=PropertyPhase.Liquid,
substance=substance,
value=10 * unit.gram / unit.mole,
uncertainty=1 * unit.gram / unit.mole,
)
density_schema = density.default_simulation_schema().workflow_schema
dielectric_schema = dielectric.default_simulation_schema().workflow_schema
density_metadata = Workflow.generate_default_metadata(
density, "smirnoff99Frosst-1.1.0.offxml", []
)
dielectric_metadata = Workflow.generate_default_metadata(
density, "smirnoff99Frosst-1.1.0.offxml", []
)
density_workflow = Workflow(density_metadata)
density_workflow.schema = density_schema
dielectric_workflow = Workflow(dielectric_metadata)
dielectric_workflow.schema = dielectric_schema
workflow_merge_function(density_workflow, dielectric_workflow)
density_workflow_graph = density_workflow.to_graph()
dielectric_workflow_graph = dielectric_workflow.to_graph()
dependants_graph_a = density_workflow_graph._protocol_graph._build_dependants_graph(
density_workflow_graph.protocols, False, apply_reduction=True
)
dependants_graph_b = (
dielectric_workflow_graph._protocol_graph._build_dependants_graph(
dielectric_workflow_graph.protocols, False, apply_reduction=True
)
)
merge_order_a = graph.topological_sort(dependants_graph_a)
merge_order_b = graph.topological_sort(dependants_graph_b)
for protocol_id_A, protocol_id_B in zip(merge_order_a, merge_order_b):
if (
protocol_id_A.find("extract_traj") < 0
and protocol_id_A.find("extract_stats") < 0
):
assert (
density_workflow.protocols[protocol_id_A].schema.json()
== dielectric_workflow.protocols[protocol_id_B].schema.json()
)
else:
assert (
density_workflow.protocols[protocol_id_A].schema.json()
!= dielectric_workflow.protocols[protocol_id_B].schema.json()
)