Skip to content

speed up getMasses by looking up each distinct element once - #2265

Open
steps-re wants to merge 2 commits into
prody:mainfrom
steps-re:perf/vectorize-getmasses
Open

speed up getMasses by looking up each distinct element once#2265
steps-re wants to merge 2 commits into
prody:mainfrom
steps-re:perf/vectorize-getmasses

Conversation

@steps-re

@steps-re steps-re commented Aug 7, 2026

Copy link
Copy Markdown

disclosing up front, same as #2262: put together with help from claude, then reproduced, measured and verified by me before opening.

two performance fixes, both in code that runs on every structure load. no behaviour change in either.

1. getMasses looks up every atom individually

getMasses in prody/utilities/misctools.py loops over every atom, calls str.capitalize() twice per atom and does a dict lookup each time. but a structure only contains a handful of distinct element symbols, so nearly all of that work is repeated. switched it to resolve each distinct symbol once and map the result back onto the atoms with unique(..., return_inverse=True).

this is on the path for every pdb, mmcif and mmtf load (pdbfile.py, ciffile.py, mmtffile.py all call it right after parsing elements), so it shows up on essentially any structure a user opens.

measured on my machine, output identical in every case:

structure before after
1,079 atoms 0.36 ms 0.03 ms 11.1x
12,750 atoms (antibody) 3.59 ms 0.27 ms 13.3x
100,000 atoms 43.4 ms 5.1 ms 8.5x
1,000,000 atoms 292 ms 45 ms 6.4x

the unknown-symbol case still returns 0.0 rather than raising, the single-string case still returns a scalar, and case-insensitive matching is unchanged.

2. max_len loop in setData

prody/atomic/atomgroup.py finds the longest chain identifier with an explicit python loop. max(map(len, array), default=0) does the same in C, and the default covers the empty case the loop handled by leaving max_len at 0. small, but it runs on every chain assignment.

verification

added TestGetMasses to prody/tests/utilities/test_misctools.py: 7 tests covering known elements, case insensitivity, unknown symbols, repeated symbols (which is what would break if the unique/inverse mapping were wrong), empty input, single-string input and array input. i checked these have teeth by deliberately breaking the mapping and confirming they fail.

ran prody/tests/atomic and prody/tests/utilities: 618 passed. prody/tests/proteins has 16 failures and 1 error on my machine, but they are identical on unpatched main, so they are pre-existing and unrelated to this change (mostly test_insty and a network-dependent mmtf test).

happy to split these into two prs if you'd rather take them separately.

mike

steps-re and others added 2 commits August 7, 2026 19:17
getMasses looped over every atom, calling str.capitalize twice and
doing a dictionary lookup each time, although a structure contains only
a handful of distinct element symbols. The lookup is now done once per
distinct symbol and mapped back onto the atoms.

getMasses runs on every PDB, mmCIF and MMTF structure load. Measured
6-13x faster with identical output: 3.59 -> 0.27 ms for a 12750 atom
antibody, 291 -> 45 ms for a million atoms.

Co-Authored-By: Claude <noreply@anthropic.com>
max(map(len, array), default=0) is equivalent and does the comparison
in C. The default also covers the empty case, which the loop handled by
leaving max_len at 0.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant