"""
JaxGB window module
===================
This module provides some well known window functions, which are
commonly applied to LISA data and therefore should be included in the
fast response.
Retrieve a window
-----------------
Use the following function to generate a window for a given size,
using its name, optional parameters or an already-generated array.
All Scipy window functions and aliases from
:mod:`scipy.signal.window.get_window` are supported, in addition to
the windows defined in this module.
.. autofunction:: get_window
Custom windows
--------------
.. autofunction:: tukey4
.. autofunction:: planck
"""
import jax
import jax.numpy as jnp
import numpy as np
import scipy.signal
[docs]
def get_window(window: str | tuple | jax.Array, size: int) -> jax.Array:
"""Return a window function of given size.
Parameters
----------
window
Either a string with the window name or a tuple consisting
of window name and window parameters. If None, a
rectangular window is returned. If array-like, return as is.
size
Number of desired points.
Returns
-------
jax.Array of shape (size,)
window.
"""
if window is None:
return jnp.ones((size))
if isinstance(window, (jnp.ndarray, np.ndarray)):
if len(window) != size:
raise ValueError(f"window length should match given size {size}")
return jnp.asarray(window)
if isinstance(window, tuple) and window[0] == "planck":
return planck(size, *window[1:])
if isinstance(window, tuple) and window[0] == "tukey4":
return tukey4(size, *window[1:])
return scipy.signal.get_window(window, size)
[docs]
def tukey4(n: int = 128, alpha: float = 0.1) -> jax.Array:
"""
Construct a tukey-taper window to the power 4 of length n.
Parameters
----------
n
Number of points in the window.
alpha
Shape parameter of the Tukey window, representing the fraction of the
window inside the cosine tapered region. Default is 0.1.
Returns
-------
jax.Array of shape (n,)
Tukey-taper window.
"""
return scipy.signal.get_window(("tukey", alpha), n) ** 4
[docs]
def planck(n: int = 128, alpha: float = 0.1) -> jax.Array:
"""
Construct a Planck-taper window of length n.
Parameters
----------
n
Number of points in the window.
alpha
Fraction of the window length to taper at each end.
Must be between 0 and 0.5. Default is 0.1.
Returns
-------
jax.Array of shape (n,)
Planck-taper window.
"""
if alpha < 1e-9 or alpha > 0.5:
raise ValueError(
f"alpha {alpha} is not within valid bounds, allowed range is [1e-9, 0.5]."
)
nrange = jnp.arange(n)
w = jnp.zeros(n)
n1 = alpha * (n - 1) # define the transition regions
n2 = (1 - alpha) * (n - 1)
# Region 1: Rising taper (0 <= n < n1)
mask1 = nrange < n1
z1 = jnp.where(
mask1,
alpha * (n - 1) / (nrange + 1e-15)
+ alpha * (n - 1) / (nrange - alpha * (n - 1) + 1e-15),
0.0,
) # avoid division by zero by using where mask is True
w = jnp.where(mask1, 1.0 / (1.0 + jnp.exp(np.clip(z1, -700.0, 700.0))), w)
# Region 2: Flat top (n1 <= n <= n2)
mask2 = (nrange >= n1) & (nrange <= n2)
w = jnp.where(mask2, 1.0, w)
# Region 3: Falling taper (n2 < n < n)
mask3 = nrange > n2
z2 = jnp.where(
mask3,
alpha * (n - 1) / (n - 1 - nrange + 1e-15)
+ alpha * (n - 1) / (n - 1 - nrange - alpha * (n - 1) + 1e-15),
0.0,
)
w = jnp.where(mask3, 1.0 / (1.0 + jnp.exp(np.clip(z2, -700.0, 700.0))), w)
return w