Separate the filtering process out from task (#80)

* add seperate filter task

* add filter class and implement

* filter code update from mac

* update test and remove unused variable

* adjust comment of function

* update based on PR comments

* remove unnecessary filter Class

---------

Co-authored-by: Xiang Chen <xchen286@ssec01.idies.jhu.edu>

Author: xiangchenjhu

bluephos/bluephos_pipeline.py

@@ -10,6 +10,14 @@
 from bluephos.tasks.optimizegeometries import OptimizeGeometriesTask
 from bluephos.tasks.smiles2sdf import Smiles2SDFTask
 from bluephos.tasks.dft import DFTTask
+from bluephos.tasks.filter_pipeline import (
+    FilterNNInTask,
+    FilterNNOutTask,
+    FilterXTBInTask,
+    FilterXTBOutTask,
+    FilterDFTInTask,
+    FilterDFTOutTask,
+)
 
 
 def ligand_pair_generator(halides_file, acids_file):
@@ -36,24 +44,32 @@ def ligand_pair_generator(halides_file, acids_file):
             yield pd.DataFrame([ligand_pair])
 
 
-def rerun_candidate_generator(input_dir, t_nn, t_ste):
+def check_column(df, column_name, condition_func, default=True):
+    """
+    Helper function to check if a column exists and apply a condition.
+
+    """
+    if column_name not in df.columns:
+        return pd.Series([default] * len(df), index=df.index)
+    else:
+        return condition_func(df[column_name])
+
+
+def rerun_candidate_generator(input_dir, t_nn, t_ste, t_dft):
     """
     Generates candidate DataFrames from parquet files in the input directory.
 
     Core Algorithm:
-    - If the absolute value of 'z' is less than t_nn,
-    - and 'ste' is None or its absolute value is less than t_ste,
-    - and 'dft_energy_diff' is None,
-    This row is then added to a new DataFrame and yielded for re-run.
-
-    Additional Context:
-    1. All valid ligand pairs should already have run through the NN process and have a 'z' score.
-    2. If a row's 'ste' is None, then it's 'dft_energy_diff' should also be None.
+    - A row is selected for re-run if:
+        1. 'z' is None or its absolute value is less than t_nn,
+        2. and 'ste' is None or its absolute value is less than t_ste,
+        3. and 'dft_energy_diff' is None.
 
     Args:
         input_dir (str): Directory containing input parquet files.
         t_nn (float): Threshold for 'z' score.
         t_ste (float): Threshold for 'ste'.
+        t_dft (float): (Optional) Threshold for 'dft_energy_diff' (not currently used)
 
     Yields:
         DataFrame: A single-row DataFrame containing candidate data.
@@ -63,10 +79,9 @@ def rerun_candidate_generator(input_dir, t_nn, t_ste):
         df["ste"] = df["ste"].replace({None: np.nan})
 
         filtered = df[
-            (df["z"].notnull())
-            & (df["z"].abs() < t_nn)
-            & ((df["ste"].isnull()) | (df["ste"].abs() < t_ste))
-            & (df["dft_energy_diff"].isna())
+            check_column(df, "z", lambda col: col.isnull() | (col.abs() < t_nn))
+            & check_column(df, "ste", lambda col: col.isnull() | (col.abs() < t_ste))
+            & check_column(df, "dft_energy_diff", lambda col: col.isna())
         ]
         for _, row in filtered.iterrows():
             yield row.to_frame().transpose()
@@ -78,15 +93,48 @@ def ligand_smiles_reader_generator(ligand_smiles):
         yield row.to_frame().transpose()
 
 
-def get_generator(ligand_smiles, halides, acids, input_dir, t_nn, t_ste):
+def get_generator(ligand_smiles, halides, acids, input_dir, t_nn, t_ste, t_dft):
     """
     Get the appropriate generator based on the input directory presence.
     """
     if ligand_smiles:
         return lambda: ligand_smiles_reader_generator(ligand_smiles)
     elif not input_dir:
         return lambda: ligand_pair_generator(halides, acids)
-    return lambda: rerun_candidate_generator(input_dir, t_nn, t_ste)
+    return lambda: rerun_candidate_generator(input_dir, t_nn, t_ste, t_dft)
+
+
+def build_pipeline_graph(input_dir: str, ligand_smiles: str):
+    """
+    Construct the pipeline graph based on input conditions.
+
+    Args:
+        input_dir (str): Directory containing input parquet files.
+        ligand_smiles (str): Path to the ligand SMILES CSV file.
+
+    Returns:
+        list: A list of task tuples representing the pipeline graph.
+    """
+    full_pipeline = [
+        (GenerateLigandTableTask, Smiles2SDFTask),  # Generate ligands, then convert SMILES to SDF
+        (Smiles2SDFTask, NNTask),  # Use SMILES to run NN prediction
+        (NNTask, FilterNNOutTask),  # NN filter "out" goes to sink
+        (NNTask, FilterNNInTask),  # NN filter "in" continues to the next task
+        (FilterNNInTask, OptimizeGeometriesTask),  # Optimize geometries for filtered ligands
+        (OptimizeGeometriesTask, FilterXTBOutTask),  # XTB filter "out" goes to sink
+        (OptimizeGeometriesTask, FilterXTBInTask),  # XTB filter "in" continues to the next task
+        (FilterXTBInTask, DFTTask),  # Run DFT calculation for filtered ligands
+        (DFTTask, FilterDFTOutTask),  # DFT filter "out" goes to sink
+        (DFTTask, FilterDFTInTask),  # DFT filter "in" could be processed further
+    ]
+
+    if ligand_smiles:
+        return full_pipeline[1:]  # from NNTask (Case 1: Input as ligand SMILES CSV file)
+
+    if input_dir:
+        return full_pipeline[2:]  # from Smiles2SDFTask (Case 2: Use parquet files for rerun)
+
+    return full_pipeline  # from GenerateLigandTableTask (Case 3: Input as halides and acids CSV files)
 
 
 def get_pipeline(
@@ -99,38 +147,19 @@ def get_pipeline(
     input_dir=None,  # Directory containing input parquet files(rerun). Defaults to None.
     dft_package="orca",  # DFT package to use. Defaults to "orca".
     xtb=True,  # Enable xTb optimize geometries task. Defaults to True.
-    t_nn=1.5,  # Threshold for 'z' score. Defaults to None
-    t_ste=1.9,  # Threshold for 'ste'. Defaults to None
+    t_nn=1.5,  # Threshold for 'z' score.
+    t_ste=1.9,  # Threshold for 'ste'.
+    t_dft=2.5,  # Threshold for 'dft'.
 ):
     """
     Set up and return the BluePhos discovery pipeline executor
     Returns:
         RayStreamGraphExecutor: An executor for the BluePhos discovery pipeline
     """
-    steps = (
-        [
-            GenerateLigandTableTask,
-            Smiles2SDFTask,
-            NNTask,
-            OptimizeGeometriesTask,
-            DFTTask,
-        ]
-        if not (input_dir or ligand_smiles)  # input as halides and acids CSV files
-        else [
-            NNTask,
-            OptimizeGeometriesTask,
-            DFTTask,
-        ]
-        if not ligand_smiles  # input as parquet files
-        else [
-            Smiles2SDFTask,
-            NNTask,
-            OptimizeGeometriesTask,
-            DFTTask,
-        ]  # input as ligand smiles CSV file
-    )
-    generator = get_generator(ligand_smiles, halides, acids, input_dir, t_nn, t_ste)
-    pipeline_executor = RayStreamGraphExecutor(graph=steps, generator=generator)
+
+    generator = get_generator(ligand_smiles, halides, acids, input_dir, t_nn, t_ste, t_dft)
+    pipeline_graph = build_pipeline_graph(input_dir, ligand_smiles)
+    pipeline_executor = RayStreamGraphExecutor(graph=pipeline_graph, generator=generator)
 
     context_dict = {
         "ligand_smiles": ligand_smiles,
@@ -143,6 +172,7 @@ def get_pipeline(
         "xtb": xtb,
         "t_nn": t_nn,
         "t_ste": t_ste,
+        "t_dft": t_dft,
     }
 
     for key, value in context_dict.items():
@@ -161,7 +191,10 @@ def get_pipeline(
     ap.add_argument("--input_dir", required=False, help="Directory containing input parquet files")
     ap.add_argument("--t_nn", type=float, required=False, default=1.5, help="Threshold for 'z' score (default: 1.5)")
     ap.add_argument("--t_ste", type=float, required=False, default=1.9, help="Threshold for 'ste' (default: 1.9)")
-    ap.add_argument("--no_xtb", action="store_false", dest="xtb", help="Disable xTB optimization (default: enabled)")
+    ap.add_argument("--t_dft", type=float, required=False, default=2.5, help="Threshold for 'dft' (default: 2.5)")
+    ap.add_argument(
+        "--no_xtb", action="store_false", dest="xtb", default=True, help="Disable xTB optimization (default: enabled)"
+    )
 
     ap.add_argument(
         "--dft_package",
@@ -186,6 +219,7 @@ def get_pipeline(
             args.xtb,
             args.t_nn,
             args.t_ste,
+            args.t_dft,
         ),
         args,
     )

bluephos/tasks/dft.py

@@ -43,12 +43,12 @@ def get_dft_calculator(dft_package, n_cpus):
 # Process each row of the DataFrame to perform DFT calculations
 def process_dataframe(row, t_ste, dft_calculator):
     mol_id = row["ligand_identifier"]
-    ste = row["ste"]
     energy_diff = row["dft_energy_diff"]
 
-    if ste is None or abs(ste) >= t_ste or energy_diff is not None:
-        logger.info(f"Skipping DFT on molecule {mol_id} based on z or t_ste conditions.")
-        return row
+    # Skip DFT processing if enery_diff already existed (re-run condition)
+    if energy_diff is not None:
+        logger.info(f"Skipping DFT on molecule {mol_id} because dft_energy_diff existed (re-run).")
+        return row  # Return the row unchanged
 
     if row["xyz"] not in ["failed", None]:
         base_name = row["ligand_identifier"]
@@ -97,5 +97,5 @@ def dft_run(df: pd.DataFrame, t_ste: float, dft_package: str) -> pd.DataFrame:
         "dft_package": "dft_package",  #  Either "ase" or "orca"
     },
     batch_size=1,
-    num_cpus=32,
+    num_cpus=8,
 )

bluephos/tasks/filter_pipeline.py

@@ -0,0 +1,62 @@
+import pandas as pd
+from dplutils.pipeline import PipelineTask
+
+
+# Helper function to create both filter_in and filter_out tasks
+def filter(df: pd.DataFrame, column: str, threshold: float, filter_in=True) -> pd.DataFrame:
+    filtered = df[column] < threshold
+    if filter_in:
+        return df[filtered]
+    return df[~filtered]
+
+
+# Dynamically create filter_in and filter_out tasks for NN
+FilterNNInTask = PipelineTask(
+    "filter_nn_in",
+    filter,
+    kwargs={"column": "z"},
+    context_kwargs={"threshold": "t_nn"},
+    # filter_in
+)
+
+FilterNNOutTask = PipelineTask(
+    "filter_nn_out",
+    filter,
+    kwargs={"column": "z", "filter_in": False},
+    context_kwargs={"threshold": "t_nn"},
+    # filter_out
+)
+
+# Dynamically create filter_in and filter_out tasks for XTB
+FilterXTBInTask = PipelineTask(
+    "filter_xtb_in",
+    filter,
+    kwargs={"column": "ste"},
+    context_kwargs={"threshold": "t_ste"},
+    # filter_in
+)
+
+FilterXTBOutTask = PipelineTask(
+    "filter_xtb_out",
+    filter,
+    kwargs={"column": "ste", "filter_in": False},
+    context_kwargs={"threshold": "t_ste"},
+    # filter_out
+)
+
+# Dynamically create filter_in and filter_out tasks for DFT
+FilterDFTInTask = PipelineTask(
+    "filter_dft_in",
+    filter,
+    kwargs={"column": "dft_energy_diff"},
+    context_kwargs={"threshold": "t_dft"},
+    # filter_in
+)
+
+FilterDFTOutTask = PipelineTask(
+    "filter_dft_out",
+    filter,
+    kwargs={"column": "dft_energy_diff", "filter_in": False},
+    context_kwargs={"threshold": "t_dft"},
+    # filter_out
+)

bluephos/tasks/optimizegeometries.py

@@ -31,17 +31,16 @@ def calculate_ste(mol):
         return None
 
 
-def optimize(row, t_nn, xtb):
+def optimize(row, xtb):
     mol_id = row["ligand_identifier"]
-    z = row["z"]
     ste = row["ste"]
 
     # Log the values of z and ste for debugging
     logger.info(f"Processing molecule {mol_id} ...")
 
-    # Skip processing based on conditions
-    if z is None or abs(z) >= t_nn or ste is not None:
-        logger.info(f"Skipping xTB optimization on molecule {mol_id} based on z or t_ste conditions.")
+    # Skip processing if ste already existed (re-run condition)
+    if ste is not None:
+        logger.info(f"Skipping xTB optimization on molecule {mol_id} because t_ste existed (re-run).")
         return row  # Return the row unchanged
 
     mol = row["structure"]
@@ -100,13 +99,13 @@ def optimize(row, t_nn, xtb):
             return row  # Return the updated row
 
 
-def optimize_geometries(df: pd.DataFrame, t_nn: float, xtb: bool) -> pd.DataFrame:
+def optimize_geometries(df: pd.DataFrame, xtb: bool) -> pd.DataFrame:
     for col in ["xyz", "ste"]:
         if col not in df.columns:
             df[col] = None
 
     # Apply the optimize function to each row
-    df = df.apply(optimize, axis=1, t_nn=t_nn, xtb=xtb)
+    df = df.apply(optimize, axis=1, xtb=xtb)
 
     return df
 
@@ -115,7 +114,6 @@ def optimize_geometries(df: pd.DataFrame, t_nn: float, xtb: bool) -> pd.DataFram
     "optimize_geometries",
     optimize_geometries,
     context_kwargs={
-        "t_nn": "t_nn",
         "xtb": "xtb",
     },
     batch_size=1,

tests/test_optimizegeometries_task.py

@@ -31,14 +31,11 @@ def test_optimize(mock_optimize_geometry, mock_octahedral_embed, setup_dataframe
     mock_octahedral_embed.return_value = None  # Does not need to return anything
     mock_optimize_geometry.return_value = None  # Does not need to return anything
 
-    # Define a mock t_nn argument
-    mock_t_nn = 1.5  # Replace with a suitable value for t_nn
-
     # Define a mock xtb argument
     mock_xtb = True
 
     # Run optimize
-    output_dataframe = optimize_geometries(setup_dataframe, mock_t_nn, mock_xtb)
+    output_dataframe = optimize_geometries(setup_dataframe, mock_xtb)
 
     # Check if XYZ data was added or set to failed
     assert output_dataframe.loc[0, "xyz"] is not None