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

Allow applying weights with relative cic paint

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-30 01:57:32 +01:00
parent b3a264ad53
commit b09580d59e
2 changed files with 28 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

5
jaxpm/painting_utils.py
View file

@ -104,8 +104,7 @@ def _scatter_chunk(carry, chunk):
spatial_shape)
# scatter
ind = tuple(ind[..., i] for i in range(spatial_ndim))
mesh = mesh.at[ind].add(val * frac)
mesh = mesh.at[ind].add(jnp.multiply(jnp.expand_dims(val, axis=-1), frac))
carry = mesh, offset, cell_size, mesh_shape
return carry, None
@ -117,11 +116,9 @@ def scatter(pmid,
val=1.,
offset=0,
cell_size=1.):
ptcl_num, spatial_ndim = pmid.shape
val = jnp.asarray(val)
mesh = jnp.asarray(mesh)
remainder, chunks = _chunk_split(ptcl_num, chunk_size, pmid, disp, val)
carry = mesh, offset, cell_size, mesh.shape
if remainder is not None: