guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-29 16:41:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Replace np.power 0.5 by np.sqrt (#43)
Some checks failed
Tests / run_tests (3.10) (push) Failing after 2m35s
Details
Tests / run_tests (3.11) (push) Failing after 2m37s
Details
Code Formatting / formatting (push) Successful in 3m1s
Details
Tests / run_tests (3.12) (push) Failing after 1m16s
Details

Browse source 
* Switch **2  by np.sqrt

* format

---------

Co-authored-by: Francois Lanusse <fr.eiffel@gmail.com>
This commit is contained in:
Wassim KABALAN 2025-06-29 10:51:45 +02:00 • committed by GitHub
parent 6693e5c725
commit 7d76573701
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 6 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
jaxpm/utils.py
View file

@ -52,7 +52,7 @@ def _initialize_pk(mesh_shape, box_shape, kedges, los):
kshapes = np.eye(len(mesh_shape), dtype=np.int32) * -2 + 1
kvec = [(2 * np.pi * m / l) * np.fft.fftfreq(m).reshape(kshape)
for m, l, kshape in zip(mesh_shape, box_shape, kshapes)]
kmesh = sum(ki**2 for ki in kvec)**0.5
kmesh = jnp.sqrt(sum(ki**2 for ki in kvec))
dig = np.digitize(kmesh.reshape(-1), kedges)
kcount = np.bincount(dig, minlength=len(kedges) + 1)