Aquila-Consortium/JaxPM_highres
0
0
Fork 1
forked from guilhem_lavaux/JaxPM
Code Pull requests Activity

fix minor issue

Browse source
This commit is contained in:
EiffL 2022-03-26 03:03:59 +01:00
parent 3f9dfa504a
commit 10b093f07f
1 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
jaxpm/nn.py
View file

@ -35,9 +35,9 @@ class NeuralSplineFourierFilter(hk.Module):
self.n_knots = n_knots
self.latent_size = latent_size
def __call__(self, k, a):
def __call__(self, x, a):
"""
k: array, scale, normalized to fftfreq default
x: array, scale, normalized to fftfreq default
a: scalar, scale factor
"""
@ -57,4 +57,4 @@ class NeuralSplineFourierFilter(hk.Module):
# Augment with repeating points
ak = jnp.concatenate([jnp.zeros((3,)), k, jnp.ones((3,))])
return _deBoorVectorized(jnp.clip(k/jnp.sqrt(3), 0, 1-1e-4), ak, w, 3)
return _deBoorVectorized(jnp.clip(x/jnp.sqrt(3), 0, 1-1e-4), ak, w, 3)