Coverage for cosmolayer/store/coarse_graining.py: 91%

1"""Merge hydrogens into their heavy-atom neighbor, shrinking a SegmentStore 

2to a united-atom store. 

3 

4``SegmentStore.coarse_grain`` builds a new store where every hydrogen 

5``Chem.RemoveHs`` would actually remove is merged into the heavy atom it's 

6bonded to: fewer atoms, same segments. Requires every molecule's stored 

7``smiles`` to carry atom-map numbers reflecting local (COSMO) atom index 

8(guaranteed for any store built after GH issue #43's fix) -- that's what 

9lets a removed hydrogen's segments be redirected to the right heavy-atom 

10neighbor. 

11""" 

12 

13from rdkit import Chem 

14 

15 

16def compute_atom_remap( 

17 mapped_smiles: str, 

18) -> tuple[dict[int, int], frozenset[int], str]: 

19 """Compute one molecule's old-to-new local atom index map and its 

20 coarse-grained, re-mapped SMILES. 

21 

22 Parses ``mapped_smiles`` with hydrogens kept, then calls 

23 ``Chem.RemoveHs`` -- which returns a *new* ``Mol``, not an in-place 

24 edit -- to find out which hydrogens RDKit's own rules actually remove. 

25 ``Chem.RemoveHs`` never reorders surviving atoms relative to each 

26 other, only deletes, so a surviving atom's position in the reduced 

27 molecule (``0, 1, 2, ...``) is exactly its rank among survivors in the 

28 original local-index order -- the new, compacted local index. A 

29 removed hydrogen's segments belong wherever its (single) heavy-atom 

30 neighbor ended up. 

31 

32 A hydrogen ``Chem.RemoveHs`` conservatively keeps (a non-default 

33 isotope, or a neighbor with non-tetrahedral stereochemistry it can't 

34 safely represent without the explicit atom) is not merged: it survives 

35 as its own atom in the coarse-grained molecule, with its own new 

36 index, same as any heavy atom. 

37 

38 Parameters 

39 ---------- 

40 mapped_smiles : str 

41 Atom-mapped SMILES for one molecule, with map numbers a clean 

42 0-based or 1-based permutation of local atom index (as 

43 ``molecules_df["smiles"]`` stores it -- see GH issue #43). 

44 

45 Returns 

46 ------- 

47 new_local_index : dict[int, int] 

48 Every original local atom index mapped to its new, compacted 

49 local index -- both surviving atoms (mapped to their own new 

50 index) and merged hydrogens (mapped to their heavy-atom 

51 neighbor's new index) have an entry. 

52 survivors : frozenset[int] 

53 Original local atom indices that are still their own atom in the 

54 coarse-grained molecule (as opposed to a merged hydrogen) -- one 

55 per new, compacted local index. 

56 new_mapped_smiles : str 

57 The coarse-grained molecule's SMILES, atom-mapped with contiguous 

58 map numbers starting at the same base (0 or 1) as the input, in 

59 the same relative order as the surviving atoms originally had. 

60 

61 Raises 

62 ------ 

63 ValueError 

64 If ``mapped_smiles`` can't be parsed, or its atoms aren't a clean 

65 0-based or 1-based permutation of local atom index (including the 

66 unmapped case, where every atom's map number is 0). 

67 """ 

68 params = Chem.SmilesParserParams() 

69 params.removeHs = False 

70 mol = Chem.MolFromSmiles(mapped_smiles, params) 

71 if mol is None: 

72 raise ValueError(f"RDKit could not parse SMILES {mapped_smiles!r}") 

73 

74 num_atoms = mol.GetNumAtoms() 

75 map_nums = {atom.GetAtomMapNum() for atom in mol.GetAtoms()} 

76 if map_nums == set(range(num_atoms)): 

77 base = 0 

78 elif map_nums == set(range(1, num_atoms + 1)): 

79 base = 1 

80 else: 

81 raise ValueError( 

82 f"SMILES {mapped_smiles!r} has no usable atom-map numbers -- " 

83 "expected a 0-based or 1-based permutation of its atom count " 

84 f"({num_atoms}), got map numbers {sorted(map_nums)}." 

85 ) 

86 

87 reduced = Chem.RemoveHs(mol) 

88 

89 new_local_index: dict[int, int] = {} 

90 survivors: set[int] = set() 

91 for new_idx, atom in enumerate(reduced.GetAtoms()): 

92 old_local = atom.GetAtomMapNum() - base 

93 new_local_index[old_local] = new_idx 

94 survivors.add(old_local) 

95 

96 surviving_map_nums = {a.GetAtomMapNum() for a in reduced.GetAtoms()} 

97 for atom in mol.GetAtoms(): 

98 if atom.GetAtomicNum() != 1 or atom.GetAtomMapNum() in surviving_map_nums: 

99 continue 

100 neighbors = atom.GetNeighbors() 

101 if len(neighbors) != 1: 

102 raise ValueError( 

103 f"expected exactly one neighbor for a hydrogen RemoveHs " 

104 f"actually removed, got {len(neighbors)} (atom map num " 

105 f"{atom.GetAtomMapNum()}, smiles {mapped_smiles!r})" 

106 ) 

107 old_local = atom.GetAtomMapNum() - base 

108 neighbor_old_local = neighbors[0].GetAtomMapNum() - base 

109 new_local_index[old_local] = new_local_index[neighbor_old_local] 

110 

111 for atom in reduced.GetAtoms(): 

112 atom.SetAtomMapNum(atom.GetIdx() + base) 

113 new_mapped_smiles = Chem.MolToSmiles(reduced) 

114 

115 return new_local_index, frozenset(survivors), new_mapped_smiles