Performance with many species about jax-md HOT 7 OPEN

Yes and no. The issue is actually the compile time and not the run time. When using species, there is a double for loop in smap.pair that gets unrolled when compiling. After compiling, however, the execution seems to be completely comparable. To see this, the following code snippet

N = 600                 # number of total particles
no_species = 20         # number of species types

packing_fraction = 0.90
box_size = box_size_at_packing_fraction_2d(packing_fraction, jnp.full((N,),1.0))
displacement, shift = space.periodic(box_size)

R_init = random.uniform(key, (N, 2), maxval=box_size)
sigma = jnp.full((N,N), 1.0)
energy_fn = jit(energy.soft_sphere_pair(displacement, species=None, sigma=sigma))

%time E = energy_fn(R_init)
%time E = energy_fn(R_init)

species = get_species(no_species, R_init.shape[0])
sigma = jnp.full((no_species,no_species), 1.0)
energy_fn = jit(energy.soft_sphere_pair(displacement, species=species, sigma=sigma))

%time E = energy_fn(R_init)
%time E = energy_fn(R_init)

spits out

CPU times: user 393 ms, sys: 11.1 ms, total: 404 ms
Wall time: 260 ms
CPU times: user 243 µs, sys: 1.12 ms, total: 1.37 ms
Wall time: 1.06 ms
CPU times: user 26.1 s, sys: 236 ms, total: 26.3 s
Wall time: 20.3 s
CPU times: user 1.63 ms, sys: 18 µs, total: 1.65 ms
Wall time: 1.31 ms

The compile time goes from .26s to 20s, but the run time only goes from 1ms to 1.3ms.

If this compile time becomes unreasonable, you can always give your parameters as N by N matrices, though obviously this has the potential to give you memory issues for large systems. A couple years ago a student of mine wrote a version of smap.pair that he claimed would compile quickly for many species, so presumably a fix is possible, but I never saw his implementation so I don't know how general it was.

from jax-md.

hi! I'm curious - is it still an issue?

from jax-md.

This is awesome, many thanks for putting time into this detailed answer!! I also checked the code last night and was surprised to see a double loop - but from your answer, it seems that this double doesn't affect the performance that much.

One of the reasons I've asked is that we're currently considering switching our simulations to jax-md (or, rather, some of our simulations) - so we're investigating pros and cons of different packages.

from jax-md.

Just out of curiosity, what are the rough estimates for numbers of species and numbers of particles you typically use?

from jax-md.

well, it's not much, to be honest, typically not more than 5. :)
With that said, the ~30x slowdown that you initially observed looked very scary, so I just wanted to understand how big of an issue it could be in the future.

from jax-md.

Here is a cleaner version of the test: https://colab.research.google.com/drive/1duxvhx1aVeaEYNdog32ltzN2flTEBhYK?usp=sharing

For both cases (with and without species), you run apply(state) once to compile, and then run again to time. Ideally this allows you to substitute your own potentials and simulation (rather than soft_spheres/minimization that I used), and include things like neighbor lists, etc. as necessary.

Importantly, I believe the compile time should not depend on the number of particles, just the number of species, so I don't think you have a lot to worry about.

from jax-md.

awesome! Many thanks - your tests look very convincing! Again, many thanks for putting your valuable time into this response - the results look very reassuring!

from jax-md.