guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-08 04:40:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add assert to make pyright happy

Browse source
This commit is contained in:
Wassim KABALAN 2024-08-02 21:21:59 +02:00
parent 75604d2396
commit ccbfee3615
1 changed files with 4 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
jaxpm/pm.py
View file

@ -19,11 +19,12 @@ def pm_forces(positions, mesh_shape=None, delta=None, r_split=0, halo_size=0):
Computes gravitational forces on particles using a PM scheme Computes gravitational forces on particles using a PM scheme
""" """
if mesh_shape is None: if mesh_shape is None:
assert(delta is not None) , "If mesh_shape is not provided, delta should be provided"
mesh_shape = delta.shape mesh_shape = delta.shape
kvec = fftk(mesh_shape) kvec = fftk(mesh_shape)
if delta is None: if delta is None:
delta_k = fft3d(cic_paint_dx(positions, halo_size=0)) delta_k = fft3d(cic_paint_dx(positions, halo_size=halo_size))
else: else:
delta_k = fft3d(delta) delta_k = fft3d(delta)
@ -32,7 +33,7 @@ def pm_forces(positions, mesh_shape=None, delta=None, r_split=0, halo_size=0):
r_split=r_split) r_split=r_split)
# Computes gravitational forces # Computes gravitational forces
forces = jnp.stack([ forces = jnp.stack([
cic_read_dx(ifft3d(gradient_kernel(kvec, i) * pot_k), halo_size=0) cic_read_dx(ifft3d(gradient_kernel(kvec, i) * pot_k), halo_size=halo_size)
for i in range(3) for i in range(3)
], ],
axis=-1) axis=-1)