MilesCranmer
commented
4 years ago Here's an example I made for a Logarithmic Halo Potential.
The integration over 1e8 timesteps for 5000 different initial conditions takes only 10 seconds in total. And this would be automatically ran on the GPU if one is available.
from jax import numpy as jnp
from jax import jit, grad, vmap, random
from jax.experimental.ode import odeint
#Set up integration:
@jit
def phi(x):
amp = 1
core = 1e-8
q = 0.5
R = x[0]
phi = x[1]
z = x[2]
return 0.5 * jnp.log(R**2 + (z/q)**2 + core**2)
#The gradient of this potential:
nacceleration = jit(grad(phi))
@jit
def odefunc(y, t):
x = y[:3]
xt = y[3:]
xtt = -nacceleration(x)
return jnp.concatenate([xt, xtt])
_odefunc = jit(vmap(odefunc, (0, None), 0))
@jit
def vodefunc(yall, t):
yall = yall.reshape(-1, 6)
out = _odefunc(yall, t)
return out.reshape(-1)
@jit
def integrate_all(y0s, ts):
y0s = y0s.reshape(-1)
yall = odeint(vodefunc, y0s, ts, mxsteps=500)
return yall.reshape(yall.shape[0], -1, 6)
#Set up initial conditions:
rng = random.PRNGKey(0)
n = 5000
y0 = jnp.stack([
random.uniform(rng, (n,), minval=0.1, maxval=5),
jnp.zeros(n),
jnp.zeros(n),
random.normal(rng+1, (n,)),
jnp.zeros(n),
random.normal(rng+2, (n,))]).T
ts = jnp.linspace(0, 10000, int(1e8))
yall = integrate_all(y0, ts)
jobovy
stale[bot]
Hi @jobovy,
I was wondering if you would consider using JAX (https://github.com/google/jax) instead of numpy for some of the galpy backend? A lot of numpy calls can be replaced from:
import numpy as nptofrom jax import numpy as np.All code is then JIT-compiled with LLVM, parallelized, and GPU-ready, without any further changes. In practice it can be faster than handwritten C code too.
The biggest unique advantage JAX has over numpy or, e.g., numba, is the operation
vmap.vmapis functionally similar to multiprocessing.Pool(...).map in that you batch some function over an array in parallel. The key difference is that vmap will actually go through your code, and add an extra axis in all your linear algebra operations. We used it extensively to build our Lagrangian Neural Networks (https://twitter.com/MilesCranmer/status/1237788581478969350?s=20) in order to efficiently batch an operation over an array. I think it would be really nice to be able to call vmap on a galpy operation, since one could then efficiently parallelize eg many evaluations in the same potential, and on a GPU.Another huge advantage is obviously the
gradoperation that lets you differentiate any function, even through an integration. That would let you do optimization by gradient descent anywhere in galpy.Cheers! Miles