Efficient calculation of bond overlap

Question

For a small library dealing with molecules, I have to calculate a so called connectivity table for the chemical bonds. I assume that there is a bond between atom i and atom j, if the distance between them is smaller than their summed radius. This function has to be fast (So it's not premature optimisation ;) ). At the moment I have two possible solutions, where give_bond_array1 is the faster one. I have two questions:

1. Why is give_bond_array1 faster? Because from the code I would expect give_bond_array2 to be faster. The whole calculation is boiled down to vectorised calls to numpy.
2. Do you have an idea for a faster implementation, or is the next step Fortran, C...?
3. Why is the subtraction step in give_bond_array2 the bottleneck? It is only n**2 number of operations compared to 3 / 2 * n * (n-1) in other lines.

I have created a file minimal_example.py with the following content:

import numpy as np

def give_bond_array1(positions, bond_radii, self_bonding_allowed=False):
    """Calculate a boolean array where ``A[i,j] is True`` indicates a
    bond between the i-th and j-th atom.
    """ 
    coords = ['x', 'y', 'z']
    radii = np.add.outer(bond_radii, bond_radii)
    squared_radii = radii ** 2
    delta = {axis: None for axis in coords}
    for i, axis in enumerate(coords):
        coord = positions[:, i]
        delta[axis] = coord - coord.reshape((len(coord), 1)) 
    squared_distances = delta['x']**2 + delta['y']**2 + delta['z']**2
    overlap = squared_radii - squared_distances
    bond_array = overlap >= 0
    if not self_bonding_allowed:
        np.fill_diagonal(bond_array, False)
    return bond_array

def give_bond_array2(positions, bond_radii, self_bonding_allowed=False):
    """Calculate a boolean array where ``A[i,j] is True`` indicates a
    bond between the i-th and j-th atom.
    """ 
    radii = np.add.outer(bond_radii, bond_radii)
    squared_radii = radii ** 2
    delta = (positions[None, :, :] - positions[:, None, :]) 
    squared_delta = delta ** 2
    squared_distances = np.sum(squared_delta, axis=2)
    overlap = squared_radii - squared_distances
    bond_array = overlap >= 0
    if not self_bonding_allowed:
        np.fill_diagonal(bond_array, False)
    return bond_array

n_atoms = 15000
positions, bond_radii = np.random.rand(n_atoms, 3), np.random.rand(n_atoms)
positions, bond_radii = [np.array(x, dtype='float32', order='F') for x in [positions, bond_radii]]

For the testing I got the following output:

In [1]: from minimal_example import *
   ...: %load_ext line_profiler
   ...:

In [2]: %lprun -f give_bond_array1 give_bond_array1(positions, bond_radii)

Timer unit: 1e-06 s

Total time: 5.72649 s
File: /home/oweser/Dropbox/Tests_and_debugging/chemcoord_test/minimal_example.py
Function: give_bond_array1 at line 3

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     3                                           def give_bond_array1(positions, bond_radii, self_bonding_allowed=False):
     4                                               """Calculate a boolean array where ``A[i,j] is True`` indicates a
     5                                               bond between the i-th and j-th atom.
     6                                               """
     7         1            3      3.0      0.0      coords = ['x', 'y', 'z']
     8         1       332504 332504.0      5.8      radii = np.add.outer(bond_radii, bond_radii)
     9         1       383706 383706.0      6.7      squared_radii = radii ** 2
    10         1           18     18.0      0.0      delta = {axis: None for axis in coords}
    11         4           32      8.0      0.0      for i, axis in enumerate(coords):
    12         3           35     11.7      0.0          coord = positions[:, i]
    13         3      1114585 371528.3     19.5          delta[axis] = coord - coord.reshape((len(coord), 1))
    14         1      3030061 3030061.0     52.9      squared_distances = delta['x']**2 + delta['y']**2 + delta['z']**2
    15         1       640631 640631.0     11.2      overlap = squared_radii - squared_distances
    16         1       224410 224410.0      3.9      bond_array = overlap >= 0
    17         1            4      4.0      0.0      if not self_bonding_allowed:
    18         1          504    504.0      0.0          np.fill_diagonal(bond_array, False)
    19         1            2      2.0      0.0      return bond_array




In [1]: from minimal_example import *
   ...: %load_ext line_profiler
   ...:

In [2]: %lprun -f give_bond_array2 give_bond_array2(positions, bond_radii)

Timer unit: 1e-06 s

Total time: 75.13 s
File: /home/oweser/Dropbox/Tests_and_debugging/chemcoord_test/minimal_example.py
Function: give_bond_array2 at line 21

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    21                                           def give_bond_array2(positions, bond_radii, self_bonding_allowed=False):
    22                                               """Calculate a boolean array where ``A[i,j] is True`` indicates a
    23                                               bond between the i-th and j-th atom.
    24                                               """
    25         1       335759 335759.0      0.4      radii = np.add.outer(bond_radii, bond_radii)
    26         1       383206 383206.0      0.5      squared_radii = radii ** 2
    27         1      1106056 1106056.0      1.5      delta = (positions[None, :, :] - positions[:, None, :])
    28         1      1578934 1578934.0      2.1      squared_delta = delta ** 2
    29         1     29164470 29164470.0     38.8      squared_distances = np.sum(squared_delta, axis=2)
    30         1     39366942 39366942.0     52.4      overlap = squared_radii - squared_distances
    31         1      1611265 1611265.0      2.1      bond_array = overlap >= 0
    32         1            3      3.0      0.0      if not self_bonding_allowed:
    33         1      1583333 1583333.0      2.1          np.fill_diagonal(bond_array, False)
    34         1            2      2.0      0.0      return bond_array

PS: I know that I should also cut my molecule into smaller pieces in order to make "divide and conquer". This is already implemented and greatly speeds up calculation. Nevertheless I would like to understand why the above functions differ so much in performance.

PPS: The number of atoms (n_atoms in the example) may safely be assumed to be less than 10^5.

Answer 1

Thanks to the timings of Gareth Rees I think I found the reason for the timing differences. First thing to note is, that I got similar timings on my better computer. But the timings on the worse computer were consistently different. I tried to answer the question on my own. If I use wrong terminology please correct me.

The main reason is that the maximum amount of used memory and the number of reshaping operations is higher in the first function and the operating system had to start swapping on the bad computer. So the finished version of the function looks like this:

def give_bond_array3(positions, bond_radii, self_bonding_allowed=False):
    """Calculate a boolean array where ``A[i,j] is True`` indicates a
    bond between the i-th and j-th atom.
    """
    squared_radii = (bond_radii + bond_radii[:, None])**2

    coord = positions[:, 0]
    squared_distances = (coord - coord[:, None])**2
    for i in range(1,3):
        coord = positions[:, i]
        squared_distances += (coord - coord[:, None])**2
    overlap = squared_radii - squared_distances
    bond_array = overlap >= 0
    if not self_bonding_allowed:
        np.fill_diagonal(bond_array, False)
    return bond_array

The main trick is that the large (n_atoms, n_atoms, 3) tensor of position differences between each pair of atom is not created at once. Instead (n_atoms, n_atoms) matrices are created for each one of the three coordinate axes and added in place.

One thing that made give_bond_array2 so especially bad was the assignment to variables like squared_radii instead of directly working with radii**2. The assigment meant there existed a reference to the object in memory, preventing the garbage collector from doing its work. Usually this would be no problem, because upon function exit the memory is freed anyway, but if you run into swapping problems during function execution, it becomes a thing.

If you just avoid the problem of variable assigment the function may be written as:

def give_bond_array4(positions, bond_radii, self_bonding_allowed=False):
    """Calculate a boolean array where ``A[i,j] is True`` indicates a
    bond between the i-th and j-th atom.
    """
    overlap = (np.add.outer(bond_radii, bond_radii)**2
               - np.sum((positions[None, :, :] - positions[:, None, :])**2, axis=2))
    bond_array = overlap >= 0
    if not self_bonding_allowed:
        np.fill_diagonal(bond_array, False)
    return bond_array

On the bad computer the timings are then like following:

1. give_bond_array3:4.40s
2. give_bond_array4: 5.63s

I explained the timing differences between give_bond_array3 and give_bond_array4 with the additional reshaping operations of which give_bond_array3 makes less use of. Additionally the coord vector is in all cases linearly aligned in memory. On the other hand the summing operation on tensor (positions[None, :, :] - positions[:, None, :])**2 along the axis=2 has to "jump" probably from address to address.

A final nice thing to note is the now visible timing difference between C and Fortran layout. The timings for give_bond_array3 are:

1. Fortran memory layout: 4.40 s
2. C memory layout: 4.85 s

Summary Don't do assignments for intermediate steps if memory is a problem.

Answer 2

I am unable to replicate your timing results. Your two functions take almost identical amounts of time when I run them (Python 3.5.3, NumPy 1.12.1):

>>> from timeit import timeit
>>> timeit(lambda:give_bond_array1(positions, bond_radii), number=1)
5.692430718801916
>>> timeit(lambda:give_bond_array2(positions, bond_radii), number=1)
5.850847749970853

So I can't tell you why the timings are so different in your case. Could it be the profiler that's at fault? Do you get the same difference in runtime when you don't use the profiler?