From d241bbb7969db9de991ad729a21aa7796743888d Mon Sep 17 00:00:00 2001 From: finlayclark Date: Wed, 24 Apr 2024 10:28:26 +0100 Subject: [PATCH] Convert ABFE workflow with system preparation to a graph --- maize/graphs/exs/biosimspace/abfe.py | 102 ++++++++++++++++----------- 1 file changed, 59 insertions(+), 43 deletions(-) diff --git a/maize/graphs/exs/biosimspace/abfe.py b/maize/graphs/exs/biosimspace/abfe.py index 979b7c3..9471df8 100644 --- a/maize/graphs/exs/biosimspace/abfe.py +++ b/maize/graphs/exs/biosimspace/abfe.py @@ -31,14 +31,14 @@ from .system_preparation import SystemPreparationBound, SystemPreparationFree __all__ = [ - "AbsoluteBindingFreeEnergySingle", - "AbsoluteBindingFreeEnergyMulti", + "AbsoluteBindingFreeEnergySingleRepeat", + "AbsoluteBindingFreeEnergyMultiRepeat", "abfe_no_prep_workflow", "abfe_with_prep_workflow", ] -class AbsoluteBindingFreeEnergySingle(Graph): +class AbsoluteBindingFreeEnergySingleRepeat(Graph): """ A class for running a single absolute binding free energy calculation using SOMD through BioSimSpace. This requires parameterised and equilibrated @@ -241,7 +241,7 @@ def build(self) -> None: self.out = self.map_port(collect_results.out) -class AbsoluteBindingFreeEnergyMulti(Graph): +class AbsoluteBindingFreeEnergyMultiRepeat(Graph): """ A class for running multiple repeat absolute binding free energy calculations using SOMD through BioSimSpace. This requires unparameterised @@ -263,7 +263,7 @@ def build(self) -> None: # Add all the repeat ABFE nodes abfe_subgraph = self.add( parallel( - AbsoluteBindingFreeEnergySingle, + AbsoluteBindingFreeEnergySingleRepeat, n_branches=self.n_repeats.value, inputs=[], constant_inputs=["inp_bound", "inp_free"], @@ -282,6 +282,53 @@ def build(self) -> None: self.out = self.map_port(accumulate_results.out, name="out") +class AbsoluteBindingFreeEnergyMultiwithPrep(Graph): + """ + A class for running multiple repeat absolute binding free energy calculations + from unparameterised input structures. + """ + + def build(self) -> None: + + # Run system preparation for each leg + sys_prep_free = self.add(SystemPreparationFree, name="SystemPreparationFree") + sys_prep_bound = self.add(SystemPreparationBound, name="SystemPreparationBound") + + # Run repeats of ABFE calculations + abfe_calc = self.add(AbsoluteBindingFreeEnergyMultiRepeat, name="AbsoluteBindingFreeEnergy") + + # Save the ABFE results + save_results = self.add(SaveAFEResults, name="SaveAFEResults") + + # Connect the nodes + self.connect(sys_prep_free.out, abfe_calc.inp_free) + self.connect(sys_prep_bound.out, abfe_calc.inp_bound) + self.connect(abfe_calc.out, save_results.inp) + + # Map inputs/ parameters for prep stages + self.combine_parameters(sys_prep_free.inp, sys_prep_bound.inp, name="lig_sdf_path") + self.combine_parameters(sys_prep_bound.protein_pdb, name="protein_pdb_path") + # Get a dict of all parameters for each system preparation stage where parameters with the same name + # Share the same key in the dict + params_prep = defaultdict(list) + params_prep = {param.name: [param] for param in sys_prep_bound.parameters.values()} + # Now add the free params, appending to the list if the key already exists + for param in sys_prep_free.parameters.values(): + params_prep[param.name].append(param) + + for param_name, param_vals in params_prep.items(): + if param_name in ["protein_pdb"]: + continue + prefix = "prep_" if "force_field" not in param_name else "" + self.combine_parameters(*param_vals, name=f"{prefix}{param_name}") + + # Map inputs/ parameters for the abfe stage + for param in abfe_calc.parameters.values(): + self.combine_parameters(param, name=f"abfe_{param.name}") + + self.combine_parameters(save_results.file, name="results_file_name") + + def get_abfe_no_prep_workflow() -> Workflow: """ A workflow to perform absolute binding free energy calculations given @@ -290,14 +337,13 @@ def get_abfe_no_prep_workflow() -> Workflow: flow = Workflow(name="absolute_binding_free_energy_no_prep") - abfe_calc = flow.add(AbsoluteBindingFreeEnergyMulti, name="AbsoluteBindingFreeEnergy") + abfe_calc = flow.add(AbsoluteBindingFreeEnergyMultiRepeat, name="AbsoluteBindingFreeEnergy") save_results = flow.add(SaveAFEResults, name="SaveAFEResults") # Connect the nodes/ subgraphs flow.connect(abfe_calc.out, save_results.inp) # Map the inputs/ parameters - flow.map(abfe_calc.n_repeats) flow.combine_parameters(abfe_calc.inp_bound, name="inp_bound") flow.combine_parameters(abfe_calc.inp_free, name="inp_free") flow.map(*abfe_calc.parameters.values()) @@ -319,43 +365,13 @@ def get_abfe_with_prep_workflow() -> Workflow: """ flow = Workflow(name="absolute_binding_free_energy") - # Run system preparation for each leg - sys_prep_free = flow.add(SystemPreparationFree, name="SystemPreparationFree") - sys_prep_bound = flow.add(SystemPreparationBound, name="SystemPreparationBound") - - # Run repeats of ABFE calculations - abfe_calc = flow.add(AbsoluteBindingFreeEnergyMulti, name="AbsoluteBindingFreeEnergy") - - # Save the ABFE results - save_results = flow.add(SaveAFEResults, name="SaveAFEResults") - - # Connect the nodes - flow.connect(sys_prep_free.out, abfe_calc.inp_free) - flow.connect(sys_prep_bound.out, abfe_calc.inp_bound) - flow.connect(abfe_calc.out, save_results.inp) - - # Map inputs/ parameters for prep stages - flow.combine_parameters(sys_prep_free.inp, sys_prep_bound.inp, name="lig_sdf_path") - flow.combine_parameters(sys_prep_bound.protein_pdb, name="protein_pdb_path") - # Get a dict of all parameters for each system preparation stage where parameters with the same name - # Share the same key in the dict - params_prep = defaultdict(list) - params_prep = {param.name: [param] for param in sys_prep_bound.parameters.values()} - # Now add the free params, appending to the list if the key already exists - for param in sys_prep_free.parameters.values(): - params_prep[param.name].append(param) - - for param_name, param_vals in params_prep.items(): - if param_name in ["protein_pdb"]: - continue - prefix = "prep_" if "force_field" not in param_name else "" - flow.combine_parameters(*param_vals, name=f"{prefix}{param_name}") - - # Map inputs/ parameters for the abfe stage - for param in abfe_calc.parameters.values(): - flow.combine_parameters(param, name=f"abfe_{param.name}") + # Run setup and repeats of ABFE calculations + abfe_with_prep = flow.add( + AbsoluteBindingFreeEnergyMultiwithPrep, name="AbsoluteBindingFreeEnergy" + ) - flow.combine_parameters(save_results.file, name="results_file_name") + # Map the inputs/ parameters + flow.map(*abfe_with_prep.parameters.values()) return flow