Skip to content

ASA calculation performance tweaks #1148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
josemduarte merged 3 commits into from
Sep 1, 2026
+178 −101
Merged

ASA calculation performance tweaks #1148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
60d03c4
Pure array solution for sphere points. About 10% faster
josemduarte Aug 24, 2026
5c8cff0
Use parallel primitive arrays for neighbors in AsaCalculator
josemduarte Aug 26, 2026
1216380
Merge branch 'master' into jd/asa-perf
josemduarte Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
257 changes: 167 additions & 90 deletions biojava-structure/src/main/java/org/biojava/nbio/structure/asa/AsaCalculator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
import org.slf4j.LoggerFactory;

import javax.vecmath.Point3d;
import javax.vecmath.Vector3d;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
Expand Down Expand Up @@ -101,12 +100,56 @@ public void run() {
}
}

static class IndexAndDistance {
final int index;
final double dist;
IndexAndDistance(int index, double dist) {
this.index = index;
this.dist = dist;
/**
* The neighbors of a single atom, as parallel primitive arrays of neighbor atom indices and their distances to
* the central atom, sorted by increasing distance.
* <p>
* Parallel primitive arrays are used rather than an array of index-distance objects because there are ~30
* neighbors per atom: for a large structure that would mean millions of small short-lived objects.
*/
static class Neighbors {

/** The neighbor atom indices, ordered by increasing distance to the central atom */
final int[] indices;
/** The distances to the central atom, in increasing order and parallel to {@link #indices} */
final double[] dists;

private Neighbors(int[] indices, double[] dists) {
this.indices = indices;
this.dists = dists;
}

/**
* Creates a Neighbors from the first count elements of the given buffers, copying them to exact-size arrays
* and sorting them by increasing distance.
* @param indicesBuffer the neighbor indices, only the first count elements are used
* @param distsBuffer the neighbor distances, only the first count elements are used
* @param count the number of neighbors
* @return the sorted neighbors
*/
static Neighbors createSorted(int[] indicesBuffer, double[] distsBuffer, int count) {
int[] indices = Arrays.copyOf(indicesBuffer, count);
double[] dists = Arrays.copyOf(distsBuffer, count);
// Sorting by closest to farthest away neighbors achieves faster runtimes when checking for occluded
// sphere sample points in calcSingleAsa. This follows the ideas exposed in
// Eisenhaber et al, J Comp Chemistry 1994 (https://onlinelibrary.wiley.com/doi/epdf/10.1002/jcc.540160303)
// This is essential for performance: it brings down the number of occlusion checks to
// an average of n_sphere_points/10 per atom, producing ~ x4 performance gain overall.
// An insertion sort is used because the arrays are small (~30 elements on average) and because it avoids
// both the boxing of a comparator-based sort and the object allocation an index-distance array would need.
for (int i = 1; i < count; i++) {
double dist = dists[i];
int index = indices[i];
int j = i - 1;
while (j >= 0 && dists[j] > dist) {
dists[j + 1] = dists[j];
indices[j + 1] = indices[j];
j--;
}
dists[j + 1] = dist;
indices[j + 1] = index;
}
return new Neighbors(indices, dists);
}
}

Expand All @@ -116,9 +159,16 @@ static class IndexAndDistance {
private final double[] radii;
private final double probe;
private final int nThreads;
private Vector3d[] spherePoints;
/**
* The sphere points to sample, as a flat array of interleaved x,y,z coordinates (thus of size 3 x nSpherePoints).
* A flat array of primitives (rather than an array of Vector3d objects) is used for performance: it keeps the
* points contiguous in memory and avoids a pointer dereference per point in the innermost loop of
* {@link #calcSingleAsa(int)}.
*/
private double[] spherePoints;
private int nSpherePoints;
private double cons;
private IndexAndDistance[][] neighborIndices;
private Neighbors[] neighbors;

private boolean useSpatialHashingForNeighbors;

Expand Down Expand Up @@ -239,7 +289,8 @@ private void initSpherePoints(int nSpherePoints) {
logger.debug("Will use {} sphere points", nSpherePoints);

// initialising the sphere points to sample
spherePoints = generateSpherePoints(nSpherePoints);
this.nSpherePoints = nSpherePoints;
this.spherePoints = generateSpherePoints(nSpherePoints);

cons = 4.0 * Math.PI / nSpherePoints;
}
Expand Down Expand Up @@ -285,10 +336,10 @@ public double[] calculateAsas() {
long start = System.currentTimeMillis();
if (useSpatialHashingForNeighbors) {
logger.debug("Will use spatial hashing to find neighbors");
neighborIndices = findNeighborIndicesSpatialHashing();
neighbors = findNeighborIndicesSpatialHashing();
} else {
logger.debug("Will not use spatial hashing to find neighbors");
neighborIndices = findNeighborIndices();
neighbors = findNeighborIndices();
}
long end = System.currentTimeMillis();
logger.debug("Took {} s to find neighbors", (end-start)/1000.0);
Expand Down Expand Up @@ -334,109 +385,126 @@ void setUseSpatialHashingForNeighbors(boolean useSpatialHashingForNeighbors) {
* Returns list of 3d coordinates of points on a unit sphere using the
* Golden Section Spiral algorithm.
* @param nSpherePoints the number of points to be used in generating the spherical dot-density
* @return the array of points as Vector3d objects
* @return a flat array of interleaved x,y,z coordinates, of size 3 x nSpherePoints
*/
private Vector3d[] generateSpherePoints(int nSpherePoints) {
Vector3d[] points = new Vector3d[nSpherePoints];
private double[] generateSpherePoints(int nSpherePoints) {
double[] points = new double[3 * nSpherePoints];
double inc = Math.PI * (3.0 - Math.sqrt(5.0));
double offset = 2.0 / nSpherePoints;
for (int k=0;k<nSpherePoints;k++) {
double y = k * offset - 1.0 + (offset / 2.0);
double r = Math.sqrt(1.0 - y*y);
double phi = k * inc;
points[k] = new Vector3d(Math.cos(phi)*r, y, Math.sin(phi)*r);
points[3*k ] = Math.cos(phi)*r;
points[3*k + 1] = y;
points[3*k + 2] = Math.sin(phi)*r;
}
return points;
}

/**
* Returns the 2-dimensional array with neighbor indices for every atom.
* @return 2-dimensional array of size: n_atoms x n_neighbors_per_atom
* Tells whether the 2 given atoms are neighbors for the purposes of the ASA calculation, i.e. whether their
* probe-inflated spheres overlap.
* @param i the index of the first atom
* @param j the index of the second atom
* @param dist the distance between the 2 atoms
* @return true if they are neighbors, false otherwise
*/
private boolean areNeighbors(int i, int j, double dist) {
return dist < radii[i] + probe + probe + radii[j];
}

/**
* Returns the neighbors of every atom, sorted by increasing distance.
* @return array of size n_atoms
*/
IndexAndDistance[][] findNeighborIndices() {
Neighbors[] findNeighborIndices() {

// looking at a typical protein case, number of neighbours are from ~10 to ~50, with an average of ~30
int initialCapacity = 60;
Neighbors[] nbs = new Neighbors[atomCoords.length];

IndexAndDistance[][] nbsIndices = new IndexAndDistance[atomCoords.length][];
// looking at a typical protein case, number of neighbours are from ~10 to ~50, with an average of ~30.
// The buffers are reused for all atoms and grown on demand, so that only the exact-size arrays kept in the
// returned Neighbors objects are allocated per atom.
int[] indicesBuffer = new int[60];
double[] distsBuffer = new double[60];

for (int k=0; k<atomCoords.length; k++) {
double radius = radii[k] + probe + probe;

List<IndexAndDistance> thisNbIndices = new ArrayList<>(initialCapacity);
int count = 0;

for (int i = 0; i < atomCoords.length; i++) {
if (i == k) continue;

double dist = atomCoords[i].distance(atomCoords[k]);

if (dist < radius + radii[i]) {
thisNbIndices.add(new IndexAndDistance(i, dist));
if (areNeighbors(k, i, dist)) {
if (count == indicesBuffer.length) {
indicesBuffer = Arrays.copyOf(indicesBuffer, count * 2);
distsBuffer = Arrays.copyOf(distsBuffer, count * 2);
}
indicesBuffer[count] = i;
distsBuffer[count] = dist;
count++;
}
}

IndexAndDistance[] indicesArray = thisNbIndices.toArray(new IndexAndDistance[0]);
nbsIndices[k] = indicesArray;
nbs[k] = Neighbors.createSorted(indicesBuffer, distsBuffer, count);
}
return nbsIndices;
return nbs;
}

/**
* Returns the 2-dimensional array with neighbor indices for every atom,
* Returns the neighbors of every atom, sorted by increasing distance,
* using spatial hashing to avoid all to all distance calculation.
* @return 2-dimensional array of size: n_atoms x n_neighbors_per_atom
* @return array of size n_atoms
*/
IndexAndDistance[][] findNeighborIndicesSpatialHashing() {

// looking at a typical protein case, number of neighbours are from ~10 to ~50, with an average of ~30
int initialCapacity = 60;
Neighbors[] findNeighborIndicesSpatialHashing() {

List<Contact> contactList = calcContacts();
Map<Integer, List<IndexAndDistance>> indices = new HashMap<>(atomCoords.length);

// A first pass to count the neighbors per atom, so that exact-size arrays can be allocated in the second
// pass. Atom indices are dense, so plain arrays are used rather than a map: that avoids boxing the indices
// and the repeated hashing, which are significant given that there are ~30 contacts per atom.
int[] counts = new int[atomCoords.length];
for (Contact contact : contactList) {
// note contacts are stored 1-way only, with j>i
int i = contact.getI();
int j = contact.getJ();

List<IndexAndDistance> iIndices;
List<IndexAndDistance> jIndices;
if (!indices.containsKey(i)) {
iIndices = new ArrayList<>(initialCapacity);
indices.put(i, iIndices);
} else {
iIndices = indices.get(i);
}
if (!indices.containsKey(j)) {
jIndices = new ArrayList<>(initialCapacity);
indices.put(j, jIndices);
} else {
jIndices = indices.get(j);
}

double radius = radii[i] + probe + probe;
double dist = contact.getDistance();
if (dist < radius + radii[j]) {
iIndices.add(new IndexAndDistance(j, dist));
jIndices.add(new IndexAndDistance(i, dist));
if (areNeighbors(i, j, contact.getDistance())) {
counts[i]++;
counts[j]++;
}
}

// convert map to array for fast access
IndexAndDistance[][] nbsIndices = new IndexAndDistance[atomCoords.length][];
for (Map.Entry<Integer, List<IndexAndDistance>> entry : indices.entrySet()) {
List<IndexAndDistance> list = entry.getValue();
IndexAndDistance[] indexAndDistances = list.toArray(new IndexAndDistance[0]);
nbsIndices[entry.getKey()] = indexAndDistances;
int[][] indices = new int[atomCoords.length][];
double[][] dists = new double[atomCoords.length][];
for (int i = 0; i < atomCoords.length; i++) {
// note that some atoms might have no neighbors at all, in which case these are empty arrays
indices[i] = new int[counts[i]];
dists[i] = new double[counts[i]];
}

// important: some atoms might have no neighbors at all: we need to initialise to empty arrays
for (int i=0; i<nbsIndices.length; i++) {
if (nbsIndices[i] == null) {
nbsIndices[i] = new IndexAndDistance[0];
int[] filled = new int[atomCoords.length];
for (Contact contact : contactList) {
int i = contact.getI();
int j = contact.getJ();
double dist = contact.getDistance();
if (areNeighbors(i, j, dist)) {
indices[i][filled[i]] = j;
dists[i][filled[i]] = dist;
filled[i]++;
indices[j][filled[j]] = i;
dists[j][filled[j]] = dist;
filled[j]++;
}
}

return nbsIndices;
Neighbors[] nbs = new Neighbors[atomCoords.length];
for (int i = 0; i < atomCoords.length; i++) {
nbs[i] = Neighbors.createSorted(indices[i], dists[i], counts[i]);
}

return nbs;
}

Point3d[] getAtomCoords() {
Expand All @@ -460,14 +528,12 @@ private List<Contact> calcContacts() {
private double calcSingleAsa(int i) {
Point3d atom_i = atomCoords[i];

int n_neighbor = neighborIndices[i].length;
IndexAndDistance[] neighbor_indices = neighborIndices[i];
// Sorting by closest to farthest away neighbors achieves faster runtimes when checking for occluded
// sphere sample points below. This follows the ideas exposed in
// Eisenhaber et al, J Comp Chemistry 1994 (https://onlinelibrary.wiley.com/doi/epdf/10.1002/jcc.540160303)
// This is essential for performance. In my tests this brings down the number of occlusion checks in loop below to
// an average of n_sphere_points/10 per atom i, producing ~ x4 performance gain overall
Arrays.sort(neighbor_indices, Comparator.comparingDouble(o -> o.dist));
// note the neighbors are already sorted by increasing distance (see Neighbors#createSorted), which is
// essential for the performance of the occlusion checks in the loop below
Neighbors nbs = neighbors[i];
int[] neighbor_indices = nbs.indices;
double[] neighbor_dists = nbs.dists;
int n_neighbor = neighbor_indices.length;

double radius_i = probe + radii[i];

Expand All @@ -476,36 +542,47 @@ private double calcSingleAsa(int i) {
int[] numDistsCalced = null;
if (logger.isDebugEnabled()) numDistsCalced = new int[n_neighbor];

// now we precalculate anything depending only on i,j in equation 3 in Eisenhaber 1994
double[] sqRadii = new double[n_neighbor];
Vector3d[] aj_minus_ais = new Vector3d[n_neighbor];
// Now we precalculate anything depending only on i,j in equation 3 in Eisenhaber 1994.
// The per-neighbor data is laid out in a single flat array as quadruplets
// [aj_minus_ai.x, aj_minus_ai.y, aj_minus_ai.z, cutoff], so that the innermost loop below is a purely
// sequential scan over contiguous primitives, with no object dereferencing. That matches the access
// pattern of the early break and is significantly faster than an array of Vector3d objects.
double[] nbData = new double[4 * n_neighbor];
for (int nbArrayInd =0; nbArrayInd<n_neighbor; nbArrayInd++) {
int j = neighbor_indices[nbArrayInd].index;
double dist = neighbor_indices[nbArrayInd].dist;
int j = neighbor_indices[nbArrayInd];
double dist = neighbor_dists[nbArrayInd];
double radius_j = radii[j] + probe;
Point3d atom_j = atomCoords[j];
int off = 4 * nbArrayInd;
nbData[off ] = atom_j.x - atom_i.x;
nbData[off + 1] = atom_j.y - atom_i.y;
nbData[off + 2] = atom_j.z - atom_i.z;
// see equation 3 in Eisenhaber 1994
sqRadii[nbArrayInd] = (dist*dist + radius_i*radius_i - radius_j*radius_j)/(2*radius_i);
Vector3d aj_minus_ai = new Vector3d(atomCoords[j]);
aj_minus_ai.sub(atom_i);
aj_minus_ais[nbArrayInd] = aj_minus_ai;
nbData[off + 3] = (dist*dist + radius_i*radius_i - radius_j*radius_j)/(2*radius_i);
}

for (Vector3d point: spherePoints){
int nbDataLength = nbData.length;

for (int p = 0; p < nSpherePoints; p++) {
double pointX = spherePoints[3*p];
double pointY = spherePoints[3*p + 1];
double pointZ = spherePoints[3*p + 2];

boolean is_accessible = true;

// note that the neighbors are sorted by distance, achieving optimal performance in this inner loop
// See Eisenhaber et al, J Comp Chemistry 1994

for (int nbArrayInd =0; nbArrayInd<n_neighbor; nbArrayInd++) {
for (int off = 0; off < nbDataLength; off += 4) {

// see equation 3 in Eisenhaber 1994. This is slightly more efficient than
// calculating distances to the actual sphere points on atom_i (which would be obtained with:
// Point3d test_point = new Point3d(point.x*radius + atom_i.x,point.y*radius + atom_i.y,point.z*radius + atom_i.z))
double dotProd = aj_minus_ais[nbArrayInd].dot(point);
double dotProd = nbData[off]*pointX + nbData[off + 1]*pointY + nbData[off + 2]*pointZ;

if (numDistsCalced!=null) numDistsCalced[nbArrayInd]++;
if (numDistsCalced!=null) numDistsCalced[off >> 2]++;

if (dotProd > sqRadii[nbArrayInd]) {
if (dotProd > nbData[off + 3]) {
is_accessible = false;
break;
}
Expand Down
Loading
Loading