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

fixing formatting

Browse source
This commit is contained in:
Francois Lanusse 2025-06-28 23:21:31 +02:00
parent 1fcfc2aef2
commit 9d0b047f06
2 changed files with 56 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

25
jaxpm/spherical.py
View file

@ -1,12 +1,16 @@
from functools import partial
import healpy as hp
import jax
import jax.numpy as jnp
import jax_healpy as jhp
import matplotlib.pyplot as plt
import jax
from functools import partial
import healpy as hp
@partial(jax.jit, static_argnames=('nside', 'fov', 'center_radec' , 'd_R' , 'box_size'))
def paint_spherical(volume, nside, fov, center_radec, observer_position, box_size, R, d_R):
@partial(jax.jit,
static_argnames=('nside', 'fov', 'center_radec', 'd_R', 'box_size'))
def paint_spherical(volume, nside, fov, center_radec, observer_position,
box_size, R, d_R):
width, height, depth = volume.shape
ra0, dec0 = center_radec
fov_width, fov_height = fov
@ -16,7 +20,9 @@ def paint_spherical(volume, nside, fov, center_radec, observer_position, box_siz
res_deg = jhp.nside2resol(nside, arcmin=True) / 60
if pixel_scale_x > res_deg or pixel_scale_y > res_deg:
print(f"WARNING Pixel scale ({pixel_scale_x:.4f} deg, {pixel_scale_y:.4f} deg) is larger than the Healpy resolution ({res_deg:.4f} deg). Increase the field of view or decrease the nside.")
print(
f"WARNING Pixel scale ({pixel_scale_x:.4f} deg, {pixel_scale_y:.4f} deg) is larger than the Healpy resolution ({res_deg:.4f} deg). Increase the field of view or decrease the nside."
)
y_idx, x_idx = jnp.indices((height, width))
ra_grid = ra0 + x_idx * pixel_scale_x
@ -24,13 +30,14 @@ def paint_spherical(volume, nside, fov, center_radec, observer_position, box_siz
ra_flat = ra_grid.flatten() * jnp.pi / 180.0
dec_flat = dec_grid.flatten() * jnp.pi / 180.0
R_s = jnp.arange(0 , d_R, 1.0) + R
R_s = jnp.arange(0, d_R, 1.0) + R
XYZ = R_s.reshape(-1, 1, 1) * jhp.ang2vec(ra_flat, dec_flat, lonlat=False)
observer_position = jnp.array(observer_position)
# Convert observer position from box units to grid units
observer_position = observer_position / jnp.array(box_size) * jnp.array(volume.shape)
observer_position = observer_position / jnp.array(box_size) * jnp.array(
volume.shape)
coords = XYZ + jnp.asarray(observer_position)[jnp.newaxis, jnp.newaxis, :]
pixels = jhp.ang2pix(nside, ra_flat, dec_flat, lonlat=False)