fix + test

test/util/test_rdkit.py

@@ -4,7 +4,16 @@
 from yarp.reaction.enum import enumerate_products
 from yarp.yarpecule.graph.adjacency import table_generator
 
-from yarp.util.rdkit import yarpecule_to_rdmol, geom_from_rdmol, rdkit_ff_opt
+from yarp.util.rdkit import (
+    yarpecule_to_rdmol,
+    geom_from_rdmol,
+    rdkit_ff_opt,
+    smiles_to_rdmol,
+    adj_from_rdmol,
+    el_from_rdmol,
+)
+from yarp.yarpecule.atom_mapping import canon_order
+from yarp.util.properties import el_mass
 
 class TestGeom:
     def test_preserved_geom(self):
@@ -68,3 +77,35 @@ def test_3hp_2_aldehyde(self):
         diff = opt_adj - target_product.adj_mat
         assert not np.all(diff == 0)
 
+class TestRDKitSMILESHelpers:
+    def test_mapped_aromatic_adjacency_matches_unmapped(self):
+        mapped_smiles = "[n:0]1([H:7])[n:4][c:1]([O:3][H:6])[c:2]([H:8])[n:5]1"
+        unmapped_smiles = "Oc1cn[nH]n1"
+
+        mapped_mol = smiles_to_rdmol(mapped_smiles)
+        unmapped_mol = smiles_to_rdmol(unmapped_smiles)
+
+        mapped_elements = [el.lower() for el in el_from_rdmol(mapped_mol)]
+        unmapped_elements = [el.lower() for el in el_from_rdmol(unmapped_mol)]
+
+        mapped_adj = adj_from_rdmol(mapped_mol)
+        unmapped_adj = adj_from_rdmol(unmapped_mol)
+
+        mapped_masses = np.array([el_mass[el] for el in mapped_elements])
+        unmapped_masses = np.array([el_mass[el] for el in unmapped_elements])
+
+        mapped_elements, mapped_adj, _ = canon_order(
+            mapped_elements,
+            mapped_adj,
+            masses=mapped_masses,
+            return_index=False,
+        )
+        unmapped_elements, unmapped_adj, _ = canon_order(
+            unmapped_elements,
+            unmapped_adj,
+            masses=unmapped_masses,
+            return_index=False,
+        )
+
+        assert mapped_elements == unmapped_elements
+        assert np.array_equal(mapped_adj, unmapped_adj)

yarp/util/rdkit.py

@@ -5,8 +5,10 @@
 from rdkit.Chem import AllChem
 from rdkit.Chem import rdchem
 from rdkit.Geometry import Point3D
+from yarp.util.properties import el_mass
 import numpy as np
 import warnings
+import re
 
 BOND_MAP = {
     1: rdchem.BondType.SINGLE,
@@ -18,6 +20,69 @@
 }
 # TODO: How to handle dative bonds for organometallics?
 
+EXPLICIT_MAPPED_H_PATTERN = re.compile(r"\[(?:\d+)?H[^\]]*:\d+[^\]]*\]")
+
+
+def has_explicit_mapped_hydrogen(smiles):
+    return EXPLICIT_MAPPED_H_PATTERN.search(smiles) is not None
+
+
+def smiles_to_rdmol(smiles, preserve_explicit_h=False):
+
+    """
+    Converts a SMILES string to an RDKIT Mol object. If hydrogen atoms are explicityly mapped then 
+    they are preserved. Otherwise, default RDKit behavior is used to remove hydrogens and then add them back in. 
+
+    """
+    preserve_h = preserve_explicit_h or has_explicit_mapped_hydrogen(smiles)
+
+    if preserve_h:
+        params = Chem.SmilesParserParams()
+        params.removeHs = False
+        mol = Chem.MolFromSmiles(smiles, params)
+
+    else:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is not None:
+            mol = Chem.AddHs(mol)
+
+    if mol is None:
+        raise ValueError(f"RDKit could not parse SMILES: {smiles}")
+
+    return mol
+
+
+def adj_from_rdmol(mol):
+    adj_mat = np.zeros((mol.GetNumAtoms(), mol.GetNumAtoms()))
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        adj_mat[i, j] = adj_mat[j, i] = 1
+    return adj_mat
+
+
+def el_from_rdmol(mol):
+    return [atom.GetSymbol() for atom in mol.GetAtoms()]
+
+
+def atom_info_from_rdmol(mol):
+    atom_info = {}
+    for i, atom in enumerate(mol.GetAtoms()):
+        element = atom.GetSymbol().lower()
+        isotope = atom.GetIsotope()
+        atom_map = int(atom.GetProp("molAtomMapNumber")) if atom.HasProp("molAtomMapNumber") else None
+
+        atom_info[i] = {
+            "atom_index": i,
+            "atom_map": atom_map,
+            "element": element,
+            "formal_charge": atom.GetFormalCharge(),
+            "mass": float(isotope) if isotope else el_mass[element],
+            "stereo": {"atom": None, "bonds": {}},
+            "aromatic_input": atom.GetIsAromatic(),
+        }
+
+    return atom_info
 
 def yarpecule_to_rdmol(elements, adj, bond_orders, atom_info=None, geo=None, sanitize=True):
     """
@@ -180,7 +245,7 @@ def rdkit_ff_opt(ypcule, lot='uff', maxiter=200):
         Level of theory used for quick optimization
         ERM: mmff94 has a tendency to reform the reactant geometry
         when used to generate initial geom of products post product enumeration
-
+    
     maxiter : int
         Maximum number of optimization steps
 
@@ -200,4 +265,4 @@ def rdkit_ff_opt(ypcule, lot='uff', maxiter=200):
 
     opt_geom = geom_from_rdmol(rdmol)
 
-    return opt_geom
+    return opt_geom

yarp/yarpecule/atom_mapping.py

@@ -203,3 +203,19 @@ def graph_seps(adj_mat_0):
         adj_mat = np.dot(adj_mat, adj_mat_0)
 
     return seps
+
+def original_to_yarp_map(ypcule):
+    """
+    Return {input_atom_map: yarp_atom_map} for atoms that had maps in the
+    user-provided partial-map SMILES.
+    """
+    out = {}
+
+    for i, info in ypcule._atom_info.items():
+        input_map = info.get("input_atom_map")
+        yarp_map = info.get("atom_map")
+
+        if input_map is not None:
+            out[input_map] = yarp_map
+
+    return out

yarp/yarpecule/graph/smiles.py

@@ -9,7 +9,7 @@
 from yarp.yarpecule.atom_mapping import canon_order
 
 
-def smiles2adjmat(smiles, verbose=False):
+def smiles2adjmat(smiles, verbose=False, reorder_mapped=True):
     """
     In-house Savoie group SMILES parser. Written in python and transparent to debug. The main motivation
     was to consistently handle protonation of radicals and atoms with formal charges. The usual SMILES
@@ -364,7 +364,8 @@ def smiles2adjmat(smiles, verbose=False):
         dupes = sorted({m for m in provided_maps if provided_maps.count(m) > 1})
         raise ValueError(f"Duplicate atom-map indices in input SMILES: {dupes}")
 
-    adjmat, atom_info = reorder_by_mappings(adjmat, atom_info)
+    if reorder_mapped:
+        adjmat, atom_info = reorder_by_mappings(adjmat, atom_info)
 
     bond_electron_mat = adjmat.copy()
     for i in atom_info:

yarp/yarpecule/input_parsers.py

@@ -4,12 +4,55 @@
 """
 from pathlib import Path
 
+import re
 import numpy as np
 from rdkit.Chem import rdmolfiles, BondType, rdchem, Atom, MolFromSmiles, AddHs, AllChem, rdmolfiles
 from yarp.util.properties import el_to_an, el_n_expand_octet, el_expand_octet, el_mass
 from yarp.yarpecule.graph.smiles import smiles2adjmat, OctetError
+from yarp.util.rdkit import (
+    adj_from_rdmol,
+    atom_info_from_rdmol,
+    el_from_rdmol,
+    geom_from_rdmol,
+    smiles_to_rdmol,
+)
 
 
+ATOM_MAP_PATTERN = re.compile(r":\d+(?=[^\]]*\])")
+EXPLICIT_H_PATTERN = re.compile(r"\[(?:\d+)?H[^\]]*\]")
+
+
+def strip_atom_maps(smiles):
+    return ATOM_MAP_PATTERN.sub("", smiles)
+
+
+def has_explicit_hydrogen_atom(smiles):
+    return EXPLICIT_H_PATTERN.search(smiles) is not None
+
+
+def prepare_partial_mapped_smiles(smiles):
+    """
+    If the SMILES is partially mapped, return an unmapped working SMILES plus
+    the original per-atom input maps in original token order.
+
+    Fully mapped and fully unmapped SMILES are left alone.
+    """
+    _, _, original_atom_info = smiles2adjmat(smiles, reorder_mapped=False)
+
+    original_maps = [
+        original_atom_info[i].get("atom_map")
+        for i in original_atom_info
+    ]
+
+    has_any_map = any(m is not None for m in original_maps)
+    has_all_maps = all(m is not None for m in original_maps)
+
+    if not has_any_map or has_all_maps:
+        return smiles, None, False
+
+    unmapped_smiles = strip_atom_maps(smiles)
+    return unmapped_smiles, original_atom_info, True
+
 def xyz_parse(xyz, read_types=False, multiple=False):
     """
     Simple wrapper function for grabbing the coordinates and elements from an xyz file.
@@ -414,8 +457,23 @@ def xyz_from_smiles(smiles, mode="yarp"):
         # Parse basics
         # NOTE: bemat is used to generate geometry via RDKit, but not returned for
         # downstream use in yarpecule - ERM
-        adj_mat, bemat, atom_info = smiles2adjmat(smiles)
+
+        parse_smiles, original_atom_info, partial_mapped_input = prepare_partial_mapped_smiles(smiles)
+
+        adj_mat, bemat, atom_info = smiles2adjmat(parse_smiles)
+
+        if partial_mapped_input:
+            if len(original_atom_info) != len(atom_info):
+                raise ValueError(
+            "Partial-map handling failed: original and unmapped SMILES produced "
+            "different atom counts."
+        )
+
+            for i in atom_info:
+                atom_info[i]["input_atom_map"] = original_atom_info[i].get("atom_map")
+
         elements = [atom_info[i]["element"] for i in atom_info]
+
         fc = [int(atom_info[i]["formal_charge"]) for i in atom_info]
         q = int(sum(fc))
 
@@ -435,7 +493,7 @@ def xyz_from_smiles(smiles, mode="yarp"):
                     "WARNING: yarp aromatic SMILES geometry fallback used RDKit "
                     f"for {smiles} after octet validation failed at atoms {violations}."
                 )
-                geo = geo_via_rdkit(smiles, atom_info)
+                geo = geo_via_rdkit(parse_smiles,atom_info,preserve_explicit_h=partial_mapped_input and has_explicit_hydrogen_atom(parse_smiles),)
                 return elements, geo, adj_mat, q, atom_info
             raise OctetError(violations)
 
@@ -481,46 +539,18 @@ def xyz_from_smiles(smiles, mode="yarp"):
 
     # RDKit branch
     else:
-        m = MolFromSmiles(smiles)  # create molecule using rdkit
-        m = AddHs(m)  # make the hydrogens explicit
-        AllChem.EmbedMolecule(m, randomSeed=0xf00d)  # create a 3D geometry
-        N_atoms = len(m.GetAtoms())  # find the number of atoms
-        elements = []  # initialize list to hold element labels
-        geo = np.zeros((N_atoms, 3))  # initialize array to hold geometry
-        q = 0  # total charge on the molecule
-        atom_info = {}
+        m = smiles_to_rdmol(smiles)
+        AllChem.EmbedMolecule(m, randomSeed=0xf00d)
 
-        # loop over atoms, save their labels, positions, and total charge
-        for i in range(N_atoms):
-            atom = m.GetAtomWithIdx(i)
-            elements += [atom.GetSymbol()]
-            coord = m.GetConformer().GetAtomPosition(i)
-            geo[i] = np.array([coord.x, coord.y, coord.z])
-            q += atom.GetFormalCharge()
-            isotope = atom.GetIsotope()
-            mass = float(isotope) if isotope else el_mass[atom.GetSymbol().lower()]
-            atom_map = None
-            if atom.HasProp("molAtomMapNumber"):
-                atom_map = int(atom.GetProp("molAtomMapNumber"))
-            atom_info[i] = {
-                "atom_index": i,
-                "atom_map": atom_map,
-                "element": atom.GetSymbol().lower(),
-                "formal_charge": atom.GetFormalCharge(),
-                "mass": mass,
-                "stereo": {"atom": None, "bonds": {}},
-                "aromatic_input": atom.GetIsAromatic(),
-            }
-
-        # Generate adjacency matrix
-        adj_mat = np.zeros((N_atoms, N_atoms))
-        for i in [(_.GetBeginAtomIdx(), _.GetEndAtomIdx()) for _ in m.GetBonds()]:
-            adj_mat[i[0], i[1]] = 1
-            adj_mat[i[1], i[0]] = 1
+        elements = el_from_rdmol(m)
+        geo = geom_from_rdmol(m)
+        adj_mat = adj_from_rdmol(m)
+        atom_info = atom_info_from_rdmol(m)
+        q = int(sum(atom_info[i]["formal_charge"] for i in atom_info))
 
     return elements, geo, adj_mat, q, atom_info
 
-def geo_via_rdkit(smiles, atom_info):
+def geo_via_rdkit(smiles, atom_info, preserve_explicit_h=False):
     """
     Generate a geometry with RDKit and align it to the atom ordering used by
     the in-house SMILES parser.
@@ -529,22 +559,15 @@ def geo_via_rdkit(smiles, atom_info):
     identifies the correct graph but its bond-electron matrix is too crude to
     survive octet validation.
     """
-    m = MolFromSmiles(smiles)
-    if m is None:
-        raise ValueError(f"RDKit could not parse SMILES: {smiles}")
-    m = AddHs(m)
+    m = smiles_to_rdmol(smiles, preserve_explicit_h=preserve_explicit_h)
     AllChem.EmbedMolecule(m, randomSeed=0xf00d)
 
-    rdkit_elements = [atom.GetSymbol().lower() for atom in m.GetAtoms()]
+    rdkit_elements = [el.lower() for el in el_from_rdmol(m)]
     yarp_elements = [atom_info[i]["element"] for i in atom_info]
 
     yarp_maps = [atom_info[i].get("atom_map", None) for i in atom_info]
-    rdkit_maps = []
-    for atom in m.GetAtoms():
-        if atom.HasProp("molAtomMapNumber"):
-            rdkit_maps.append(int(atom.GetProp("molAtomMapNumber")))
-        else:
-            rdkit_maps.append(None)
+    rdkit_atom_info = atom_info_from_rdmol(m)
+    rdkit_maps = [rdkit_atom_info[i]["atom_map"] for i in rdkit_atom_info]
 
     if all(_ is not None for _ in yarp_maps):
         rdkit_by_map = {atom_map: idx for idx, atom_map in enumerate(rdkit_maps) if atom_map is not None}
@@ -557,9 +580,9 @@ def geo_via_rdkit(smiles, atom_info):
     else:
         raise ValueError("RDKit fallback atom order did not match the in-house parser ordering.")
 
+    rdkit_geo = geom_from_rdmol(m)
     geo = np.zeros((len(order), 3))
     for i, rdkit_idx in enumerate(order):
-        coord = m.GetConformer().GetAtomPosition(rdkit_idx)
-        geo[i] = np.array([coord.x, coord.y, coord.z])
+        geo[i] = rdkit_geo[rdkit_idx]
 
-    return geo
+    return geo

yarp/yarpecule/yarpecule.py

@@ -303,6 +303,7 @@ def _read_structure(self, mol, mode, strict=False, atom_info=None):
             normalized_atom_info[i] = {
                 "atom_index": i,
                 "atom_map": record.get("atom_map", None),
+                "input_atom_map": record.get("input_atom_map", None),
                 "element": self._elements[i],
                 "formal_charge": record.get("formal_charge", None),
                 "mass": record.get("mass", el_mass[self._elements[i]]),