Merge pull request #12 from DifferentiableUniverseInitiative/PGD

PGD

jaxpm/kernels.py

@@ -98,3 +98,29 @@ def cic_compensation(kvec):
   kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
   wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
   return wts
+
+def PGD_kernel(kvec, kl, ks):
+  """
+  Computes the PGD kernel
+  Parameters:
+  -----------
+  kvec: array
+    Array of k values in Fourier space
+  kl: float
+    initial long range scale parameter
+  ks: float
+    initial dhort range scale parameter
+  Returns:
+  --------
+  v: array
+    kernel
+  """
+  kk = sum(ki**2 for ki in kvec)
+  kl2 = kl**2
+  ks4 = ks**4
+  mask = (kk == 0).nonzero()
+  kk[mask] = 1
+  v = jnp.exp(-kl2 / kk) * jnp.exp(-kk**2 / ks4)
+  imask = (~(kk == 0)).astype(int)
+  v *= imask
+  return v

jaxpm/pm.py

@@ -3,7 +3,7 @@
 
 import jax_cosmo as jc
 
-from jaxpm.kernels import fftk, gradient_kernel, laplace_kernel, longrange_kernel
+from jaxpm.kernels import fftk, gradient_kernel, laplace_kernel, longrange_kernel, PGD_kernel
 from jaxpm.painting import cic_paint, cic_read
 from jaxpm.growth import growth_factor, growth_rate, dGfa
 
@@ -70,3 +70,27 @@ def nbody_ode(state, a, cosmo):
         return dpos, dvel
 
     return nbody_ode
+
+
+def pgd_correction(pos, params):
+    """
+    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
+    """
+    kvec = fftk(mesh_shape)
+
+    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)
+
+    pot_k_pgd=(delta_k * laplace_kernel(kvec))*PGD_range
+
+    forces_pgd= jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k_pgd), pos) 
+                      for i in range(3)],axis=-1)
+
+    dpos_pgd = forces_pgd*alpha
+
+    return dpos_pgd