Optimize Atom Mapping procedure

[kfir4444]

Atom Mapping Performance Improvements — optimize_am branch

Overview

Reduce runtime of the atom mapping procedure, utilzing algorithmic and programmatic solutions.

---

Change 1 — arc/species/zmat.py: Collinear triplet guard in determine_d_atoms()

Problem

zmat_nerf implements the SN-NeRF atom placement algorithm: given three already-placed reference atoms A, B, C, it places a new atom D at bond length r from C, bond angle θ(B,C,D), and dihedral φ(A,B,C,D). This is the inner loop of _add_nth_atom_to_coords in zmat.py. However, this doesn't work if co-linearity exist.

determine_d_atoms() could select a reference atom triplet (B, C, D) that is nearly
collinear. _add_nth_atom_to_coords() later computes cross(B−A, normalize(C−B)) for
NeRF placement; a near-zero cross product produces NaN coordinates and causes mapping
failures on linear-geometry intermediates.

Fix

Added _nerf_cross_sq() — a double-precision function computing
|cross(B−A, normalize(C−B))|². A threshold _NERF_CROSS_SQ_MIN = 0.01
(≈3.8° from linear at 1.5 Å bond) gates triplet acceptance:

if _nerf_cross_sq(coords, ...) >= _NERF_CROSS_SQ_MIN:
    d_atoms.append(i)
    break

The same guard is applied to the [n, 2, 1, 0] fallback path, with a scan
for a valid 4th reference atom when the fallback itself is collinear.

---

Change 2 — arc/species/_zmat_kernels.py + arc/species/vectors.py: C kernel wiring

One key bottleneck of atom mapping comes from vector computations:
The cost decomposition for C₈H₁₈ species to species mapping (cProfile, 0.770 s total, 1,284,756 calls)

Component	Calls	Cumul. time	Fraction
fix_dihedrals_by_backbone_mapping	2	0.764 s	99.2%
set_dihedral	20	0.651 s	84.5%
xyz_to_zmat	60	0.514 s	66.8%
_add_nth_atom_to_zmat	1,560	0.481 s	62.5%
calculate_param	10,357	0.377 s	48.9%
np.cross	5,644	0.191 s	24.8%
get_angle	7,597	0.190 s	24.7%
get_dihedral	1,438	0.153 s	19.9%
zmat_to_coords	60	0.128 s	16.6%
numpy.moveaxis	16,932	0.110 s	14.3%
order_atoms (inside copy)	4	0.091 s	11.8%
numpy.asarray	29,025	0.059 s	7.7%

Key optimizations of the vector coputations imperative for fast atom mapping. This is achieved by compiling kernels in C for these types of computations, and calling them via the _zmat_kernels.py (new file)

ctypes loader for the pre-compiled _zmat_c_kernels.so. Exposes:

• zmat_r_f32 — interatomic distance
• zmat_a_f32 — bond angle
• zmat_d_f32 — dihedral angle
• zmat_nerf — NeRF atom placement (SN-NeRF formula)

Sets available = True only when the shared library loads without error; all callers
gate on this flag so the code degrades gracefully when the .so is absent, to avoid compiling errors.

vectors.py

Fast paths in calculate_distance and calculate_dihedral_angle:

if _ck_available:
    return _ck.zmat_r_f32(c0[0], c0[1], c0[2], c1[0], c1[1], c1[2])

calculate_angle was left as numpy — the C kernel uses float32 which caused a
precision failure in test_get_zmat_param_value at the ~0.001° level. This can be debated offline.

Also added apply_rodrigues_rotation() used by Change 3. (source: rotation)

---

Change 3 — arc/mapping/engine.py: Rodrigues rotation replaces set_dihedral() round-trip

Problem

fix_dihedrals_by_backbone_mapping() corrected each torsion angle by calling
set_dihedral(), which performs a full xyz → zmat → modify → zmat → xyz round-trip.
This is O(N²) per torsion and O(N³) overall for large molecules.

Fix

Direct Cartesian rotation via Rodrigues' formula:

1. Identify downstream atoms using BFS from the far-side pivot (_downstream_atoms_adj).
2. Rotate only those atoms: R(u,δ) = I cosδ + (1−cosδ)u⊗u + sinδ[u]×
3. No zmat round-trip — O(N) per torsion, O(N²) overall.

apply_rodrigues_rotation(coords, axis_origin=coords[pivot_i],
                         axis_unit=axis_unit, angle_rad=delta,
                         indices=downstream)

Additional changes in this file:

• identify_superimposable_candidates: capped at _MAX_BACKBONE_CANDIDATES = 4
  (was unbounded) and deduplicates inline, removing the prune_identical_dicts post-pass. No apparent performence loss.

---

Change 4 — arc/species/converter.py: Fast path in order_atoms_in_mol_list()

Problem

For species with pre-ordered mol_list, order_atoms_in_mol_list ran VF2
for every resonance structure even when atom IDs already matched positionally.

Fix

if (len(mol.atoms) == len(ref_mol.atoms)
        and all(m.id == r.id and m.id != -1
                for m, r in zip(mol.atoms, ref_mol.atoms))):
    continue

Also changed order_mol_by_atom_map to use reordered.update(sort_atoms=False),
preventing implicit reordering after explicit atom positioning.

---

Change 5 — arc/family/family.py: Three optimizations in the family-matching pipeline

5a. Adjacency-list CPI cache in determine_possible_reaction_products_from_family

check_product_isomorphism (CPI) was called once per product_list from
generate_products. For large families (e.g. R_Addition on C10H8+H generates
2400 product_lists), most entries are graph-duplicates — same adjacency structure,
different atom-label mappings. The isomorphism result depends only on the graph,
so it is cached by adjacency-list key:

adj_key = tuple(m.to_adjacency_list() for m in template_mols)
if adj_key in iso_cache:
    if not iso_cache[adj_key]: continue
else:
    result = check_product_isomorphism(products=template_mols, p_species=p_species)
    iso_cache[adj_key] = result
    if not result: continue

For reaction C10H9 ↔ C10H8+H: 2400 CPI calls → 150 unique per family (16× reduction).

5b. Shared fwd_iso_cache/rev_iso_cache across family calls

When rmg_family_set='all' is requested, 112 families are iterated, many duplicated
(same label appears as both RMG and ARC copy). With a shared cache, the second occurrence
of any family hits entirely from cache (0 CPI evaluations).

For C16H11 isomerization: CPI calls 794 → 406, time 19.55s → 9.47s.

flip_reaction() is now called once before the loop instead of once per family.
The pre-build is wrapped in try/except to preserve the original graceful-skip
behaviour when to_smiles() fails (e.g. openbabel mocked in tests).

5c. mol_list search limit for large molecules

For molecules with more than 2 resonance structures AND more than 20 atoms,
generate_unimolecular_products limits the search to the first 2 Kekulé forms:

if len(mol_list) > 2 and len(mol_list[0].atoms) > 20:
    mol_list = mol_list[:2]

Additional resonance structures produce near-duplicate product graphs already
deduplicated by the adjacency-list cache at the CPI level. The VF2 cost of
generating them (3–4 s for a 27-atom polycyclic per family) is eliminated.

For C16H11, 27 atoms, 7 resonance structures: generate_products
for Intra_R_Add_Exocyclic drops from ~3s to ~0.9s per call.

---

Change 6 — arc/species/species.py: skip_conformers on scissors()/_scissors()

Added skip_conformers: bool = False parameter propagated to the generate_conformers
calls inside _scissors(). Allows mapping-layer callers to get the cut species
structure without triggering conformer generation.

---