merge hugos LPT2 code

jaxpm/kernels.py

@@ -1,5 +1,3 @@
-from enum import Enum
-
 import jax.numpy as jnp
 import jax_cosmo as jc
 import numpy as np
@@ -45,16 +43,18 @@ def interpolate_power_spectrum(input, k, pk, sharding=None):
 def gradient_kernel(kvec, direction, order=1):
     """
     Computes the gradient kernel in the requested direction
-    Parameters:
+    
+    Parameters
     -----------
-    kvec: array
+    kvec: list
-      Array of k values in Fourier space
+        List of wave-vectors in Fourier space
     direction: int
-      Index of the direction in which to take the gradient
+        Index of the direction in which to take the gradient
-    Returns:
+
+    Returns
     --------
     wts: array
-      Complex kernel
+        Complex kernel values
     """
     if order == 0:
         wts = 1j * kvec[direction]
@@ -69,37 +69,43 @@ def gradient_kernel(kvec, direction, order=1):
         return wts
 
 
-def laplace_kernel(kvec):
+def invlaplace_kernel(kvec):
     """
-    Compute the Laplace kernel from a given K vector
+    Compute the inverse Laplace kernel
-    Parameters:
+
+    Parameters
     -----------
-    kvec: array
+    kvec: list
-      Array of k values in Fourier space
+        List of wave-vectors
-    Returns:
+
+    Returns
     --------
     wts: array
-      Complex kernel
+        Complex kernel values
     """
     kk = sum(ki**2 for ki in kvec)
-    wts = jnp.where(kk == 0, 1., 1. / kk)
+    kk_nozeros = jnp.where(kk==0, 1, kk) 
-    return wts
+    return - jnp.where(kk==0, 0, 1 / kk_nozeros)
 
 
 def longrange_kernel(kvec, r_split):
     """
-  Computes a long range kernel
+    Computes a long range kernel
-  Parameters:
+
-  -----------
+    Parameters
-  kvec: array
+    -----------
-    Array of k values in Fourier space
+    kvec: list
-  r_split: float
+        List of wave-vectors
+    r_split: float
+        Splitting radius
+        
+    Returns
+    --------
+    wts: array
+        Complex kernel values
+    
     TODO: @modichirag add documentation
-  Returns:
+    """
-  --------
-  wts: array
-    kernel
-  """
     if r_split != 0:
         kk = sum(ki**2 for ki in kvec)
         return np.exp(-kk * r_split**2)
@@ -109,15 +115,21 @@ def longrange_kernel(kvec, r_split):
 
 def cic_compensation(kvec):
     """
-  Computes cic compensation kernel.
+    Computes cic compensation kernel.
-  Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
+    Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
-  Itself based on equation 18 (with p=2) of
+    Itself based on equation 18 (with p=2) of
-        `Jing et al 2005 <https://arxiv.org/abs/astro-ph/0409240>`_
+          [Jing et al 2005](https://arxiv.org/abs/astro-ph/0409240)
-  Args:
+
-    kvec: array of k values in Fourier space
+    Parameters:
-  Returns:
+    -----------
-    v: array of kernel
+    kvec: list
-  """
+        List of wave-vectors
+        
+    Returns:
+    --------
+    wts: array
+        Complex kernel values
+    """
     kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
     wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
     return wts
@@ -125,20 +137,22 @@ def cic_compensation(kvec):
 
 def PGD_kernel(kvec, kl, ks):
     """
-  Computes the PGD kernel
+    Computes the PGD kernel
-  Parameters:
+
-  -----------
+    Parameters:
-  kvec: array
+    -----------
-    Array of k values in Fourier space
+    kvec: list
-  kl: float
+        List of wave-vectors
-    initial long range scale parameter
+    kl: float
-  ks: float
+        Initial long range scale parameter
-    initial dhort range scale parameter
+    ks: float
-  Returns:
+        Initial dhort range scale parameter
-  --------
+
-  v: array
+    Returns:
-    kernel
+    --------
-  """
+    v: array
+        Complex kernel values
+    """
     kk = sum(ki**2 for ki in kvec)
     kl2 = kl**2
     ks4 = ks**4

jaxpm/pm.py

@@ -9,7 +9,7 @@ from jaxpm.distributed import (autoshmap, fft3d, get_local_shape, ifft3d,
                                normal_field)
 from jaxpm.growth import (dGf2a, dGfa, growth_factor, growth_factor_second,
                           growth_rate, growth_rate_second)
-from jaxpm.kernels import (PGD_kernel, fftk, gradient_kernel, laplace_kernel,
+from jaxpm.kernels import (PGD_kernel, fftk, gradient_kernel, invlaplace_kernel,
                            longrange_kernel)
 from jaxpm.painting import cic_paint, cic_paint_dx, cic_read, cic_read_dx
 
@@ -29,18 +29,20 @@ def pm_forces(positions,
         mesh_shape = delta.shape
 
     if delta is None:
-        delta_k = fft3d(
+        field = cic_paint_dx(positions, halo_size=halo_size, sharding=sharding)
-            cic_paint_dx(positions, halo_size=halo_size, sharding=sharding))
+        delta_k = fft3d(field)
-    else:
+    elif jnp.isrealobj(delta):
         delta_k = fft3d(delta)
+    else:
+        delta_k = delta
 
     kvec = fftk(delta_k)
     # Computes gravitational potential
-    pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,
+    pot_k = delta_k * invlaplace_kernel(kvec) * longrange_kernel(kvec,
                                                               r_split=r_split)
     # Computes gravitational forces
     forces = jnp.stack([
-        cic_read_dx(ifft3d(gradient_kernel(kvec, i) * pot_k),
+        cic_read_dx(ifft3d( - gradient_kernel(kvec, i) * pot_k),
                     halo_size=halo_size,
                     sharding=sharding) for i in range(3)
     ],
@@ -49,44 +51,10 @@ def pm_forces(positions,
     return forces
 
 
-def lpt2_source(mesh_size, initial_conditions):
+def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None,order=1):
-
-    kvec = fftk(mesh_size)
-    # TODO : this has already been done for LPT1, we should reuse it
-    delta_k = fft3d(initial_conditions)
-
-    source = jnp.zeros_like(delta_k)
-
-    D1 = [1, 2, 0]
-    D2 = [2, 0, 1]
-
-    # laplace_kernel should be actually inv laplace_kernel
-    # adding a minus sign here that will be negated when computing forces
-    # because F = -grad(phi)
-    # and phi = -laplace_kernel(delta_k)
-    pot_k = delta_k * laplace_kernel(delta_k)
-
-    nabla_i_nabla_i = [
-        ifft3d(gradient_kernel(kvec, i)**2 * pot_k) for i in range(3)
-    ]
-    # for diagonal terms
-    source += nabla_i_nabla_i[D1[0]] * nabla_i_nabla_i[D2[0]]
-    source += nabla_i_nabla_i[D1[1]] * nabla_i_nabla_i[D2[1]]
-    source += nabla_i_nabla_i[D1[2]] * nabla_i_nabla_i[D2[2]]
-
-    # off diag terms
-    for i in range(3):
-        nabla_i_nabla_j = gradient_kernel(kvec, D1[i]) * gradient_kernel(
-            kvec, D2[i])
-        phi = ifft3d(nabla_i_nabla_j * pot_k)
-        source -= phi**2
-
-    return source
-
-
-def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None):
     """
-    Computes first order LPT displacement
+    Computes first and second order LPT displacement and momentum, 
+    e.g. Eq. 2 and 3 [Jenkins2010](https://arxiv.org/pdf/0910.0258)
     """
     gpu_mesh = sharding.mesh if sharding is not None else None
     spec = sharding.spec if sharding is not None else P()
@@ -99,48 +67,48 @@ def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None):
       out_specs=spec)()  # yapf: disable
 
 
+    a = jnp.atleast_1d(a)
+    E = jnp.sqrt(jc.background.Esqr(cosmo, a)) 
+    delta_k = fft3d(initial_conditions)
     initial_force = pm_forces(displacement,
-                              delta=initial_conditions,
+                              delta=delta_k,
                               halo_size=halo_size,
                               sharding=sharding)
-    a = jnp.atleast_1d(a)
     dx = growth_factor(cosmo, a) * initial_force
-    p = a**2 * growth_rate(cosmo, a) * jnp.sqrt(jc.background.Esqr(cosmo,
+    p = a**2 * growth_rate(cosmo, a) * E * dx
-                                                                   a)) * dx
+    f = a**2 * E * dGfa(cosmo,a) * initial_force
-    f = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGfa(cosmo,
+    if order == 2:
-                                                             a) * initial_force
+        kvec = fftk(delta_k)
-    return dx, p, f
+        pot_k = delta_k * invlaplace_kernel(kvec)
 
+        delta2 = 0
+        shear_acc = 0
+        # for i, ki in enumerate(kvec):
+        for i in range(3):
+            # Add products of diagonal terms = 0 + s11*s00 + s22*(s11+s00)...
+            # shear_ii = jnp.fft.irfftn(- ki**2 * pot_k)
+            nabla_i_nabla_i = gradient_kernel(kvec, i)**2
+            shear_ii = jnp.fft.irfftn(nabla_i_nabla_i * pot_k)
+            delta2 += shear_ii * shear_acc 
+            shear_acc += shear_ii
 
-# @Credit Hugo Simon https://github.com/hsimonfroy/montecosmo
+            # for kj in kvec[i+1:]:
-def lpt2(cosmo, initial_conditions, dx, p, f, a, halo_size=0):
+            for j in range(i+1, 3):
+                # Substract squared strict-up-triangle terms
+                # delta2 -= jnp.fft.irfftn(- ki * kj * pot_k)**2
+                nabla_i_nabla_j = gradient_kernel(kvec, i) * gradient_kernel(kvec, j)
+                delta2 -= jnp.fft.irfftn(nabla_i_nabla_j * pot_k)**2
         
-    mesh_size = initial_conditions.shape
+        delta_k2 = fft3d(delta2)
-    local_mesh_shape = (*get_local_shape(initial_conditions.shape), 3)
+        init_force2 = pm_forces(displacement, delta=delta_k2,halo_size=halo_size,sharding=sharding)
-    # TODO
+        # NOTE: growth_factor_second is renormalized: - D2 = 3/7 * growth_factor_second
-    # Displacements have been created in the previous step
+        dx2 = 3/7 * growth_factor_second(cosmo, a) * init_force2
-    # find a way to reuse them
+        p2 = a**2 * growth_rate_second(cosmo, a) * E * dx2
-    displacement = autoshmap(
+        f2 = a**2 * E * dGf2a(cosmo, a) * init_force2
-      partial(jnp.zeros, shape=(local_mesh_shape), dtype='float32'),
-      in_specs=(),
-      out_specs=P('x', 'y'))()  # yapf: disable
 
-    lpt2_delta = lpt2_source(mesh_size, initial_conditions)
+        dx += dx2
-    delta2_k = fft3d(lpt2_delta)
+        p  += p2
-
+        f  += f2
-    lpt2_forces = pm_forces(displacement,
-                            mesh_size,
-                            delta_k=delta2_k,
-                            halo_size=halo_size)
-    dx2 = 3 / 7 * growth_factor_second(cosmo, a) * lpt2_forces
-    p2 = a**2 * growth_rate_second(cosmo, a) * jnp.sqrt(
-        jc.background.Esqr(cosmo, a)) * dx2
-    f2 = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGf2a(cosmo,
-                                                               a) * lpt2_forces
-
-    dx += dx2
-    p += p2
-    f += f2
 
     return dx, p, f
 
@@ -185,10 +153,33 @@ def make_ode_fn(mesh_shape, halo_size=0, sharding=None):
 
     return nbody_ode
 
+def get_ode_fn(cosmo, mesh_shape, halo_size=0, sharding=None):
+
+    def nbody_ode(a, state, args):
+        """
+        State is an array [position, velocities]
+
+        Compatible with [Diffrax API](https://docs.kidger.site/diffrax/)
+        """
+        pos, vel = state
+        forces = pm_forces(pos, mesh_shape, halo_size=halo_size, sharding=sharding) * 1.5 * cosmo.Omega_m
+
+        # Computes the update of position (drift)
+        dpos = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * vel
+        
+        # Computes the update of velocity (kick)
+        dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces
+
+        return jnp.stack([dpos, dvel])
+
+    return nbody_ode
+
 
 def pgd_correction(pos, mesh_shape, params):
     """
-    improve the short-range interactions of PM-Nbody simulations with potential gradient descent method, based on https://arxiv.org/abs/1804.00671
+    improve the short-range interactions of PM-Nbody simulations with potential gradient descent method, 
+    based on https://arxiv.org/abs/1804.00671
+
     args:
       pos: particle positions [npart, 3]
       params: [alpha, kl, ks] pgd parameters
@@ -197,24 +188,20 @@ def pgd_correction(pos, mesh_shape, params):
     delta = cic_paint(jnp.zeros(mesh_shape), pos)
     alpha, kl, ks = params
     delta_k = jnp.fft.rfftn(delta)
-    PGD_range = PGD_kernel(kvec, kl, ks)
+    PGD_range=PGD_kernel(kvec, kl, ks)
     
-    pot_k_pgd = (delta_k * laplace_kernel(kvec)) * PGD_range
+    pot_k_pgd=(delta_k * invlaplace_kernel(kvec))*PGD_range
 
-    forces_pgd = jnp.stack([
+    forces_pgd= jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k_pgd), pos) 
-        cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i) * pot_k_pgd), pos)
+                      for i in range(3)],axis=-1)
-        for i in range(3)
-    ],
-                           axis=-1)
     
-    dpos_pgd = forces_pgd * alpha
+    dpos_pgd = forces_pgd*alpha
    
     return dpos_pgd
 
 
 def make_neural_ode_fn(model, mesh_shape):
-
+    def neural_nbody_ode(state, a, cosmo:Cosmology, params):
-    def neural_nbody_ode(state, a, cosmo, params):
         """
         state is a tuple (position, velocities)
         """
@@ -226,19 +213,15 @@ def make_neural_ode_fn(model, mesh_shape):
         delta_k = jnp.fft.rfftn(delta)
 
         # Computes gravitational potential
-        pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,
+        pot_k = delta_k * invlaplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)
-                                                                  r_split=0)
 
         # Apply a correction filter
-        kk = jnp.sqrt(sum((ki / jnp.pi)**2 for ki in kvec))
+        kk = jnp.sqrt(sum((ki/jnp.pi)**2 for ki in kvec))
-        pot_k = pot_k * (1. + model.apply(params, kk, jnp.atleast_1d(a)))
+        pot_k = pot_k *(1. + model.apply(params, kk, jnp.atleast_1d(a)))
 
         # Computes gravitational forces
-        forces = jnp.stack([
+        forces = jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k), pos) 
-            cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i) * pot_k), pos)
+                          for i in range(3)],axis=-1)
-            for i in range(3)
-        ],
-                           axis=-1)
 
         forces = forces * 1.5 * cosmo.Omega_m
 
@@ -249,5 +232,4 @@ def make_neural_ode_fn(model, mesh_shape):
         dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces
 
         return dpos, dvel
-
     return neural_nbody_ode