From 85e1c65951a2f8541ddf045f57d10f09abbd253c Mon Sep 17 00:00:00 2001 From: EiffL Date: Thu, 28 Apr 2022 00:21:46 +0200 Subject: [PATCH] adding hamiltonian gnn demo --- jaxpm/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/jaxpm/utils.py b/jaxpm/utils.py index b541c65..312f73f 100644 --- a/jaxpm/utils.py +++ b/jaxpm/utils.py @@ -66,10 +66,10 @@ def power_spectrum(field, kmin=5, dk=0.5, boxsize=False): #calculating powerspectra real = jnp.real(pk).reshape([-1]) imag = jnp.imag(pk).reshape([-1]) - - Psum = jnp.bincount(dig, weights=(W.flatten() * imag), minlength=xsum.size) * 1j - Psum += jnp.bincount(dig, weights=(W.flatten() * real), minlength=xsum.size) - + + Psum = jnp.bincount(dig, weights=(W.flatten() * imag), length=xsum.size) * 1j + Psum += jnp.bincount(dig, weights=(W.flatten() * real), length=xsum.size) + P = ((Psum / Nsum)[1:-1] * boxsize.prod()).astype('float32') #normalization for powerspectra @@ -78,4 +78,4 @@ def power_spectrum(field, kmin=5, dk=0.5, boxsize=False): #find central values of each bin kbins = kedges[:-1] + (kedges[1:] - kedges[:-1]) / 2 - return kbins, P / norm \ No newline at end of file + return kbins, P / norm