mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-06-14 18:01:11 +00:00
Compare commits
5 commits
f391a79211
...
7623e60581
Author | SHA1 | Date | |
---|---|---|---|
|
7623e60581 | ||
|
d874790543 | ||
|
187cf5c4ba | ||
|
67a80e1041 | ||
|
e1a8134b8e |
4 changed files with 5 additions and 2 deletions
2
.github/workflows/tests.yml
vendored
2
.github/workflows/tests.yml
vendored
|
@ -35,6 +35,8 @@ jobs:
|
|||
pip install pytest
|
||||
pip install diffrax
|
||||
pip install .
|
||||
echo "numpy version installed:"
|
||||
python -c "import numpy; print(numpy.__version__)"
|
||||
|
||||
- name: Run Single Device Tests
|
||||
run: |
|
||||
|
|
|
@ -166,7 +166,7 @@ def uniform_particles(mesh_shape, sharding=None):
|
|||
axis=-1)
|
||||
|
||||
|
||||
def normal_field(seed , shape, sharding=None, dtype='float32'):
|
||||
def normal_field(seed, shape, sharding=None, dtype=float):
|
||||
"""Generate a Gaussian random field with the given power spectrum."""
|
||||
gpu_mesh = sharding.mesh if sharding is not None else None
|
||||
if gpu_mesh is not None and not (gpu_mesh.empty):
|
||||
|
|
|
@ -131,7 +131,7 @@ def linear_field(mesh_shape, box_size, pk, seed, sharding=None):
|
|||
Generate initial conditions.
|
||||
"""
|
||||
# Initialize a random field with one slice on each gpu
|
||||
field = normal_field(seed=seed , shape=mesh_shape, sharding=sharding)
|
||||
field = normal_field(seed=seed, shape=mesh_shape, sharding=sharding)
|
||||
field = fft3d(field)
|
||||
kvec = fftk(field)
|
||||
kmesh = sum((kk / box_size[i] * mesh_shape[i])**2
|
||||
|
|
|
@ -1,3 +1,4 @@
|
|||
pfft-python @ git+https://github.com/MP-Gadget/pfft-python
|
||||
pmesh @ git+https://github.com/MP-Gadget/pmesh
|
||||
fastpm @ git+https://github.com/ASKabalan/fastpm-python
|
||||
numpy==2.2.6
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue