guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-23 18:10:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

adding hamiltonian gnn demo

Browse source
This commit is contained in:
EiffL 2022-04-28 00:21:46 +02:00
parent 1795319e4d
commit d0a15e8c78
1 changed files with 5 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
jaxpm/utils.py
View file

@ -66,10 +66,10 @@ def power_spectrum(field, kmin=5, dk=0.5, boxsize=False):
#calculating powerspectra #calculating powerspectra
real = jnp.real(pk).reshape([-1]) real = jnp.real(pk).reshape([-1])
imag = jnp.imag(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() * imag), length=xsum.size) * 1j
Psum += jnp.bincount(dig, weights=(W.flatten() * real), minlength=xsum.size) Psum += jnp.bincount(dig, weights=(W.flatten() * real), length=xsum.size)
P = ((Psum / Nsum)[1:-1] * boxsize.prod()).astype('float32') P = ((Psum / Nsum)[1:-1] * boxsize.prod()).astype('float32')
#normalization for powerspectra #normalization for powerspectra
@ -78,4 +78,4 @@ def power_spectrum(field, kmin=5, dk=0.5, boxsize=False):
#find central values of each bin #find central values of each bin
kbins = kedges[:-1] + (kedges[1:] - kedges[:-1]) / 2 kbins = kedges[:-1] + (kedges[1:] - kedges[:-1]) / 2
return kbins, P / norm return kbins, P / norm