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

Merge remote-tracking branch 'origin/main' into 41-spherical-lensing

Browse source
This commit is contained in:
Wassim Kabalan 2025-06-29 11:32:30 +02:00
commit 09d82283f5
4 changed files with 6 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
jaxpm/growth.py
View file

@ -28,7 +28,7 @@ def E(cosmo, a):
where :math:`f(a)` is the Dark Energy evolution parameter computed where :math:`f(a)` is the Dark Energy evolution parameter computed
by :py:meth:`.f_de`. by :py:meth:`.f_de`.
""" """
return np.power(Esqr(cosmo, a), 0.5) return np.sqrt(Esqr(cosmo, a))
def df_de(cosmo, a, epsilon=1e-5): def df_de(cosmo, a, epsilon=1e-5):

2
jaxpm/pm.py
View file

@ -139,7 +139,7 @@ def linear_field(mesh_shape, box_size, pk, seed, sharding=None):
pkmesh = pk(kmesh) * (mesh_shape[0] * mesh_shape[1] * mesh_shape[2]) / ( pkmesh = pk(kmesh) * (mesh_shape[0] * mesh_shape[1] * mesh_shape[2]) / (
box_size[0] * box_size[1] * box_size[2]) box_size[0] * box_size[1] * box_size[2])
field = field * (pkmesh)**0.5 field = field * jnp.sqrt(pkmesh)
field = ifft3d(field) field = ifft3d(field)
return field return field

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 kshapes = np.eye(len(mesh_shape), dtype=np.int32) * -2 + 1
kvec = [(2 * np.pi * m / l) * np.fft.fftfreq(m).reshape(kshape) kvec = [(2 * np.pi * m / l) * np.fft.fftfreq(m).reshape(kshape)
for m, l, kshape in zip(mesh_shape, box_shape, kshapes)] 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) dig = np.digitize(kmesh.reshape(-1), kedges)
kcount = np.bincount(dig, minlength=len(kedges) + 1) kcount = np.bincount(dig, minlength=len(kedges) + 1)

5
tests/conftest.py
View file

@ -96,8 +96,9 @@ def fpm_initial_conditions(cosmo, particle_mesh):
whitec = particle_mesh.generate_whitenoise(42, whitec = particle_mesh.generate_whitenoise(42,
type='complex', type='complex',
unitary=False) unitary=False)
lineark = whitec.apply(lambda k, v: pk_fn(sum(ki**2 for ki in k)**0.5)**0.5 lineark = whitec.apply(lambda k, v: jnp.sqrt(
* v * (1 / v.BoxSize).prod()**0.5) pk_fn(jnp.sqrt(sum(ki**2 for ki in k)))) * v * jnp.sqrt(
(1 / v.BoxSize).prod()))
init_mesh = lineark.c2r().value # XXX init_mesh = lineark.c2r().value # XXX
return lineark, grid, init_mesh return lineark, grid, init_mesh