Merge remote-tracking branch 'upstream/main' into ASKabalan/jaxdecomp_proto

jaxpm/pm.py

@@ -172,7 +172,7 @@ def make_ode_fn(mesh_shape, halo_size=0):
     return nbody_ode
 
 
-def pgd_correction(pos, params):
+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
     args:
@@ -180,20 +180,51 @@ def pgd_correction(pos, params):
       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
+    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
+
+
+def make_neural_ode_fn(model, mesh_shape):
+    def neural_nbody_ode(state, a, cosmo, params):
+        """
+        state is a tuple (position, velocities)
+        """
+        pos, vel = state
+        kvec = fftk(mesh_shape)
+
+        delta = cic_paint(jnp.zeros(mesh_shape), pos)
+
+        delta_k = jnp.fft.rfftn(delta)
+
+        # Computes gravitational potential
+        pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)
+
+        # Apply a correction filter
+        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)))
+
+        # Computes gravitational forces
+        forces = jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) 
+                          for i in range(3)],axis=-1)
+
+        forces = forces * 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 dpos, dvel
+    return neural_nbody_ode