add-isomer-option (#92)

* add-isomer-option * format * update test case * format * restore --------- Co-authored-by: Xiang Chen <[email protected]>
ssec-jhu · Dec 5, 2024 · de5ed3b · de5ed3b
1 parent 93ed163
commit de5ed3b
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/bluephos/bluephos_pipeline.py b/bluephos/bluephos_pipeline.py
@@ -152,6 +152,7 @@ def get_pipeline(
     t_nn=1.5,  # Threshold for 'z' score.
     t_ste=1.9,  # Threshold for 'ste'.
     t_dft=2.5,  # Threshold for 'dft'.
+    isomer="fac",  # default isomer
 ):
     """
     Set up and return the BluePhos discovery pipeline executor
@@ -175,6 +176,7 @@ def get_pipeline(
         "t_nn": t_nn,
         "t_ste": t_ste,
         "t_dft": t_dft,
+        "isomer": isomer,
     }
 
     for key, value in context_dict.items():
@@ -205,6 +207,12 @@ def get_pipeline(
         choices=["orca", "ase"],
         help="DFT package to use (default: orca)",
     )
+    ap.add_argument(
+        "--isomer",
+        required=False,
+        default="fac",
+        help="Set the isomer type for the OptimizeGeometriesTask (e.g., mer, fac)",
+    )
     args = ap.parse_args()
 
     # Run the pipeline with the provided arguments
@@ -222,6 +230,7 @@ def get_pipeline(
             args.t_nn,
             args.t_ste,
             args.t_dft,
+            args.isomer,
         ),
         args,
     )
diff --git a/bluephos/tasks/optimizegeometries.py b/bluephos/tasks/optimizegeometries.py
@@ -34,12 +34,12 @@ def calculate_ste(mol):
         return None
 
 
-def optimize(row, xtb):
+def optimize(row, xtb, isomer):
     mol_id = row["ligand_identifier"]
     ste = row["ste"]
 
-    # Log the values of z and ste for debugging
-    logger.info(f"Processing molecule {mol_id} ...")
+    # Log the value of isomer for debugging
+    logger.info(f"Processing molecule {mol_id} with isomer '{isomer}'...")
 
     # Skip processing if ste already existed (re-run condition)
     if ste is not None:
@@ -56,7 +56,6 @@ def optimize(row, xtb):
         return row  # Return the updated row
 
     mol = AddHs(Chem.MolFromMolBlock(mol))
-    isomer = "fac"  # Example isomer, can be dynamically set if needed
     mol_id = mol.GetProp("_Name") + f"_{isomer}"
     mol.SetProp("_Name", mol_id)
 
@@ -105,13 +104,13 @@ def optimize(row, xtb):
             return row  # Return the updated row
 
 
-def optimize_geometries(df: pd.DataFrame, xtb: bool) -> pd.DataFrame:
+def optimize_geometries(df: pd.DataFrame, xtb: bool, isomer: str) -> pd.DataFrame:
     for col in ["xyz", "ste", "xtb_walltime", "xyz_len"]:
         if col not in df.columns:
             df[col] = None
 
     # Apply the optimize function to each row
-    df = df.apply(optimize, axis=1, xtb=xtb)
+    df = df.apply(optimize, axis=1, xtb=xtb, isomer=isomer)
 
     return df
 
@@ -121,6 +120,7 @@ def optimize_geometries(df: pd.DataFrame, xtb: bool) -> pd.DataFrame:
     optimize_geometries,
     context_kwargs={
         "xtb": "xtb",
+        "isomer": "isomer",
     },
     batch_size=1,
     num_cpus=1,

diff --git a/tests/test_optimizegeometries_task.py b/tests/test_optimizegeometries_task.py
@@ -34,8 +34,11 @@ def test_optimize(mock_optimize_geometry, mock_octahedral_embed, setup_dataframe
     # Define a mock xtb argument
     mock_xtb = True
 
+    # Define a mock isomer argument
+    mock_isomer = "fac"
+
     # Run optimize
-    output_dataframe = optimize_geometries(setup_dataframe, mock_xtb)
+    output_dataframe = optimize_geometries(setup_dataframe, mock_xtb, mock_isomer)
 
     # Check if XYZ data was added or set to failed
     assert output_dataframe.loc[0, "xyz"] is not None