pywavelet 0.1.1__tar.gz → 0.2.0__tar.gz

{pywavelet-0.1.1 → pywavelet-0.2.0}/CHANGELOG.rst

@@ -5,16 +5,54 @@ CHANGELOG
 =========
 
 
+.. _changelog-v0.2.0:
+
+v0.2.0 (2025-01-23)
+===================
+
+Features
+--------
+
+* feat: add jax as optional backend (`264613e`_)
+
+.. _264613e: https://github.com/pywavelet/pywavelet/commit/264613e5a58042641eb6814530dab36bb54b3371
+
+
+.. _changelog-v0.1.2:
+
+v0.1.2 (2025-01-20)
+===================
+
+Bug Fixes
+---------
+
+* fix: adjust the WaveletMask repr (`50c6201`_)
+
+Chores
+------
+
+* chore(release): 0.1.2 (`7b64515`_)
+
+.. _50c6201: https://github.com/pywavelet/pywavelet/commit/50c6201efb7689dd9757a5e4c6047d241015cb96
+.. _7b64515: https://github.com/pywavelet/pywavelet/commit/7b64515a2bf2418719f68cb6b15f1c204938408d
+
+
 .. _changelog-v0.1.1:
 
 v0.1.1 (2025-01-16)
 ===================
 
+Chores
+------
+
+* chore(release): 0.1.1 (`bd15453`_)
+
 Unknown
 -------
 
 * Merge branch 'main' of github.com:avivajpeyi/pywavelet into main (`69eefa2`_)
 
+.. _bd15453: https://github.com/pywavelet/pywavelet/commit/bd15453e028705548232b802b2d21bbebd307ca7
 .. _69eefa2: https://github.com/pywavelet/pywavelet/commit/69eefa29b7873c30fcb74ad1e051eb20101a277a
 
 

{pywavelet-0.1.1 → pywavelet-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.1.1
+Version: 0.2.0
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/
@@ -21,6 +21,8 @@ Requires-Dist: tqdm
 Requires-Dist: rich
 Requires-Dist: rocket-fft
 Requires-Dist: astropy>=5.2.1
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.1.0; extra == "dev"

{pywavelet-0.1.1 → pywavelet-0.2.0}/README.rst

@@ -23,6 +23,7 @@ pywavelet is available on PyPI and can be installed with `pip <https://pip.pypa.
     $ pip install pywavelet
 
 For developers
+--------------
 
 First set up a conda environment with the latest version of python.
 
@@ -35,7 +36,8 @@ First set up a conda environment with the latest version of python.
     $ pip install -e ".[dev]"
     $ pre-commit install
 
-## Test code
+Test code
+---------
 
 Locate directory /tests from root directory. run 
 

{pywavelet-0.1.1 → pywavelet-0.2.0}/docs/example.ipynb

@@ -19,8 +19,8 @@
      "hide-cell"
     ],
     "ExecuteTime": {
-     "end_time": "2024-10-14T02:14:53.298769Z",
-     "start_time": "2024-10-14T02:14:53.269666Z"
+     "end_time": "2025-01-23T03:42:44.922293Z",
+     "start_time": "2025-01-23T03:42:43.279788Z"
     }
    },
    "source": [
@@ -29,7 +29,7 @@
     "from typing import List\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from pywavelet.transforms.types import TimeSeries\n",
+    "from pywavelet.types import TimeSeries\n",
     "from pywavelet.transforms import from_time_to_wavelet, from_wavelet_to_time\n",
     "\n",
     "\n",
@@ -67,6 +67,7 @@
     "    )\n",
     "    ax.set_xlabel(\"Residuals\")\n",
     "    ax.set_ylabel(\"Count\")\n",
+    "    ax.ticklabel_format(axis=\"x\", style=\"sci\", scilimits=(0, 0))\n",
     "    return ax\n",
     "\n",
     "# Sizes\n",
@@ -103,7 +104,7 @@
     "plt.close()"
    ],
    "outputs": [],
-   "execution_count": 39
+   "execution_count": 3
   },
   {
    "cell_type": "markdown",
@@ -125,14 +126,14 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-10-14T02:15:20.552385Z",
-     "start_time": "2024-10-14T02:15:19.877799Z"
+     "end_time": "2025-01-23T03:42:42.572318Z",
+     "start_time": "2025-01-23T03:42:39.773430Z"
     }
    },
    "cell_type": "code",
    "source": [
     "import numpy as np\n",
-    "from pywavelet.transforms.types import FrequencySeries\n",
+    "from pywavelet.types import FrequencySeries\n",
     "from pywavelet.transforms import from_freq_to_wavelet, from_wavelet_to_freq\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
@@ -162,7 +163,7 @@
     "plt.close()"
    ],
    "outputs": [],
-   "execution_count": 44
+   "execution_count": 2
   },
   {
    "metadata": {},
@@ -171,13 +172,6 @@
     "![](roundtrip_freq.png)\n",
     "\n"
    ]
-  },
-  {
-   "metadata": {},
-   "cell_type": "code",
-   "outputs": [],
-   "execution_count": null,
-   "source": ""
   }
  ],
  "metadata": {

{pywavelet-0.1.1 → pywavelet-0.2.0}/pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+jax = ["jax"]
 dev = [
     "pytest>=6.0",
     "pytest-cov>=4.1.0",

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

pywavelet-0.2.0/src/pywavelet/backend.py

@@ -0,0 +1,26 @@
+import os
+
+try:
+    import jax
+
+    jax_available = True
+
+
+except ImportError:
+    jax_available = False
+
+use_jax = jax_available and os.getenv("PYWAVELET_JAX", "0") == "1"
+
+if use_jax:
+    import jax.numpy as xp  # type: ignore
+    from jax.scipy.fft import fft, ifft, rfft, irfft, rfftfreq  # type: ignore
+    from jax.scipy.special import betainc  # type: ignore
+
+
+else:
+    import numpy as xp  # type: ignore
+    from numpy.fft import fft, ifft, rfft, irfft, rfftfreq  # type: ignore
+    from scipy.special import betainc # type: ignore
+
+
+PI = xp.pi

pywavelet-0.2.0/src/pywavelet/transforms/__init__.py

@@ -0,0 +1,17 @@
+from .numpy import (
+    from_wavelet_to_time,
+    from_wavelet_to_freq,
+    from_time_to_wavelet,
+    from_freq_to_wavelet,
+)
+from .phi_computer import phi_vec, phitilde_vec_norm, phitilde_vec
+
+__all__ = [
+    "from_wavelet_to_time",
+    "from_wavelet_to_freq",
+    "from_time_to_wavelet",
+    "from_freq_to_wavelet",
+    "phitilde_vec_norm",
+    "phi_vec",
+    "phitilde_vec",
+]

pywavelet-0.2.0/src/pywavelet/transforms/jax/forward/__init__.py

@@ -0,0 +1,6 @@
+from .main import from_freq_to_wavelet, from_time_to_wavelet
+from ....logger import logger
+
+logger.warning("JAX SUBPACKAGE NOT YET TESTED")
+
+__all__ = ["from_time_to_wavelet", "from_freq_to_wavelet"]

pywavelet-0.2.0/src/pywavelet/transforms/jax/forward/from_freq.py

@@ -0,0 +1,56 @@
+import jax
+import jax.numpy as jnp
+from functools import partial
+from jax import jit
+from jax.numpy.fft import ifft
+
+@partial(jit, static_argnames=('Nf', 'Nt'))
+def transform_wavelet_freq_helper(
+        data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray
+) -> jnp.ndarray:
+    # Initially all wrk being done in time-rws, freq-cols
+    wave = jnp.zeros((Nt, Nf))
+    f_bins = jnp.arange(Nf)
+
+    i_base = Nt // 2
+    jj_base = f_bins * Nt // 2
+
+    initial_values = jnp.where(
+        (f_bins == 0) | (f_bins == Nf),
+        phif[0] * data[f_bins * Nt // 2] / 2.0,
+        phif[0] * data[f_bins * Nt // 2]
+    )
+
+    DX = jnp.zeros((Nf, Nt), dtype=jnp.complex64)
+    DX = DX.at[:, Nt // 2].set(initial_values)
+
+    j_range = jnp.arange(1 - Nt // 2, Nt // 2)
+    j = jnp.abs(j_range)
+    i = i_base + j_range
+
+    cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)
+    cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)
+    cond3 = j[None, :] == 0
+
+    jj = jj_base[:, None] + j_range[None, :]
+    val = jnp.where(cond1 | cond2, 0.0, phif[j] * data[jj])
+    DX = DX.at[:, i].set(jnp.where(cond3, DX[:, i], val))
+
+    # Vectorized ifft
+    DX_trans = ifft(DX, axis=1)
+
+    # Vectorized __fill_wave_2_jax
+    n_range = jnp.arange(Nt)
+    cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1
+    cond2 = jnp.expand_dims(f_bins % 2 == 1, axis=-1) # shape: (Nf, 1)
+
+    real_part = jnp.where(cond2, -jnp.imag(DX_trans), jnp.real(DX_trans))
+    imag_part = jnp.where(cond2, jnp.real(DX_trans), jnp.imag(DX_trans))
+
+    wave = jnp.where(cond1, imag_part.T, real_part.T)
+
+    ## Special cases for f_bin 0 and Nf
+    wave = wave.at[::2, 0].set(jnp.real(DX_trans[0, ::2] * jnp.sqrt(2)))
+    wave = wave.at[1::2, -1].set(jnp.real(DX_trans[-1, ::2] * jnp.sqrt(2)))
+
+    return wave.T

pywavelet-0.2.0/src/pywavelet/transforms/jax/forward/from_time.py

@@ -0,0 +1,51 @@
+import jax
+import jax.numpy as jnp
+from jax import jit
+from jax.numpy.fft import rfft
+from functools import partial
+
+@partial(jit, static_argnames=('Nf', 'Nt', 'mult'))
+def transform_wavelet_time_helper(
+    data: jnp.ndarray, phi: jnp.ndarray, Nf: int, Nt: int,  mult: int
+) -> jnp.ndarray:
+    """Helper function to do the wavelet transform in the time domain using JAX"""
+    # Define constants
+    ND = Nf * Nt
+    K = mult * 2 * Nf
+
+    # Pad the data with K extra values
+    data_pad = jnp.concatenate((data, data[:K]))
+
+    # Generate time bin indices
+    time_bins = jnp.arange(Nt)
+    jj_base = (time_bins[:, None] * Nf - K // 2) % ND
+    jj = (jj_base + jnp.arange(K)[None, :]) % ND
+
+    # Apply the window (phi) to the data
+    wdata = data_pad[jj] * phi[None, :]
+
+    # Perform FFT on the windowed data
+    wdata_trans = rfft(wdata, axis=1)
+
+    # Initialize the wavelet transform result
+    wave = jnp.zeros((Nt, Nf))
+
+    # Handle m=0 case for even time bins
+    even_mask = (time_bins % 2 == 0) & (time_bins < Nt - 1)
+    even_indices = jnp.nonzero(even_mask, size=even_mask.shape[0])[0]
+
+    # Update wave for m=0 using even time bins
+    wave = wave.at[even_indices, 0].set(jnp.real(wdata_trans[even_indices, 0]) / jnp.sqrt(2))
+    wave = wave.at[even_indices + 1, 0].set(jnp.real(wdata_trans[even_indices, Nf * mult]) / jnp.sqrt(2))
+
+    # Handle other cases (j > 0) using vectorized operations
+    j_range = jnp.arange(1, Nf)
+    odd_condition = ((time_bins[:, None] + j_range[None, :]) % 2 == 1)
+
+    wave = wave.at[:, 1:].set(
+        jnp.where(odd_condition,
+                  -jnp.imag(wdata_trans[:, j_range * mult]),
+                  jnp.real(wdata_trans[:, j_range * mult]))
+    )
+
+    return wave.T

pywavelet-0.2.0/src/pywavelet/transforms/jax/forward/main.py

@@ -0,0 +1,110 @@
+from typing import Union
+
+import jax.numpy as jnp
+
+from ....logger import logger
+from ....types import FrequencySeries, TimeSeries, Wavelet
+from ....types.wavelet_bins import _get_bins, _preprocess_bins
+from ...phi_computer import phi_vec, phitilde_vec_norm
+from .from_freq import transform_wavelet_freq_helper
+from .from_time import transform_wavelet_time_helper
+
+
+def from_time_to_wavelet(
+    timeseries: TimeSeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+    mult: int = 32,
+    **kwargs,
+) -> Wavelet:
+    """Transforms time-domain data to wavelet-domain data.
+
+    Warning: there can be significant leakage if mult is too small and the
+    transform is only approximately exact if mult=Nt/2
+
+    Parameters
+    ----------
+    timeseries : TimeSeries
+        Time domain data
+    Nf : int
+        Number of frequency bins
+    Nt : int
+        Number of time bins
+    nx : float, optional
+        Number of standard deviations for the phi_vec, by default 4.
+    mult : int, optional
+        Number of time bins to use for the wavelet transform, by default 32
+    **kwargs:
+        Additional keyword arguments passed to the Wavelet.from_data constructor.
+
+    Returns
+    -------
+    Wavelet
+        Wavelet domain data
+
+    """
+    Nf, Nt = _preprocess_bins(timeseries, Nf, Nt)
+    dt = timeseries.dt
+    t_bins, f_bins = _get_bins(timeseries, Nf, Nt)
+
+    ND = Nf * Nt
+
+    if len(timeseries) != ND:
+        logger.warning(
+            f"len(freqseries)={len(timeseries)} != Nf*Nt={ND}. Truncating to freqseries[:{ND}]"
+        )
+        timeseries = timeseries[:ND]
+    if mult > Nt / 2:
+        logger.warning(
+            f"mult={mult} is too large for Nt={Nt}. This may lead to bogus results."
+        )
+
+    mult = min(mult, Nt // 2)  # make sure K isn't bigger than ND
+    phi = jnp.array(phi_vec(Nf, d=nx, q=mult))
+    wave = transform_wavelet_time_helper(timeseries.data, Nf=Nf, Nt=Nt, phi=phi, mult=mult)
+    return Wavelet(
+        wave* jnp.sqrt(2), time=t_bins, freq=f_bins
+    )
+
+
+def from_freq_to_wavelet(
+    freqseries: FrequencySeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+    **kwargs,
+) -> Wavelet:
+    """Transforms frequency-domain data to wavelet-domain data.
+
+    Parameters
+    ----------
+    freqseries : FrequencySeries
+        Frequency domain data
+    Nf : int
+        Number of frequency bins
+    Nt : int
+        Number of time bins
+    nx : float, optional
+        Number of standard deviations for the phi_vec, by default 4.
+    **kwargs:
+        Additional keyword arguments passed to the Wavelet.from_data constructor.
+
+    Returns
+    -------
+    Wavelet
+        Wavelet domain data
+
+    """
+    Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
+    t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
+    phif = jnp.array(phitilde_vec_norm(Nf, Nt,  d=nx))
+    wave =  transform_wavelet_freq_helper(
+        freqseries.data, Nf=Nf, Nt=Nt, phif=phif
+    )
+
+    return Wavelet(
+        (2 / Nf) * wave * jnp.sqrt(2),
+        time=t_bins,
+        freq=f_bins
+    )

{pywavelet-0.1.1/src/pywavelet/transforms → pywavelet-0.2.0/src/pywavelet/transforms/numpy}/__init__.py

@@ -1,9 +1,10 @@
 from .forward import from_freq_to_wavelet, from_time_to_wavelet
 from .inverse import from_wavelet_to_freq, from_wavelet_to_time
 
+
 __all__ = [
     "from_wavelet_to_time",
     "from_wavelet_to_freq",
     "from_time_to_wavelet",
     "from_freq_to_wavelet",
-]
+]

{pywavelet-0.1.1/src/pywavelet/transforms → pywavelet-0.2.0/src/pywavelet/transforms/numpy}/forward/from_freq.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 from numba import njit
-from numpy import fft
 
 
 def transform_wavelet_freq_helper(

{pywavelet-0.1.1/src/pywavelet/transforms → pywavelet-0.2.0/src/pywavelet/transforms/numpy}/forward/main.py

@@ -2,10 +2,10 @@ from typing import Union
 
 import numpy as np
 
-from ...logger import logger
-from ...types import FrequencySeries, TimeSeries, Wavelet
-from ...types.wavelet_bins import _get_bins, _preprocess_bins
-from ..phi_computer import phi_vec, phitilde_vec_norm
+from ....logger import logger
+from ....types import FrequencySeries, TimeSeries, Wavelet
+from ....types.wavelet_bins import _get_bins, _preprocess_bins
+from ...phi_computer import phi_vec, phitilde_vec_norm
 from .from_freq import transform_wavelet_freq_helper
 from .from_time import transform_wavelet_time_helper
 
@@ -72,7 +72,7 @@ def from_time_to_wavelet(
         )
 
     mult = min(mult, Nt // 2)  # Ensure mult is not larger than ND/2
-    phi = phi_vec(Nf, dt=dt, d=nx, q=mult)
+    phi = phi_vec(Nf, d=nx, q=mult)
     wave = transform_wavelet_time_helper(timeseries.data, Nf, Nt, phi, mult).T
     return Wavelet(wave * np.sqrt(2), time=t_bins, freq=f_bins)
 
@@ -113,8 +113,7 @@ def from_freq_to_wavelet(
     """
     Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
     t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
-    dt = freqseries.dt
-    phif = phitilde_vec_norm(Nf, Nt, dt=dt, d=nx)
+    phif = phitilde_vec_norm(Nf, Nt,  d=nx)
     wave = transform_wavelet_freq_helper(freqseries.data, Nf, Nt, phif)
 
     return Wavelet((2 / Nf) * wave.T * np.sqrt(2), time=t_bins, freq=f_bins)

{pywavelet-0.1.1/src/pywavelet/transforms → pywavelet-0.2.0/src/pywavelet/transforms/numpy}/inverse/main.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from ...transforms.phi_computer import phi_vec, phitilde_vec_norm
-from ...types import FrequencySeries, TimeSeries, Wavelet
+from ...phi_computer import phi_vec, phitilde_vec_norm
+from ....types import FrequencySeries, TimeSeries, Wavelet
 from .to_freq import inverse_wavelet_freq_helper_fast
 from .to_time import inverse_wavelet_time_helper_fast
 
@@ -46,7 +46,7 @@ def from_wavelet_to_time(
     """
 
     mult = min(mult, wave_in.Nt // 2)  # Ensure mult is not larger than ND/2
-    phi = phi_vec(wave_in.Nf, d=nx, q=mult, dt=dt) / 2
+    phi = phi_vec(wave_in.Nf, d=nx, q=mult) / 2
     h_t = inverse_wavelet_time_helper_fast(
         wave_in.data.T, phi, wave_in.Nf, wave_in.Nt, mult
     )
@@ -84,7 +84,7 @@ def from_wavelet_to_freq(
     to ensure the proper backwards transformation.
     """
 
-    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt, dt=dt, d=nx)
+    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt,  d=nx)
     freq_data = inverse_wavelet_freq_helper_fast(
         wave_in.data, phif, wave_in.Nf, wave_in.Nt
     )

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/transforms/phi_computer.py

@@ -1,13 +1,10 @@
-import numpy as np
-from numpy import fft
-from scipy.special import betainc
+from ..backend import xp, PI, betainc, ifft
 
-PI = np.pi
 
 
 def phitilde_vec(
-    omega: np.ndarray, Nf: int, dt: float, d: float = 4.0
-) -> np.ndarray:
+    omega: xp.ndarray, Nf: int, d: float = 4.0
+) -> xp.ndarray:
     """Compute phi_tilde(omega_i) array, nx is filter steepness, defaults to 4.
 
     Eq 11 of https://arxiv.org/pdf/2009.00043.pdf (Cornish et al. 2020)
@@ -22,7 +19,7 @@ def phitilde_vec(
 
     Parameters
     ----------
-    ω : np.ndarray
+    ω : xp.ndarray
         Array of angular frequencies
     Nf : int
         Number of frequency bins
@@ -31,35 +28,35 @@ def phitilde_vec(
 
     Returns
     -------
-    np.ndarray
+    xp.ndarray
         Array of phi_tilde(omega_i) values
 
     """
     dF = 1.0 / (2 * Nf)  # NOTE: missing 1/dt?
     dOmega = 2 * PI * dF  # Near Eq 10 # 2 pi times DF
-    inverse_sqrt_dOmega = 1.0 / np.sqrt(dOmega)
+    inverse_sqrt_dOmega = 1.0 / xp.sqrt(dOmega)
 
     A = dOmega / 4
     B = dOmega - 2 * A  # Cannot have B \leq 0.
     if B <= 0:
         raise ValueError("B must be greater than 0")
 
-    phi = np.zeros(omega.size)
-    mask = (A <= np.abs(omega)) & (np.abs(omega) < A + B)  # Minor changes
+    phi = xp.zeros(omega.size)
+    mask = (A <= xp.abs(omega)) & (xp.abs(omega) < A + B)  # Minor changes
     vd = (PI / 2.0) * __nu_d(omega[mask], A, B, d=d)  # different from paper
-    phi[mask] = inverse_sqrt_dOmega * np.cos(vd)
-    phi[np.abs(omega) < A] = inverse_sqrt_dOmega
+    phi[mask] = inverse_sqrt_dOmega * xp.cos(vd)
+    phi[xp.abs(omega) < A] = inverse_sqrt_dOmega
     return phi
 
 
 def __nu_d(
-    omega: np.ndarray, A: float, B: float, d: float = 4.0
-) -> np.ndarray:
+    omega: xp.ndarray, A: float, B: float, d: float = 4.0
+) -> xp.ndarray:
     """Compute the normalized incomplete beta function.
 
     Parameters
     ----------
-    ω : np.ndarray
+    ω : xp.ndarray
         Array of angular frequencies
     A : float
         Lower bound for the beta function
@@ -70,29 +67,29 @@ def __nu_d(
 
     Returns
     -------
-    np.ndarray
+    xp.ndarray
         Array of ν_d values
 
     scipy.special.betainc
     https://docs.scipy.org/doc/scipy-1.7.1/reference/reference/generated/scipy.special.betainc.html
 
     """
-    x = (np.abs(omega) - A) / B
+    x = (xp.abs(omega) - A) / B
     return betainc(d, d, x) / betainc(d, d, 1)
 
 
-def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.ndarray:
+def phitilde_vec_norm(Nf: int, Nt: int,  d: float) -> xp.ndarray:
     """Normalize phitilde for inverse frequency domain transform."""
 
     # Calculate the frequency values
     ND = Nf * Nt
-    omegas = 2 * np.pi / ND * np.arange(0, Nt // 2 + 1)
+    omegas = 2 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)
 
     # Calculate the unnormalized phitilde (u_phit)
-    u_phit = phitilde_vec(omegas, Nf, dt, d)
+    u_phit = phitilde_vec(omegas, Nf,  d)
 
     # Normalize the phitilde
-    normalising_factor = np.pi ** (-1 / 2)  # Ollie's normalising factor
+    normalising_factor = PI ** (-1 / 2)  # Ollie's normalising factor
 
     # Notes: this is the overall normalising factor that is different from Cornish's paper
     # It is the only way I can force this code to be consistent with our work in the
@@ -117,15 +114,15 @@ def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.ndarray:
     return u_phit / (normalising_factor)
 
 
-def phi_vec(Nf: int, dt, d: float = 4.0, q: int = 16) -> np.ndarray:
+def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> xp.ndarray:
     """get time domain phi as fourier transform of phitilde_vec"""
-    insDOM = 1.0 / np.sqrt(PI / Nf)
+    insDOM = 1.0 / xp.sqrt(PI / Nf)
     K = q * 2 * Nf
-    half_K = q * Nf  # np.int64(K/2)
+    half_K = q * Nf  # xp.int64(K/2)
 
     dom = 2 * PI / K  # max frequency is K/2*dom = pi/dt = OM
 
-    DX = np.zeros(K, dtype=np.complex128)
+    DX = xp.zeros(K, dtype=xp.complex128)
 
     # zero frequency
     DX[0] = insDOM
@@ -133,20 +130,20 @@ def phi_vec(Nf: int, dt, d: float = 4.0, q: int = 16) -> np.ndarray:
     DX = DX.copy()
     # postive frequencies
     DX[1 : half_K + 1] = phitilde_vec(
-        dom * np.arange(1, half_K + 1), Nf, dt, d
+        dom * xp.arange(1, half_K + 1), Nf, d
     )
     # negative frequencies
     DX[half_K + 1 :] = phitilde_vec(
-        -dom * np.arange(half_K - 1, 0, -1), Nf, dt, d
+        -dom * xp.arange(half_K - 1, 0, -1), Nf,  d
     )
-    DX = K * fft.ifft(DX, K)
+    DX = K * ifft(DX, K)
 
-    phi = np.zeros(K)
-    phi[0:half_K] = np.real(DX[half_K:K])
-    phi[half_K:] = np.real(DX[0:half_K])
+    phi = xp.zeros(K)
+    phi[0:half_K] = xp.real(DX[half_K:K])
+    phi[half_K:] = xp.real(DX[0:half_K])
 
-    nrm = np.sqrt(K / dom)  # *np.linalg.norm(phi)
+    nrm = xp.sqrt(K / dom)  # *xp.linalg.norm(phi)
 
-    fac = np.sqrt(2.0) / nrm
+    fac = xp.sqrt(2.0) / nrm
     phi *= fac
     return phi

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/types/common.py

@@ -1,9 +1,6 @@
-from typing import Tuple, Union
-
-import numpy as xp
-from numpy.fft import fft, irfft, rfft, rfftfreq  # type: ignore
-
+from typing import Tuple, Union, Callable
 from ..logger import logger
+from ..backend import xp
 
 
 def _len_check(d):
@@ -11,7 +8,7 @@ def _len_check(d):
         logger.warning(f"Data length {len(d)} is suggested to be a power of 2")
 
 
-def is_documented_by(original):
+def is_documented_by(original:Callable):
     def wrapper(target):
         target.__doc__ = original.__doc__
         return target

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/types/frequencyseries.py

@@ -2,7 +2,8 @@ from typing import Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 
-from .common import fmt_pow2, fmt_time, irfft, is_documented_by, xp
+from ..backend import xp, irfft
+from .common import fmt_pow2, fmt_time, is_documented_by
 from .plotting import plot_freqseries, plot_periodogram
 
 __all__ = ["FrequencySeries"]
@@ -226,7 +227,7 @@ class FrequencySeries:
         Wavelet
             The corresponding wavelet.
         """
-        from ..transforms.forward import from_freq_to_wavelet
+        from ..transforms import from_freq_to_wavelet
 
         return from_freq_to_wavelet(self, Nf=Nf, Nt=Nt, nx=nx)
 

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/types/timeseries.py

@@ -10,10 +10,9 @@ from .common import (
     fmt_time,
     fmt_timerange,
     is_documented_by,
-    rfft,
-    rfftfreq,
-    xp,
 )
+from ..backend import xp, rfftfreq, rfft
+
 from .plotting import plot_spectrogram, plot_timeseries
 
 __all__ = ["TimeSeries"]

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/types/wavelet.py

@@ -1,7 +1,6 @@
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 
 import matplotlib.pyplot as plt
-import numpy as np
 
 from .common import fmt_timerange, is_documented_by, xp
 from .plotting import plot_wavelet_grid, plot_wavelet_trend
@@ -286,7 +285,7 @@ class Wavelet:
         TimeSeries
             A `TimeSeries` object representing the time-domain signal.
         """
-        from ..transforms.inverse import from_wavelet_to_time
+        from ..transforms import from_wavelet_to_time
 
         return from_wavelet_to_time(self, dt=self.delta_t, nx=nx, mult=mult)
 
@@ -299,7 +298,7 @@ class Wavelet:
         FrequencySeries
             A `FrequencySeries` object representing the frequency-domain signal.
         """
-        from ..transforms.inverse import from_wavelet_to_freq
+        from ..transforms import from_wavelet_to_freq
 
         return from_wavelet_to_freq(self, dt=self.delta_t, nx=nx)
 
@@ -315,8 +314,8 @@ class Wavelet:
 
         frange = ",".join([f"{f:.2e}" for f in (self.freq[0], self.freq[-1])])
         trange = fmt_timerange((self.t0, self.tend))
-        Nfpow2 = int(np.log2(self.shape[0]))
-        Ntpow2 = int(np.log2(self.shape[1]))
+        Nfpow2 = int(xp.log2(self.shape[0]))
+        Ntpow2 = int(xp.log2(self.shape[1]))
         shapef = f"NfxNf=[2^{Nfpow2}, 2^{Ntpow2}]"
         return f"Wavelet({shapef}, [{frange}]Hz, {trange})"
 
@@ -441,14 +440,16 @@ class WaveletMask(Wavelet):
         return self.data
 
     def __repr__(self):
-        return f"WaveletMask({self.mask.shape}, {fmt_timerange(self.time)}, {self.freq})"
+        rpr = super().__repr__()
+        rpr = rpr.replace("Wavelet", "WaveletMask")
+        return rpr
 
     @classmethod
     def from_frange(
         cls, time_grid: xp.ndarray, freq_grid: xp.ndarray, frange: List[float]
     ):
         self = cls.zeros_from_grid(time_grid, freq_grid)
-        self.mask[
+        self.data[
             (freq_grid >= frange[0]) & (freq_grid <= frange[1]), :
         ] = True
         return self

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet/types/wavelet_bins.py

@@ -1,6 +1,6 @@
 from typing import Tuple, Union
 
-import numpy as np
+from ..backend import xp
 
 from .frequencyseries import FrequencySeries
 from .timeseries import TimeSeries
@@ -33,7 +33,7 @@ def _get_bins(
     data: Union[TimeSeries, FrequencySeries],
     Nf: Union[int, None] = None,
     Nt: Union[int, None] = None,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple[xp.ndarray, xp.ndarray]:
     T = data.duration
     t_bins, f_bins = compute_bins(Nf, Nt, T)
 
@@ -46,12 +46,12 @@ def _get_bins(
     return t_bins, f_bins
 
 
-def compute_bins(Nf: int, Nt: int, T: float) -> Tuple[np.ndarray, np.ndarray]:
+def compute_bins(Nf: int, Nt: int, T: float) -> Tuple[xp.ndarray, xp.ndarray]:
     """Get the bins for the wavelet transform
     Eq 4-6 in Wavelets paper
     """
     delta_T = T / Nt
     delta_F = 1 / (2 * delta_T)
-    t_bins = np.arange(0, Nt) * delta_T
-    f_bins = np.arange(0, Nf) * delta_F
+    t_bins = xp.arange(0, Nt) * delta_T
+    f_bins = xp.arange(0, Nf) * delta_F
     return t_bins, f_bins

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.1.1
+Version: 0.2.0
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/
@@ -21,6 +21,8 @@ Requires-Dist: tqdm
 Requires-Dist: rich
 Requires-Dist: rocket-fft
 Requires-Dist: astropy>=5.2.1
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.1.0; extra == "dev"

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ docs/roundtrip_time.png
 docs/_static/demo.gif
 src/pywavelet/__init__.py
 src/pywavelet/_version.py
+src/pywavelet/backend.py
 src/pywavelet/logger.py
 src/pywavelet/utils.py
 src/pywavelet.egg-info/PKG-INFO
@@ -27,14 +28,20 @@ src/pywavelet.egg-info/requires.txt
 src/pywavelet.egg-info/top_level.txt
 src/pywavelet/transforms/__init__.py
 src/pywavelet/transforms/phi_computer.py
-src/pywavelet/transforms/forward/__init__.py
-src/pywavelet/transforms/forward/from_freq.py
-src/pywavelet/transforms/forward/from_time.py
-src/pywavelet/transforms/forward/main.py
-src/pywavelet/transforms/inverse/__init__.py
-src/pywavelet/transforms/inverse/main.py
-src/pywavelet/transforms/inverse/to_freq.py
-src/pywavelet/transforms/inverse/to_time.py
+src/pywavelet/transforms/jax/__init__.py
+src/pywavelet/transforms/jax/forward/__init__.py
+src/pywavelet/transforms/jax/forward/from_freq.py
+src/pywavelet/transforms/jax/forward/from_time.py
+src/pywavelet/transforms/jax/forward/main.py
+src/pywavelet/transforms/numpy/__init__.py
+src/pywavelet/transforms/numpy/forward/__init__.py
+src/pywavelet/transforms/numpy/forward/from_freq.py
+src/pywavelet/transforms/numpy/forward/from_time.py
+src/pywavelet/transforms/numpy/forward/main.py
+src/pywavelet/transforms/numpy/inverse/__init__.py
+src/pywavelet/transforms/numpy/inverse/main.py
+src/pywavelet/transforms/numpy/inverse/to_freq.py
+src/pywavelet/transforms/numpy/inverse/to_time.py
 src/pywavelet/types/__init__.py
 src/pywavelet/types/common.py
 src/pywavelet/types/frequencyseries.py

{pywavelet-0.1.1 → pywavelet-0.2.0}/src/pywavelet.egg-info/requires.txt

@@ -18,3 +18,6 @@ isort
 mypy
 jupyter-book
 GitPython
+
+[jax]
+jax

{pywavelet-0.1.1 → pywavelet-0.2.0}/tests/test_mask.py

@@ -17,7 +17,7 @@ def test_mask(plot_dir):
     assert np.isclose(compute_likelihood(d, h, psd, mask), 0)
 
     mask1 = WaveletMask.from_frange(h.time, h.freq, [f0 + 0.5, f0 + 1.5])
-    # assert np.isclose(compute_likelihood(d, h, psd, mask1), 0) == False
+    # assert np.isclose(compute_likelihood(d, h, psd_analysis, mask1), 0) == False
 
     # plt the 3 differnet datasets
     fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)

{pywavelet-0.1.1 → pywavelet-0.2.0}/tests/test_phi.py

@@ -1,7 +1,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from pywavelet.transforms.phi_computer import (
+from pywavelet.transforms import (
     phi_vec,
     phitilde_vec,
     phitilde_vec_norm,
@@ -16,10 +16,10 @@ def test_phi(plot_dir):
     delT = Nf * dt
     delF = 1.0 / (2 * Nf * dt)
 
-    phi = phi_vec(Nf=Nf, dt=dt, d=d, q=q) / delT
+    phi = phi_vec(Nf=Nf, d=d, q=q) / delT
     t = np.linspace(-q, q, len(phi))
     f = np.linspace(-1, 1, 1000)
-    phitilde = phitilde_vec(f, Nf=Nf, d=d, dt=dt) / delF
+    phitilde = phitilde_vec(f, Nf=Nf, d=d) / delF
 
     fig, ax = plt.subplots(
         2,