Crest reproducibility/joint optimization update

yarp/reaction/conf_sampling/joint_opt.py

@@ -1,34 +1,115 @@
 import copy
+import os
 import numpy as np
-from openbabel import openbabel as ob
+from openbabel import pybel
+from rdkit.Chem import AllChem
+
 from yarp.util.write_files import mol_write_yp
+from yarp.util.rdkit import yarpecule_to_rdmol, geom_from_rdmol
+from yarp.yarpecule.graph.adjacency import table_generator
+
 
-def ob_joint_optimize(conformer, target_bem, ff_name="uff"):
+def ob_joint_optimize(conformer, target_bem, lot="uff"):
     """
-    Applies constraints based on the target BEM and runs an OpenBabel FF optimization.
-    Returns a NEW conformer object with the biased geometry.
+    Biases a conformer's geometry toward a target BEM via low-level force
+    field optimization. Returns a NEW conformer object with the biased
+    geometry, or None if no optimizer could produce a geometry consistent
+    with the target BEM's connectivity.
+
+    Mirrors the RDKit-first / Open Babel fallback pattern used by
+    yarp.reaction.generate_rxns.quick_geom_opt: RDKit is tried first, and its
+    result is only kept if the resulting geometry's connectivity matches
+    target_bem; otherwise Open Babel is tried as a fallback, and if that also
+    fails to reproduce the target connectivity, None is returned so the
+    caller can skip the pair.
     """
-    obMol = build_ob_mol(conformer.elements, conformer.geo, target_bem)
-    ff = ob.OBForceField.FindForceField(ff_name) or ob.OBForceField.FindForceField("uff")
-    if ff is None:
-        raise RuntimeError(f"OpenBabel force field not found: {ff_name} or uff")
+    target_adj = bondmat_to_adjmat(target_bem)
+
+    # First, attempt to bias the geometry with RDKit
+    rd_opt_g = rdkit_joint_opt(conformer, target_bem, target_adj, lot=lot)
+
+    # Check if optimization reproduced the target connectivity
+    if rd_opt_g is not None:
+        rd_adj = table_generator(conformer.elements, rd_opt_g)
+        rd_diff = rd_adj - target_adj
+    if rd_opt_g is None or not np.all(rd_diff == 0):
+        # If RDKit generated a garbage geom (or failed outright), try Open Babel
+        ob_opt_g = obabel_joint_opt(conformer, target_bem, target_adj, ff_name=lot)
 
-    if not ff.Setup(obMol):
-        ff = ob.OBForceField.FindForceField("uff")
-        if ff is None or not ff.Setup(obMol):
-            raise RuntimeError("Failed to set up OpenBabel force field for joint optimization")
+        # If Open Babel fails too, we return None
+        if ob_opt_g is None:
+            return None
+        ob_adj = table_generator(conformer.elements, ob_opt_g)
+        ob_diff = ob_adj - target_adj
+        if not np.all(ob_diff == 0):
+            return None
 
-    ff.ConjugateGradients(500)
-    ff.GetCoordinates(obMol)
+        opt_geo = ob_opt_g
+
+    # Otherwise, if RDKit gave a valid geom, use that one
+    else:
+        opt_geo = rd_opt_g
 
     biased_conf = copy.deepcopy(conformer)
-    biased_conf.geo = np.array([[obMol.GetAtom(i).GetX(), obMol.GetAtom(i).GetY(), obMol.GetAtom(i).GetZ()]
-                                for i in range(1, obMol.NumAtoms() + 1)])
+    biased_conf.geo = opt_geo
     biased_conf.type = f"biased_{conformer.type}"
 
     return biased_conf
 
 
+def rdkit_joint_opt(conformer, target_bem, target_adj, lot="uff", maxiter=200):
+    """
+    Attempt to bias conformer geometry toward target_bem using RDKit.
+
+    Returns the optimized geometry (N x 3 array), or None if RDKit could not
+    build/optimize a mol from the imposed (possibly non-physical, mid-reaction)
+    target bonding.
+    """
+    try:
+        rdmol = yarpecule_to_rdmol(elements=conformer.elements, adj=target_adj,
+                                    bond_orders=target_bem, geo=conformer.geo)
+
+        if lot == "uff":
+            AllChem.UFFOptimizeMolecule(rdmol, maxIters=maxiter, ignoreInterfragInteractions=False)
+        elif lot == "mmff94":
+            AllChem.MMFFOptimizeMolecule(rdmol, maxIters=maxiter, ignoreInterfragInteractions=False)
+
+        return geom_from_rdmol(rdmol)
+    except (ValueError, RuntimeError):
+        return None
+
+
+def obabel_joint_opt(conformer, target_bem, target_adj, ff_name="uff", maxiter=500):
+    """
+    Attempt to bias conformer geometry toward target_bem using Open Babel.
+
+    Mirrors quick_geom_opt's Open Babel fallback (yarp.util.obabel.obabel_ff_opt):
+    writes a temporary mol file with the imposed target bonding and optimizes
+    it via pybel's `localopt`, so both fallbacks behave identically instead of
+    diverging on which part of the Open Babel API they happen to call.
+
+    Returns the optimized geometry (N x 3 array), or None if Open Babel could
+    not set up/optimize a force field for the imposed target bonding.
+    """
+    mol_file = '.tmp_joint.mol'
+    try:
+        mol_write_yp(mol_file, conformer.elements, conformer.geo, target_bem, target_adj)
+
+        mol = next(pybel.readfile("mol", mol_file))
+        mol.localopt(forcefield=ff_name, steps=maxiter)
+
+        opt_geo = np.zeros_like(conformer.geo)
+        for count_i, i in enumerate(opt_geo):
+            opt_geo[count_i] = mol.atoms[count_i].coords
+
+        return opt_geo
+    except (ValueError, RuntimeError):
+        return None
+    finally:
+        if os.path.exists(mol_file):
+            os.remove(mol_file)
+
+
 def bondmat_to_adjmat(bond_mat):
     adj_mat = copy.deepcopy(bond_mat)
     for count_i, i in enumerate(bond_mat):
@@ -38,28 +119,3 @@ def bondmat_to_adjmat(bond_mat):
             if count_i == count_j:
                 adj_mat[count_i][count_i] = 0.0
     return adj_mat
-
-
-def build_ob_mol(elements, coords, bond_mat):
-    """
-    Builds an OBMol object using explicit BEM bonding info.
-    """
-    adj_mat = bondmat_to_adjmat(bond_mat)
-    obMol = ob.OBMol()
-    conv = ob.OBConversion()
-    conv.SetInFormat("mol")
-
-    import os
-    import tempfile
-
-    tmp_file = None
-    try:
-        with tempfile.NamedTemporaryFile(suffix=".mol", delete=False) as tmp:
-            tmp_file = tmp.name
-        mol_write_yp(tmp_file, elements, coords, bond_mat, adj_mat)
-        conv.ReadFile(obMol, tmp_file)
-    finally:
-        if tmp_file is not None:
-            os.unlink(tmp_file)
-
-    return obMol

yarp/reaction/conf_sampling/select_pairs.py

@@ -23,20 +23,43 @@ def select_gsm_pairs(rxn, config):
     # --- STEP A: Apply Joint Optimization (Biasing) ---
     lot = config.bias_lot
     mode = config.joint_opt.lower()
+    verbose = getattr(config, "verbose", False)
 
     biased_r = r_confs
     biased_p = p_confs
     # biased_p: product geometries guided by reactant geometries
     # biased_r: reactant geometries guided by product geometries
+    #
+    # ob_joint_optimize returns None when neither optimizer can produce a
+    # geometry consistent with the target BEM (mirrors quick_geom_opt in
+    # generate_rxns.py); such conformers are dropped, keeping r_confs/p_confs
+    # aligned with their biased counterparts.
     if mode in ['dual', 'r_only']:
-        biased_p = [ob_joint_optimize(c, rxn.reactant.paired_bem, lot) for c in r_confs]
+        kept_r_confs, biased_p = [], []
+        for c in r_confs:
+            b = ob_joint_optimize(c, rxn.reactant.paired_bem, lot)
+            if b is None:
+                if verbose:
+                    print(f"  + SKIPPED! Unable to bias reactant conformer toward product BEM")
+                continue
+            kept_r_confs.append(c)
+            biased_p.append(b)
+        r_confs = kept_r_confs
     if mode in ['dual', 'p_only']:
-        biased_r = [ob_joint_optimize(c, rxn.product.paired_bem, lot) for c in p_confs]
+        kept_p_confs, biased_r = [], []
+        for c in p_confs:
+            b = ob_joint_optimize(c, rxn.product.paired_bem, lot)
+            if b is None:
+                if verbose:
+                    print(f"  + SKIPPED! Unable to bias product conformer toward reactant BEM")
+                continue
+            kept_p_confs.append(c)
+            biased_r.append(b)
+        p_confs = kept_p_confs
 
     # Default to "conformation-poor" model, unless an excess of conformers is present
     n_conf = config.n_conf
     total_conf = len(r_confs) + len(p_confs)
-    verbose = getattr(config, "verbose", False)
     if verbose:
         print("Total number of reactant + product conformers generated = ", total_conf)
 

yarp/reaction/conformer.py

@@ -2,7 +2,7 @@
 Definition of the conformer object class.
 """
 import numpy as np
-from yarp.util.write_files import xyz_generate_string
+from yarp.util.write_files import xyz_generate_string, xyz_write
 
 class conformer:
     """
@@ -109,3 +109,15 @@ def to_xyz_string(self):
             return xyz_generate_string(elements=self.elements, geo=self.geo)
         else:
             return None
+
+    def export_geometry(self, filename):
+        """
+        Export the conformer's geometry to an xyz file.
+        This shouldn't ever change any of the attributes of the conformer.
+
+        Parameters
+        ----------
+        filename : str
+            The name of the file to export the geometry to.
+        """
+        xyz_write(filename, self.elements, self.geo, comment=self.type)

yarp/reaction/external/calc_base.py

@@ -23,14 +23,17 @@ def set_scratch_dir(self, path: Path):
         self.scratch_dir = path
         self.scratch_dir.mkdir(parents=True, exist_ok=True)
 
-    def get_container_prefix(self, image_name: str, work_dir: str, *, apptainer_run: bool = False) -> str:
+    def get_container_prefix(self, image_name: str, work_dir: str, *, apptainer_run: bool = False, env_vars: dict = None) -> str:
         """
         The universal toggle for container execution.
         Maps the scratch directory to /work inside the container.
         Automatically checks for and pulls missing images.
 
         apptainer_run: If True, use ``apptainer run`` (OCI ENTRYPOINT/CMD, like ``docker run``).
         If False, use ``apptainer exec`` (requires an explicit executable before flags).
+
+        env_vars: Optional dict of environment variables to inject into the container
+        (e.g. {"OMP_NUM_THREADS": 4}). Passed via -e for Docker and --env for Apptainer.
         """
         if self.job_manager.container == "docker":
             # 1. Check if the image exists locally.
@@ -46,7 +49,10 @@ def get_container_prefix(self, image_name: str, work_dir: str, *, apptainer_run:
                 print(f"Docker image '{image_name}' not found locally. Pulling from registry...")
                 subprocess.run(["docker", "pull", "--platform", "linux/amd64", image_name], check=True)
 
-            return f"docker run --platform linux/amd64 --rm -v {work_dir}:/work -w /work {image_name}"
+            env_flags = ""
+            if env_vars:
+                env_flags = " ".join(f"-e {k}={v}" for k, v in env_vars.items()) + " "
+            return f"docker run --platform linux/amd64 --rm {env_flags}-v {work_dir}:/work -w /work {image_name}"
 
         elif self.job_manager.container == "apptainer":
             # Sanitize the image name so it works as a safe, flat filename
@@ -95,8 +101,11 @@ def get_container_prefix(self, image_name: str, work_dir: str, *, apptainer_run:
                         )
                     raise RuntimeError(msg) from None
 
+            env_flags = ""
+            if env_vars:
+                env_flags = " ".join(f"--env {k}={v}" for k, v in env_vars.items()) + " "
             verb = "run" if apptainer_run else "exec"
-            return f"apptainer {verb} -e --bind {work_dir}:/work --pwd /work {sif_path}"
+            return f"apptainer {verb} -e {env_flags}--bind {work_dir}:/work --pwd /work {sif_path}"
 
         elif self.job_manager.container == "singularity":
             # Sanitize the image name so it works as a safe, flat filename
@@ -130,8 +139,11 @@ def get_container_prefix(self, image_name: str, work_dir: str, *, apptainer_run:
                     check=True,
                 )
 
+            env_flags = ""
+            if env_vars:
+                env_flags = " ".join(f"--env {k}={v}" for k, v in env_vars.items()) + " "
             verb = "run" if apptainer_run else "exec"
-            return f"singularity {verb} -e --bind {work_dir}:/work --pwd /work {sif_path}"
+            return f"singularity {verb} -e {env_flags}--bind {work_dir}:/work --pwd /work {sif_path}"
 
         else:
             raise ValueError(f"Unsupported container runner: {self.job_manager.container}")

yarp/reaction/external/conf_gen.py

@@ -37,8 +37,16 @@ def write_submission_script(self) -> Path:
         """Write the bash script that the JobManager will execute."""
         script_path = self.scratch_dir / "run_crest_cmd.sh"
 
-        # Construct the core command
-        prefix = self.get_container_prefix(self.image_name, self.scratch_dir)
+        # When a seed is set, pin OMP/MKL threads inside the container so xTB's
+        # geometry optimizations use the same thread count as CREST's -T setting.
+        env_vars = None
+        if self.config.seed is not None:
+            env_vars = {
+                "OMP_NUM_THREADS": self.config.n_cpus,
+                "MKL_NUM_THREADS": self.config.n_cpus,
+            }
+
+        prefix = self.get_container_prefix(self.image_name, self.scratch_dir, env_vars=env_vars)
         crest_cmd = self._get_crest_command()
         full_command = f"{prefix} {crest_cmd}"
 
@@ -47,7 +55,6 @@ def write_submission_script(self) -> Path:
                 f.write("#!/bin/bash\n")
                 self.write_scheduler_headers(f)
                 f.write(f"cd {self.scratch_dir}\n")
-                # Pipe output to a log file so it doesn't flood the local terminal
                 f.write(f"{full_command} > crest_run.log 2> crest_run.err\n")
         except PermissionError:
             raise PermissionError(
@@ -138,6 +145,9 @@ def _get_crest_command(self):
         # molecular descriptors
         cmd += f" --chrg {self.config.charge} --uhf {self.config.n_unpaired_electrons}"
 
+        if self.config.seed is not None:
+            cmd += f" --seed {self.config.seed}"
+
         # conformer generation thresholds (ERM: expand this later, if needed)
         # ERM: no current way to cap CREST outputs at a set number of generated conformers!
         # You can damp down via adjusting the energy window threshold, but that's it

yarp/reaction/external/irc_val.py

@@ -184,9 +184,11 @@ def write_submission_script(self):
         # 1. The INNER Script (Runs inside the Docker container)
         # ---------------------------------------------------------
         inner_script_path = self.scratch_dir / "run_pysis_inner.sh"
-        
+
         with open(inner_script_path, "w") as f:
             f.write("#!/bin/bash\n\n")
+            f.write(f"export OMP_NUM_THREADS={self.config.n_cpus}\n")
+            f.write(f"export MKL_NUM_THREADS={self.config.n_cpus}\n\n")
             f.write("echo 'Starting YARP serial IRC execution...'\n\n")
 
             for i in range(1, self._get_num_runs() + 1):

yarp/reaction/external/min_opt.py

@@ -80,8 +80,11 @@ def write_submission_script(self) -> Path:
         """Write the bash script that the JobManager will execute."""
         script_path = self.scratch_dir / "run_pysis_rpopt.sh"
 
-        # Construct the core command
-        prefix = self.get_container_prefix(self.image_name, self.scratch_dir)
+        env_vars = {
+            "OMP_NUM_THREADS": self.config.n_cpus,
+            "MKL_NUM_THREADS": self.config.n_cpus,
+        }
+        prefix = self.get_container_prefix(self.image_name, self.scratch_dir, env_vars=env_vars)
         pysis_cmd = "pysis min_opt.yaml > min_opt.log 2> min_opt.err"
         full_command = f"{prefix} {pysis_cmd}"