mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-06-29 08:31:11 +00:00
57 lines
2 KiB
Python
57 lines
2 KiB
Python
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
|
|
|
|
|
|
@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
|
|
|
|
pixel_scale_x = fov_width / width
|
|
pixel_scale_y = fov_height / height
|
|
|
|
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."
|
|
)
|
|
|
|
y_idx, x_idx = jnp.indices((height, width))
|
|
ra_grid = ra0 + x_idx * pixel_scale_x
|
|
dec_grid = dec0 + y_idx * pixel_scale_y
|
|
|
|
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
|
|
|
|
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)
|
|
|
|
coords = XYZ + jnp.asarray(observer_position)[jnp.newaxis, jnp.newaxis, :]
|
|
|
|
pixels = jhp.ang2pix(nside, ra_flat, dec_flat, lonlat=False)
|
|
|
|
npix = jhp.nside2npix(nside)
|
|
|
|
@partial(jax.vmap, in_axes=(0, None, None))
|
|
def interpolate_volume(coords, volume, pixels):
|
|
voxels = jax.scipy.ndimage.map_coordinates(volume, coords.T, order=1)
|
|
sums = jnp.bincount(pixels, weights=voxels, length=npix)
|
|
return sums
|
|
|
|
sum_map = interpolate_volume(coords, volume, pixels).sum(axis=0)
|
|
counts = jnp.bincount(pixels, length=npix)
|
|
sum_map = jnp.where(counts > 0, sum_map / counts, jhp.UNSEEN)
|
|
|
|
return sum_map
|