pywavelet 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

pywavelet/_version.py

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

pywavelet/transforms/jax/inverse/main.py

@@ -0,0 +1,69 @@
+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 .to_freq import inverse_wavelet_freq_helper
+# from .inverse_wavelet_time_funcs import inverse_wavelet_time_helper
+
+
+def from_wavelet_to_time(
+    wave_in: Wavelet,
+    dt: float,
+    nx: float = 4.0,
+    mult: int = 32,
+) -> TimeSeries:
+    """Inverse wavelet transform to time domain.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        input wavelet
+    dt : float
+        time step
+    nx : float, optional
+        parameter for phi_vec, by default 4.0
+    mult : int, optional
+        parameter for phi_vec, by default 32
+
+    Returns
+    -------
+    TimeSeries
+        Time domain signal
+    """
+    # Can we just do this?
+    freq = from_wavelet_to_freq(wave_in, dt=dt, nx=nx)
+    return freq.to_timeseries()
+
+
+def from_wavelet_to_freq(
+    wave_in: Wavelet, dt: float, nx=4.0
+) -> FrequencySeries:
+    """Inverse wavelet transform to frequency domain.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        input wavelet
+    dt : float
+        time step
+    nx : float, optional
+        parameter for phitilde_vec_norm, by default 4.0
+
+    Returns
+    -------
+    FrequencySeries
+        Frequency domain signal
+
+    """
+    phif = jnp.array(phitilde_vec_norm(wave_in.Nf, wave_in.Nt, dt=dt, d=nx))
+    freq_data = inverse_wavelet_freq_helper(
+        wave_in.data, phif=phif, Nf=wave_in.Nf, Nt=wave_in.Nt
+    )
+
+    freq_data *= 2 ** (
+        -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)

pywavelet/transforms/jax/inverse/to_freq.py

@@ -0,0 +1,70 @@
+import jax
+import jax.numpy as jnp
+from jax import jit
+from jax.numpy.fft import fft
+
+from functools import partial
+
+
+@partial(jit, static_argnames=('Nf', 'Nt'))
+def inverse_wavelet_freq_helper(
+        wave_in: jnp.ndarray, phif: jnp.ndarray, Nf: int, Nt: int
+) -> jnp.ndarray:
+    """JAX vectorized function for inverse_wavelet_freq"""
+    wave_in = wave_in.T
+    ND = Nf * Nt
+
+    m_range = jnp.arange(Nf + 1)
+    prefactor2s = jnp.zeros((Nf + 1, Nt), dtype=jnp.complex128)
+
+    n_range = jnp.arange(Nt)
+
+    # m == 0 case
+    prefactor2s = prefactor2s.at[0].set(2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt, 0])
+
+    # m == Nf case
+    prefactor2s = prefactor2s.at[Nf].set(2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt + 1, 0])
+
+    # Other m cases
+    m_mid = m_range[1:Nf]
+    n_grid, m_grid = jnp.meshgrid(n_range, m_mid)
+    val = wave_in[n_grid, m_grid]
+    mult2 = jnp.where((n_grid + m_grid) % 2, -1j, 1)
+    prefactor2s = prefactor2s.at[1:Nf].set(mult2 * val)
+
+    # Vectorized FFT
+    fft_prefactor2s = fft(prefactor2s, axis=1)
+
+    # Vectorized __unpack_wave_inverse
+    ## TODO: Check with Giorgio
+    # ND or ND // 2 + 1?
+    # https://github.com/pywavelet/pywavelet/blob/63151a47cde9edc14f1e7e0bf17f554e78ad257c/src/pywavelet/transforms/from_wavelets/inverse_wavelet_freq_funcs.py
+    res = jnp.zeros(ND, dtype=jnp.complex128)
+
+    # m == 0 or m == Nf cases
+    i_ind_range = jnp.arange(Nt // 2)
+    i_0 = jnp.abs(i_ind_range)
+    i_Nf = jnp.abs(Nf * Nt // 2 - i_ind_range)
+    ind3_0 = (2 * i_0) % Nt
+    ind3_Nf = (2 * i_Nf) % Nt
+
+    res = res.at[i_0].add(fft_prefactor2s[0, ind3_0] * phif[i_ind_range])
+    res = res.at[i_Nf].add(fft_prefactor2s[Nf, ind3_Nf] * phif[i_ind_range])
+
+    # Special case for m == Nf
+    res = res.at[Nf * Nt // 2].add(fft_prefactor2s[Nf, 0] * phif[Nt // 2])
+
+    # Other m cases
+    m_mid = m_range[1:Nf]
+    i_ind_range = jnp.arange(Nt // 2 + 1)
+    m_grid, i_ind_grid = jnp.meshgrid(m_mid, i_ind_range)
+
+    i1 = Nt // 2 * m_grid - i_ind_grid
+    i2 = Nt // 2 * m_grid + i_ind_grid
+    ind31 = (Nt // 2 * m_grid - i_ind_grid) % Nt
+    ind32 = (Nt // 2 * m_grid + i_ind_grid) % Nt
+
+    res = res.at[i1].add(fft_prefactor2s[m_grid, ind31] * phif[i_ind_grid])
+    res = res.at[i2].add(fft_prefactor2s[m_grid, ind32] * phif[i_ind_grid])
+
+    return res

{pywavelet-0.2.0.dist-info → pywavelet-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.0
+Version: 0.2.1
 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.0.dist-info → pywavelet-0.2.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pywavelet/__init__.py,sha256=zcK3Qj4wTrGZF1rU3aT6yA9LvliAOD4DVOY7gNfHhCI,53
-pywavelet/_version.py,sha256=H-qsvrxCpdhaQzyddR-yajEqI71hPxLa4KxzpP3uS1g,411
+pywavelet/_version.py,sha256=MxUhzLJIZQfEpDTTcKSxciTGrMLd5v2VmMlHa2HGeo0,411
 pywavelet/backend.py,sha256=k4pDi6f4cwNY6HsUIx1xfuga9f2wLnFr_FIb7Fs1Mds,553
 pywavelet/logger.py,sha256=DyKC-pJ_N9GlVeXL00E1D8hUd8GceBg-pnn7g1YPKcM,391
 pywavelet/utils.py,sha256=l47C643nGlV9q4a0G7wtKzuas0Ou4En2e1FTATCgwlw,1907
@@ -10,6 +10,9 @@ pywavelet/transforms/jax/forward/__init__.py,sha256=Ki2RJCfkE9Zy59mqT3oEtGK9Ro9k
 pywavelet/transforms/jax/forward/from_freq.py,sha256=PsUC7RfrN6pRWWkMSXYHk9z5lxCXW3DfF0m-Rd1GOBE,1785
 pywavelet/transforms/jax/forward/from_time.py,sha256=xNeoZq54B6Gi3TdTTYLr_euaFeJcwpms-lSyCG53AdI,1726
 pywavelet/transforms/jax/forward/main.py,sha256=mm0R4m0pXcnzZB0jCckAc4ynG8STH5mldCmHyyU_PGo,3091
+pywavelet/transforms/jax/inverse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pywavelet/transforms/jax/inverse/main.py,sha256=ZK8NyfMI6oFYMKcALatETWnCystH0LWjEYInvnMMmh0,1714
+pywavelet/transforms/jax/inverse/to_freq.py,sha256=ASNARcDBJQr4EizAP_77e5ai36iPwP6hzfvwGbZQ6BM,2295
 pywavelet/transforms/numpy/__init__.py,sha256=qFLpGpW3VJSbDp2JpD0Gx7PdwDjH-wrW_aO84ASkIgA,255
 pywavelet/transforms/numpy/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
 pywavelet/transforms/numpy/forward/from_freq.py,sha256=JmJyjrNSb64WnpP50VZRt0BICP64iZJP5QAZTZoexkw,2675
@@ -26,7 +29,7 @@ pywavelet/types/plotting.py,sha256=JNDxeP-fB8U09E90J-rVT-h5yCGA_tGRHtctbgINiRo,1
 pywavelet/types/timeseries.py,sha256=u35bIqFo3QdlQRBEu6maeWA7DePS11LQ6WMiLjZPcWo,9456
 pywavelet/types/wavelet.py,sha256=el48oyAfwtSw2tCQLUb85F9lKr0qMSRJPUmAUU8TS50,12552
 pywavelet/types/wavelet_bins.py,sha256=GoQGKeZlPc-KbYY7LoxAhB-HI4diHpPcTABBXRfUTLA,1459
-pywavelet-0.2.0.dist-info/METADATA,sha256=rg9LNZxrykv39lKIGuX65v7WzuY93_D0oY268mMe5iw,1362
-pywavelet-0.2.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-pywavelet-0.2.0.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
-pywavelet-0.2.0.dist-info/RECORD,,
+pywavelet-0.2.1.dist-info/METADATA,sha256=pVblWF5CqZ_hFU9J8_bAENKh5HqJnIdSy8y1RR3TkVU,1362
+pywavelet-0.2.1.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+pywavelet-0.2.1.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
+pywavelet-0.2.1.dist-info/RECORD,,