From 6b6b4141958410b3444c02497ed34e28ad6d43bd 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 85b2f4f097077b0d14e5b0a41d51464801cbe466 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 77827fcf44bc0c117fbb3dde9f982f575ddd0452 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 dd71d359e01838de9215629bd4a647ef5fcb2ef7 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 3df0c05d293e24c74219d2299e7397acf9f9da11 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)