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_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_paprika_schema(
cls,
existing_schema: SimulationSchema = None,
n_solvent_molecules: int = 2500,
n_thermalization_steps: int = 50000,
n_equilibration_steps: int = 200000,
n_production_steps: int = 2500000,
dt_thermalization: unit.Quantity = 1.0 * unit.femtosecond,
dt_equilibration: unit.Quantity = 2.0 * unit.femtosecond,
dt_production: unit.Quantity = 2.0 * unit.femtosecond,
debug: bool = False,
):
"""Returns the default calculation schema to use when estimating
a host-guest binding affinity measurement with an APR calculation
using the ``paprika`` package.
Notes
-----
* This schema requires additional metadata to be able to estimate
each metadata. This metadata is automatically generated for properties
loaded from the ``taproom`` package using the ``TaproomDataSet`` object.
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.
n_solvent_molecules
The number of solvent molecules to add to the box.
n_thermalization_steps
The number of thermalization simulations steps to perform.
Sample generated during this step will be discarded.
n_equilibration_steps
The number of equilibration simulations steps to perform.
Sample generated during this step will be discarded.
n_production_steps
The number of production simulations steps to perform.
Sample generated during this step will be used in the final
free energy calculation.
dt_thermalization
The integration timestep during thermalization
dt_equilibration
The integration timestep during equilibration
dt_production
The integration timestep during production
debug
Whether to return a debug schema. This is nearly identical
to the default schema, albeit with significantly less
solvent molecules (10), all simulations run in NVT and much
shorter simulation runs (500 steps). If True, the other input
arguments will be ignored.
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)
# Initialize the protocols which will serve as templates for those
# used in the actual workflows.
solvation_template = cls._paprika_default_solvation_protocol(
n_solvent_molecules=n_solvent_molecules)
(
minimization_template,
*simulation_templates,
) = cls._paprika_default_simulation_protocols(
n_thermalization_steps=n_thermalization_steps,
n_equilibration_steps=n_equilibration_steps,
n_production_steps=n_production_steps,
dt_thermalization=dt_thermalization,
dt_equilibration=dt_equilibration,
dt_production=dt_production,
)
if debug:
solvation_template.max_molecules = 10
solvation_template.mass_density = 0.01 * unit.grams / unit.milliliters
for simulation_template in simulation_templates:
simulation_template.ensemble = Ensemble.NVT
simulation_template.steps_per_iteration = 500
simulation_template.output_frequency = 50
# Set up a replicator which will perform the attach-pull calculation for
# each of the guest orientations
orientation_replicator = ProtocolReplicator("orientation_replicator")
orientation_replicator.template_values = ProtocolPath(
"guest_orientations", "global")
restraint_schemas = {
"static":
ProtocolPath(
f"guest_orientations[{orientation_replicator.placeholder_id}]."
f"static_restraints",
"global",
),
"conformational":
ProtocolPath(
f"guest_orientations[{orientation_replicator.placeholder_id}]."
f"conformational_restraints",
"global",
),
"guest":
ProtocolPath("guest_restraints", "global"),
"wall":
ProtocolPath("wall_restraints", "global"),
"symmetry":
ProtocolPath("symmetry_restraints", "global"),
}
# Build the protocols to compute the attach and pull free energies.
(
attach_pull_protocols,
attach_pull_replicators,
attach_free_energy,
pull_free_energy,
reference_work,
) = cls._paprika_build_attach_pull_protocols(
orientation_replicator,
restraint_schemas,
solvation_template,
minimization_template,
*simulation_templates,
)
# Build the protocols to compute the release free energies.
(
release_protocols,
release_replicator,
release_free_energy,
) = cls._paprika_build_release_protocols(
orientation_replicator,
restraint_schemas,
solvation_template,
minimization_template,
*simulation_templates,
)
# Compute the symmetry correction.
symmetry_correction = ComputeSymmetryCorrection("symmetry_correction")
symmetry_correction.n_microstates = ProtocolPath(
"n_guest_microstates", "global")
symmetry_correction.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
# Sum together the free energies of the individual orientations
orientation_free_energy = miscellaneous.AddValues(
f"orientation_free_energy_{orientation_replicator.placeholder_id}")
orientation_free_energy.values = [
attach_free_energy,
pull_free_energy,
reference_work,
release_free_energy,
ProtocolPath("result", symmetry_correction.id),
]
# Finally, combine all of the values together
total_free_energy = analysis.AverageFreeEnergies("total_free_energy")
total_free_energy.values = ProtocolPath("result",
orientation_free_energy.id)
total_free_energy.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
calculation_schema.workflow_schema = WorkflowSchema()
calculation_schema.workflow_schema.protocol_schemas = [
*(protocol.schema for protocol in attach_pull_protocols),
*(protocol.schema for protocol in release_protocols),
symmetry_correction.schema,
orientation_free_energy.schema,
total_free_energy.schema,
]
calculation_schema.workflow_schema.protocol_replicators = [
orientation_replicator,
*attach_pull_replicators,
release_replicator,
]
# Define where the final value comes from.
calculation_schema.workflow_schema.final_value_source = ProtocolPath(
"result", total_free_energy.id)
return calculation_schema
def _paprika_build_release_protocols(
cls,
orientation_replicator: ProtocolReplicator,
restraint_schemas: Dict[str, ProtocolPath],
solvation_template: coordinates.SolvateExistingStructure,
minimization_template: openmm.OpenMMEnergyMinimisation,
thermalization_template: openmm.OpenMMSimulation,
equilibration_template: openmm.OpenMMSimulation,
production_template: openmm.OpenMMSimulation,
):
# Define a replicator to set up each release window
release_replicator = ProtocolReplicator("release_replicator")
release_replicator.template_values = ProtocolPath(
"release_windows_indices", "global")
orientation_placeholder = orientation_replicator.placeholder_id
release_replicator_id = (f"{release_replicator.placeholder_id}_"
f"{orientation_placeholder}")
# Filter out only the solvent substance to help with the solvation step.
filter_solvent = miscellaneous.FilterSubstanceByRole(
"host-filter_solvent")
filter_solvent.input_substance = ProtocolPath("host_substance",
"global")
filter_solvent.component_roles = [Component.Role.Solvent]
# Construct a set of coordinates for a host molecule correctly
# aligned to the z-axis.
align_coordinates = PrepareReleaseCoordinates(
"release_align_coordinates")
align_coordinates.substance = ProtocolPath("host_substance", "global")
align_coordinates.complex_file_path = ProtocolPath(
"host_coordinate_path", "global")
solvate_coordinates = copy.deepcopy(solvation_template)
solvate_coordinates.id = "release_solvate_coordinates"
solvate_coordinates.substance = ProtocolPath("filtered_substance",
filter_solvent.id)
solvate_coordinates.solute_coordinate_file = ProtocolPath(
"output_coordinate_path", align_coordinates.id)
# Apply the force field parameters. This only needs to be done for one
# of the windows.
apply_parameters = forcefield.BaseBuildSystem(
"release_apply_parameters")
apply_parameters.force_field_path = ProtocolPath(
"force_field_path", "global")
apply_parameters.substance = ProtocolPath("host_substance", "global")
apply_parameters.coordinate_file_path = ProtocolPath(
"coordinate_file_path", solvate_coordinates.id)
# Add the dummy atoms.
add_dummy_atoms = AddDummyAtoms("release_add_dummy_atoms")
add_dummy_atoms.substance = ProtocolPath("host_substance", "global")
add_dummy_atoms.input_coordinate_path = ProtocolPath(
"coordinate_file_path",
solvate_coordinates.id,
)
add_dummy_atoms.input_system = ProtocolPath("parameterized_system",
apply_parameters.id)
add_dummy_atoms.offset = ProtocolPath("dummy_atom_offset", "global")
# Apply the restraints files
generate_restraints = GenerateReleaseRestraints(
f"release_generate_restraints_{orientation_placeholder}")
generate_restraints.host_coordinate_path = ProtocolPath(
"output_coordinate_path", add_dummy_atoms.id)
generate_restraints.release_lambdas = ProtocolPath(
"release_lambdas", "global")
generate_restraints.restraint_schemas = restraint_schemas
apply_restraints = ApplyRestraints(
f"release_apply_restraints_{release_replicator_id}")
apply_restraints.restraints_path = ProtocolPath(
"restraints_path", generate_restraints.id)
apply_restraints.phase = "release"
apply_restraints.window_index = ReplicatorValue(release_replicator.id)
apply_restraints.input_system = ProtocolPath("output_system",
add_dummy_atoms.id)
# Setup the simulations for the release phase.
(
release_minimization,
release_thermalization,
release_equilibration,
release_production,
) = cls._paprika_build_simulation_protocols(
ProtocolPath("output_coordinate_path", add_dummy_atoms.id),
ProtocolPath("output_system", apply_restraints.id),
"release",
release_replicator_id,
minimization_template,
thermalization_template,
equilibration_template,
production_template,
)
# Analyze the release phase.
analyze_release_phase = AnalyzeAPRPhase(
f"analyze_release_phase_{orientation_placeholder}")
analyze_release_phase.topology_path = ProtocolPath(
"output_coordinate_path", add_dummy_atoms.id)
analyze_release_phase.trajectory_paths = ProtocolPath(
"trajectory_file_path", release_production.id)
analyze_release_phase.phase = "release"
analyze_release_phase.restraints_path = ProtocolPath(
"restraints_path", generate_restraints.id)
# Return the full list of protocols which make up the release parts
# of a host-guest APR calculation.
protocols = [
filter_solvent,
align_coordinates,
solvate_coordinates,
apply_parameters,
add_dummy_atoms,
generate_restraints,
apply_restraints,
release_minimization,
release_thermalization,
release_equilibration,
release_production,
analyze_release_phase,
]
return (
protocols,
release_replicator,
ProtocolPath("result", analyze_release_phase.id),
)
def _paprika_build_attach_pull_protocols(
cls,
orientation_replicator: ProtocolReplicator,
restraint_schemas: Dict[str, ProtocolPath],
solvation_template: coordinates.SolvateExistingStructure,
minimization_template: openmm.OpenMMEnergyMinimisation,
thermalization_template: openmm.OpenMMSimulation,
equilibration_template: openmm.OpenMMSimulation,
production_template: openmm.OpenMMSimulation,
):
# Define a replicator to set and solvate up the coordinates for each pull window
orientation_placeholder = orientation_replicator.placeholder_id
pull_replicator = ProtocolReplicator(
f"pull_replicator_{orientation_placeholder}")
pull_replicator.template_values = ProtocolPath("pull_windows_indices",
"global")
pull_replicator_id = (f"{pull_replicator.placeholder_id}_"
f"{orientation_placeholder}")
attach_replicator = ProtocolReplicator(
f"attach_replicator_{orientation_placeholder}")
attach_replicator.template_values = ProtocolPath(
"attach_windows_indices", "global")
attach_replicator_id = (f"{attach_replicator.placeholder_id}_"
f"{orientation_placeholder}")
# Filter out only the solvent substance to help with the solvation step.
filter_solvent = miscellaneous.FilterSubstanceByRole(
"host-guest-filter_solvent")
filter_solvent.input_substance = ProtocolPath("substance", "global")
filter_solvent.component_roles = [Component.Role.Solvent]
# Define the protocols which will set and solvate up the coordinates for each
# pull window
align_coordinates = PreparePullCoordinates(
f"pull_align_coordinates_{pull_replicator_id}")
align_coordinates.substance = ProtocolPath("substance", "global")
align_coordinates.complex_file_path = ProtocolPath(
f"guest_orientations[{orientation_placeholder}].coordinate_path",
"global")
align_coordinates.guest_orientation_mask = ProtocolPath(
"guest_orientation_mask", "global")
align_coordinates.pull_window_index = ReplicatorValue(
pull_replicator.id)
align_coordinates.pull_distance = ProtocolPath("pull_distance",
"global")
align_coordinates.n_pull_windows = ProtocolPath(
"n_pull_windows", "global")
solvate_coordinates = copy.deepcopy(solvation_template)
solvate_coordinates.id = f"pull_solvate_coordinates_{pull_replicator_id}"
solvate_coordinates.substance = ProtocolPath("filtered_substance",
filter_solvent.id)
solvate_coordinates.solute_coordinate_file = ProtocolPath(
"output_coordinate_path", align_coordinates.id)
# Apply the force field parameters. This only needs to be done once.
apply_parameters = forcefield.BuildSmirnoffSystem(
f"pull_apply_parameters_{orientation_placeholder}")
apply_parameters.force_field_path = ProtocolPath(
"force_field_path", "global")
apply_parameters.substance = ProtocolPath("substance", "global")
apply_parameters.coordinate_file_path = ProtocolPath(
"coordinate_file_path",
f"pull_solvate_coordinates_0_{orientation_placeholder}",
)
# Add the dummy atoms.
add_dummy_atoms = AddDummyAtoms(
f"pull_add_dummy_atoms_{pull_replicator_id}")
add_dummy_atoms.substance = ProtocolPath("substance", "global")
add_dummy_atoms.input_coordinate_path = ProtocolPath(
"coordinate_file_path", solvate_coordinates.id)
add_dummy_atoms.input_system = ProtocolPath("parameterized_system",
apply_parameters.id)
add_dummy_atoms.offset = ProtocolPath("dummy_atom_offset", "global")
attach_coordinate_path = ProtocolPath(
"output_coordinate_path",
f"pull_add_dummy_atoms_0_{orientation_placeholder}",
)
attach_system = ProtocolPath(
"output_system",
f"pull_add_dummy_atoms_0_{orientation_placeholder}")
# Apply the attach restraints
generate_attach_restraints = GenerateAttachRestraints(
f"attach_generate_restraints_{orientation_placeholder}")
generate_attach_restraints.complex_coordinate_path = attach_coordinate_path
generate_attach_restraints.attach_lambdas = ProtocolPath(
"attach_lambdas", "global")
generate_attach_restraints.restraint_schemas = restraint_schemas
apply_attach_restraints = ApplyRestraints(
f"attach_apply_restraints_{attach_replicator_id}")
apply_attach_restraints.restraints_path = ProtocolPath(
"restraints_path", generate_attach_restraints.id)
apply_attach_restraints.phase = "attach"
apply_attach_restraints.window_index = ReplicatorValue(
attach_replicator.id)
apply_attach_restraints.input_system = attach_system
# Apply the pull restraints
generate_pull_restraints = GeneratePullRestraints(
f"pull_generate_restraints_{orientation_placeholder}")
generate_pull_restraints.complex_coordinate_path = attach_coordinate_path
generate_pull_restraints.attach_lambdas = ProtocolPath(
"attach_lambdas", "global")
generate_pull_restraints.n_pull_windows = ProtocolPath(
"n_pull_windows", "global")
generate_pull_restraints.restraint_schemas = restraint_schemas
apply_pull_restraints = ApplyRestraints(
f"pull_apply_restraints_{pull_replicator_id}")
apply_pull_restraints.restraints_path = ProtocolPath(
"restraints_path", generate_pull_restraints.id)
apply_pull_restraints.phase = "pull"
apply_pull_restraints.window_index = ReplicatorValue(
pull_replicator.id)
apply_pull_restraints.input_system = ProtocolPath(
"output_system", add_dummy_atoms.id)
# Setup the simulations for the attach and pull phases.
(
attach_minimization,
attach_thermalization,
attach_equilibration,
attach_production,
) = cls._paprika_build_simulation_protocols(
attach_coordinate_path,
ProtocolPath("output_system", apply_attach_restraints.id),
"attach",
attach_replicator_id,
minimization_template,
thermalization_template,
equilibration_template,
production_template,
)
(
pull_minimization,
pull_thermalization,
pull_equilibration,
pull_production,
) = cls._paprika_build_simulation_protocols(
ProtocolPath("output_coordinate_path", add_dummy_atoms.id),
ProtocolPath("output_system", apply_pull_restraints.id),
"pull",
pull_replicator_id,
minimization_template,
thermalization_template,
equilibration_template,
production_template,
)
# Analyze the attach phase.
attach_free_energy = AnalyzeAPRPhase(
f"analyze_attach_phase_{orientation_placeholder}")
attach_free_energy.topology_path = attach_coordinate_path
attach_free_energy.trajectory_paths = ProtocolPath(
"trajectory_file_path", attach_production.id)
attach_free_energy.phase = "attach"
attach_free_energy.restraints_path = ProtocolPath(
"restraints_path", generate_attach_restraints.id)
# Analyze the pull phase.
pull_free_energy = AnalyzeAPRPhase(
f"analyze_pull_phase_{orientation_placeholder}")
pull_free_energy.topology_path = attach_coordinate_path
pull_free_energy.trajectory_paths = ProtocolPath(
"trajectory_file_path", pull_production.id)
pull_free_energy.phase = "pull"
pull_free_energy.restraints_path = ProtocolPath(
"restraints_path", generate_pull_restraints.id)
reference_state_work = ComputeReferenceWork(
f"pull_reference_work_{orientation_placeholder}")
reference_state_work.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
reference_state_work.restraints_path = ProtocolPath(
"restraints_path", generate_pull_restraints.id)
# Return the full list of protocols which make up the attach and pull parts
# of a host-guest APR calculation.
protocols = [
filter_solvent,
align_coordinates,
solvate_coordinates,
apply_parameters,
add_dummy_atoms,
generate_attach_restraints,
apply_attach_restraints,
generate_pull_restraints,
apply_pull_restraints,
attach_minimization,
attach_thermalization,
attach_equilibration,
attach_production,
pull_minimization,
pull_thermalization,
pull_equilibration,
pull_production,
attach_free_energy,
pull_free_energy,
reference_state_work,
]
protocol_replicators = [pull_replicator, attach_replicator]
return (
protocols,
protocol_replicators,
ProtocolPath("result", attach_free_energy.id),
ProtocolPath("result", pull_free_energy.id),
ProtocolPath("result", reference_state_work.id),
)
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_simulation_schema(
cls,
absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_molecules=1000,
) -> SimulationSchema:
"""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 to use for the fully mixed system.
(
mixture_protocols,
mixture_value,
mixture_stored_data,
) = generate_simulation_protocols(
analysis.AverageObservable("extract_observable_mixture"),
use_target_uncertainty,
id_suffix="_mixture",
n_molecules=n_molecules,
)
# Specify the average observable which should be estimated.
mixture_protocols.analysis_protocol.observable = ProtocolPath(
f"observables[{cls._observable_type().value}]",
mixture_protocols.production_simulation.id,
)
(
mixture_protocols.analysis_protocol.divisor,
mixture_n_molar_molecules,
) = cls._n_molecules_divisor(
ProtocolPath("output_number_of_molecules",
mixture_protocols.build_coordinates.id),
"_mixture",
)
# Define the protocols to use for each component, creating a replicator that
# will copy these for each component in the mixture substance.
component_replicator = ProtocolReplicator("component_replicator")
component_replicator.template_values = ProtocolPath(
"components", "global")
component_substance = ReplicatorValue(component_replicator.id)
component_protocols, _, component_stored_data = generate_simulation_protocols(
analysis.AverageObservable(
f"extract_observable_component_{component_replicator.placeholder_id}"
),
use_target_uncertainty,
id_suffix=f"_component_{component_replicator.placeholder_id}",
n_molecules=n_molecules,
)
# Make sure the protocols point to the correct substance.
component_protocols.build_coordinates.substance = component_substance
# Specify the average observable which should be estimated.
component_protocols.analysis_protocol.observable = ProtocolPath(
f"observables[{cls._observable_type().value}]",
component_protocols.production_simulation.id,
)
(
component_protocols.analysis_protocol.divisor,
component_n_molar_molecules,
) = cls._n_molecules_divisor(
ProtocolPath("output_number_of_molecules",
component_protocols.build_coordinates.id),
f"_component_{component_replicator.placeholder_id}",
)
# Weight the component value by the mole fraction.
weight_by_mole_fraction = miscellaneous.WeightByMoleFraction(
f"weight_by_mole_fraction_{component_replicator.placeholder_id}")
weight_by_mole_fraction.value = ProtocolPath(
"value", component_protocols.analysis_protocol.id)
weight_by_mole_fraction.full_substance = ProtocolPath(
"substance", "global")
weight_by_mole_fraction.component = component_substance
component_protocols.converge_uncertainty.add_protocols(
weight_by_mole_fraction)
# Make sure the convergence criteria is set to use the per component
# uncertainty target.
if use_target_uncertainty:
component_protocols.converge_uncertainty.conditions[
0].right_hand_value = ProtocolPath("per_component_uncertainty",
"global")
# Finally, set up the protocols which will be responsible for adding together
# the component observables, and subtracting these from the mixture system value.
add_component_observables = miscellaneous.AddValues(
"add_component_observables")
add_component_observables.values = ProtocolPath(
"weighted_value",
component_protocols.converge_uncertainty.id,
weight_by_mole_fraction.id,
)
calculate_excess_observable = miscellaneous.SubtractValues(
"calculate_excess_observable")
calculate_excess_observable.value_b = mixture_value
calculate_excess_observable.value_a = ProtocolPath(
"result", add_component_observables.id)
# Build the final workflow schema
schema = WorkflowSchema()
schema.protocol_schemas = [
component_protocols.build_coordinates.schema,
component_protocols.assign_parameters.schema,
component_protocols.energy_minimisation.schema,
component_protocols.equilibration_simulation.schema,
component_protocols.converge_uncertainty.schema,
component_protocols.decorrelate_trajectory.schema,
component_protocols.decorrelate_observables.schema,
mixture_protocols.build_coordinates.schema,
mixture_protocols.assign_parameters.schema,
mixture_protocols.energy_minimisation.schema,
mixture_protocols.equilibration_simulation.schema,
mixture_protocols.converge_uncertainty.schema,
mixture_protocols.decorrelate_trajectory.schema,
mixture_protocols.decorrelate_observables.schema,
add_component_observables.schema,
calculate_excess_observable.schema,
]
if component_n_molar_molecules is not None:
schema.protocol_schemas.append(component_n_molar_molecules.schema)
if mixture_n_molar_molecules is not None:
schema.protocol_schemas.append(mixture_n_molar_molecules.schema)
schema.protocol_replicators = [component_replicator]
schema.final_value_source = ProtocolPath(
"result", calculate_excess_observable.id)
schema.outputs_to_store = {
"full_system":
mixture_stored_data,
f"component_{component_replicator.placeholder_id}":
component_stored_data,
}
calculation_schema.workflow_schema = schema
return calculation_schema
def generate_base_reweighting_protocols(
statistical_inefficiency: S,
reweight_observable: T,
replicator_id: str = "data_replicator",
id_suffix: str = "",
) -> Tuple[ReweightingProtocols[S, T], ProtocolReplicator]:
"""Constructs a set of protocols which, when combined in a workflow schema, may be
executed to reweight a set of cached simulation data to estimate the average
value of an observable.
Parameters
----------
statistical_inefficiency
The protocol which will be used to compute the statistical inefficiency and
equilibration time of the observable of interest. This information will be
used to decorrelate the cached data prior to reweighting.
reweight_observable
The MBAR reweighting protocol to use to reweight the observable to the target
state. This method will automatically set the reduced potentials on the
object.
replicator_id: str
The id to use for the cached data replicator.
id_suffix: str
A string suffix to append to each of the protocol ids.
Returns
-------
The protocols to add to the workflow, a reference to the average value of the
estimated observable (an ``Observable`` object), and the replicator which will
clone the workflow for each piece of cached simulation data.
"""
# Create the replicator which will apply these protocol once for each piece of
# cached simulation data.
data_replicator = ProtocolReplicator(replicator_id=replicator_id)
data_replicator.template_values = ProtocolPath("full_system_data", "global")
# Validate the inputs.
assert isinstance(statistical_inefficiency, analysis.BaseAverageObservable)
assert data_replicator.placeholder_id in statistical_inefficiency.id
assert data_replicator.placeholder_id not in reweight_observable.id
replicator_suffix = f"_{data_replicator.placeholder_id}{id_suffix}"
# Unpack all the of the stored data.
unpack_stored_data = storage.UnpackStoredSimulationData(
"unpack_data{}".format(replicator_suffix)
)
unpack_stored_data.simulation_data_path = ReplicatorValue(replicator_id)
# Join the individual trajectories together.
join_trajectories = reweighting.ConcatenateTrajectories(
f"join_trajectories{id_suffix}"
)
join_trajectories.input_coordinate_paths = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id
)
join_trajectories.input_trajectory_paths = ProtocolPath(
"trajectory_file_path", unpack_stored_data.id
)
join_observables = reweighting.ConcatenateObservables(
f"join_observables{id_suffix}"
)
join_observables.input_observables = ProtocolPath(
"observables", unpack_stored_data.id
)
# Calculate the reduced potentials for each of the reference states.
build_reference_system = forcefield.BaseBuildSystem(
f"build_system{replicator_suffix}"
)
build_reference_system.force_field_path = ProtocolPath(
"force_field_path", unpack_stored_data.id
)
build_reference_system.coordinate_file_path = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id
)
build_reference_system.substance = ProtocolPath("substance", unpack_stored_data.id)
reduced_reference_potential = openmm.OpenMMEvaluateEnergies(
f"reduced_potential{replicator_suffix}"
)
reduced_reference_potential.parameterized_system = ProtocolPath(
"parameterized_system", build_reference_system.id
)
reduced_reference_potential.thermodynamic_state = ProtocolPath(
"thermodynamic_state", unpack_stored_data.id
)
reduced_reference_potential.coordinate_file_path = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id
)
reduced_reference_potential.trajectory_file_path = ProtocolPath(
"output_trajectory_path", join_trajectories.id
)
# Calculate the reduced potential of the target state.
build_target_system = forcefield.BaseBuildSystem(f"build_system_target{id_suffix}")
build_target_system.force_field_path = ProtocolPath("force_field_path", "global")
build_target_system.substance = ProtocolPath("substance", "global")
build_target_system.coordinate_file_path = ProtocolPath(
"output_coordinate_path", join_trajectories.id
)
reduced_target_potential = openmm.OpenMMEvaluateEnergies(
f"reduced_potential_target{id_suffix}"
)
reduced_target_potential.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
reduced_target_potential.parameterized_system = ProtocolPath(
"parameterized_system", build_target_system.id
)
reduced_target_potential.coordinate_file_path = ProtocolPath(
"output_coordinate_path", join_trajectories.id
)
reduced_target_potential.trajectory_file_path = ProtocolPath(
"output_trajectory_path", join_trajectories.id
)
reduced_target_potential.gradient_parameters = ProtocolPath(
"parameter_gradient_keys", "global"
)
# Compute the observable gradients.
zero_gradients = gradients.ZeroGradients(f"zero_gradients{id_suffix}")
zero_gradients.force_field_path = ProtocolPath("force_field_path", "global")
zero_gradients.gradient_parameters = ProtocolPath(
"parameter_gradient_keys", "global"
)
# Decorrelate the target potentials and observables.
if not isinstance(statistical_inefficiency, analysis.BaseAverageObservable):
raise NotImplementedError()
decorrelate_target_potential = analysis.DecorrelateObservables(
f"decorrelate_target_potential{id_suffix}"
)
decorrelate_target_potential.time_series_statistics = ProtocolPath(
"time_series_statistics", statistical_inefficiency.id
)
decorrelate_target_potential.input_observables = ProtocolPath(
"output_observables", reduced_target_potential.id
)
decorrelate_observable = analysis.DecorrelateObservables(
f"decorrelate_observable{id_suffix}"
)
decorrelate_observable.time_series_statistics = ProtocolPath(
"time_series_statistics", statistical_inefficiency.id
)
decorrelate_observable.input_observables = ProtocolPath(
"output_observables", zero_gradients.id
)
# Decorrelate the reference potentials. Due to a quirk of how workflow replicators
# work the time series statistics need to be passed via a dummy protocol first.
#
# Because the `statistical_inefficiency` and `decorrelate_reference_potential`
# protocols are replicated by the same replicator the `time_series_statistics`
# input of `decorrelate_reference_potential_X` will take its value from
# the `time_series_statistics` output of `statistical_inefficiency_X` rather than
# as a list of of [statistical_inefficiency_0.time_series_statistics...
# statistical_inefficiency_N.time_series_statistics]. Passing the statistics via
# an un-replicated intermediate resolves this.
replicate_statistics = miscellaneous.DummyProtocol(
f"replicated_statistics{id_suffix}"
)
replicate_statistics.input_value = ProtocolPath(
"time_series_statistics", statistical_inefficiency.id
)
decorrelate_reference_potential = analysis.DecorrelateObservables(
f"decorrelate_reference_potential{replicator_suffix}"
)
decorrelate_reference_potential.time_series_statistics = ProtocolPath(
"output_value", replicate_statistics.id
)
decorrelate_reference_potential.input_observables = ProtocolPath(
"output_observables", reduced_reference_potential.id
)
# Finally, apply MBAR to get the reweighted value.
reweight_observable.reference_reduced_potentials = ProtocolPath(
"output_observables[ReducedPotential]", decorrelate_reference_potential.id
)
reweight_observable.target_reduced_potentials = ProtocolPath(
"output_observables[ReducedPotential]", decorrelate_target_potential.id
)
reweight_observable.observable = ProtocolPath(
"output_observables", decorrelate_observable.id
)
reweight_observable.frame_counts = ProtocolPath(
"time_series_statistics.n_uncorrelated_points", statistical_inefficiency.id
)
protocols = ReweightingProtocols(
unpack_stored_data,
#
join_trajectories,
join_observables,
#
build_reference_system,
reduced_reference_potential,
#
build_target_system,
reduced_target_potential,
#
statistical_inefficiency,
replicate_statistics,
#
decorrelate_reference_potential,
decorrelate_target_potential,
#
decorrelate_observable,
zero_gradients,
#
reweight_observable,
)
return protocols, data_replicator
def _default_reweighting_schema(
cls,
observable_type: ObservableType,
absolute_tolerance: unit.Quantity = UNDEFINED,
relative_tolerance: float = UNDEFINED,
n_effective_samples: int = 50,
) -> ReweightingSchema:
"""Returns the default calculation schema to use when estimating this class of
property by re-weighting cached simulation data.
This internal implementation allows re-weighting a different observable than
may be specified by the `_observable_type` class property.
Parameters
----------
absolute_tolerance
The absolute tolerance to estimate the property to within.
relative_tolerance
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_effective_samples
The minimum number of effective samples to require when
reweighting the cached simulation data.
Returns
-------
The default re-weighting calculation schema.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = ReweightingSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
# Set up the storage queries
calculation_schema.storage_queries = cls._default_reweighting_storage_query(
)
# Define the protocols which will re-weight the observable computed for the
# fully mixed system.
mixture_protocols, mixture_data_replicator = generate_reweighting_protocols(
observable_type,
"mixture_data_replicator",
"_mixture",
)
mixture_protocols.reweight_observable.required_effective_samples = (
n_effective_samples)
divide_by_mixture_molecules = miscellaneous.DivideValue(
"divide_by_mixture_molecules")
divide_by_mixture_molecules.value = ProtocolPath(
"value", mixture_protocols.reweight_observable.id)
(
divide_by_mixture_molecules.divisor,
mixture_n_molar_molecules,
) = cls._n_molecules_divisor(
ProtocolPath(
"total_number_of_molecules",
mixture_protocols.unpack_stored_data.id.replace(
mixture_data_replicator.placeholder_id, "0"),
),
"_mixture",
)
# Define the protocols to use for each component, creating a replicator that
# will copy these for each component in the full substance.
component_replicator = ProtocolReplicator("component_replicator")
component_replicator.template_values = ProtocolPath(
"components", "global")
component_protocols, component_data_replicator = generate_reweighting_protocols(
observable_type,
f"component_{component_replicator.placeholder_id}_data_replicator",
f"_component_{component_replicator.placeholder_id}",
)
component_protocols.reweight_observable.required_effective_samples = (
n_effective_samples)
component_data_replicator.template_values = ProtocolPath(
f"component_data[$({component_replicator.id})]", "global")
divide_by_component_molecules = miscellaneous.DivideValue(
f"divide_by_component_{component_replicator.placeholder_id}_molecules"
)
divide_by_component_molecules.value = ProtocolPath(
"value", component_protocols.reweight_observable.id)
(
divide_by_component_molecules.divisor,
component_n_molar_molecules,
) = cls._n_molecules_divisor(
ProtocolPath(
"total_number_of_molecules",
component_protocols.unpack_stored_data.id.replace(
component_data_replicator.placeholder_id, "0"),
),
f"_component_{component_replicator.placeholder_id}",
)
# Make sure the protocols point to the correct substance.
component_substance = ReplicatorValue(component_replicator.id)
component_protocols.build_reference_system.substance = component_substance
component_protocols.build_target_system.substance = component_substance
# Weight the component value by the mole fraction.
weight_by_mole_fraction = miscellaneous.WeightByMoleFraction(
f"weight_by_mole_fraction_{component_replicator.placeholder_id}")
weight_by_mole_fraction.value = ProtocolPath(
"result", divide_by_component_molecules.id)
weight_by_mole_fraction.full_substance = ProtocolPath(
"substance", "global")
weight_by_mole_fraction.component = component_substance
# Finally, set up the protocols which will be responsible for adding together
# the component observables, and subtracting these from the full system value.
add_component_observables = miscellaneous.AddValues(
"add_component_observables")
add_component_observables.values = ProtocolPath(
"weighted_value",
weight_by_mole_fraction.id,
)
calculate_excess_observable = miscellaneous.SubtractValues(
"calculate_excess_observable")
calculate_excess_observable.value_b = ProtocolPath(
"result", divide_by_mixture_molecules.id)
calculate_excess_observable.value_a = ProtocolPath(
"result", add_component_observables.id)
# Build the final workflow schema
schema = WorkflowSchema()
schema.protocol_schemas = [
*[x.schema for x in mixture_protocols if x is not None],
divide_by_mixture_molecules.schema,
*[x.schema for x in component_protocols if x is not None],
divide_by_component_molecules.schema,
weight_by_mole_fraction.schema,
add_component_observables.schema,
calculate_excess_observable.schema,
]
if component_n_molar_molecules is not None:
schema.protocol_schemas.append(component_n_molar_molecules.schema)
if mixture_n_molar_molecules is not None:
schema.protocol_schemas.append(mixture_n_molar_molecules.schema)
schema.protocol_replicators = [
mixture_data_replicator,
component_replicator,
component_data_replicator,
]
schema.final_value_source = ProtocolPath(
"result", calculate_excess_observable.id)
calculation_schema.workflow_schema = schema
return calculation_schema
def _apply_replicator_to_replicators(replicator, schema, template_values):
"""Applies a replicator to any replicators which depend upon
it (e.g. replicators with ids similar to `other_id_$(replicator.id)`).
Parameters
----------
replicator: ProtocolReplicator
The replicator being applied.
schema: WorkflowSchema
The workflow schema to which the replicator belongs.
template_values: List of Any
The values which the replicator is applying.
"""
# Look over all of the replicators left to apply and update them
# to point to the newly replicated protocols where appropriate.
new_indices = [str(index) for index in range(len(template_values))]
replicators = []
for original_replicator in schema.protocol_replicators:
# Check whether this replicator will be replicated.
if replicator.placeholder_id not in original_replicator.id:
replicators.append(original_replicator)
continue
# Create the replicated replicators
for template_index in new_indices:
replicator_id = original_replicator.id.replace(
replicator.placeholder_id, template_index)
new_replicator = ProtocolReplicator(replicator_id)
new_replicator.template_values = original_replicator.template_values
# Make sure to replace any reference to the applied replicator
# with the actual index.
if isinstance(new_replicator.template_values, ProtocolPath):
updated_path = new_replicator.template_values.full_path.replace(
replicator.placeholder_id, template_index)
new_replicator.template_values = ProtocolPath.from_string(
updated_path)
elif isinstance(new_replicator.template_values, list):
updated_values = []
for template_value in new_replicator.template_values:
if not isinstance(template_value, ProtocolPath):
updated_values.append(template_value)
continue
updated_path = template_value.full_path.replace(
replicator.placeholder_id, template_index)
updated_values.append(
ProtocolPath.from_string(updated_path))
new_replicator.template_values = updated_values
replicators.append(new_replicator)
schema.protocol_replicators = replicators
def generate_base_reweighting_protocols(
analysis_protocol,
mbar_protocol,
replicator_id="data_repl",
id_suffix="",
):
"""Constructs a set of protocols which, when combined in a workflow schema,
may be executed to reweight a set of existing data to estimate a particular
property. The reweighted observable of interest will be calculated by
following the passed in `analysis_protocol`.
Parameters
----------
analysis_protocol: AveragePropertyProtocol
The protocol which will take input from the stored data,
and generate a set of observables to reweight.
mbar_protocol: BaseReweightingProtocol
A template mbar reweighting protocol, which has it's reference
observables already set. This method will automatically set the
reduced potentials on this object.
replicator_id: str
The id to use for the data replicator.
id_suffix: str
A string suffix to append to each of the protocol ids.
Returns
-------
BaseReweightingProtocols:
A named tuple of the protocol which should form the bulk of
a property estimation workflow.
ProtocolReplicator:
A replicator which will clone the workflow for each piece of
stored data.
"""
assert isinstance(analysis_protocol, analysis.AveragePropertyProtocol)
assert f"$({replicator_id})" in analysis_protocol.id
assert f"$({replicator_id})" not in mbar_protocol.id
replicator_suffix = "_$({}){}".format(replicator_id, id_suffix)
# Unpack all the of the stored data.
unpack_stored_data = storage.UnpackStoredSimulationData(
"unpack_data{}".format(replicator_suffix))
unpack_stored_data.simulation_data_path = ReplicatorValue(replicator_id)
# The autocorrelation time of each of the stored files will be calculated for this property
# using the passed in analysis protocol.
if isinstance(analysis_protocol, analysis.ExtractAverageStatistic):
analysis_protocol.statistics_path = ProtocolPath(
"statistics_file_path", unpack_stored_data.id)
elif isinstance(analysis_protocol, analysis.AverageTrajectoryProperty):
analysis_protocol.input_coordinate_file = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id)
analysis_protocol.trajectory_path = ProtocolPath(
"trajectory_file_path", unpack_stored_data.id)
# Decorrelate the frames of the stored trajectory and statistics arrays.
decorrelate_statistics = analysis.ExtractUncorrelatedStatisticsData(
"decorrelate_stats{}".format(replicator_suffix))
decorrelate_statistics.statistical_inefficiency = ProtocolPath(
"statistical_inefficiency", analysis_protocol.id)
decorrelate_statistics.equilibration_index = ProtocolPath(
"equilibration_index", analysis_protocol.id)
decorrelate_statistics.input_statistics_path = ProtocolPath(
"statistics_file_path", unpack_stored_data.id)
decorrelate_trajectory = analysis.ExtractUncorrelatedTrajectoryData(
"decorrelate_traj{}".format(replicator_suffix))
decorrelate_trajectory.statistical_inefficiency = ProtocolPath(
"statistical_inefficiency", analysis_protocol.id)
decorrelate_trajectory.equilibration_index = ProtocolPath(
"equilibration_index", analysis_protocol.id)
decorrelate_trajectory.input_coordinate_file = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id)
decorrelate_trajectory.input_trajectory_path = ProtocolPath(
"trajectory_file_path", unpack_stored_data.id)
# Stitch together all of the trajectories
join_trajectories = reweighting.ConcatenateTrajectories("concat_traj" +
id_suffix)
join_trajectories.input_coordinate_paths = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id)
join_trajectories.input_trajectory_paths = ProtocolPath(
"output_trajectory_path", decorrelate_trajectory.id)
join_statistics = reweighting.ConcatenateStatistics("concat_stats" +
id_suffix)
join_statistics.input_statistics_paths = ProtocolPath(
"output_statistics_path", decorrelate_statistics.id)
# Calculate the reduced potentials for each of the reference states.
build_reference_system = forcefield.BaseBuildSystem(
"build_system{}".format(replicator_suffix))
build_reference_system.force_field_path = ProtocolPath(
"force_field_path", unpack_stored_data.id)
build_reference_system.substance = ProtocolPath("substance",
unpack_stored_data.id)
build_reference_system.coordinate_file_path = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id)
reduced_reference_potential = openmm.OpenMMReducedPotentials(
"reduced_potential{}".format(replicator_suffix))
reduced_reference_potential.system_path = ProtocolPath(
"system_path", build_reference_system.id)
reduced_reference_potential.thermodynamic_state = ProtocolPath(
"thermodynamic_state", unpack_stored_data.id)
reduced_reference_potential.coordinate_file_path = ProtocolPath(
"coordinate_file_path", unpack_stored_data.id)
reduced_reference_potential.trajectory_file_path = ProtocolPath(
"output_trajectory_path", join_trajectories.id)
reduced_reference_potential.kinetic_energies_path = ProtocolPath(
"output_statistics_path", join_statistics.id)
# Calculate the reduced potential of the target state.
build_target_system = forcefield.BaseBuildSystem("build_system_target" +
id_suffix)
build_target_system.force_field_path = ProtocolPath(
"force_field_path", "global")
build_target_system.substance = ProtocolPath("substance", "global")
build_target_system.coordinate_file_path = ProtocolPath(
"output_coordinate_path", join_trajectories.id)
reduced_target_potential = openmm.OpenMMReducedPotentials(
"reduced_potential_target" + id_suffix)
reduced_target_potential.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
reduced_target_potential.system_path = ProtocolPath(
"system_path", build_target_system.id)
reduced_target_potential.coordinate_file_path = ProtocolPath(
"output_coordinate_path", join_trajectories.id)
reduced_target_potential.trajectory_file_path = ProtocolPath(
"output_trajectory_path", join_trajectories.id)
reduced_target_potential.kinetic_energies_path = ProtocolPath(
"output_statistics_path", join_statistics.id)
# Finally, apply MBAR to get the reweighted value.
mbar_protocol.reference_reduced_potentials = ProtocolPath(
"statistics_file_path", reduced_reference_potential.id)
mbar_protocol.target_reduced_potentials = ProtocolPath(
"statistics_file_path", reduced_target_potential.id)
if (isinstance(mbar_protocol, reweighting.ReweightStatistics)
and mbar_protocol.statistics_type != ObservableType.PotentialEnergy
and mbar_protocol.statistics_type != ObservableType.TotalEnergy
and mbar_protocol.statistics_type != ObservableType.Enthalpy and
mbar_protocol.statistics_type != ObservableType.ReducedPotential):
mbar_protocol.statistics_paths = ProtocolPath(
"output_statistics_path", decorrelate_statistics.id)
elif isinstance(mbar_protocol, reweighting.ReweightStatistics):
mbar_protocol.statistics_paths = [
ProtocolPath("statistics_file_path", reduced_target_potential.id)
]
mbar_protocol.frame_counts = ProtocolPath(
"number_of_uncorrelated_samples", decorrelate_statistics.id)
base_protocols = BaseReweightingProtocols(
unpack_stored_data,
analysis_protocol,
decorrelate_statistics,
decorrelate_trajectory,
join_trajectories,
join_statistics,
build_reference_system,
reduced_reference_potential,
build_target_system,
reduced_target_potential,
mbar_protocol,
)
# Create the replicator object.
component_replicator = ProtocolReplicator(replicator_id=replicator_id)
component_replicator.template_values = ProtocolPath(
"full_system_data", "global")
return base_protocols, component_replicator
def generate_gradient_protocol_group(
template_reweighting_protocol,
force_field_path,
coordinate_file_path,
trajectory_file_path,
statistics_file_path,
replicator_id="repl",
substance_source=None,
id_suffix="",
enable_pbc=True,
effective_sample_indices=None,
):
"""Constructs a set of protocols which, when combined in a workflow schema,
may be executed to reweight a set of existing data to estimate a particular
property. The reweighted observable of interest will be calculated by
following the passed in `analysis_protocol`.
Parameters
----------
template_reweighting_protocol: BaseMBARProtocol
A template protocol which will be used to reweight the observable of
interest to small perturbations to the parameter of interest. These
will then be used to calculate the finite difference gradient.
The template *must* have it's `reference_reduced_potentials` input set.
The `target_reduced_potentials` input will be set automatically by this
function.
In the case that the template is of type `ReweightStatistics` and the
observable is an energy, the statistics path will automatically be pointed
to the energies evaluated using the perturbed parameter as opposed to the
energy measured during the reference simulation.
force_field_path: ProtocolPath
The path to the force field parameters which the observables are being
estimated at.
coordinate_file_path: ProtocolPath
A path to the initial coordinates of the simulation trajectory which
was used to estimate the observable of interest.
trajectory_file_path: ProtocolPath
A path to the simulation trajectory which was used
to estimate the observable of interest.
statistics_file_path: ProtocolPath, optional
A path to the statistics which were generated alongside
the trajectory passed to the `trajectory_file_path`. These
should have been generated using the passed `force_field_path`.
replicator_id: str
A unique id which will be used for the protocol replicator which will
replicate this group for every parameter of interest.
substance_source: PlaceholderValue, optional
An optional protocol path to the substance whose gradient
is being estimated. If None, the global property substance
is used.
id_suffix: str
An optional string to append to the end of each of the
protocol ids.
enable_pbc: bool
If true, periodic boundary conditions are employed when recalculating
the reduced potentials.
effective_sample_indices: ProtocolPath, optional
A placeholder variable which in future will ensure that only samples
with a non-zero weight are included in the gradient calculation.
Returns
-------
ProtocolGroup
The protocol group which will estimate the gradient of
an observable with respect to one parameter.
ProtocolReplicator
The replicator which will copy the gradient group for
every parameter of interest.
ProtocolPath
A protocol path which points to the final gradient value.
"""
assert isinstance(template_reweighting_protocol,
reweighting.BaseMBARProtocol)
assert template_reweighting_protocol.reference_reduced_potentials is not None
assert template_reweighting_protocol.reference_reduced_potentials != UNDEFINED
id_suffix = f"_$({replicator_id}){id_suffix}"
# Set values of the optional parameters.
substance_source = (ProtocolPath("substance", "global")
if substance_source is None else substance_source)
effective_sample_indices = (effective_sample_indices if
effective_sample_indices is not None else [])
# Define the protocol which will evaluate the reduced potentials of the
# reference, forward and reverse states using only a subset of the full
# force field.
reduced_potentials = openmm.OpenMMGradientPotentials(
f"gradient_reduced_potentials{id_suffix}")
reduced_potentials.substance = substance_source
reduced_potentials.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global")
reduced_potentials.force_field_path = force_field_path
reduced_potentials.statistics_path = statistics_file_path
reduced_potentials.trajectory_file_path = trajectory_file_path
reduced_potentials.coordinate_file_path = coordinate_file_path
reduced_potentials.parameter_key = ReplicatorValue(replicator_id)
reduced_potentials.enable_pbc = enable_pbc
reduced_potentials.effective_sample_indices = effective_sample_indices
# Set up the protocols which will actually reweight the value of the
# observable to the forward and reverse states.
template_reweighting_protocol.bootstrap_iterations = 1
template_reweighting_protocol.required_effective_samples = 0
# We need to make sure we use the observable evaluated at the target state
# if the observable depends on the parameter being reweighted.
use_target_state_energies = isinstance(
template_reweighting_protocol, reweighting.ReweightStatistics) and (
template_reweighting_protocol.statistics_type
== ObservableType.PotentialEnergy
or template_reweighting_protocol.statistics_type
== ObservableType.ReducedPotential
or template_reweighting_protocol.statistics_type
== ObservableType.TotalEnergy
or template_reweighting_protocol.statistics_type
== ObservableType.Enthalpy)
template_reweighting_schema = template_reweighting_protocol.schema
# Create the reweighting protocols from the template schema.
reverse_mbar_schema = copy.deepcopy(template_reweighting_schema)
reverse_mbar_schema.id = f"reverse_reweight{id_suffix}"
reverse_mbar = registered_workflow_protocols[reverse_mbar_schema.type](
reverse_mbar_schema.id)
reverse_mbar.schema = reverse_mbar_schema
reverse_mbar.target_reduced_potentials = ProtocolPath(
"reverse_potentials_path", reduced_potentials.id)
forward_mbar_schema = copy.deepcopy(template_reweighting_schema)
forward_mbar_schema.id = f"forward_reweight{id_suffix}"
forward_mbar = registered_workflow_protocols[forward_mbar_schema.type](
forward_mbar_schema.id)
forward_mbar.schema = forward_mbar_schema
forward_mbar.target_reduced_potentials = ProtocolPath(
"forward_potentials_path", reduced_potentials.id)
if use_target_state_energies:
reverse_mbar.statistics_paths = [
ProtocolPath("reverse_potentials_path", reduced_potentials.id)
]
forward_mbar.statistics_paths = [
ProtocolPath("forward_potentials_path", reduced_potentials.id)
]
# Set up the protocol which will actually evaluate the parameter gradient
# using the central difference method.
central_difference = gradients.CentralDifferenceGradient(
f"central_difference{id_suffix}")
central_difference.parameter_key = ReplicatorValue(replicator_id)
central_difference.reverse_observable_value = ProtocolPath(
"value", reverse_mbar.id)
central_difference.forward_observable_value = ProtocolPath(
"value", forward_mbar.id)
central_difference.reverse_parameter_value = ProtocolPath(
"reverse_parameter_value", reduced_potentials.id)
central_difference.forward_parameter_value = ProtocolPath(
"forward_parameter_value", reduced_potentials.id)
# Assemble all of the protocols into a convenient group.
gradient_group = groups.ProtocolGroup(f"gradient_group{id_suffix}")
gradient_group.add_protocols(reduced_potentials, reverse_mbar,
forward_mbar, central_difference)
# Create the replicator which will copy the group for each parameter gradient
# which will be calculated.
parameter_replicator = ProtocolReplicator(replicator_id=replicator_id)
parameter_replicator.template_values = ProtocolPath(
"parameter_gradient_keys", "global")
return (
gradient_group,
parameter_replicator,
ProtocolPath("gradient", gradient_group.id, central_difference.id),
)
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
# Set up the storage queries
calculation_schema.storage_queries = (
ExcessMolarVolume._default_reweighting_storage_query()
)
# Set up a replicator that will re-run the component reweighting workflow for each
# component in the system.
component_replicator = ProtocolReplicator(replicator_id="component_replicator")
component_replicator.template_values = ProtocolPath("components", "global")
gradient_replicator = ProtocolReplicator("gradient")
gradient_replicator.template_values = ProtocolPath(
"parameter_gradient_keys", "global"
)
# Set up the protocols which will reweight data for the full system.
full_data_replicator_id = "full_data_replicator"
(
full_protocols,
full_volume,
full_data_replicator,
full_gradient_group,
full_gradient_source,
) = ExcessMolarVolume._get_reweighting_protocols(
"_full",
gradient_replicator.id,
full_data_replicator_id,
n_effective_samples=n_effective_samples,
)
# Set up the protocols which will reweight data for each component.
component_data_replicator_id = (
f"component_{component_replicator.placeholder_id}_data_replicator"
)
(
component_protocols,
component_volumes,
component_data_replicator,
component_gradient_group,
component_gradient_source,
) = ExcessMolarVolume._get_reweighting_protocols(
"_component",
gradient_replicator.id,
component_data_replicator_id,
replicator_id=component_replicator.id,
weight_by_mole_fraction=True,
substance_reference=ReplicatorValue(component_replicator.id),
n_effective_samples=n_effective_samples,
)
# Make sure the replicator is only replicating over component data.
component_data_replicator.template_values = ProtocolPath(
f"component_data[$({component_replicator.id})]", "global"
)
add_component_molar_volumes = miscellaneous.AddValues(
"add_component_molar_volumes"
)
add_component_molar_volumes.values = component_volumes
calculate_excess_volume = miscellaneous.SubtractValues(
"calculate_excess_potential"
)
calculate_excess_volume.value_b = full_volume
calculate_excess_volume.value_a = ProtocolPath(
"result", add_component_molar_volumes.id
)
# Combine the gradients.
add_component_gradients = miscellaneous.AddValues(
f"add_component_gradients" f"_{gradient_replicator.placeholder_id}"
)
add_component_gradients.values = component_gradient_source
combine_gradients = miscellaneous.SubtractValues(
f"combine_gradients_{gradient_replicator.placeholder_id}"
)
combine_gradients.value_b = full_gradient_source
combine_gradients.value_a = ProtocolPath("result", add_component_gradients.id)
# Build the final workflow schema.
schema = WorkflowSchema()
schema.protocol_schemas = [
*(x.schema for x in full_protocols),
*(x.schema for x in component_protocols),
add_component_molar_volumes.schema,
calculate_excess_volume.schema,
full_gradient_group.schema,
component_gradient_group.schema,
add_component_gradients.schema,
combine_gradients.schema,
]
schema.protocol_replicators = [
full_data_replicator,
component_replicator,
component_data_replicator,
gradient_replicator,
]
schema.gradients_sources = [ProtocolPath("result", combine_gradients.id)]
schema.final_value_source = ProtocolPath("result", calculate_excess_volume.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
use_target_uncertainty = (
absolute_tolerance != UNDEFINED or relative_tolerance != UNDEFINED
)
# Define the id of the replicator which will clone the gradient protocols
# for each gradient key to be estimated.
gradient_replicator_id = "gradient_replicator"
# Set up a workflow to calculate the molar volume of the full, mixed system.
(
full_system_protocols,
full_system_molar_molecules,
full_system_volume,
full_output,
full_system_gradient_group,
full_system_gradient_replicator,
full_system_gradient,
) = ExcessMolarVolume._get_simulation_protocols(
"_full",
gradient_replicator_id,
use_target_uncertainty=use_target_uncertainty,
n_molecules=n_molecules,
)
# Set up a general workflow for calculating the molar volume of one of the system components.
component_replicator_id = "component_replicator"
component_substance = ReplicatorValue(component_replicator_id)
# Make sure to weight by the mole fractions of the actual full system as these may be slightly
# different to the mole fractions of the measure property due to rounding.
full_substance = ProtocolPath(
"output_substance", full_system_protocols.build_coordinates.id
)
(
component_protocols,
component_molar_molecules,
component_volumes,
component_output,
component_gradient_group,
component_gradient_replicator,
component_gradient,
) = ExcessMolarVolume._get_simulation_protocols(
"_component",
gradient_replicator_id,
replicator_id=component_replicator_id,
weight_by_mole_fraction=True,
component_substance_reference=component_substance,
full_substance_reference=full_substance,
use_target_uncertainty=use_target_uncertainty,
n_molecules=n_molecules,
)
# Finally, set up the protocols which will be responsible for adding together
# the component molar volumes, and subtracting these from the mixed system molar volume.
add_component_molar_volumes = miscellaneous.AddValues(
"add_component_molar_volumes"
)
add_component_molar_volumes.values = component_volumes
calculate_excess_volume = miscellaneous.SubtractValues(
"calculate_excess_volume"
)
calculate_excess_volume.value_b = full_system_volume
calculate_excess_volume.value_a = ProtocolPath(
"result", add_component_molar_volumes.id
)
# Create the replicator object which defines how the pure component
# molar volume estimation protocols will be replicated for each component.
component_replicator = ProtocolReplicator(replicator_id=component_replicator_id)
component_replicator.template_values = ProtocolPath("components", "global")
# Combine the gradients.
add_component_gradients = miscellaneous.AddValues(
f"add_component_gradients" f"_$({gradient_replicator_id})"
)
add_component_gradients.values = component_gradient
combine_gradients = miscellaneous.SubtractValues(
f"combine_gradients_$({gradient_replicator_id})"
)
combine_gradients.value_b = full_system_gradient
combine_gradients.value_a = ProtocolPath("result", add_component_gradients.id)
# Combine the gradient replicators.
gradient_replicator = ProtocolReplicator(replicator_id=gradient_replicator_id)
gradient_replicator.template_values = ProtocolPath(
"parameter_gradient_keys", "global"
)
# Build the final workflow schema
schema = WorkflowSchema()
schema.protocol_schemas = [
component_protocols.build_coordinates.schema,
component_protocols.assign_parameters.schema,
component_protocols.energy_minimisation.schema,
component_protocols.equilibration_simulation.schema,
component_protocols.converge_uncertainty.schema,
component_molar_molecules.schema,
full_system_protocols.build_coordinates.schema,
full_system_protocols.assign_parameters.schema,
full_system_protocols.energy_minimisation.schema,
full_system_protocols.equilibration_simulation.schema,
full_system_protocols.converge_uncertainty.schema,
full_system_molar_molecules.schema,
component_protocols.extract_uncorrelated_trajectory.schema,
component_protocols.extract_uncorrelated_statistics.schema,
full_system_protocols.extract_uncorrelated_trajectory.schema,
full_system_protocols.extract_uncorrelated_statistics.schema,
add_component_molar_volumes.schema,
calculate_excess_volume.schema,
component_gradient_group.schema,
full_system_gradient_group.schema,
add_component_gradients.schema,
combine_gradients.schema,
]
schema.protocol_replicators = [gradient_replicator, component_replicator]
# Finally, tell the schemas where to look for its final values.
schema.gradients_sources = [ProtocolPath("result", combine_gradients.id)]
schema.final_value_source = ProtocolPath("result", calculate_excess_volume.id)
schema.outputs_to_store = {
"full_system": full_output,
f"component_$({component_replicator_id})": component_output,
}
calculation_schema.workflow_schema = schema
return calculation_schema