pywavelet 0.2.2__tar.gz → 0.2.4__tar.gz

{pywavelet-0.2.2 → pywavelet-0.2.4}/.github/workflows/ci.yml

@@ -28,6 +28,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install -e .[dev]
+          python -m pip install -e .[jax]
           pre-commit install
 
       - name: pre-commit

{pywavelet-0.2.2 → pywavelet-0.2.4}/CHANGELOG.rst

@@ -5,16 +5,60 @@ CHANGELOG
 =========
 
 
+.. _changelog-v0.2.4:
+
+v0.2.4 (2025-01-24)
+===================
+
+Unknown
+-------
+
+* Merge branch 'main' of github.com:avivajpeyi/pywavelet into main (`d2c84d9`_)
+
+.. _d2c84d9: https://github.com/pywavelet/pywavelet/commit/d2c84d980b1701baf99e40ba6191cbd9336cfa59
+
+
+.. _changelog-v0.2.3:
+
+v0.2.3 (2025-01-24)
+===================
+
+Bug Fixes
+---------
+
+* fix: add masks to filter out gaps (`26fe40a`_)
+
+* fix: add backend check for os.environ (`98c0818`_)
+
+* fix: add test for jax (`1940394`_)
+
+Chores
+------
+
+* chore(release): 0.2.3 (`d067461`_)
+
+.. _26fe40a: https://github.com/pywavelet/pywavelet/commit/26fe40ace80d5f9d598e1efeba2f8ca4a6f1043b
+.. _98c0818: https://github.com/pywavelet/pywavelet/commit/98c0818078190d829a23734f932f1f93c9932167
+.. _1940394: https://github.com/pywavelet/pywavelet/commit/194039437a3a9b3ada303d101b4e2573ab7d0afd
+.. _d067461: https://github.com/pywavelet/pywavelet/commit/d0674615df328774a0d80eb224b5c503fbd8f332
+
+
 .. _changelog-v0.2.2:
 
 v0.2.2 (2025-01-23)
 ===================
 
+Chores
+------
+
+* chore(release): 0.2.2 (`eed5d68`_)
+
 Unknown
 -------
 
 * Merge branch 'main' of github.com:pywavelet/pywavelet (`e8e2115`_)
 
+.. _eed5d68: https://github.com/pywavelet/pywavelet/commit/eed5d6864276fc5f90c4866749903e3e358df5ca
 .. _e8e2115: https://github.com/pywavelet/pywavelet/commit/e8e2115e797a5001f236ff027a14ef226151dcc1
 
 

{pywavelet-0.2.2 → pywavelet-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.2
+Version: 0.2.4
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.2'
-__version_tuple__ = version_tuple = (0, 2, 2)
+__version__ = version = '0.2.4'
+__version_tuple__ = version_tuple = (0, 2, 4)

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet/backend.py

@@ -1,5 +1,7 @@
 import os
 
+from .logger import logger
+
 try:
     import jax
 
@@ -13,14 +15,17 @@ 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.numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
     from jax.scipy.special import betainc  # type: ignore
 
+    logger.info("Using JAX backend")
 
 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
+    from numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
+    from scipy.special import betainc  # type: ignore
+
+    logger.info("Using NumPy+numba backend")
 
 
-PI = xp.pi
+PI = xp.pi

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

@@ -0,0 +1,28 @@
+from ..backend import use_jax
+
+if use_jax:
+    from .jax import (
+        from_freq_to_wavelet,
+        from_time_to_wavelet,
+        from_wavelet_to_freq,
+        from_wavelet_to_time,
+    )
+else:
+    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, phitilde_vec_norm
+
+__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.4/src/pywavelet/transforms/jax/__init__.py

@@ -0,0 +1,12 @@
+from ...logger import logger
+from .forward import from_freq_to_wavelet, from_time_to_wavelet
+from .inverse import from_wavelet_to_freq, from_wavelet_to_time
+
+logger.warning("JAX SUBPACKAGE NOT FULLY TESTED")
+
+__all__ = [
+    "from_wavelet_to_time",
+    "from_wavelet_to_freq",
+    "from_time_to_wavelet",
+    "from_freq_to_wavelet",
+]

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

@@ -0,0 +1,97 @@
+from functools import partial
+
+import jax.numpy as jnp
+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:
+    """
+    Transforms input data from the frequency domain to the wavelet domain using a
+    pre-computed wavelet filter (`phif`) and performs an efficient inverse FFT.
+
+    Parameters:
+    - data (jnp.ndarray): 1D array representing the input data in the frequency domain.
+    - Nf (int): Number of frequency bins.
+    - Nt (int): Number of time bins. (Note: Nt * Nf == len(data))
+    - phif (jnp.ndarray): Pre-computed wavelet filter for frequency components.
+
+    Returns:
+    - wave (jnp.ndarray): 2D array of wavelet-transformed data with shape (Nf, Nt).
+    """
+
+    # Initialize the wavelet output array with zeros (time-rows, frequency-columns)
+    wave = jnp.zeros((Nt, Nf))
+    f_bins = jnp.arange(Nf)  # Frequency bin indices
+
+    # Compute base indices for time (i_base) and frequency (jj_base)
+    i_base = Nt // 2
+    jj_base = f_bins * Nt // 2
+
+    # Set initial values for the center of the transformation
+    initial_values = jnp.where(
+        (f_bins == 0)
+        | (f_bins == Nf),  # Edge cases: DC (f=0) and Nyquist (f=Nf)
+        phif[0] * data[f_bins * Nt // 2] / 2.0,  # Adjust for symmetry
+        phif[0] * data[f_bins * Nt // 2],
+    )
+
+    # Initialize a 2D array to store intermediate FFT input values
+    DX = jnp.zeros(
+        (Nf, Nt), dtype=jnp.complex64
+    )  # TODO: Check dtype -- is complex64 sufficient?
+    DX = DX.at[:, Nt // 2].set(
+        initial_values
+    )  # Set initial values at the center of the transformation (2 sided FFT)
+
+    # Compute time indices for all offsets around the midpoint
+    j_range = jnp.arange(
+        1 - Nt // 2, Nt // 2
+    )  # Time offsets (centered around zero)
+    j = jnp.abs(j_range)  # Absolute offset indices
+    i = i_base + j_range  # Time indices in output array
+
+    # Compute conditions for edge cases
+    cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)  # Nyquist
+    cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)  # DC
+    cond3 = j[None, :] == 0  # Center of the transformation (no offset)
+
+    # Compute frequency indices for the input data
+    jj = jj_base[:, None] + j_range[None, :]  # Frequency offsets
+    val = jnp.where(
+        cond1 | cond2, 0.0, phif[j] * data[jj]
+    )  # Wavelet filter application
+    DX = DX.at[:, i].set(
+        jnp.where(cond3, DX[:, i], val)
+    )  # Update DX with computed values
+    # At this point, DX contains the data FFT'd with the wavelet filter
+    # (each row is a frequency bin, each column is a time bin)
+
+    # Perform the inverse FFT along the time dimension
+    DX_trans = ifft(DX, axis=1)
+
+    # Fill the wavelet output array based on the inverse FFT results
+    n_range = jnp.arange(Nt)  # Time indices
+    cond1 = (
+        n_range[:, None] + f_bins[None, :]
+    ) % 2 == 1  # Odd/even alternation
+    cond2 = jnp.expand_dims(f_bins % 2 == 1, axis=-1)  # Odd frequency bins
+
+    # Assign real and imaginary parts based on conditions
+    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 frequency bins 0 (DC) and Nf (Nyquist)
+    wave = wave.at[::2, 0].set(
+        jnp.real(DX_trans[0, ::2] * jnp.sqrt(2))
+    )  # DC component
+    wave = wave.at[1::2, -1].set(
+        jnp.real(DX_trans[-1, ::2] * jnp.sqrt(2))
+    )  # Nyquist component
+
+    # Return the wavelet-transformed array (transposed for freq-major layout)
+    return wave.T  # (Nt, Nf) -> (Nf, Nt)

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet/transforms/jax/inverse/main.py

@@ -1,10 +1,9 @@
 import jax.numpy as jnp
 from jax.numpy.fft import rfftfreq
 
-from ...phi_computer import phi_vec, phitilde_vec_norm
 from ....types import FrequencySeries, TimeSeries, Wavelet
+from ...phi_computer import phi_vec, phitilde_vec_norm
 from .to_freq import inverse_wavelet_freq_helper
-# from .inverse_wavelet_time_funcs import inverse_wavelet_time_helper
 
 
 def from_wavelet_to_time(
@@ -65,5 +64,5 @@ def from_wavelet_to_freq(
         -1 / 2
     )  # Normalise to get the proper backwards transformation
 
-    freqs = rfftfreq(wave_in.ND*2, d=dt)[1:]
-    return FrequencySeries(data=freq_data, freq=freqs)
+    freqs = rfftfreq(wave_in.ND * 2, d=dt)[1:]
+    return FrequencySeries(data=freq_data, freq=freqs)

{pywavelet-0.2.2/src/pywavelet/transforms/jax → pywavelet-0.2.4/src/pywavelet/transforms/numpy}/forward/__init__.py

@@ -1,6 +1,3 @@
 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.4/src/pywavelet/transforms/numpy/inverse/__init__.py

@@ -0,0 +1,3 @@
+from .main import from_wavelet_to_freq, from_wavelet_to_time
+
+__all__ = ["from_wavelet_to_time", "from_wavelet_to_freq"]

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet/types/wavelet.py

@@ -1,6 +1,7 @@
 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
@@ -343,14 +344,18 @@ class Wavelet:
 
     def __mul__(self, other):
         """Element-wise multiplication of two Wavelet objects."""
-        if isinstance(other, Wavelet):
-            return Wavelet(
-                data=self.data * other.data, time=self.time, freq=self.freq
-            )
+        if isinstance(other, WaveletMask):
+            data = self.data.copy()
+            data[~other.mask] = np.nan
+            return Wavelet(data=data, time=self.time, freq=self.freq)
         elif isinstance(other, float):
             return Wavelet(
                 data=self.data * other, time=self.time, freq=self.freq
             )
+        elif isinstance(other, WaveletMask):
+            return Wavelet(
+                data=self.data * other.data, time=self.time, freq=self.freq
+            )
 
     def __truediv__(self, other):
         """Element-wise division of two Wavelet objects."""
@@ -445,11 +450,40 @@ class WaveletMask(Wavelet):
         return rpr
 
     @classmethod
-    def from_frange(
-        cls, time_grid: xp.ndarray, freq_grid: xp.ndarray, frange: List[float]
+    def from_restrictions(
+        cls,
+        time_grid: xp.ndarray,
+        freq_grid: xp.ndarray,
+        frange: List[float],
+        tgaps: List[Tuple[float, float]] = [],
     ):
+        """
+        Create a WaveletMask object from restrictions on time and frequency.
+
+        Parameters
+        ----------
+        time_grid : xp.ndarray
+            Array of time points.
+        freq_grid : xp.ndarray
+            Array of corresponding frequency points.
+        frange : List[float]
+            Frequency range to include.
+        tgaps : List[Tuple[float, float]]
+            List of time gaps to exclude.
+
+        Returns
+        -------
+        WaveletMask
+            A WaveletMask object with the specified restrictions.
+        """
         self = cls.zeros_from_grid(time_grid, freq_grid)
         self.data[
             (freq_grid >= frange[0]) & (freq_grid <= frange[1]), :
         ] = True
+
+        for tgap in tgaps:
+            self.data[
+                :, (time_grid >= tgap[0]) & (time_grid <= tgap[1])
+            ] = False
+        self.data = self.data.astype(bool)
         return self

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.2
+Version: 0.2.4
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/

{pywavelet-0.2.2 → pywavelet-0.2.4}/src/pywavelet.egg-info/SOURCES.txt

@@ -53,6 +53,7 @@ src/pywavelet/types/timeseries.py
 src/pywavelet/types/wavelet.py
 src/pywavelet/types/wavelet_bins.py
 tests/conftest.py
+tests/test_jax.py
 tests/test_lnl.py
 tests/test_mask.py
 tests/test_phi.py

pywavelet-0.2.4/tests/test_jax.py

@@ -0,0 +1,105 @@
+import importlib
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+from conftest import monochromatic_wnm
+
+from pywavelet.transforms import from_freq_to_wavelet, from_time_to_wavelet
+from pywavelet.types import FrequencySeries, TimeSeries
+from pywavelet.types.wavelet_bins import compute_bins
+from pywavelet.utils import compute_snr, evolutionary_psd_from_stationary_psd
+
+
+def test_toy_model_snr(plot_dir):
+    f0 = 20
+    dt = 0.0125
+    A = 2
+    Nt = 128
+    Nf = 256
+    ND = Nt * Nf
+    t = np.arange(0, ND) * dt
+    PSD_AMP = 1
+
+    ########################################
+    # Part1: Analytical SNR calculation
+    #######################################
+
+    # Eq 21
+    y = A * np.sin(2 * np.pi * f0 * t)  # Signal waveform we wish to test
+    signal_timeseries = TimeSeries(y, t)
+    signal_freq = signal_timeseries.to_frequencyseries()
+    psd_freq = FrequencySeries(
+        PSD_AMP * np.ones(len(signal_freq)), signal_freq.freq
+    )
+    snr = signal_freq.optimal_snr(psd_freq)
+
+    ########################################
+    # Part2: Wavelet domain (numpy)
+    ########################################
+
+    signal_wavelet = from_freq_to_wavelet(signal_freq, Nf=Nf, Nt=Nt)
+    psd_wavelet = evolutionary_psd_from_stationary_psd(
+        psd=psd_freq.data,
+        psd_f=psd_freq.freq,
+        f_grid=signal_wavelet.freq,
+        t_grid=signal_wavelet.time,
+        dt=dt,
+    )
+    wdm_snr = compute_snr(signal_wavelet, signal_wavelet, psd_wavelet)
+    assert np.isclose(snr, wdm_snr, atol=0.5), f"{snr}!={wdm_snr}"
+
+    ########################################
+    # Part3: Wavelet domain (jax)
+    ########################################
+
+    from pywavelet.transforms.jax import (
+        from_freq_to_wavelet as jax_from_freq_to_wavelet,
+    )
+
+    signal_wavelet_jax = jax_from_freq_to_wavelet(signal_freq, Nf=Nf, Nt=Nt)
+    psd_wavelet_jax = evolutionary_psd_from_stationary_psd(
+        psd=psd_freq.data,
+        psd_f=psd_freq.freq,
+        f_grid=signal_wavelet_jax.freq,
+        t_grid=signal_wavelet_jax.time,
+        dt=dt,
+    )
+    wdm_snr_jax = compute_snr(
+        signal_wavelet_jax, signal_wavelet_jax, psd_wavelet_jax
+    )
+    assert np.isclose(snr, wdm_snr_jax, atol=0.5), f"{snr}!={wdm_snr_jax}"
+
+    wdm_diff = signal_wavelet - signal_wavelet_jax
+
+    ########################################
+    # Part4: Plot
+    ########################################
+
+    fig, ax = plt.subplots(1, 3, figsize=(15, 6))
+    signal_wavelet.plot(ax=ax[0])
+    signal_wavelet_jax.plot(ax=ax[1])
+    wdm_diff.plot(ax=ax[2])
+    ax[0].set_title(f"Numpy SNR={wdm_snr:.2f}")
+    ax[1].set_title(f"Jax SNR={wdm_snr_jax:.2f}")
+    ax[2].set_title("Difference")
+    plt.tight_layout()
+    plt.savefig(f"{plot_dir}/jax_vs_np.png")
+
+
+def test_backend_loader():
+    # temporarily set os.environ["PYWAVELET_JAX"] = "1"
+
+    import pywavelet.backend
+
+    os.environ["PYWAVELET_JAX"] = "1"
+    importlib.reload(pywavelet.backend)
+    from pywavelet.backend import use_jax
+
+    assert use_jax
+    os.environ["PYWAVELET_JAX"] = "0"
+
+    importlib.reload(pywavelet.backend)
+    from pywavelet.backend import use_jax
+
+    assert not use_jax

{pywavelet-0.2.2 → pywavelet-0.2.4}/tests/test_mask.py

@@ -13,25 +13,26 @@ def test_mask(plot_dir):
     psd = Wavelet(np.ones((h.Nf, h.Nt)), h.time, h.freq)
     assert compute_likelihood(d, h, psd) == 0
 
-    mask = WaveletMask.from_frange(h.time, h.freq, [f0 - 0.5, f0 + 0.5])
+    mask = WaveletMask.from_restrictions(h.time, h.freq, [f0 - 0.5, f0 + 0.5])
+    dmasked = d * mask
     assert np.isclose(compute_likelihood(d, h, psd, mask), 0)
+    assert np.isclose(compute_likelihood(dmasked, h, psd), 0)
+    # number of nans in dmasked.data
+    assert np.isnan(dmasked.data).sum() != 0
 
-    mask1 = WaveletMask.from_frange(h.time, h.freq, [f0 + 0.5, f0 + 1.5])
+    mask1 = WaveletMask.from_restrictions(
+        h.time, h.freq, [f0 + 0.5, f0 + 1.5], tgaps=[[1.7 * 60, 3.4 * 60]]
+    )
+    dmasked1 = d * mask1
     # 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)
-    kwgs = dict(
-        cmap="viridis",
-        aspect="auto",
-        origin="lower",
-        extent=[d.time[0], d.time[-1], d.freq[0], d.freq[-1]],
-    )
-    axes[0].imshow(d.data, **kwgs)
+    d.plot(ax=axes[0])
     axes[0].set_title("data")
-    axes[1].imshow(d.data * mask.mask, **kwgs)
+    (d * mask).plot(ax=axes[1])
     axes[1].set_title("data*mask[f0-0.5, f0+0.5]")
-    axes[2].imshow(d.data * mask1.mask, **kwgs)
+    (d * mask1).plot(ax=axes[2])
     axes[2].set_title("data*mask[f0+0.5, f0+1.5]")
     plt.tight_layout()
     fig.savefig(f"{plot_dir}/test_mask.png")

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

@@ -1,17 +0,0 @@
-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.2/src/pywavelet/transforms/jax/forward/from_freq.py

@@ -1,56 +0,0 @@
-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