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

add interpolate operator

Browse source
This commit is contained in:
Wassim KABALAN 2024-07-09 02:34:24 +02:00
parent c99b2c3ef4
commit e506dc2a7e
1 changed files with 8 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
jaxpm/ops.py
View file

@ -1,5 +1,8 @@
import jax.numpy as jnp
import jax_cosmo as jc
import numpy as np import numpy as np
from _src.spmd_config import pm_operators
from jaxpm._src.spmd_config import pm_operators
def fftn(arr): def fftn(arr):
@ -32,3 +35,7 @@ def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
def generate_initial_positions(shape): def generate_initial_positions(shape):
return pm_operators.generate_initial_positions(shape) return pm_operators.generate_initial_positions(shape)
def interpolate_ic(kfield, kk, cosmo: jc.Cosmology, box_size):
return pm_operators.interpolate_ic(kfield, kk, cosmo, box_size)