galpy 2.0.0b1 is the first pre-release of galpy 2.0, which adds JAX and PyTorch as array backends to galpy's NumPy and C implementation. It is based on v1.12.1 and includes every change through that release. Because it is a pre-release, pip installs it only with --pre:
pip install --pre galpy # NumPy only, as before
pip install --pre "galpy[jax]" # with the JAX backend
pip install --pre "galpy[torch]" # with the PyTorch backend
galpy 2.x requires Python ≥ 3.11. galpy 2.0 is not yet on conda-forge.
Array backends. Potentials, orbit integration, action-angle coordinates, distribution functions, and stream models accept JAX arrays and PyTorch tensors, compute with that framework, and return arrays of the same kind. galpy calculations can therefore be differentiated with automatic differentiation (jax.grad, torch.autograd), compiled (jax.jit, torch.compile), vectorized (jax.vmap), and run on a GPU. With NumPy inputs, galpy behaves as before. The backend is chosen by the type of the arrays passed in, by the galpy.backend.use("jax") / use("torch") context manager, or by the new [backend] section of the configuration file. Start with the quick-start tutorial; each section of the documentation has a backend tutorial of its own (potentials, orbits, action-angle, distribution functions, streamdf, streamspraydf).
- Orbits integrate inside the backend with the differentiable adaptive-step integrators
method='diffrax'(JAX) andmethod='torchode'/'torchdiffeq'(PyTorch), with constant-step in-backend leapfrog and symplectic integrators for many orbits at once, and through galpy's C integrators, which return the derivative with respect to the initial condition from their variational equations. - The Staeckel, adiabatic, spherical, and isochrone-approximation actions, frequencies, and angles, and the inverse transformation of
actionAngleVerticalInverse, are differentiable with respect to the phase-space coordinates and the parameters of the potential. - The spherical, disk, and stream distribution functions are differentiable with respect to their parameters and those of the potential, including the stream track. Sampling takes an explicit random key (
key=) and draws by inverse transformation, which makes the samples differentiable with respect to the parameters of the distribution.
Deliberate changes to the NumPy results (otherwise byte-for-byte identical to v1.12.1): the King model solves its Poisson equation to a relative tolerance of 1e-10 (changes at the 1e-5 to 1e-3 level); the velocity sampler of quasiisothermaldf, the samplers of the two-dimensional disk DFs, and the sampler of streamdf draw by inverse transformation rather than by rejection (different draws, same distribution, and the two-dimensional disk DF draws no longer have a few-per-cent bias in the mean azimuthal velocity); the forces of EinastoPotential and PowerSphericalPotentialwCutoff were rewritten for accuracy at small radii (differences of at most 5e-15 relative).
Removed (deprecated for removal after v1.13): lists of potentials (combine potentials with +); the Rnorm=, Vnorm=, and custom_transform= keywords and the calc_stream_lb method of streamdf; the leading= keyword of fardal15spraydf and chen24spraydf (use tail=). Orbit.integrate requires the times to be an array.
See HISTORY.txt for the full list of changes and the v2 documentation.