From d3026f7849479b61ae4dbf748a6fd69eb01c71e8 Mon Sep 17 00:00:00 2001 From: denise lanzieri Date: Tue, 17 May 2022 15:28:30 +0200 Subject: [PATCH 1/5] 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 From b7376313463fca6efc760a13044c154731f8c154 Mon Sep 17 00:00:00 2001 From: Denise Lanzieri <72620117+dlanzieri@users.noreply.github.com> Date: Wed, 18 May 2022 09:55:31 +0200 Subject: [PATCH 2/5] Update jaxpm/pm.py Co-authored-by: Francois Lanusse --- jaxpm/pm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/jaxpm/pm.py b/jaxpm/pm.py index 31872dc..d587b23 100644 --- a/jaxpm/pm.py +++ b/jaxpm/pm.py @@ -87,7 +87,6 @@ def pgd_correction(pos, cosmo, params): 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 From 8266f113696de606c10d04d55c691013575d66b2 Mon Sep 17 00:00:00 2001 From: Denise Lanzieri <72620117+dlanzieri@users.noreply.github.com> Date: Wed, 18 May 2022 09:58:34 +0200 Subject: [PATCH 3/5] Update jaxpm/pm.py Co-authored-by: Francois Lanusse --- jaxpm/pm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxpm/pm.py b/jaxpm/pm.py index d587b23..32a9691 100644 --- a/jaxpm/pm.py +++ b/jaxpm/pm.py @@ -72,7 +72,7 @@ def make_ode_fn(mesh_shape): return nbody_ode -def pgd_correction(pos, cosmo, params): +def pgd_correction(pos, params): """ improve the short-range interactions of PM-Nbody simulations with potential gradient descent method """ From 782c9dbc6c601600f8ffa52241daf50ce406e176 Mon Sep 17 00:00:00 2001 From: Denise Lanzieri <72620117+dlanzieri@users.noreply.github.com> Date: Wed, 18 May 2022 09:58:46 +0200 Subject: [PATCH 4/5] Update jaxpm/pm.py Co-authored-by: Francois Lanusse --- jaxpm/pm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxpm/pm.py b/jaxpm/pm.py index 32a9691..7aac781 100644 --- a/jaxpm/pm.py +++ b/jaxpm/pm.py @@ -74,7 +74,7 @@ def make_ode_fn(mesh_shape): def pgd_correction(pos, params): """ - improve the short-range interactions of PM-Nbody simulations with potential gradient descent method + improve the short-range interactions of PM-Nbody simulations with potential gradient descent method, based on https://arxiv.org/abs/1804.00671 """ kvec = fftk(mesh_shape) From 5394abbed1f492a8840f13045ce0556c719686f5 Mon Sep 17 00:00:00 2001 From: Denise Lanzieri <72620117+dlanzieri@users.noreply.github.com> Date: Wed, 18 May 2022 09:59:59 +0200 Subject: [PATCH 5/5] Update jaxpm/pm.py Co-authored-by: Francois Lanusse --- jaxpm/pm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/jaxpm/pm.py b/jaxpm/pm.py index 7aac781..d9870f7 100644 --- a/jaxpm/pm.py +++ b/jaxpm/pm.py @@ -75,6 +75,9 @@ def make_ode_fn(mesh_shape): 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)