fix up a few things

OpenFreeEnergy · Jan 7, 2025 · 45b8025 · 45b8025
1 parent 020ba71
commit 45b8025
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/openfecli/commands/plan_rbfe_network.py b/openfecli/commands/plan_rbfe_network.py
@@ -50,13 +50,15 @@ def plan_rbfe_network_main(
     from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import (
         RBFEAlchemicalNetworkPlanner,
     )
+    from openfe.protocols.openmm_utils.charge_generation import assign_offmol_partial_charges
+    from openfe import SmallMoleculeComponent
 
     write("assigning partial charges -- this may be slow")
 
     charged_small_molecules = []
     for smc in small_molecules:
         offmol = smc.to_openff()
-        charge_generation.assign_offmol_partial_charges(
+        assign_offmol_partial_charges(
             offmol=offmol,
             overwrite=False,
             method=partial_charge_settings.partial_charge_method,

diff --git a/openfecli/commands/plan_rhfe_network.py b/openfecli/commands/plan_rhfe_network.py
@@ -13,7 +13,7 @@
 
 def plan_rhfe_network_main(
     mapper, mapping_scorer, ligand_network_planner, small_molecules,
-    solvent,
+    solvent, partial_charge_settings,
 ):
     """Utility method to plan a relative hydration free energy network.
 
@@ -29,6 +29,9 @@ def plan_rhfe_network_main(
         molecules of the system
     solvent : SolventComponent
         Solvent component used for solvation
+    partial_charge_settings : OpenFFPartialChargeSettings
+        how to assign partial charges to the input ligands
+        (if they don't already have partial charges).
 
     Returns
     -------
@@ -39,6 +42,21 @@ def plan_rhfe_network_main(
     from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import (
         RHFEAlchemicalNetworkPlanner
     )
+    from openfe.protocols.openmm_utils.charge_generation import assign_offmol_partial_charges
+    from openfe import SmallMoleculeComponent
+
+    charged_small_molecules = []
+    for smc im small_molecules:
+        offmol = smc.to_openff()
+        assign_offmol_partial_charges(
+            offmol=offmol,
+            overwrite=False,
+            method=partial_charge_settings.partial_charge_method,
+            toolkit_backend=partial_charge_settings.off_toolkit_backend,
+            generate_n_conformers=partial_charge_settings.number_of_conformers,
+            nagl_model=partial_charge_settings.nagl_model
+        )
+        charged_small_molecules.append(SmallMoleculeComponent.from_openff(offmol))
 
     network_planner = RHFEAlchemicalNetworkPlanner(
         mappers=mapper,