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

adds fix to make code jittablel

Browse source
This commit is contained in:
EiffL 2022-05-17 23:37:55 +02:00
parent 5dc239927f
commit 5108e56ee8
2 changed files with 7 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

6
jaxpm/lensing.py
View file

@ -26,11 +26,9 @@ def density_plane(positions,
xy = xy / nx * plane_resolution
# Selecting only particles that fall inside the volume of interest
mask = (d > (center - width / 2)) & (d <= (center + width / 2))
xy = xy[mask]
weight = jnp.where((d > (center - width / 2)) & (d <= (center + width / 2)), 1., 0.)
# Painting density plane
density_plane = cic_paint_2d(jnp.zeros([plane_resolution, plane_resolution]), xy)
density_plane = cic_paint_2d(jnp.zeros([plane_resolution, plane_resolution]), xy, weight)
# Apply density normalization
density_plane = density_plane / ((nx / plane_resolution) *