Preserve R/S and E/Z stereochemistry in polymer converters

Extends PR #1's stereochemistry fix (previously molecule + markush only) to the
polymer path, which until now silently dropped chirality and double-bond
geometry on round-trip.

- polymer_to_mermaid: emit absolute CIP R/S (_CIPCode) in atom ids instead of
  order-dependent CHI_TETRAHEDRAL_CW/CCW (same bug fixed in the molecule converter).
- mermaid_to_psmiles: parse _R/_S id suffixes and ===|E|/===|Z| markers, then
  restore them after sanitize — chirality via brute-forced CIP matching, and
  double-bond geometry via SetDoubleBondNeighborDirections so SMILES emits / and \.
- tests/test_polymer_stereochemistry.py: PLA R/S, cis/trans polybutadiene,
  combined chiral+E/Z, and plain-polymer regression coverage.

Verified: chiral and E/Z PSMILES now round-trip to canonical-identical SMILES,
and cis is distinguished from trans.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: YZY-stack

molecode/polymer/mermaid_to_psmiles.py

@@ -93,6 +93,70 @@ def parse_element(label: str) -> tuple[str, int]:
 
 _NODE_RE = re.compile(r'^\s+(\w+)\[([^\]]+)\]\s*$')
 
+# Absolute CIP chirality suffix on an atom id, e.g. ``B0_C2_R`` → "R".
+# The preceding char is always a digit (the per-element counter), so this never
+# collides with an element symbol such as sulfur (``B0_S1``).
+_CHIRAL_RE = re.compile(r'\d_([RS])$')
+
+
+# ── Stereochemistry restoration (after sanitize / full topology) ──────────────
+
+def _restore_chirality(mol: RWMol, idx_to_cip: dict[int, str]) -> None:
+    """Restore absolute CIP R/S by brute-forcing the chiral tag.
+
+    CW/CCW is order-dependent, so we try each tag, recompute CIP, and keep the
+    one whose ``_CIPCode`` matches the desired absolute label (mirrors the
+    small-molecule converter).
+    """
+    if not idx_to_cip:
+        return
+    for idx, desired in idx_to_cip.items():
+        atom = mol.GetAtomWithIdx(idx)
+        matched = False
+        for tag in (Chem.ChiralType.CHI_TETRAHEDRAL_CW,
+                    Chem.ChiralType.CHI_TETRAHEDRAL_CCW):
+            atom.SetChiralTag(tag)
+            try:
+                Chem.AssignStereochemistry(mol, cleanIt=True, force=True)
+            except Exception:
+                continue
+            if atom.HasProp('_CIPCode') and atom.GetProp('_CIPCode') == desired:
+                matched = True
+                break
+        if not matched:
+            atom.SetChiralTag(Chem.ChiralType.CHI_UNSPECIFIED)
+    try:
+        Chem.AssignStereochemistry(mol, cleanIt=False, force=True)
+    except Exception:
+        pass
+
+
+def _restore_double_bond_stereo(mol: RWMol, records: list) -> None:
+    """Restore ===|E| / ===|Z| double-bond configuration.
+
+    Sets the double-bond stereo + reference atoms, then derives the neighbouring
+    single-bond directions so the SMILES writer actually emits ``/`` and ``\``.
+    """
+    applied = False
+    for idx1, idx2, stereo in records:
+        bond = mol.GetBondBetweenAtoms(idx1, idx2)
+        if bond is None:
+            continue
+        n1 = [n.GetIdx() for n in mol.GetAtomWithIdx(idx1).GetNeighbors()
+              if n.GetIdx() != idx2]
+        n2 = [n.GetIdx() for n in mol.GetAtomWithIdx(idx2).GetNeighbors()
+              if n.GetIdx() != idx1]
+        if n1 and n2:
+            bond.SetStereoAtoms(n1[0], n2[0])
+            bond.SetStereo(Chem.BondStereo.STEREOE if stereo == 'E'
+                           else Chem.BondStereo.STEREOZ)
+            applied = True
+    if applied:
+        try:
+            Chem.SetDoubleBondNeighborDirections(mol)
+        except Exception:
+            pass
+
 
 # ── Main converter ────────────────────────────────────────────────────────────
 
@@ -107,6 +171,7 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
 
     # 1. Collect node declarations (skip TL / TR terminus nodes)
     nodes: dict[str, tuple[str, int]] = {}  # nid → (element, formal_charge)
+    chirality: dict[str, str] = {}          # nid → "R" / "S"
     for line in lines:
         m = _NODE_RE.match(line)
         if m:
@@ -117,10 +182,13 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
                 symbol, charge = parse_element(label)
                 nodes[nid] = (symbol, charge)
             except ValueError:
-                pass
+                continue
+            cm = _CHIRAL_RE.search(nid)
+            if cm:
+                chirality[nid] = cm.group(1)
 
     # 2. Collect edges and identify attachment nodes
-    edges: list[tuple[str, Chem.BondType, str]] = []
+    edges: list[tuple[str, Chem.BondType, str, Optional[str]]] = []
     entry_nodes: set[str] = set()
     exit_nodes:  set[str] = set()
 
@@ -130,6 +198,7 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
             continue
         a, bond_sym, b = m.group(1), m.group(2), m.group(3)
         bt = _BOND_MAP.get(bond_sym, Chem.BondType.SINGLE)
+        ez = 'E' if bond_sym == '===|E|' else ('Z' if bond_sym == '===|Z|' else None)
 
         if a == "TL":
             entry_nodes.add(b)
@@ -141,7 +210,7 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
             exit_nodes.add(b)
         else:
             if a in nodes and b in nodes:
-                edges.append((a, bt, b))
+                edges.append((a, bt, b, ez))
 
     if not entry_nodes or not exit_nodes:
         return None
@@ -166,12 +235,15 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
     ap_left  = rw.AddAtom(Chem.Atom(0))
     ap_right = rw.AddAtom(Chem.Atom(0))
 
-    for a, bt, b in edges:
+    stereo_records: list[tuple[int, int, str]] = []  # (idx1, idx2, 'E'/'Z')
+    for a, bt, b, ez in edges:
         if a in nid_to_idx and b in nid_to_idx:
             try:
                 rw.AddBond(nid_to_idx[a], nid_to_idx[b], bt)
             except Exception:
                 return None
+            if ez:
+                stereo_records.append((nid_to_idx[a], nid_to_idx[b], ez))
 
     if entry_nid not in nid_to_idx or exit_nid not in nid_to_idx:
         return None
@@ -191,6 +263,12 @@ def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
         except Exception:
             return None
 
+    # 5. Restore stereochemistry (R/S and E/Z) after the full topology exists.
+    _restore_chirality(
+        rw, {nid_to_idx[n]: cip for n, cip in chirality.items() if n in nid_to_idx}
+    )
+    _restore_double_bond_stereo(rw, stereo_records)
+
     return Chem.MolToSmiles(rw)
 
 

molecode/polymer/polymer_to_mermaid.py

@@ -122,12 +122,15 @@ def _new_node_id(self, atom: Chem.Atom) -> str:
         symbol = atom.GetSymbol()
         self._element_counter[symbol] += 1
         cnt = self._element_counter[symbol]
-        chiral = atom.GetChiralTag()
+        # 使用 RDKit 计算的绝对 CIP 构型（_CIPCode），而不是把
+        # CHI_TETRAHEDRAL_CW/CCW 直接当 R/S —— CW/CCW 依赖原子顺序，
+        # 只有 _CIPCode 才是可序列化的绝对 R/S 标签。
+        # AssignStereochemistry 已在 convert() 中调用，确保 _CIPCode 存在。
         suffix = ""
-        if chiral == Chem.ChiralType.CHI_TETRAHEDRAL_CW:
-            suffix = "_R"
-        elif chiral == Chem.ChiralType.CHI_TETRAHEDRAL_CCW:
-            suffix = "_S"
+        if atom.HasProp("_CIPCode"):
+            cip = atom.GetProp("_CIPCode")
+            if cip in ("R", "S"):
+                suffix = f"_{cip}"
         return f"{self.prefix}_{symbol}{cnt}{suffix}"
 
     def convert(self) -> Tuple[List[str], str, str]:

tests/test_polymer_stereochemistry.py

@@ -0,0 +1,58 @@
+"""Polymer round-trips preserve R/S chirality and E/Z double-bond stereo.
+
+Mirrors tests/test_molecule_stereochemistry.py but for the polymer converters,
+which encode a repeat unit (PSMILES with two ``*`` attachment points) as a
+Mermaid graph and back. Before the CIP fix the polymer path dropped both
+chirality and E/Z; these tests pin the corrected behaviour.
+"""
+
+from rdkit import Chem
+
+from molecode.polymer import polymer_to_mermaid, mermaid_to_psmiles
+
+
+def _round_trip(psmiles: str, n: int = 10) -> str:
+    graph = polymer_to_mermaid(psmiles, n=n, name="Test")
+    back = mermaid_to_psmiles(graph)
+    assert back is not None
+    return Chem.CanonSmiles(back)
+
+
+def test_r_chirality_round_trip_uses_absolute_cip():
+    ps = "*OC(=O)[C@@H](C)*"          # PLA-like, one stereocentre
+    graph = polymer_to_mermaid(ps, n=10, name="Test")
+    assert "_R" in graph or "_S" in graph
+    assert _round_trip(ps) == Chem.CanonSmiles(ps)
+
+
+def test_s_chirality_round_trip_uses_absolute_cip():
+    ps = "*OC(=O)[C@H](C)*"
+    assert _round_trip(ps) == Chem.CanonSmiles(ps)
+
+
+def test_trans_double_bond_round_trip():
+    ps = "*C/C=C/C*"
+    graph = polymer_to_mermaid(ps, n=5, name="Test")
+    assert "===|E|" in graph
+    assert _round_trip(ps, n=5) == Chem.CanonSmiles(ps)
+
+
+def test_cis_double_bond_round_trip():
+    ps = "*C/C=C\\C*"
+    graph = polymer_to_mermaid(ps, n=5, name="Test")
+    assert "===|Z|" in graph
+    assert _round_trip(ps, n=5) == Chem.CanonSmiles(ps)
+
+
+def test_cis_and_trans_are_distinguished():
+    assert _round_trip("*C/C=C/C*", n=5) != _round_trip("*C/C=C\\C*", n=5)
+
+
+def test_chirality_and_ez_combined():
+    ps = "*O[C@@H](C)/C=C/C*"
+    assert _round_trip(ps) == Chem.CanonSmiles(ps)
+
+
+def test_plain_polymers_still_round_trip():
+    for ps in ("*CC*", "*CC(C)*", "*CCO*", "*NCCCCCC(=O)*", "*CC(c1ccccc1)*"):
+        assert _round_trip(ps) == Chem.CanonSmiles(ps)