perf: query the KD-tree once per structure in _identify_bonds

[JustinKyleKirkland]

Description

_identify_bonds loops over every atom in Python and issues one KDTree.query call per atom, then builds a list comprehension of cutoff distances for each one. On an 80k atom receptor that is 80k round trips through scipy. Profiling PDBQTReceptor(pdbqt_string) on an 80k atom input puts it at 1.738 s of a 2.441 s total, so roughly 70 percent of the time spent building a receptor is in this one function.

cKDTree.query already accepts the whole coordinate array at once. This changes the function to make a single query for the entire structure with workers=-1, and turns the radius lookup and cutoff comparison into array operations. Only the final dict assembly stays in Python:

distances, indices = KDTree.query(positions, k=k, workers=-1)
neighbors = indices[:, 1:]
r_cov = np.array([covalent_radius[autodock4_atom_types_elements[t]] for t in atom_types])
optimal_distances = bond_allowance_factor * (r_cov[:, None] + r_cov[neighbors])
is_bonded = distances[:, 1:] < optimal_distances

The function is duplicated verbatim in molecule_pdbqt.py and receptor_pdbqt.py, so both copies get the same change and remain identical to each other. Deduplicating them is worth doing separately and is not attempted here.

No linked issue.

Timing

_identify_bonds on synthetic structures at protein-like density, same machine, best of a warmed run:

atoms	develop	this PR	speedup
2,000	0.030 s	0.004 s	6.9x
10,000	0.160 s	0.022 s	7.3x
20,000	0.345 s	0.041 s	8.4x
40,000	0.752 s	0.081 s	9.3x
80,000	1.579 s	0.165 s	9.6x

End to end, the whole PDBQTReceptor(pdbqt_string) call, which also includes parsing and get_atom_indices_by_residue that this PR does not touch:

atoms	develop	this PR	speedup
2,000	0.045 s	0.017 s	2.7x
20,000	0.480 s	0.148 s	3.2x
80,000	2.265 s	0.755 s	3.0x

Measured on Python 3.13.11, NumPy 2.5.1, SciPy 1.18.0, 8 cores. Batching alone accounts for most of it and workers=-1 contributes roughly another 3.8x on the query itself, so the gain will be smaller on a single core machine but the batching part holds regardless.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Other (please describe): performance, with one small behaviour change noted below

Checklist:

Please go through the following checklist before submitting the PR:

• Written a description of the feature or bug fix above.
• Ran pytest locally and all tests have passed.
• If applicable, documentation was updated with new feature.
• Docstring included for any new classes/functions/methods, or for significantly modifed existing functions.
• (Optional) new test added to account for new feature.

Additional Information

Output is unchanged. I checked the new implementation against a verbatim copy of the current one on 40 random structures spanning 2 to 400 atoms and varied density and composition, on 20 non-contiguous index subsets of the kind the annotation lists produce, and on all five pdbqt files in test/rdkitmol_from_docking_data. Every bond dict came back identical.

One behaviour change is worth flagging. A single atom input previously raised TypeError: 'int' object is not subscriptable in receptor_pdbqt, because query with k=1 returns scalars rather than arrays and the slicing could not handle them. It now returns no bonds. molecule_pdbqt already returned no bonds for that case through its own early return, so this makes the two agree, but it does turn a crash into a quiet empty result.

workers= was added in SciPy 1.6, released December 2020. requirements.txt previously pinned no versions at all, so this adds scipy>=1.6 as the one bound this change actually needs, rather than trying to pin everything.

Adds test/identify_bonds_test.py with 7 tests, run against both copies of the function: two and three atom geometries with known bonding, a methane-like case, non-contiguous indices, single atom and empty input, and agreement with a brute force reference over 20 random structures. Six of the seven also pass against develop, which is the point, they characterise existing behaviour rather than the new implementation. The seventh is the single atom case described above. Full suite goes from 94 passed / 4 skipped to 101 passed / 4 skipped.

[joanimato]

Very good, thanks!