Comments (7)
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.
Related Issues (20)
- Cannot import 'FunctionalFullyConnectedTensorProduct' from 'e3nn_jax' HOT 1
- test_nequip_silicon in energy_test.py is broken.
- test_nve_2d_neighbor_list_multi_atom_species in rigid_body_test.py is broken HOT 1
- documentation not compiling 2
- Elasticity calculations involving rigid bodies
- Out Of Memory issue during neighbor list generation
- Inconsistence of periodic and periodic_general
- equivariant_neural_networks notebook is broken
- Question about npt simulation HOT 2
- AttributeError: module 'jax.random' has no attribute 'KeyArray' HOT 1
- Neighbor lists are broken for rigid bodies.
- PME Energies
- AttributeError: module 'jax' has no attribute 'linear_util'
- ImportError in the example notebooks HOT 1
- Run sample notebooks as of 13 April 2024. HOT 2
- Potential Definition Discrepancy with Stress Calculation
- future directions to improve jax-md's performance?
- Error when running flocking.ipynb HOT 1
- How to generate equivalent NequIP model in JAX-MD
- test_coulomb_cubeions fails
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax-md.