Gravity Operators

gravity

Gravitational Poisson solver using FFT spectral methods.

Solves the cosmological Poisson equation in conformal time:

\[\nabla^2 \phi = 4\pi G \, a^2 \, \delta\rho\]

where \(\delta\rho = \rho - \bar{\rho}\) is the overdensity. In Fourier space (periodic BCs):

\[\widetilde{\phi}(k) = -\frac{4\pi G \, a^2 \, \widetilde{\delta\rho}(k)}{k^2} \quad (k \neq 0), \qquad \widetilde{\phi}(0) = 0\]

(gauge choice: \(\langle \phi \rangle = 0\)).

The gravitational acceleration \(g = -\nabla\phi\) in Fourier space:

\[\widetilde{g}_i(k) = -i k_i \widetilde{\phi}(k)\]

Performance design

Factory functions (make_gravitational_poisson_solver, make_gravitational_force) precompute k-space arrays at construction time and embed them as compile-time constants in the returned JIT-compiled callable. This means:

• No retracing when the returned function is called inside jax.jit.
• Safe for jax.vmap over batch dimensions (all inputs are arrays).
• Safe for jax.pmap across devices.
• The k=0 singularity is resolved with jnp.where (no Python branching), so XLA sees a single, branch-free computation graph.
• Real-valued FFT (rfftn / irfftn) halves the memory footprint and reduces FFT work by ~2× compared to the full complex transform.

The mean density is removed automatically: the k=0 mode is set to zero, so the solver accepts either \(\rho\) (total density) or \(\delta\rho\) (overdensity) — the result is identical because the homogeneous background contributes only to the k=0 mode.

References

• Hockney & Eastwood, "Computer Simulation Using Particles" (1981), Ch. 6
• Springel, "GADGET-2" (2005), MNRAS 364, 1105, Appendix A
• Dodelson & Schmidt, "Modern Cosmology" (2020), Eq. (5.54)

make_gravitational_poisson_solver(grid_shape, dx)

Return a JIT-compiled gravitational Poisson solver for a fixed grid.

Builds and caches the k-space inverse-Laplacian kernel once. The returned callable accepts only JAX arrays (plus Python/JAX scalars for a and M_pl), making it fully composable with jax.jit, jax.vmap, and jax.pmap.

Parameters

grid_shape: Spatial dimensions, e.g. (N,), (Nx, Ny), (Nx, Ny, Nz). dx: Uniform lattice spacing.

Returns

solve : Callable solve(rho, a, M_pl) -> phi

- ``rho``  – density or overdensity field, shape ``grid_shape``.
  The k=0 mode is automatically zeroed (mean removed).
- ``a``    – scale factor (scalar JAX array or Python float).
- ``M_pl`` – reduced Planck mass; ``G = 1 / M_pl²``.
- ``phi``  – gravitational potential, shape ``grid_shape``.

Examples

.. code-block:: python

solve = make_gravitational_poisson_solver((128, 128, 128), dx=1.0)
phi = solve(delta_rho, a=cosmo.a, M_pl=cosmo.M_pl)

# vmap over a batch of density fields:
batched_solve = jax.vmap(solve, in_axes=(0, None, None))
phi_batch = batched_solve(delta_rho_batch, a, M_pl)

Source code in jaxlatt/operators/gravity.py

def make_gravitational_poisson_solver(grid_shape: tuple[int, ...], dx: float) -> Callable:
    """
    Return a JIT-compiled gravitational Poisson solver for a fixed grid.

    Builds and caches the k-space inverse-Laplacian kernel once. The
    returned callable accepts only JAX arrays (plus Python/JAX scalars for
    ``a`` and ``M_pl``), making it fully composable with ``jax.jit``,
    ``jax.vmap``, and ``jax.pmap``.

    Parameters
    ----------
    grid_shape:
        Spatial dimensions, e.g. ``(N,)``, ``(Nx, Ny)``, ``(Nx, Ny, Nz)``.
    dx:
        Uniform lattice spacing.

    Returns
    -------
    solve : Callable
        ``solve(rho, a, M_pl) -> phi``

        - ``rho``  – density or overdensity field, shape ``grid_shape``.
          The k=0 mode is automatically zeroed (mean removed).
        - ``a``    – scale factor (scalar JAX array or Python float).
        - ``M_pl`` – reduced Planck mass; ``G = 1 / M_pl²``.
        - ``phi``  – gravitational potential, shape ``grid_shape``.

    Examples
    --------
    .. code-block:: python

        solve = make_gravitational_poisson_solver((128, 128, 128), dx=1.0)
        phi = solve(delta_rho, a=cosmo.a, M_pl=cosmo.M_pl)

        # vmap over a batch of density fields:
        batched_solve = jax.vmap(solve, in_axes=(0, None, None))
        phi_batch = batched_solve(delta_rho_batch, a, M_pl)
    """
    ndim = len(grid_shape)

    if ndim == 3:
        nx, ny, nz = grid_shape
        _, _, _, k_sq = _kvecs_3d(nx, ny, nz, dx)

        @jit
        def solve(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfftn(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            return jnp.fft.irfftn(phi_k, s=(nx, ny, nz))

    elif ndim == 2:
        nx, ny = grid_shape
        _, _, k_sq = _kvecs_2d(nx, ny, dx)

        @jit
        def solve(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfft2(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            return jnp.fft.irfft2(phi_k, s=(nx, ny))

    elif ndim == 1:
        (n,) = grid_shape
        _, k_sq = _kvecs_1d(n, dx)

        @jit
        def solve(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfft(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            return jnp.fft.irfft(phi_k, n=n)

    else:
        raise ValueError(f"grid_shape must be 1D, 2D or 3D, got ndim={ndim}")

    return solve

make_gravitational_force(grid_shape, dx)

Return a JIT-compiled gravitational acceleration solver for a fixed grid.

Computes \(g = -\nabla\phi\) from the density field in a single FFT round-trip by multiplying phi_k by \(-i k_i\) for each spatial direction, avoiding a separate gradient pass after the Poisson solve.

The returned array has shape (ndim, *grid_shape) with g[i] being the acceleration along axis i. This layout matches the PM particle-mesh interface convention.

Parameters

grid_shape: Spatial dimensions. dx: Uniform lattice spacing.

Returns

force_fn : Callable force_fn(rho, a, M_pl) -> Array[ndim, *grid_shape]

Examples

.. code-block:: python

force_fn = make_gravitational_force((128, 128, 128), dx=1.0)
g = force_fn(delta_rho, a=cosmo.a, M_pl=cosmo.M_pl)
# g[0], g[1], g[2]  ←→  gx, gy, gz

# Parallel across devices:
pforce = jax.pmap(force_fn, in_axes=(0, None, None))
g = pforce(delta_rho_sharded, a, M_pl)

Source code in jaxlatt/operators/gravity.py

def make_gravitational_force(grid_shape: tuple[int, ...], dx: float) -> Callable:
    """
    Return a JIT-compiled gravitational acceleration solver for a fixed grid.

    Computes $g = -\\nabla\\phi$ from the density field in a single FFT round-trip by
    multiplying phi_k by $-i k_i$ for each spatial direction, avoiding a
    separate gradient pass after the Poisson solve.

    The returned array has shape ``(ndim, *grid_shape)`` with ``g[i]``
    being the acceleration along axis ``i``. This layout matches the PM
    particle-mesh interface convention.

    Parameters
    ----------
    grid_shape:
        Spatial dimensions.
    dx:
        Uniform lattice spacing.

    Returns
    -------
    force_fn : Callable
        ``force_fn(rho, a, M_pl) -> Array[ndim, *grid_shape]``

    Examples
    --------
    .. code-block:: python

        force_fn = make_gravitational_force((128, 128, 128), dx=1.0)
        g = force_fn(delta_rho, a=cosmo.a, M_pl=cosmo.M_pl)
        # g[0], g[1], g[2]  ←→  gx, gy, gz

        # Parallel across devices:
        pforce = jax.pmap(force_fn, in_axes=(0, None, None))
        g = pforce(delta_rho_sharded, a, M_pl)
    """
    ndim = len(grid_shape)

    if ndim == 3:
        nx, ny, nz = grid_shape
        KX, KY, KZ, k_sq = _kvecs_3d(nx, ny, nz, dx)
        s = (nx, ny, nz)

        @jit
        def force_fn(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfftn(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            # g_i = -∂φ/∂x_i  →  g̃_i(k) = -i k_i φ̃(k)
            gx = jnp.fft.irfftn(-1j * KX * phi_k, s=s)
            gy = jnp.fft.irfftn(-1j * KY * phi_k, s=s)
            gz = jnp.fft.irfftn(-1j * KZ * phi_k, s=s)
            return jnp.stack([gx, gy, gz])

    elif ndim == 2:
        nx, ny = grid_shape
        KX, KY, k_sq = _kvecs_2d(nx, ny, dx)
        s = (nx, ny)

        @jit
        def force_fn(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfft2(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            gx = jnp.fft.irfft2(-1j * KX * phi_k, s=s)
            gy = jnp.fft.irfft2(-1j * KY * phi_k, s=s)
            return jnp.stack([gx, gy])

    elif ndim == 1:
        (n,) = grid_shape
        k, k_sq = _kvecs_1d(n, dx)

        @jit
        def force_fn(rho: Array, a: float, M_pl: float) -> Array:
            G = 1.0 / (M_pl * M_pl)
            coeff = -4.0 * jnp.pi * G * a * a
            rho_k = jnp.fft.rfft(rho)
            phi_k = jnp.where(k_sq > 0.0, coeff * rho_k / k_sq, 0.0 + 0.0j)
            gx = jnp.fft.irfft(-1j * k * phi_k, n=n)
            return jnp.stack([gx])

    else:
        raise ValueError(f"grid_shape must be 1D, 2D or 3D, got ndim={ndim}")

    return force_fn

gravitational_potential(rho, dx, a, M_pl)

Compute gravitational potential \(\phi\) from density field (1D/2D/3D).

Solves \(\nabla^2 \phi = 4\pi G \, a^2 \, \delta\rho\) via FFT. The mean of rho is removed automatically (k=0 mode is zeroed).

.. note:: This rebuilds the k-space kernel on every call. Use :func:make_gravitational_poisson_solver in tight loops or inside jax.jit/jax.vmap.

Parameters

rho: Density or overdensity field, shape (N,), (Nx, Ny), or (Nx, Ny, Nz). dx: Lattice spacing. a: Scale factor. M_pl: Reduced Planck mass.

Returns

phi : Array Gravitational potential, same shape as rho.

Source code in jaxlatt/operators/gravity.py

def gravitational_potential(
    rho: Array,
    dx: float,
    a: float,
    M_pl: float,
) -> Array:
    """
    Compute gravitational potential $\\phi$ from density field (1D/2D/3D).

    Solves $\\nabla^2 \\phi = 4\\pi G \\, a^2 \\, \\delta\\rho$ via FFT. The mean of ``rho`` is removed
    automatically (k=0 mode is zeroed).

    .. note::
        This rebuilds the k-space kernel on every call. Use
        :func:`make_gravitational_poisson_solver` in tight loops or inside
        ``jax.jit``/``jax.vmap``.

    Parameters
    ----------
    rho:
        Density or overdensity field, shape ``(N,)``, ``(Nx, Ny)``, or
        ``(Nx, Ny, Nz)``.
    dx:
        Lattice spacing.
    a:
        Scale factor.
    M_pl:
        Reduced Planck mass.

    Returns
    -------
    phi : Array
        Gravitational potential, same shape as ``rho``.
    """
    return make_gravitational_poisson_solver(rho.shape, dx)(rho, a, M_pl)

gravitational_force(rho, dx, a, M_pl)

Compute gravitational acceleration \(g = -\nabla\phi\) from density field (1D/2D/3D).

Returns an array of shape (ndim, *rho.shape), where result[i] is the acceleration component along axis i.

.. note:: This rebuilds the k-space kernel on every call. Use :func:make_gravitational_force in tight loops or inside jax.jit/jax.vmap.

Parameters

rho: Density or overdensity field. dx: Lattice spacing. a: Scale factor. M_pl: Reduced Planck mass.

Returns

g : Array Gravitational acceleration, shape (ndim, *rho.shape).

Source code in jaxlatt/operators/gravity.py

def gravitational_force(
    rho: Array,
    dx: float,
    a: float,
    M_pl: float,
) -> Array:
    """
    Compute gravitational acceleration $g = -\\nabla\\phi$ from density field (1D/2D/3D).

    Returns an array of shape ``(ndim, *rho.shape)``, where ``result[i]``
    is the acceleration component along axis ``i``.

    .. note::
        This rebuilds the k-space kernel on every call. Use
        :func:`make_gravitational_force` in tight loops or inside
        ``jax.jit``/``jax.vmap``.

    Parameters
    ----------
    rho:
        Density or overdensity field.
    dx:
        Lattice spacing.
    a:
        Scale factor.
    M_pl:
        Reduced Planck mass.

    Returns
    -------
    g : Array
        Gravitational acceleration, shape ``(ndim, *rho.shape)``.
    """
    return make_gravitational_force(rho.shape, dx)(rho, a, M_pl)