"around" selection is slow and memory-intensive

[mnmelo]

As of 0.15 performing an around search can be quite expensive, both compute-wise as memory-wise. This crops up more obviously for systems with many atoms, and where the reference for around is itself also large.

This was tested in a system with ~640k atoms, where the reference has 1130 atoms (select_atoms around 15 name ATNAME). It took at least 15 seconds. Doubling the system size (and reference too) busts my 16Gb of memory.

From the discussion in #383, and later the addition of OpenMP (#529), I was under the impression our implementation was quite fast. But for comparison, for the same system and machine VMD does the same selection operation with little memory usage spike, and virtually instantaneously. Granted, VMD doesn't do PBC but is it enough to account for this?

I took a look at their code and from what I could understand they implement several optimizations some of which I thought we were already doing since #529:

• gridded search;
• parallelization (threading);
• pre-filtering of the search space (for compact references it's easy to define a parallelepiped that contains it with a padding equal to the cutoff, and you're sure your contacts will all be inside that volume);
• shortcutting (when an atom being searched is found to be in contact, no need to calculate further distances from it to the other reference atoms).

I admit I'm not familiar with the internals of distance_array and contact_matrix, but my question is: is there room for optimization or am I hitting MDAnalysis' limits here?

[kain88-de]

How does the situation look with older versions of MDAnalysis?

I don't doubt that the KDTree is memory intensive, if I remember the algorithm it has to calculate a full distance matrix. @richardjgowers has some prototype code for the gridded search.

[mnmelo]

Ah, ok, I remember reading something about @richardjgowers' gridded stuff, and thought that was already in place with OpenMP. I'll check with older versions and post here.

[richardjgowers]

So iirc, PBC conditions have to use a distance_array, so it will do a len(system) * len(reference) search. I imagine KDTree is much better than this.

The code for the around selection is all here if you want to try and optimise it.
https://github.com/MDAnalysis/mdanalysis/blob/develop/package/MDAnalysis/core/Selection.py#L242

And yeah, I played around with gridded search before, here I think it mostly works... I was having some trouble making it competitive as brute force methods for smaller system sizes, probably because of a mistake I made in the Cython code.

[richardjgowers]

@mnmelo I'm getting about a 10x speedup for a 7000x 7000 search when the max distance is 15% of the box size with the CellGrid code, if you gave it a spin with your data and it works we could look at including it in MDA.

[mnmelo]

While I definitely think this should make its way to MDA, I really can't spend time on it right now. I actually opened this issue on behalf of a colleague who bumped into memory problems when doing this kind of selection...

[richardjgowers]

https://github.com/maxscheurer/pycontact/blob/develop/PyContact/core/aroundPatch.py

Here's an implementation of a better around search for MDAnalysis that @maxscheurer put together. From a quick glance at the C level stuff, I think it's taken from VMD, which means we probably can't do the same because of license things(?), but it's definitely worth looking at and benchmarking against in any case.

[richardjgowers]

@ayushsuhane this is something that you could look at

[ayushsuhane]

Sure. The discussion here looks interesting. While it definitely seems the gridded structures are better for distance evaluations for selection of atoms (see here and here). Although both methods (CellGrid and Neighbour Search used in FATSLiM utilize similar data structures, it seems that they have efficient function implementations for different purposes. Richard's cellgrid is faster to evaluate contacts, while FATSLiM implementation is faster for selections. However, I believe both have scope for optimization (could be similar to the VMD implementation which @mnmelo has referred to). I would think a cython implementation of gridded search would be very beneficial for MDAnalysis. Any views and/or suggestions to proceed?