2lpt, get_ode, invlaplace, docstrings

jaxpm/kernels.py

@@ -3,8 +3,9 @@ import numpy as np
 
 
 def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
-    """ Return k_vector given a shape (nc, nc, nc) and box_size
-  """
+    """ 
+    Return wave-vectors for a given shape
+    """
     k = []
     for d in range(len(shape)):
         kd = np.fft.fftfreq(shape[d])
@@ -22,18 +23,20 @@ def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
 
 def gradient_kernel(kvec, direction, order=1):
     """
-  Computes the gradient kernel in the requested direction
-  Parameters:
-  -----------
-  kvec: array
-    Array of k values in Fourier space
-  direction: int
-    Index of the direction in which to take the gradient
-  Returns:
-  --------
-  wts: array
-    Complex kernel
-  """
+    Computes the gradient kernel in the requested direction
+    
+    Parameters
+    -----------
+    kvec: list
+        List of wave-vectors in Fourier space
+    direction: int
+        Index of the direction in which to take the gradient
+
+    Returns
+    --------
+    wts: array
+        Complex kernel values
+    """
     if order == 0:
         wts = 1j * kvec[direction]
         wts = jnp.squeeze(wts)
@@ -47,41 +50,43 @@ def gradient_kernel(kvec, direction, order=1):
         return wts
 
 
-def laplace_kernel(kvec):
+def invlaplace_kernel(kvec):
+    """
+    Compute the inverse Laplace kernel
+
+    Parameters
+    -----------
+    kvec: list
+        List of wave-vectors
+
+    Returns
+    --------
+    wts: array
+        Complex kernel values
     """
-  Compute the Laplace kernel from a given K vector
-  Parameters:
-  -----------
-  kvec: array
-    Array of k values in Fourier space
-  Returns:
-  --------
-  wts: array
-    Complex kernel
-  """
     kk = sum(ki**2 for ki in kvec)
-    mask = (kk == 0).nonzero()
-    kk[mask] = 1
-    wts = 1. / kk
-    imask = (~(kk == 0)).astype(int)
-    wts *= imask
-    return wts
+    kk_nozeros = jnp.where(kk==0, 1, kk) 
+    return - jnp.where(kk==0, 0, 1 / kk_nozeros)
 
 
 def longrange_kernel(kvec, r_split):
     """
-  Computes a long range kernel
-  Parameters:
-  -----------
-  kvec: array
-    Array of k values in Fourier space
-  r_split: float
+    Computes a long range kernel
+
+    Parameters
+    -----------
+    kvec: list
+        List of wave-vectors
+    r_split: float
+        Splitting radius
+        
+    Returns
+    --------
+    wts: array
+        Complex kernel values
+    
     TODO: @modichirag add documentation
-  Returns:
-  --------
-  wts: array
-    kernel
-  """
+    """
     if r_split != 0:
         kk = sum(ki**2 for ki in kvec)
         return np.exp(-kk * r_split**2)
@@ -91,15 +96,21 @@ def longrange_kernel(kvec, r_split):
 
 def cic_compensation(kvec):
     """
-  Computes cic compensation kernel.
-  Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
-  Itself based on equation 18 (with p=2) of
-        `Jing et al 2005 <https://arxiv.org/abs/astro-ph/0409240>`_
-  Args:
-    kvec: array of k values in Fourier space
-  Returns:
-    v: array of kernel
-  """
+    Computes cic compensation kernel.
+    Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
+    Itself based on equation 18 (with p=2) of
+          [Jing et al 2005](https://arxiv.org/abs/astro-ph/0409240)
+
+    Parameters:
+    -----------
+    kvec: list
+        List of wave-vectors
+        
+    Returns:
+    --------
+    wts: array
+        Complex kernel values
+    """
     kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
     wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
     return wts
@@ -107,20 +118,22 @@ def cic_compensation(kvec):
 
 def PGD_kernel(kvec, kl, ks):
     """
-  Computes the PGD kernel
-  Parameters:
-  -----------
-  kvec: array
-    Array of k values in Fourier space
-  kl: float
-    initial long range scale parameter
-  ks: float
-    initial dhort range scale parameter
-  Returns:
-  --------
-  v: array
-    kernel
-  """
+    Computes the PGD kernel
+
+    Parameters:
+    -----------
+    kvec: list
+        List of wave-vectors
+    kl: float
+        Initial long range scale parameter
+    ks: float
+        Initial dhort range scale parameter
+
+    Returns:
+    --------
+    v: array
+        Complex kernel values
+    """
     kk = sum(ki**2 for ki in kvec)
     kl2 = kl**2
     ks4 = ks**4