From 6b6b4141958410b3444c02497ed34e28ad6d43bd Mon Sep 17 00:00:00 2001
From: denise lanzieri <dl264294@dappce88.extra.cea.fr>
Date: Tue, 17 May 2022 15:28:30 +0200
Subject: [PATCH] PGD

---
 jaxpm/kernels.py | 26 ++++++++++++++++++++++++++
 jaxpm/pm.py      | 24 +++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/jaxpm/kernels.py b/jaxpm/kernels.py
index c3306c6..97d34dd 100644
--- a/jaxpm/kernels.py
+++ b/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
\ No newline at end of file
diff --git a/jaxpm/pm.py b/jaxpm/pm.py
index 3f39c9c..31872dc 100644
--- a/jaxpm/pm.py
+++ b/jaxpm/pm.py
@@ -3,7 +3,7 @@ import jax.numpy as jnp
 
 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,25 @@ def make_ode_fn(mesh_shape):
         return dpos, dvel
 
     return nbody_ode
+
+
+def pgd_correction(pos, cosmo, params):
+    """
+    improve the short-range interactions of PM-Nbody simulations with potential gradient descent method
+    """
+    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)
+    forces_pgd = forces_pgd * 1.5 * cosmo.Omega_m
+    
+    dpos_pgd = forces_pgd*alpha
+   
+    return dpos_pgd
\ No newline at end of file