forked from guilhem_lavaux/JaxPM
adding hamiltonian gnn demo
This commit is contained in:
parent
fdb25fb661
commit
85e1c65951
1 changed files with 5 additions and 5 deletions
|
@ -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
|
||||
return kbins, P / norm
|
||||
|
|
Loading…
Add table
Reference in a new issue