nx_chem_morgan.nx source
↩ module page · 379 lines · 14255 B
1// nx_chem_morgan.nx -- C2.3c milestone: Morgan extended-connectivity
2// canonical atom ranking.
3//
4// Implements Weisfeiler-Lehman (1-WL) partition refinement to compute
5// a canonical rank per atom that is invariant under graph isomorphism.
6// Two MolGraphs representing the same molecule produce the same set
7// of canonical ranks (modulo atom-index relabeling), so canonical
8// SMILES emit using these ranks yields byte-identical output for
9// graph-equivalent inputs.
10//
11// Algorithm (after Weininger 1989 + modern WL formulation):
12// 1. Initial invariant per atom: packed (z, degree, charge,
13// isotope, h_count, aromaticity) into i64.
14// 2. Sort by initial invariant, assign DENSE class indices 0..k-1.
15// All atoms with same packed invariant get same class.
16// 3. Refinement step: for each atom, signature =
17// (current_class, sorted_neighbor_classes). Sort by signature,
18// assign new dense class indices.
19// 4. Iterate refinement until class count stable (or MAX_ITER).
20// 5. Final rank = position in sorted (final_class, atom_index) order.
21//
22// EXCEED axis hit (landscape doc):
23// E1 -- bit-reproducible canonical form across versions/platforms.
24// The Morgan algorithm here is content-addressed via versioned
25// initial-invariant packing + deterministic signature sort.
26// RDKit's BackwardsIncompatibleChanges.html documents
27// cross-release canonical drift; this substrate's canonical is
28// pinned to nx_chem_morgan_canonical_rank_v1 semantics
29// (hash-pinned at the function level by source identity).
30//
31// Honest gaps (deferred):
32// - Bond orders not yet used in invariants (so cyclohexane vs
33// cyclohexene atoms could collide on initial invariant; one WL
34// refinement step resolves it via neighbor distinctness)
35// - Stereo not yet folded into canonicalization (C2.3d)
36// - Tie-breaking for true graph automorphisms uses (final_class,
37// atom_idx); not a true automorphism canonicalizer but adequate
38// for graph-equivalent input alternatives
39
40import "nx_chem.nx"
41import "nx_chem_molecule.nx"
42
43const NX_MORGAN_MAX_DEG: nx_int = 7 // max neighbors per atom (organic typical <= 4)
44const NX_MORGAN_MAX_ITER: nx_int = 32 // safety cap for refinement loop
45const NX_MORGAN_SIG_LEN: nx_int = 8 // 1 (own class) + 7 (neighbor classes) = 8 nx_int per row
46
47// =================================================================
48// Pack initial invariant for one atom into i64.
49// Layout (LSB first):
50// bits 0-7 : z (atomic number, 0..118)
51// bits 8-11 : degree (0..15; organic typical <= 4)
52// bits 12-15 : charge + 8 (encoding -8..+7)
53// bits 16-19 : h_count + 2 (encoding -2..+13; -1 sentinel as 1)
54// bits 20-27 : isotope (0..255)
55// bit 28 : aromaticity flag (any bit of aromaticity tuple set)
56// =================================================================
57func nx_chem_morgan_initial_invariant(m: *MolGraph, atom_idx: nx_int) -> nx_int {
58 let a: *Atom = ((m.atoms as nx_int) + (atom_idx * NX_ATOM_BYTES)) as *Atom
59 // compute degree
60 var deg: nx_int = 0
61 var bi: nx_int = 0
62 while bi < m.n_bonds {
63 let b: *Bond = ((m.bonds as nx_int) + (bi * NX_BOND_BYTES)) as *Bond
64 if b.a == atom_idx { deg = deg + 1 }
65 if b.b == atom_idx { deg = deg + 1 }
66 bi = bi + 1
67 }
68 var v: nx_int = a.z & 0xff
69 if deg > 15 { deg = 15 }
70 v = v | ((deg & 0xf) << 8)
71 var c: nx_int = a.charge + 8
72 if c < 0 { c = 0 }
73 if c > 15 { c = 15 }
74 v = v | ((c & 0xf) << 12)
75 var h: nx_int = a.h_count + 2
76 if h < 0 { h = 0 }
77 if h > 15 { h = 15 }
78 v = v | ((h & 0xf) << 16)
79 var iso: nx_int = a.isotope
80 if iso < 0 { iso = 0 }
81 if iso > 255 { iso = 255 }
82 v = v | ((iso & 0xff) << 20)
83 var arom: nx_int = 0
84 if a.aromaticity != 0 { arom = 1 }
85 v = v | ((arom & 0x1) << 28)
86 return v
87}
88
89// =================================================================
90// Get neighbor atom indices of atom_idx; writes into out (size >= MAX_DEG).
91// Returns count of neighbors (0..MAX_DEG).
92// =================================================================
93func nx_chem_morgan_get_neighbors(m: *MolGraph, atom_idx: nx_int, out: *nx_int) -> nx_int {
94 var count: nx_int = 0
95 var bi: nx_int = 0
96 while bi < m.n_bonds {
97 let b: *Bond = ((m.bonds as nx_int) + (bi * NX_BOND_BYTES)) as *Bond
98 if b.a == atom_idx {
99 if count < NX_MORGAN_MAX_DEG {
100 out[count] = b.b
101 count = count + 1
102 }
103 }
104 if b.b == atom_idx {
105 if count < NX_MORGAN_MAX_DEG {
106 out[count] = b.a
107 count = count + 1
108 }
109 }
110 bi = bi + 1
111 }
112 return count
113}
114
115// =================================================================
116// Insertion sort an nx_int array ascending.
117// =================================================================
118func nx_chem_morgan_sort_asc(arr: *nx_int, n: nx_int) -> nx_int {
119 var i: nx_int = 1
120 while i < n {
121 let key: nx_int = arr[i]
122 var j: nx_int = i - 1
123 var done: nx_int = 0
124 while done == 0 {
125 if j < 0 { done = 1 }
126 else {
127 if arr[j] <= key { done = 1 }
128 else {
129 arr[j + 1] = arr[j]
130 j = j - 1
131 }
132 }
133 }
134 arr[j + 1] = key
135 i = i + 1
136 }
137 return 0
138}
139
140// =================================================================
141// Compute signature row for atom_idx based on current classes.
142// sig_row format: [own_class, sorted_neighbor_class_0, ..., -1 pad].
143// =================================================================
144func nx_chem_morgan_signature(
145 m: *MolGraph, atom_idx: nx_int, class_curr: *nx_int, sig_row: *nx_int
146) -> nx_int {
147 let neighbors: *nx_int = (sys_mmap((NX_MORGAN_MAX_DEG * 8) as i64)) as *nx_int
148 let neigh_classes: *nx_int = (sys_mmap((NX_MORGAN_MAX_DEG * 8) as i64)) as *nx_int
149 let n_neigh: nx_int = nx_chem_morgan_get_neighbors(m, atom_idx, neighbors)
150 var ni: nx_int = 0
151 while ni < n_neigh {
152 neigh_classes[ni] = class_curr[neighbors[ni]]
153 ni = ni + 1
154 }
155 let _s: nx_int = nx_chem_morgan_sort_asc(neigh_classes, n_neigh)
156 sig_row[0] = class_curr[atom_idx]
157 var k: nx_int = 0
158 while k < n_neigh {
159 sig_row[1 + k] = neigh_classes[k]
160 k = k + 1
161 }
162 var pad: nx_int = n_neigh
163 while pad < NX_MORGAN_MAX_DEG {
164 sig_row[1 + pad] = -1
165 pad = pad + 1
166 }
167 return 0
168}
169
170// =================================================================
171// Lexicographic compare of two signature rows.
172// Returns -1 if a<b, 0 if equal, 1 if a>b.
173// =================================================================
174func nx_chem_morgan_sig_cmp(sig_a: *nx_int, sig_b: *nx_int) -> nx_int {
175 var i: nx_int = 0
176 while i < NX_MORGAN_SIG_LEN {
177 if sig_a[i] < sig_b[i] { return -1 }
178 if sig_a[i] > sig_b[i] { return 1 }
179 i = i + 1
180 }
181 return 0
182}
183
184// =================================================================
185// Sort indices array such that signatures[indices[*]] are in
186// ascending lexicographic order. Insertion sort; signatures rows are
187// at offset (indices[k] * NX_MORGAN_SIG_LEN * 8) bytes into sigs.
188// =================================================================
189func nx_chem_morgan_sort_indices_by_sig(indices: *nx_int, n: nx_int, sigs: *nx_int) -> nx_int {
190 var i: nx_int = 1
191 while i < n {
192 let key: nx_int = indices[i]
193 let key_sig: *nx_int = ((sigs as nx_int) + (key * NX_MORGAN_SIG_LEN * 8)) as *nx_int
194 var j: nx_int = i - 1
195 var done: nx_int = 0
196 while done == 0 {
197 if j < 0 { done = 1 }
198 else {
199 let j_sig: *nx_int = ((sigs as nx_int) + (indices[j] * NX_MORGAN_SIG_LEN * 8)) as *nx_int
200 let cmp: nx_int = nx_chem_morgan_sig_cmp(j_sig, key_sig)
201 if cmp > 0 {
202 indices[j + 1] = indices[j]
203 j = j - 1
204 }
205 else { done = 1 }
206 }
207 }
208 indices[j + 1] = key
209 i = i + 1
210 }
211 return 0
212}
213
214// =================================================================
215// One refinement step: compute signatures from class_curr, sort
216// atom indices by signature, assign dense new class indices.
217// Returns number of distinct classes after this refinement.
218// =================================================================
219func nx_chem_morgan_refine(
220 m: *MolGraph, class_curr: *nx_int, class_next: *nx_int,
221 sigs: *nx_int, indices: *nx_int
222) -> nx_int {
223 var i: nx_int = 0
224 while i < m.n_atoms {
225 let row: *nx_int = ((sigs as nx_int) + (i * NX_MORGAN_SIG_LEN * 8)) as *nx_int
226 let _s: nx_int = nx_chem_morgan_signature(m, i, class_curr, row)
227 indices[i] = i
228 i = i + 1
229 }
230 let _o: nx_int = nx_chem_morgan_sort_indices_by_sig(indices, m.n_atoms, sigs)
231 var cur_class: nx_int = 0
232 if m.n_atoms > 0 {
233 class_next[indices[0]] = 0
234 }
235 var k: nx_int = 1
236 while k < m.n_atoms {
237 let prev_idx: nx_int = indices[k - 1]
238 let curr_idx: nx_int = indices[k]
239 let prev_sig: *nx_int = ((sigs as nx_int) + (prev_idx * NX_MORGAN_SIG_LEN * 8)) as *nx_int
240 let curr_sig: *nx_int = ((sigs as nx_int) + (curr_idx * NX_MORGAN_SIG_LEN * 8)) as *nx_int
241 let cmp: nx_int = nx_chem_morgan_sig_cmp(prev_sig, curr_sig)
242 if cmp != 0 { cur_class = cur_class + 1 }
243 class_next[curr_idx] = cur_class
244 k = k + 1
245 }
246 if m.n_atoms == 0 { return 0 }
247 return cur_class + 1
248}
249
250// =================================================================
251// Convert packed initial invariants to dense class indices.
252// Same logic as morgan_refine but using raw invariants instead of
253// signatures (no neighbor consideration).
254// =================================================================
255func nx_chem_morgan_densify_initial(
256 invariants: *nx_int, n: nx_int, class_out: *nx_int, indices: *nx_int
257) -> nx_int {
258 var i: nx_int = 0
259 while i < n {
260 indices[i] = i
261 i = i + 1
262 }
263 // Insertion sort indices by invariant value
264 var j: nx_int = 1
265 while j < n {
266 let key: nx_int = indices[j]
267 let key_v: nx_int = invariants[key]
268 var k: nx_int = j - 1
269 var done: nx_int = 0
270 while done == 0 {
271 if k < 0 { done = 1 }
272 else {
273 if invariants[indices[k]] <= key_v { done = 1 }
274 else {
275 indices[k + 1] = indices[k]
276 k = k - 1
277 }
278 }
279 }
280 indices[k + 1] = key
281 j = j + 1
282 }
283 // Assign dense classes
284 var cur: nx_int = 0
285 if n > 0 { class_out[indices[0]] = 0 }
286 var p: nx_int = 1
287 while p < n {
288 let prev_v: nx_int = invariants[indices[p - 1]]
289 let curr_v: nx_int = invariants[indices[p]]
290 if prev_v != curr_v { cur = cur + 1 }
291 class_out[indices[p]] = cur
292 p = p + 1
293 }
294 if n == 0 { return 0 }
295 return cur + 1
296}
297
298// =================================================================
299// Main entry: compute canonical rank for each atom in m.
300// rank_out[i] = rank of atom i in canonical order, where rank 0 is
301// the canonical-DFS root (lowest in the final sorted order).
302// =================================================================
303func nx_chem_morgan_canonical_rank(m: *MolGraph, rank_out: *nx_int) -> nx_int {
304 if m.n_atoms == 0 { return 0 }
305 let invariants: *nx_int = (sys_mmap((m.n_atoms * 8) as i64)) as *nx_int
306 let class_a: *nx_int = (sys_mmap((m.n_atoms * 8) as i64)) as *nx_int
307 let class_b: *nx_int = (sys_mmap((m.n_atoms * 8) as i64)) as *nx_int
308 let sigs: *nx_int = (sys_mmap((m.n_atoms * NX_MORGAN_SIG_LEN * 8) as i64)) as *nx_int
309 let indices: *nx_int = (sys_mmap((m.n_atoms * 8) as i64)) as *nx_int
310 // Step 1: initial packed invariants
311 var i: nx_int = 0
312 while i < m.n_atoms {
313 invariants[i] = nx_chem_morgan_initial_invariant(m, i)
314 i = i + 1
315 }
316 // Step 2: densify initial invariants into class_a
317 var class_count: nx_int = nx_chem_morgan_densify_initial(invariants, m.n_atoms, class_a, indices)
318 // Step 3: iterative refinement
319 var iter: nx_int = 0
320 var done: nx_int = 0
321 while iter < NX_MORGAN_MAX_ITER {
322 if done == 0 {
323 let new_count: nx_int = nx_chem_morgan_refine(m, class_a, class_b, sigs, indices)
324 if new_count == class_count {
325 done = 1
326 }
327 else {
328 class_count = new_count
329 var c: nx_int = 0
330 while c < m.n_atoms {
331 class_a[c] = class_b[c]
332 c = c + 1
333 }
334 }
335 }
336 iter = iter + 1
337 }
338 // Step 4: final sort to determine rank.
339 // After convergence, class_a holds the canonical class per atom.
340 // Sort atoms by (class, atom_idx) ascending and use sort position as rank.
341 var r: nx_int = 0
342 while r < m.n_atoms {
343 indices[r] = r
344 r = r + 1
345 }
346 // Insertion sort indices by (class_a[indices[k]], indices[k])
347 var s: nx_int = 1
348 while s < m.n_atoms {
349 let key: nx_int = indices[s]
350 let key_c: nx_int = class_a[key]
351 var t: nx_int = s - 1
352 var done: nx_int = 0
353 while done == 0 {
354 if t < 0 { done = 1 }
355 else {
356 let prev_c: nx_int = class_a[indices[t]]
357 var shift: nx_int = 0
358 if prev_c > key_c { shift = 1 }
359 if prev_c == key_c {
360 if indices[t] > key { shift = 1 }
361 }
362 if shift == 1 {
363 indices[t + 1] = indices[t]
364 t = t - 1
365 }
366 else { done = 1 }
367 }
368 }
369 indices[t + 1] = key
370 s = s + 1
371 }
372 // Assign rank
373 var q: nx_int = 0
374 while q < m.n_atoms {
375 rank_out[indices[q]] = q
376 q = q + 1
377 }
378 return class_count
379}