ENH: Change `get_hbond_map` to use KDTree for distance search #5182

BradyAJohnston · 2025-12-22T10:05:21Z

Changes made in this Pull Request:

This updates the get_hbond_map function to use MDAnalysis.lib.distances.capped_distance to speed up distance calculations. Nothing changes in the API, but rather than computing a pairwise distance matrix of every residue against every other residue, we can first construct a KDTree for both the N and O atoms and use that for finding relevant pairs within a certain cutoff distance. The returned value from get_hbond_map remains the same but we have avoided having to compute distances and energies between residues that would never form an HBond to begin with.

In my benchmarks running just the get_hbond_map function this produces a speedup from 2-5x, with that difference continuing to increase with size of the protein system:

`develop`

Scaling benchmark:
Residues     Time (ms)       Time/Res (μs)
------------------------------------------------------------
50           0.175           3.505
100          0.559           5.593
150          1.201           8.007
200          2.537           12.686
250          3.966           15.864
300          5.882           19.608
============================================================

This PR

Scaling benchmark:
Residues     Time (ms)       Time/Res (μs)
------------------------------------------------------------
50           0.077           1.541
100          0.164           1.638
150          0.297           1.980
200          0.490           2.450
250          0.712           2.849
300          1.161           3.870
============================================================

PR Checklist

Issue raised/referenced?
Tests updated/added?
Documentation updated/added?
package/CHANGELOG file updated?
Is your name in package/AUTHORS? (If it is not, add it!)

Developers Certificate of Origin

I certify that I can submit this code contribution as described in the Developer Certificate of Origin, under the MDAnalysis LICENSE.

📚 Documentation preview 📚: https://mdanalysis--5182.org.readthedocs.build/en/5182/

github-actions

Hello there first time contributor! Welcome to the MDAnalysis community! We ask that all contributors abide by our Code of Conduct and that first time contributors introduce themselves on GitHub Discussions so we can get to know you. You can learn more about participating here. Please also add yourself to package/AUTHORS as part of this PR.

codecov · 2025-12-22T10:25:56Z

Codecov Report

❌ Patch coverage is 89.65517% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.71%. Comparing base (9fb8b0a) to head (4faa47c).
⚠️ Report is 1 commits behind head on develop.

Files with missing lines	Patch %	Lines
package/MDAnalysis/analysis/dssp/dssp.py	50.00%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #5182      +/-   ##
===========================================
- Coverage    92.72%   92.71%   -0.02%     
===========================================
  Files          180      180              
  Lines        22472    22480       +8     
  Branches      3188     3191       +3     
===========================================
+ Hits         20838    20842       +4     
- Misses        1176     1179       +3     
- Partials       458      459       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

package/CHANGELOG

IAlibay · 2025-12-22T19:57:31Z

package/MDAnalysis/analysis/dssp/pydssp_numpy.py

-    hmap = np.tile(h_1, (1, 1, n)).reshape(n, n, 3)
+    # We can use KDTree to find candidate pairs within cutoff distance without having to
+    # compute a full pairwise distance matrix, much faster especially for larger proteins
+    n_kdtree = KDTree(n_atoms)


I know the original code was vendored from somewhere else, so it explains the heavy reliance on linalg.norm and the lack of PBC awareness - however, now that we're making changes, should we be using one of our lib.distances calls (maybe capped_distance) instead?

I'm not at all experienced with lib.distances are you talking about for the KDTree / potential pair generation or for the actual distance computation? If just for computing distances then on the KDTree generation we would already be discarding some potential across-boundary bonds.

We should be able to provide simple periodic checking with KDTree you can specify a boxsize parameter. But this wouldn't take into account the more complex boundaries we might come across?

Okay I just dug in and found capped_distance which should do exactly what we are after with better periodic boundary checking.

The performance isn't as good but also as expected with more checks. Still a large improvement over current.

============================================================ Scaling benchmark: Residues Time (ms) Time/Res (μs) ------------------------------------------------------------ 50 0.142 2.835 100 0.291 2.906 150 0.489 3.257 200 0.736 3.682 250 1.056 4.226 300 1.629 5.429 ============================================================

I've currently updated the pair search to use capped_distances - but the actual distance computation still isn't taking into account periodic boundaries. Is there a periodic boundary aware distance calculation that doesn't do "everything vs everything" which distance_array seems to do? Ideally we would avoid repeatedly calling capped_distance as we have already constructed the KDTree once and can just use those results for computation later on, but still taking into account periodic boundary crossing.

Pair-wise distances between two equal groups of atoms can be calculated with MDAnalysis.lib.distances.calc_bonds.

marinegor · 2025-12-23T17:55:48Z

@BradyAJohnston wow, that's super impressive results, thanks! Some comments:

but the actual distance computation still isn't taking into account periodic boundaries. Is there a periodic boundary aware distance calculation that doesn't do "everything vs everything" which distance_array seems to do?

Just to make sure I understand this correctly: the current implementation is plain wrong when the boundary condition is crossed, since it just computes the wrong distance, right? I mean this part:

    # compute distances and energy as previously but only for the potential pairs
    # still returning the same energy matrix that would have otherwise been made
    o_indices = pairs[:, 1]
    n_indices = pairs[:, 0]
    d_on = np.linalg.norm(o_atoms[o_indices] - n_atoms[n_indices], axis=-1)
    d_ch = np.linalg.norm(c_atoms[o_indices] - h_1[n_indices], axis=-1)
    d_oh = np.linalg.norm(o_atoms[o_indices] - h_1[n_indices], axis=-1)
    d_cn = np.linalg.norm(c_atoms[o_indices] - n_atoms[n_indices], axis=-1)

so we're looking for periodic contition aware np.linalg.norm replacement? Or can we just put capped_distance or smth instead of -, like

    d_on = np.linalg.norm(capped_distance(o_atoms[o_indices], n_atoms[n_indices]))
    ...

When defaulting to self._trajectory.ts.dimensions, I'd perhaps consider changing / removing it if the structure comes from RCSB -- cryoEM structures default to some hardcoded unit cell which can even be too small, so it might break things.

BradyAJohnston · 2025-12-23T18:09:10Z

Hi @marinegor yes if DSSP is currently run on a protein that has secondary structure over the periodic cell then it would incorrectly not assign it.

We could just have multiple capped_distance calls which would solve the issue with periodic crossing, but each time we would be re-building a KDTree when ideally we can just re-use the pairs from the first capped_distance call which gives the relevant pairs to check already.

If there was just a periodic_distance(c_atoms[o_indices] - h_1[n_indices]) that we can call then we compute the distance and allow for periodic crossing.

In the current setup while we do detect periodic bonds to check - the results end up being wrong like you say.

marinegor · 2025-12-23T18:24:05Z

If there was just a periodic_distance(c_atoms[o_indices] - h_1[n_indices]) that we can call then we compute the distance and allow for periodic crossing.

perhaps there are some functions in distopia (namely, distopia.distances_ortho that you could use directly, or am I wrong: https://www.mdanalysis.org/distopia/api)

In the current setup while we do detect periodic bonds to check - the results end up being wrong like you say.

do you see this being a separate issue, perhaps if you don't find a quick solution to use among distopia library? I feel like this PR doesn't have to be blocked by this discussion.

IAlibay · 2025-12-23T18:34:14Z

If there was just a periodic_distance(c_atoms[o_indices] - h_1[n_indices]) that we can call then we compute the distance and allow for periodic crossing.

@BradyAJohnston have a look at minimize_vectors https://github.com/MDAnalysis/mdanalysis/blob/develop/package/MDAnalysis/lib/distances.py#L2106

IAlibay · 2025-12-23T18:36:17Z

If there was just a periodic_distance(c_atoms[o_indices] - h_1[n_indices]) that we can call then we compute the distance and allow for periodic crossing.

perhaps there are some functions in distopia (namely, distopia.distances_ortho that you could use directly, or am I wrong: https://www.mdanalysis.org/distopia/api)

Please don't try to use distopia directly, it's meant to be an (optional for now) backend replacement.

BradyAJohnston · 2025-12-23T18:41:43Z

Thanks for the suggestion @IAlibay - this would then just become as such:

    d_on = np.linalg.norm(minimize_vectors(o_atoms[o_indices] - n_atoms[n_indices], box=box), axis=-1)
    d_ch = np.linalg.norm(minimize_vectors(c_atoms[o_indices] - h_1[n_indices], box=box), axis=-1)
    d_oh = np.linalg.norm(minimize_vectors(o_atoms[o_indices] - h_1[n_indices], box=box), axis=-1)
    d_cn = np.linalg.norm(minimize_vectors(c_atoms[o_indices] - n_atoms[n_indices], box=box), axis=-1)

The usage of np.linalg.norm should now be fine we are taking into account the periodic crossing?

marinegor · 2025-12-23T18:47:35Z

The usage of np.linalg.norm should now be fine we are taking into account the periodic crossing?

I guess so, yeah -- it's just a norm calculation, there should not be any problem with that as long as the vector itself is boundary-corrected.

As for cryoEM default, the placeholder unit cell is (1, 1, 1, 90, 90, 90), so you can explicitly compare with that and e.g. default to no box calculations if it's set.

UPD: added this remark to changes request.

marinegor · 2025-12-23T18:48:49Z

package/MDAnalysis/analysis/dssp/dssp.py

+        dssp = assign(
+            coords,
+            donor_mask=self._donor_mask,
+            box=self._trajectory.ts.dimensions,


if dimensions are (1, 1, 1, 90, 90, 90) (cryoEM placeholder), just set to None.

This shouldn't be done here, but rather when loading the PDB - there is an a bad edge case where you might actually just have that unit box.

Indeed we do this already: https://github.com/MDAnalysis/mdanalysis/blob/develop/package/MDAnalysis/coordinates/PDB.py#L228-L231

On further thought, I wouldn't be completely against this check - maybe a warning would be ok too if we don't want to be too strict

Not 100% with how this should be handled. If I am raising a warning in _single_frame won't that warning be raised a bunch of times for every frame?

def _single_frame(self): coords = self._get_coords() DEFAULT_BOX_DIMENSIONS = (1, 1, 1, 90, 90, 90) box = self._trajectory.ts.dimensions if box == DEFAULT_BOX_DIMENSIONS: box = None warnings.warn( f"Default box dimensions {DEFAULT_BOX_DIMENSIONS} are treated as no box" ) dssp = assign( coords, donor_mask=self._donor_mask, box=self._trajectory.ts.dimensions, ) self.results.dssp_ndarray.append(dssp)

@BradyAJohnston - how about making that warning happen in init? The atomgroups are set there, so it would be best to do it when the class is created.

IAlibay · 2025-12-23T18:54:50Z

    d_on = np.linalg.norm(minimize_vectors(o_atoms[o_indices] - n_atoms[n_indices], box=box), axis=-1)
    d_ch = np.linalg.norm(minimize_vectors(c_atoms[o_indices] - h_1[n_indices], box=box), axis=-1)
    d_oh = np.linalg.norm(minimize_vectors(o_atoms[o_indices] - h_1[n_indices], box=box), axis=-1)
    d_cn = np.linalg.norm(minimize_vectors(c_atoms[o_indices] - n_atoms[n_indices], box=box), axis=-1)

The usage of np.linalg.norm should now be fine we are taking into account the periodic crossing?

Yep, don't trust me fully (I'm sleep deprived), but that sounds reasonable.

marinegor · 2025-12-23T19:27:45Z

One thing, it probably shouldn't happen in `_single_frame` but rather in `__init__`, otherwise i'm ok with a warning.

BradyAJohnston · 2025-12-29T17:41:59Z

package/MDAnalysis/analysis/dssp/dssp.py

+        if np.array_equal(self._box, (1, 1, 1, 90, 90, 90)):
+            self._box = None
+            warnings.warn(
+                "Box dimensions are (1, 1, 1, 90, 90, 90), not using periodic boundary conditions in DSSP calculations"


Not 100% on what the warning should be - but this is how we are imagining it? @IAlibay

Yes that looks about right. With regards to the warning, your message is close enough to what I was thinking, i.e. something like "Default unit cells of 1 A cubed found. This likely comes from a cryo-em file. To avoid usses, we are disabling periodic boundary conditions in the DSSP calculation".

A warning is currently already raised twice: once when loading the universe and again when running the DSSP (not sure why it is happening for the run(). As the unit cell dimensions are already set to None it's not triggering the new warning. In most cases the unit cell wouldn't make it "all the way" to DSSP but I'll change it to just be the same warning message.

In [9]: p = Path("/Users/brady/Downloads/9N9M.pdb") In [10]: u = mda.Universe(p) /Users/brady/git/mdanalysis/package/MDAnalysis/coordinates/PDB.py:479: UserWarning: 1 A^3 CRYST1 record, this is usually a placeholder. Unit cell dimensions will be set to None. warnings.warn( In [11]: a = dssp.DSSP(u) In [12]: a.run() /Users/brady/git/mdanalysis/package/MDAnalysis/coordinates/PDB.py:479: UserWarning: 1 A^3 CRYST1 record, this is usually a placeholder. Unit cell dimensions will be set to None. warnings.warn( /Users/brady/git/mdanalysis/package/MDAnalysis/analysis/base.py:562: UserWarning: Reader has no dt information, set to 1.0 ps self.times[idx] = ts.time Out[12]: <MDAnalysis.analysis.dssp.dssp.DSSP at 0x110d4c910>

Yeah this is what I was trying to convey earlier in the conversation when I was saying that the PDB reader takes care of that edge case already. I'm not sure what the best approach is here, thoughts @marinegor ?

orbeckst · 2026-01-06T21:55:28Z

package/MDAnalysis/analysis/dssp/pydssp_numpy.py

+            return np.linalg.norm(x[o_indices] - y[n_indices], axis=-1)
+        else:
+            return np.linalg.norm(
+                minimize_vectors(x[o_indices] - y[n_indices], box=box),


or use directly MDAnalysis.lib.distances.calc_bonds

return calc_bonds(x[o_indices], y[n_indices], box=box)

Ah this is good to know! (and yeah I missed it from naming). Will update to use that instead of this combo.

orbeckst

Good idea to use our distance functions. (see comments on using calc_bonds, whose name betrays its flexibility and thus may not be widely know as the the function for PBC-aware distances.

I would add an explicit backend keyword argument to all calls of our distance functions (and a way for the user to also set the backend when invoking the class) so that it becomes possible to enable parallel/distopia backends.

orbeckst · 2026-01-06T22:05:42Z

@marinegor may I assign you as the PR shepherd? If you don't have the bandwidth, please unassign yourself.

BradyAJohnston added 6 commits December 20, 2025 20:08

initial port to KDTree based HBond map

8f25085

change distance_matrix -> query_ball_tree

a3f9ac4

tweak comments

617e7de

tweak commments and layout

234ad87

run black

93906df

add name to AUTHORS

f578aa9

github-actions bot reviewed Dec 22, 2025

View reviewed changes

BradyAJohnston requested a review from marinegor December 22, 2025 10:09

add changelog and return comment

d831ba2

BradyAJohnston mentioned this pull request Dec 22, 2025

Add DSSP support for trajectories BradyAJohnston/MolecularNodes#1072

Merged

BradyAJohnston added 2 commits December 22, 2025 11:50

add test for donor_mask

86a045d

run black

26c4ca8

IAlibay requested changes Dec 22, 2025

View reviewed changes

BradyAJohnston added 3 commits December 23, 2025 13:11

change KDTree -> capped_distances

2105bd6

run black

e81d2e2

move changelog entry

bdaadab

marinegor requested changes Dec 23, 2025

View reviewed changes

BradyAJohnston added 2 commits December 29, 2025 16:40

use before distance calculations

f0b2e23

warning for init of default box dimensions

4faa47c

BradyAJohnston commented Dec 29, 2025

View reviewed changes

orbeckst reviewed Jan 6, 2026

View reviewed changes

orbeckst assigned marinegor Jan 6, 2026

ENH: Change get_hbond_map to use KDTree for distance search #5182

Are you sure you want to change the base?

ENH: Change get_hbond_map to use KDTree for distance search #5182

Conversation

BradyAJohnston commented Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

develop

This PR

PR Checklist

Developers Certificate of Origin

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradyAJohnston Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

marinegor commented Dec 23, 2025

Uh oh!

BradyAJohnston commented Dec 23, 2025

Uh oh!

marinegor commented Dec 23, 2025

Uh oh!

IAlibay commented Dec 23, 2025

Uh oh!

IAlibay commented Dec 23, 2025

Uh oh!

BradyAJohnston commented Dec 23, 2025

Uh oh!

marinegor commented Dec 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IAlibay commented Dec 23, 2025

Uh oh!

marinegor commented Dec 23, 2025 via email

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ENH: Change `get_hbond_map` to use KDTree for distance search #5182

ENH: Change `get_hbond_map` to use KDTree for distance search #5182

BradyAJohnston commented Dec 22, 2025 •

edited

Loading

`develop`

codecov bot commented Dec 22, 2025 •

edited

Loading

BradyAJohnston Dec 22, 2025 •

edited

Loading

marinegor commented Dec 23, 2025 •

edited

Loading