nullcal 0.2.0__py3-none-any.whl

nullcal/time_frequency_transform/inverse_wavelet_freq_funcs.py

@@ -0,0 +1,42 @@
+"""JAX kernels for inverse WDM transforms in the frequency domain."""
+
+from __future__ import annotations
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp  # noqa: E402
+
+
+def inverse_wavelet_freq_helper_fast(wave_in, phif, n_f, n_t):
+    """Return the inverse frequency-domain transform."""
+    wave_in = jnp.asarray(wave_in)
+    phif = jnp.asarray(phif, dtype=wave_in.dtype)
+    rows = jnp.arange(n_t)
+    modes = jnp.arange(1, n_f)[:, None]
+    interior = wave_in[:, 1:n_f].T
+    factors = jnp.where((rows[None, :] + modes) % 2 == 1, -1j, 1.0)
+
+    prefactors = jnp.zeros((n_f + 1, n_t), dtype=jnp.result_type(wave_in, 1j))
+    prefactors = prefactors.at[0].set(wave_in[(2 * rows) % n_t, 0] / jnp.sqrt(2.0))
+    prefactors = prefactors.at[n_f].set(wave_in[(2 * rows) % n_t + 1, 0] / jnp.sqrt(2.0))
+    prefactors = prefactors.at[1:n_f].set(interior * factors)
+    transformed = jnp.fft.fft(prefactors, axis=1)
+
+    res = jnp.zeros(n_f * n_t // 2 + 1, dtype=transformed.dtype)
+    half = n_t // 2
+
+    lower_offsets = jnp.arange(half)
+    res = res.at[lower_offsets].add(transformed[0, (2 * lower_offsets) % n_t] * phif[lower_offsets])
+
+    upper_offsets = jnp.arange(half + 1)
+    upper_indices = n_f * half - upper_offsets
+    res = res.at[upper_indices].add(transformed[n_f, (-2 * upper_offsets) % n_t] * phif[upper_offsets])
+
+    offsets = jnp.arange(1 - half, half)
+    centres = jnp.arange(1, n_f)[:, None] * half
+    indices = centres + offsets[None, :]
+    fft_indices = indices % n_t
+    contributions = transformed[1:n_f][jnp.arange(n_f - 1)[:, None], fft_indices] * phif[jnp.abs(offsets)][None, :]
+    return res.at[indices].add(contributions)

nullcal/time_frequency_transform/inverse_wavelet_time_funcs.py

@@ -0,0 +1,49 @@
+"""JAX kernels for inverse WDM transforms in the time domain."""
+
+from __future__ import annotations
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp  # noqa: E402
+
+
+def _compact_packets(wave_in, n_f, n_t):
+    """Pack adjacent WDM time rows into complex FFT inputs."""
+    wave_in = jnp.asarray(wave_in)
+    even_rows = jnp.arange(0, n_t, 2)
+    even = wave_in[even_rows]
+    odd = wave_in[even_rows + 1]
+    packets = jnp.zeros((n_t // 2, 2 * n_f), dtype=jnp.result_type(wave_in, 1j))
+    packets = packets.at[:, 0].set(jnp.sqrt(2.0) * even[:, 0])
+    packets = packets.at[:, n_f].set(jnp.sqrt(2.0) * odd[:, 0])
+
+    even_modes = jnp.arange(2, n_f - 1, 2)
+    packets = packets.at[:, even_modes].set(even[:, even_modes] - odd[:, even_modes])
+    packets = packets.at[:, 2 * n_f - even_modes].set(even[:, even_modes] + odd[:, even_modes])
+
+    odd_modes = jnp.arange(1, n_f, 2)
+    packets = packets.at[:, odd_modes].set(1j * (even[:, odd_modes] - odd[:, odd_modes]))
+    return packets.at[:, 2 * n_f - odd_modes].set(-1j * (even[:, odd_modes] + odd[:, odd_modes]))
+
+
+def inverse_wavelet_time_helper_fast(wave_in, phi, n_f, n_t, mult):
+    """Return the inverse time-domain transform using batched JAX FFTs."""
+    wave_in = jnp.asarray(wave_in)
+    phi = jnp.asarray(phi, dtype=wave_in.dtype)
+    n_d = n_f * n_t
+    k_cutoff = mult * 2 * n_f
+    transformed = jnp.fft.fft(_compact_packets(wave_in, n_f, n_t), axis=1)
+
+    pair_rows = jnp.arange(0, n_t, 2)[:, None]
+    offsets = jnp.arange(k_cutoff)[None, :]
+    fft_indices = (-k_cutoff // 2 + pair_rows * n_f + n_d + offsets) % (2 * n_f)
+    result_indices = (-k_cutoff // 2 + pair_rows * n_f + offsets) % n_d
+    real_values = jnp.take_along_axis(jnp.real(transformed), fft_indices, axis=1)
+    imag_indices = (fft_indices + n_f) % (2 * n_f)
+    imag_values = jnp.take_along_axis(jnp.imag(transformed), imag_indices, axis=1)
+
+    result = jnp.zeros(n_d, dtype=wave_in.dtype)
+    result = result.at[result_indices].add(phi[None, :] * real_values)
+    return result.at[(result_indices + n_f) % n_d].add(phi[None, :] * imag_values)

nullcal/time_frequency_transform/stft.py

@@ -0,0 +1,45 @@
+"""
+Short-time Fourier transform.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def stft(data: np.ndarray, sampling_frequency: float, frequency_resolution: float, window: np.ndarray):
+    """Short-time Fourier transform.
+
+    Args:
+        data (np.ndarray): Data.
+        sampling_frequency (float): Sampling frequency in Hz.
+        frequency_resolution (float): Frequency resolution in Jz.
+        window (np.ndarray, optional): Window function. Defaults to None.
+
+    Raises:
+        ValueError: data must be real-valued.
+
+    Returns:
+        np.ndarray: Short-time Fourier transform of data.
+    """
+    data = np.asarray(data)
+    if not np.isrealobj(data):
+        raise ValueError("data must be real-valued")
+
+    # Compute number of frames
+    n_samples = data.shape[0]
+    window_size = int(sampling_frequency / frequency_resolution)
+    n_frames = int(n_samples // window_size)
+
+    # Reshape signal into frames
+    # Shape: (n_frames, window_size)
+    frames = data.reshape(n_frames, window_size)
+
+    # Apply window function
+    windowed_frames = frames * window
+
+    # Compute rFFT for each frame
+    # Shape: (n_frames, window_size // 2 + 1), complex-valued
+    r_stft_matrix = np.fft.rfft(windowed_frames, n=window_size, axis=-1) / sampling_frequency
+
+    return r_stft_matrix

nullcal/time_frequency_transform/transform_freq_funcs.py

@@ -0,0 +1,180 @@
+"""helper functions for transform_freq"""
+
+from __future__ import annotations
+
+import jax
+import numpy as np
+import scipy.special
+
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp  # noqa: E402
+
+
+def phitilde_vec(om, n_f, nx=4.0):
+    """Compute phitilde, om i array, nx is filter steepness, defaults to 4.
+
+    Args:
+        om (1D numpy array): om.
+        n_f (int): Number of frequency bins.
+        nx (float, optional): Filter steepness. Defaults to 4.
+
+    Returns:
+        1D numpy array: z.
+    """
+    omega = np.pi
+    d_omega = omega / n_f
+    ins_d_omega = 1.0 / np.sqrt(d_omega)
+    b = omega / (2 * n_f)
+    a = (d_omega - b) / 2
+    z = np.zeros(om.size)
+
+    mask = (np.abs(om) >= a) & (np.abs(om) < a + b)
+
+    x = (np.abs(om[mask]) - a) / b
+    y = scipy.special.betainc(nx, nx, x)
+    z[mask] = ins_d_omega * np.cos(np.pi / 2.0 * y)
+
+    z[np.abs(om) < a] = ins_d_omega
+    return z
+
+
+def phitilde_vec_norm(n_f, n_t, nx):
+    """Normalize phitilde as needed for inverse frequency domain transform.
+
+    Args:
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float): Filter steepness
+
+    Returns:
+        phif: Wavelet.
+    """
+    n_d = n_f * n_t
+    oms = 2 * np.pi / n_d * np.arange(0, n_t // 2 + 1)
+    phif = phitilde_vec(oms, n_f, nx)
+    # nrm should be 1
+    nrm = np.sqrt((2 * np.sum(phif[1:] ** 2) + phif[0] ** 2) * 2 * np.pi / n_d)
+    nrm /= np.pi ** (3 / 2) / np.pi
+    phif /= nrm
+    return phif
+
+
+def tukey(data, alpha, n_sample):
+    """Apply Tukey window function to data.
+
+    Args:
+        data (1D numpy array): Data.
+        alpha (float): Rolling parameter.
+        n_sample (int): Length of data
+    """
+    imin = int(alpha * (n_sample - 1) / 2)
+    imax = int((n_sample - 1) * (1 - alpha / 2))
+    n_win = n_sample - imax
+    values = jnp.asarray(data)
+    indices = jnp.arange(n_sample)
+    safe_imin = max(imin, 1)
+    rising_values = 0.5 * (1.0 + jnp.cos(jnp.pi * (indices / safe_imin - 1.0)))
+    falling_values = 0.5 * (1.0 + jnp.cos(jnp.pi / n_win * (indices - imax)))
+    window = jnp.where(indices < imin, rising_values, 1.0)
+    window = jnp.where(indices > imax, falling_values, window)
+    result = values * window
+    if isinstance(data, np.ndarray):
+        data[...] = np.asarray(result)
+        return None
+    return result
+
+
+def _frequency_packets(data, n_f, n_t, phif):
+    """Build every windowed frequency packet in one JAX gather."""
+    data = jnp.asarray(data)
+    phif = jnp.asarray(phif, dtype=data.real.dtype)
+    modes = jnp.arange(n_f + 1)[:, None]
+    offsets = jnp.arange(1 - n_t // 2, n_t // 2)[None, :]
+    frequency_indices = modes * (n_t // 2) + offsets
+    safe_indices = jnp.clip(frequency_indices, 0, data.size - 1)
+    valid = ~(((modes == 0) & (offsets < 0)) | ((modes == n_f) & (offsets > 0)))
+    values = phif[jnp.abs(offsets)] * data[safe_indices]
+
+    packets = jnp.zeros((n_f + 1, n_t), dtype=data.dtype)
+    packets = packets.at[:, 1:].set(jnp.where(valid, values, 0.0))
+    endpoint_scale = jnp.where((jnp.arange(n_f + 1) == 0) | (jnp.arange(n_f + 1) == n_f), 0.5, 1.0)
+    centres = phif[0] * data[jnp.arange(n_f + 1) * (n_t // 2)] * endpoint_scale
+    return packets.at[:, n_t // 2].set(centres)
+
+
+def _unpack_frequency_packets(transformed, n_f, n_t, *, quadrature=False, frequency_filter=None):
+    """Map packet FFT components to the real WDM coefficient array."""
+    rows = jnp.arange(n_t)[:, None]
+    modes = jnp.arange(1, n_f)[None, :]
+    interior = transformed[1:n_f].T
+    parity = (rows + modes) % 2 == 1
+    odd_mode = modes % 2 == 1
+    if quadrature:
+        values = jnp.where(parity, jnp.real(interior), jnp.where(odd_mode, -jnp.imag(interior), jnp.imag(interior)))
+    else:
+        values = jnp.where(parity, jnp.where(odd_mode, -jnp.imag(interior), jnp.imag(interior)), jnp.real(interior))
+
+    if frequency_filter is not None:
+        selected = jnp.asarray(frequency_filter, dtype=bool)
+        values = jnp.where(selected[None, 1:n_f], values, 0.0)
+    else:
+        selected = jnp.ones(n_f + 1, dtype=bool)
+
+    wave = jnp.zeros((n_t, n_f), dtype=transformed.real.dtype)
+    wave = wave.at[:, 1:].set(values)
+    packet_indices = jnp.arange(1 if quadrature else 0, n_t, 2)
+    wave = wave.at[::2, 0].set(jnp.where(selected[0], jnp.real(transformed[0, packet_indices]) * jnp.sqrt(2.0), 0.0))
+    wave = wave.at[1::2, 0].set(
+        jnp.where(selected[n_f], jnp.real(transformed[n_f, packet_indices]) * jnp.sqrt(2.0), 0.0)
+    )
+    return wave
+
+
+def transform_wavelet_freq_helper(data, n_f, n_t, phif):
+    """Helper to do the wavelet transform using the fast wavelet domain transform.
+
+    Args:
+        data (1D numpy array): Data.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        phif (1D numpy array): Wavelet.
+
+    Returns:
+        2D numpy array: Data in wavelet domain.
+    """
+    packets = _frequency_packets(data, n_f, n_t, phif)
+    return _unpack_frequency_packets(jnp.fft.ifft(packets, axis=1), n_f, n_t)
+
+
+def transform_wavelet_freq_partial_helper(data, n_f, n_t, phif, frequency_filter):
+    """Helper to do the wavelet transform using the fast wavelet domain transform.
+
+    Args:
+        data (1D numpy array): Data.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        phif (1D numpy array): Wavelet.
+        frequency_filter (1D numpy numpy): An array to indicate which frequency bins to evaluate.
+
+    Returns:
+        2D numpy array: Data in wavelet domain.
+    """
+    packets = _frequency_packets(data, n_f, n_t, phif)
+    return _unpack_frequency_packets(jnp.fft.ifft(packets, axis=1), n_f, n_t, frequency_filter=frequency_filter)
+
+
+def transform_wavelet_freq_quadrature_helper(data, n_f, n_t, phif):
+    """Helper to do the wavelet transform using the fast wavelet domain quadrature transform.
+
+    Args:
+        data (1D numpy array): Data.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        phif (1D numpy array): Wavelet.
+
+    Returns:
+        2D numpy array: Data in wavelet domain.
+    """
+    packets = _frequency_packets(data, n_f, n_t, phif)
+    return _unpack_frequency_packets(jnp.fft.ifft(packets, axis=1), n_f, n_t, quadrature=True)

nullcal/time_frequency_transform/transform_time_funcs.py

@@ -0,0 +1,60 @@
+"""JAX kernels for the direct time-domain WDM transform."""
+
+from __future__ import annotations
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import jax.numpy as jnp  # noqa: E402
+import numpy as np  # noqa: E402
+
+from .transform_freq_funcs import phitilde_vec  # noqa: E402
+
+
+def transform_wavelet_time_helper(data, n_f, n_t, phi, mult):
+    """Transform a time series into WDM coefficients with batched JAX FFTs.
+
+    Computation uses the common floating dtype of ``data`` and ``phi``: two
+    float32 inputs stay float32, while float64 or integer data paired with the
+    float64 window use float64.
+    """
+    data = jnp.asarray(data)
+    phi = jnp.asarray(phi)
+    working_dtype = jnp.result_type(data, phi)
+    data = data.astype(working_dtype)
+    phi = phi.astype(working_dtype)
+    n_d = n_f * n_t
+    n_k = mult * 2 * n_f
+    rows = jnp.arange(n_t)[:, None]
+    offsets = jnp.arange(n_k)[None, :]
+    sample_indices = (rows * n_f - n_k // 2 + offsets) % n_d
+    packets = data[sample_indices] * phi[None, :]
+    transformed = jnp.fft.rfft(packets, axis=1)
+
+    modes = jnp.arange(1, n_f)[None, :]
+    selected = transformed[:, (jnp.arange(1, n_f) * mult)]
+    interior = jnp.where((rows + modes) % 2 == 1, -jnp.imag(selected), jnp.real(selected))
+
+    wave = jnp.zeros((n_t, n_f), dtype=transformed.real.dtype)
+    wave = wave.at[:, 1:].set(interior)
+    wave = wave.at[::2, 0].set(jnp.real(transformed[::2, 0]) / jnp.sqrt(2.0))
+    return wave.at[1::2, 0].set(jnp.real(transformed[::2, n_f * mult]) / jnp.sqrt(2.0))
+
+
+def phi_vec(n_f, nx=4.0, mult=16):
+    """Return the time-domain WDM window."""
+    omega = np.pi
+    d_omega = omega / n_f
+    inverse_sqrt_d_omega = 1.0 / np.sqrt(d_omega)
+    k_cutoff = mult * 2 * n_f
+    half_k_cutoff = mult * n_f
+    dom = 2 * np.pi / k_cutoff
+
+    spectrum = np.zeros(k_cutoff, dtype=np.complex128)
+    spectrum[0] = inverse_sqrt_d_omega
+    spectrum[1 : half_k_cutoff + 1] = phitilde_vec(dom * np.arange(1, half_k_cutoff + 1), n_f, nx)
+    spectrum[half_k_cutoff + 1 :] = phitilde_vec(-dom * np.arange(half_k_cutoff - 1, 0, -1), n_f, nx)
+    transformed = k_cutoff * jnp.fft.ifft(jnp.asarray(spectrum), k_cutoff)
+    phi = jnp.concatenate((jnp.real(transformed[half_k_cutoff:]), jnp.real(transformed[:half_k_cutoff])))
+    return phi * (jnp.sqrt(2.0) / jnp.sqrt(k_cutoff / dom))

nullcal/time_frequency_transform/utils.py

@@ -0,0 +1,21 @@
+"""
+Utility functions for the time-frequency transform.
+"""
+
+from __future__ import annotations
+
+
+def get_shape_of_wavelet_transform(t_length: int, sampling_frequency: float, frequency_resolution: float) -> tuple:
+    """Get the shape of wavelet transform.
+
+    Args:
+        t_length (int): Length of the time series.
+        sampling_frequency (float): Sampling frequency in Hz.
+        frequency_resolution (float): Frequency resolution in Hz.
+
+    Returns:
+        tuple: n_t, n_f. Number of time points and frequency points.
+    """
+    n_f = int(sampling_frequency / 2 / frequency_resolution)
+    n_t = int(t_length / n_f)
+    return n_t, n_f

nullcal/time_frequency_transform/wavelet_transforms.py

@@ -0,0 +1,249 @@
+"""helper functions for transform_time.py"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from .inverse_wavelet_freq_funcs import inverse_wavelet_freq_helper_fast
+from .inverse_wavelet_time_funcs import inverse_wavelet_time_helper_fast
+from .transform_freq_funcs import (
+    phitilde_vec_norm,
+    transform_wavelet_freq_helper,
+    transform_wavelet_freq_quadrature_helper,
+)
+from .transform_time_funcs import phi_vec, transform_wavelet_time_helper
+from .utils import get_shape_of_wavelet_transform
+
+
+def inverse_wavelet_time(wave_in, n_f, n_t, nx=4.0, mult=32) -> jax.Array:
+    """Fast inverse wavelet transform to time domain.
+
+    Args:
+        wave_in (2D numpy array): Data in wavelet domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of the filter. Defaults to 4..
+        mult (int, optional): mult. Defaults to 32.
+
+    Returns:
+        jax.Array: Data in time domain.
+    """
+    mult = min(mult, n_t // 2)  # make sure K isn't bigger than n_d
+    phi = phi_vec(n_f, nx=nx, mult=mult) / 2
+    output = inverse_wavelet_time_helper_fast(wave_in, phi, n_f, n_t, mult)
+    return output / jnp.sqrt(output.shape[0])
+
+
+def inverse_wavelet_freq_time(wave_in, n_f, n_t, nx=4.0) -> jax.Array:
+    """Inverse wavelet transform to time domain via Fourier transform of frequency domain.
+
+    Args:
+        wave_in (2D numpy array): Data in wavelet domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in time domain.
+    """
+    res_f = inverse_wavelet_freq(wave_in, n_f, n_t, nx)
+    return jnp.fft.irfft(res_f)
+
+
+def inverse_wavelet_freq(wave_in, n_f, n_t, nx=4.0) -> jax.Array:
+    """Inverse wavelet transform to frequency domain signal.
+
+    Args:
+        wave_in (2D numpy array): Data in wavelet domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in frequency domain.
+    """
+    phif = phitilde_vec_norm(n_f, n_t, nx)
+    output = inverse_wavelet_freq_helper_fast(wave_in, phif, n_f, n_t)
+    return output / jnp.sqrt((output.shape[0] - 1) * 2)
+
+
+def transform_wavelet_time(data, n_f, n_t, nx=4.0, mult=32) -> jax.Array:
+    """Do the wavelet transform in the time domain,
+    note there can be significant leakage if mult is too small and the
+    transform is only approximately exact if mult=n_t/2.
+
+    Args:
+        data (1D numpy array): Data in time domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+        mult (int, optional): mult. Defaults to 32.
+
+    Returns:
+        jax.Array: Data in wavelet domain.
+    """
+    mult = min(mult, n_t // 2)  # make sure K isn't bigger than n_d
+    phi = phi_vec(n_f, nx, mult)
+    wave = transform_wavelet_time_helper(data, n_f, n_t, phi, mult) * jnp.sqrt(data.shape[0])
+
+    return wave
+
+
+def transform_wavelet_freq_time(data, n_f, n_t, nx=4.0) -> jax.Array:
+    """Transform time domain data into wavelet domain via FFT and then frequency transform.
+
+    Args:
+        data (1D numpy array): Data in time domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in wavelet domain.
+    """
+    data_fft = jnp.fft.rfft(data)
+
+    return transform_wavelet_freq(data_fft, n_f, n_t, nx)
+
+
+def transform_wavelet_freq_time_quadrature(data, n_f, n_t, nx=4.0) -> jax.Array:
+    """Transform time domain data into wavelet quadrature domain via FFT and then frequency transform.
+
+    Args:
+        data (1D numpy array): Data in time domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in wavelet domain.
+    """
+    data_fft = jnp.fft.rfft(data)
+
+    return transform_wavelet_freq_quadrature(data_fft, n_f, n_t, nx)
+
+
+def transform_wavelet_freq_quadrature(data, n_f, n_t, nx=4.0) -> jax.Array:
+    """Do the wavelet quadrature transform using the fast wavelet domain transform.
+
+    Args:
+        data (1D complex numpy array): Data in frequency domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in wavelet domain.
+    """
+    phif = 2 / n_f * phitilde_vec_norm(n_f, n_t, nx)
+    return transform_wavelet_freq_quadrature_helper(data, n_f, n_t, phif) * jnp.sqrt((data.shape[0] - 1) * 2)
+
+
+def transform_wavelet_freq(data, n_f, n_t, nx=4.0) -> jax.Array:
+    """Do the wavelet transform using the fast wavelet domain transform.
+
+    Args:
+        data (1D complex numpy array): Data in frequency domain.
+        n_f (int): Number of frequency bins.
+        n_t (int): Number of time bins.
+        nx (float, optional): Steepness of filter. Defaults to 4..
+
+    Returns:
+        jax.Array: Data in wavelet domain.
+    """
+    phif = 2 / n_f * phitilde_vec_norm(n_f, n_t, nx)
+    return transform_wavelet_freq_helper(data, n_f, n_t, phif) * jnp.sqrt((data.shape[0] - 1) * 2)
+
+
+class WaveletTransform:
+    """A class to handle wavelet transform."""
+
+    def __init__(self, duration: float, sampling_frequency: float, frequency_resolution: float, nx: float = 4.0):
+        """Wavelet transform.
+
+        Args:
+            duration (float): Duration in second.
+            sampling_frequency (float): Sampling frequency in Hz.
+            frequency_resolution (float): Frequency resolution in Hz.
+            nx (float): Sharpness of wavelet.
+                Defaults to 4.0.
+        """
+        self.duration = duration
+        self.sampling_frequency = sampling_frequency
+        self.frequency_resolution = frequency_resolution
+        self.nx = nx
+        self._t_length = int(duration * sampling_frequency)
+        self._f_length = self._t_length // 2 + 1
+        self.shape = get_shape_of_wavelet_transform(
+            t_length=int(duration * sampling_frequency),
+            sampling_frequency=sampling_frequency,
+            frequency_resolution=frequency_resolution,
+        )
+
+    def frequency_to_wavelet(self, frequency_domain_data: np.ndarray) -> jax.Array:
+        """Transform from frequency domain to wavelet domain.
+
+        Args:
+            frequency_domain_data (np.ndarray): Frequency-domain data.
+
+        Raises:
+            ValueError: The length of frequency_domain_data does not match
+                the expected length.
+
+        Returns:
+            jax.Array: Wavelet-domain data.
+        """
+        # Check whether the length of frequency-domain data
+        # matches the expected length.
+        if len(frequency_domain_data) != self._f_length:
+            raise ValueError(
+                f"The length of frequency_domain_data: {len(frequency_domain_data)}"
+                f"does not match the expected length: {self._f_length}."
+            )
+        return transform_wavelet_freq(data=frequency_domain_data, n_f=self.shape[1], n_t=self.shape[0])
+
+    def wavelet_to_frequency(self, wavelet_domain_data: np.ndarray) -> jax.Array:
+        """Transform from wavelet domain to frequency domain.
+
+        Args:
+            wavelet_domain_data (np.ndarray): Wavelet-domain data.
+
+        Raises:
+            ValueError: The shape of wavelet_domain_data does not match
+                the expected shape.
+
+        Returns:
+            jax.Array: Frequency-domain data.
+        """
+        # Check whether the shape of wavelet-domain strain
+        # matches the expected shape.
+        if wavelet_domain_data.shape != self.shape:
+            raise ValueError(
+                f"The shape of wavelet_domain_data: {wavelet_domain_data.shape}"
+                f"does not match the expected shape: {self.shape}."
+            )
+        return inverse_wavelet_freq(wave_in=wavelet_domain_data, n_f=self.shape[1], n_t=self.shape[0], nx=self.nx)
+
+    def frequency_to_wavelet_quadrature(self, frequency_domain_data: np.ndarray) -> jax.Array:
+        """Transform from frequency domain to wavelet quadrature domain.
+
+        Args:
+            frequency_domain_data (np.ndarray): Frequency-domain data.
+
+        Raises:
+            ValueError: The length of frequency_domain_data does not match
+                the expected length.
+
+        Returns:
+            jax.Array: Wavelet quadrature domain data.
+        """
+        # Check whether the length of frequency-domain data
+        # matches the expected length.
+        if len(frequency_domain_data) != self._f_length:
+            raise ValueError(
+                f"The length of frequency_domain_data: {len(frequency_domain_data)}"
+                f"does not match the expected length: {self._f_length}."
+            )
+        return transform_wavelet_freq_quadrature(data=frequency_domain_data, n_f=self.shape[1], n_t=self.shape[0])