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_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_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 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: 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_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(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