# WDM Wavelet — API Usage Guide

**Reference paper**: V. Necula et al., *J. Phys.: Conf. Ser.* 363, 012032 (2012)
https://doi.org/10.1088/1742-6596/363/1/012032

---

## Contents

1. [Two-layer API design](#1-two-layer-api-design)
2. [High-level class API — `WDM`](#2-high-level-class-api--wdm)
3. [Core pure-function API](#3-core-pure-function-api)
4. [Filter construction](#4-filter-construction)
5. [Time-delay amplitude API](#5-time-delay-amplitude-api)
6. [Embedding in JAX JIT contexts](#6-embedding-in-jax-jit-contexts)
7. [Embedding in Numba JIT contexts](#7-embedding-in-numba-jit-contexts)
8. [Data types and precision](#8-data-types-and-precision)
9. [Quick-reference table](#9-quick-reference-table)

---

## 1. Two-layer API design

The package exposes two complementary layers:

| Layer | Where | Returns | JIT-embeddable |
|---|---|---|---|
| **Class methods** | `wdm_wavelet.WDM` | `TimeFrequencyMap`, `gwpy.TimeSeries` | No (user-facing) |
| **Pure functions** | `wdm_wavelet.core.*` | Raw `jax.Array` / `numpy.ndarray` | Yes (`@jax.jit`, `jax.vmap`) |

Use the class API for interactive or scripted analysis.  Use the pure functions when the transform must be called from inside a compiled pipeline.

---

## 2. High-level class API — `WDM`

### 2.1 Construction

```python
from wdm_wavelet import WDM

wdm = WDM(
    M=32,           # int — number of frequency bands (Nyquist index).
                    #   TF map has M+1 layers: DC(0), interior(1…M-1), Nyquist(M).
    K=64,           # int — Meyer window sharpness. Larger K → steeper roll-off, longer filter.
    beta_order=6,   # int — beta polynomial order (selects Fourier coefficient table).
    precision=10,   # int — decimal digits of reconstruction accuracy.
                    #   Controls auto-computed filter length.
    backend="jax",  # str — compute backend, locked at construction. Currently only "jax".
    # Optional:
    filter=None,    # array-like or None — custom filter taps. None = auto-compute.
    m_H=None,       # int or None — number of filter taps to use. Inferred when None.
)
```

**Notes**
- `backend` is **immutable** after construction.  No per-call `backend=` parameter exists.
- When `filter=None`, the filter is computed then truncated to the minimum length satisfying reconstruction error < `10**(-precision)`.
- Custom `filter` must satisfy `(m_H - 1) % (2*M) == 0`.
- `WDM` is registered as a JAX pytree leaf, so it can be passed to `jax.jit`-compiled functions directly.

---

### 2.2 Forward transform — `WDM.t2w`

```python
tf_map = wdm.t2w(
    strain,          # array-like | gwpy.TimeSeries | pycbc.TimeSeries
                     #   Input time series.
    sample_rate=None,# float — required when strain is a plain array.
    t0=0,            # float — start time in seconds (used for metadata only).
    MM=-1,           # int — decimation mode.
                     #   MM < 0 → full quadrature output, stride = M (default, MM=-1).
                     #   MM > 0 → power-only (no 90° phase), stride = MM.
)
# Returns: TimeFrequencyMap
#   .data        — complex128 array, shape (M+1, n_time)
#                  real part = 00-phase coefficients
#                  imag part = 90-phase coefficients (zero when MM >= 0)
#   .df          — frequency resolution in Hz
#   .dt          — time resolution in seconds
#   .t0          — start time in seconds
#   .n_freq      — M+1 (number of frequency layers)
#   .wdm_params  — dict with M, K, beta_order, precision, has_quadrature, t2w_mm
```

---

### 2.3 Inverse transforms — `WDM.w2t` and `WDM.w2tQ`

```python
strain_out = wdm.w2t(tf_map)
# Reconstructs time series from 00-phase coefficients.
# Input:  TimeFrequencyMap produced by wdm.t2w
# Output: gwpy.TimeSeries with sample_rate and t0 preserved

strain_out = wdm.w2tQ(tf_map)
# Reconstructs time series from 90-phase (quadrature) coefficients.
# Requires tf_map produced with MM < 0 (has_quadrature=True).
# Raises ValueError if map was produced in power-only mode.
```

---

### 2.4 JAX-native class wrappers (flat-array interface)

For use inside `@jax.jit`-decorated class methods, three wrappers pre-bind `self.filter`:

```python
m_L, aligned_length, tf_map = wdm.t2w_jax(signal, MM=-1)
# signal: jax.Array, shape (N,)
# MM: static int (see core API below)
# Returns same as t2w_jax pure function

reconstructed = wdm.w2t_jax(tf_map_flattened, m_layer, output_length=None)
reconstructed = wdm.w2tQ_jax(tf_map_flattened, m_layer, output_length=None)
# tf_map_flattened: jax.Array — 1-D flattened coefficient map
# m_layer: M+1 (static int)
# output_length: optional static int to truncate result
```

---

### 2.5 Time-delay amplitudes

```python
# Step 1 — build TD filters (call once per wdm instance)
wdm.set_td_filter(
    n_coeffs=4,  # int — half-width of the symmetric TD filter kernel.
    L=1,         # int — upsampling factor. Fundamental delay step = tau / L.
)

# Step 2 — query delayed amplitudes
amp = wdm.get_pixel_amplitude(
    tf_map,   # TimeFrequencyMap with both quadratures (has_quadrature=True)
    m,        # int — frequency layer index (0 … M)
    n,        # int — time bin index
    dT,       # float — time delay in units of the fundamental delay step (1/L of tau)
    quad=False,  # bool — False → 00-phase, True → 90-phase result
)
# Returns: float — delayed amplitude for pixel (m, n) at delay dT

amp_value = wdm.get_td_amp(
    tf_map,
    pixel_index,  # int — linear pixel index = n * (M+1) + m
    delay_index,  # int — delay index in [-M*L, M*L]
    mode="p",     # str — "p" (power), "s" (scaled amp), "a" (raw amp)
)

amp_vector = wdm.get_td_vec(
    tf_map,
    pixel_index,  # int — linear pixel index
    K,            # int — half-range; returns delays [-K, ..., K]
    mode="p",     # str — same as get_td_amp
)
# Returns: numpy.ndarray, shape (2*K+1,)
```

---

## 3. Core pure-function API

These functions return **raw JAX arrays only** and are safe to call inside `@jax.jit` / `jax.vmap` pipelines.

### 3.1 Forward transform — `t2w_jax`

```python
from wdm_wavelet.core.t2w import t2w_jax

m_L, aligned_length, tf_map = t2w_jax(
    M,            # int (static) — number of frequency bands.
    m_H,          # int (static) — filter half-length. Must satisfy (m_H-1) % (2*M) == 0.
    signal,       # array-like, shape (N,) — input time series (float64).
    filter_taps,  # array-like, shape (>= m_H,) — WDM prototype filter coefficients.
    MM,           # int — decimation mode (see below).
)
```

**`MM` parameter**

| Value | Behaviour |
|---|---|
| `MM < 0` | Full quadrature output. Stride = M. Returns both 00 and 90 coefficients. |
| `MM > 0` | Power-only (no 90° phase). Stride = MM. |
| `MM == 0` | Treated as `MM = M` internally (power-only, natural stride). |

**Returns**

| Name | Type | Description |
|---|---|---|
| `m_L` | `int` | Filter half-length used (`m_H` for quadrature, `0` for power-only). |
| `aligned_length` | `int` | Signal length after block alignment to `MM`. |
| `tf_map` | `jax.Array`, shape `(2, n_time, M+1)` | `tf_map[0]` = 00-phase (or power); `tf_map[1]` = 90-phase (zeros when power-only). |

where `n_time = aligned_length // MM`.

---

### 3.2 Inverse transforms — `w2t_jax` and `w2tQ_jax`

```python
from wdm_wavelet.core.w2t import w2t_jax, w2tQ_jax

signal_out = w2t_jax(
    tf_map,         # array-like — flattened TF coefficients, shape (2 * n_time * (M+1),).
                    #   Typically produced by tf_map.reshape(-1) or tf_map.ravel().
    m_layer,        # int (static) — number of frequency layers = M + 1.
    filter_taps,    # array-like, shape (m_H,) — WDM prototype filter coefficients.
    output_length=None,  # int or None — if given, truncate output to this many samples.
)
# Returns: jax.Array, shape (output_length,) or (M * n_time,)

signal_out = w2tQ_jax(
    tf_map,         # array-like — must contain BOTH quadratures: shape (2 * n_time * (M+1),).
                    #   Only the second half (90-phase) is used.
    m_layer,        # int (static)
    filter_taps,    # array-like, shape (m_H,)
    output_length=None,
)
# Returns: jax.Array, shape (output_length,) or (M * n_time,)
```

**Flattening convention for `tf_map`**

```
flat = tf_map.reshape(-1)
# Layout:  [00-phase block | 90-phase block]
# 00-phase block: shape (n_time, M+1) → indices 0 … n_time*(M+1)-1
# 90-phase block: shape (n_time, M+1) → indices n_time*(M+1) … 2*n_time*(M+1)-1
```

`w2t_jax` reads the **first** block; `w2tQ_jax` reads the **second** block.

---

## 4. Filter construction

```python
from wdm_wavelet.core.get_filter import get_filter, get_filter_jax

# Auto-dispatch (JAX preferred, Numba fallback, plain Python last resort)
filter_taps = get_filter(
    M,           # int — number of frequency bands.
    K,           # int — Meyer window sharpness.
    beta_order,  # int — beta polynomial order.
    n,           # int — number of taps to compute (upper bound; typically 3e5).
    Cos,         # array — Fourier coefficient table (from load_fourier_coeff()).
    CosSize,     # array — number of active coefficients per order.
)
# Returns: numpy.ndarray, shape (n,)

# JAX-only version (preferred; JIT-compiled with lax.fori_loop)
filter_taps = get_filter_jax(M, K, beta_order, n, Cos, CosSize)
# Returns: numpy.ndarray, shape (n,)
```

**Loading Fourier coefficient tables**

```python
from wdm_wavelet.fourier_coeff import load_fourier_coeff

Cos, Cos2, SinCos, CosSize, Cos2Size, SinCosSize = load_fourier_coeff()
# Cos / CosSize    — used by get_filter (forward transform filter)
# Cos2 / Cos2Size  — used by set_td_filter (TD filter 1, same-band)
# SinCos / SinCosSize — used by set_td_filter (TD filter 2, cross-band)
```

`WDM.__post_init__` calls `get_filter` automatically when no custom filter is provided.

---

## 5. Time-delay amplitude API

Time-delay amplitudes implement Eq. (27)-(32) from the paper.  They require `wdm.set_td_filter()` to have been called first.

### Delay index convention

The fundamental delay unit is `tau / L` where:
- `tau = 1 / sample_rate` — sample period
- `L` — upsampling factor set in `set_td_filter`

The full integer delay range is `[-M*L, ..., M*L]`, giving `2*M*L + 1` distinct delays.

### `mode` parameter for `get_td_amp` / `get_td_vec`

| `mode` | Returns |
|---|---|
| `"p"` | Power = amplitude² |
| `"s"` | Scaled amplitude (signed) |
| `"a"` | Raw (unscaled) complex amplitude |

---

## 6. Embedding in JAX JIT contexts

The core functions are designed to be called inside `@jax.jit`-compiled functions.

### 6.1 Direct embedding

```python
import jax
import jax.numpy as jnp
from wdm_wavelet.core.t2w import t2w_jax
from wdm_wavelet.core.w2t import w2t_jax

# Pre-compute static values outside JIT
M = 32
m_H = wdm.m_H
filter_taps = wdm.filter  # jax.Array (already float64)
MM = -1                   # must be a Python int (static)
m_layer = M + 1

@jax.jit
def forward_pass(signal):
    m_L, aligned_len, tf = t2w_jax(M, m_H, signal, filter_taps, MM)
    return tf  # jax.Array, shape (2, n_time, M+1)

@jax.jit
def round_trip(signal):
    _, aligned_len, tf = t2w_jax(M, m_H, signal, filter_taps, MM)
    flat = tf.reshape(-1)
    return w2t_jax(flat, m_layer, filter_taps, output_length=int(signal.shape[0]))
```

### 6.2 Using `WDM` as a pytree inside JIT

`WDM` is registered as a JAX pytree (via `jax.tree_util.register_dataclass`).  This means you can pass a `WDM` instance as an argument to a `@jax.jit` function and JAX will correctly trace through it.

```python
@jax.jit
def pipeline(wdm_instance, signal):
    # wdm_instance.filter is the only dynamic (data) leaf
    # M, K, beta_order, precision, backend, m_H are static (meta) leaves
    _, _, tf = t2w_jax(wdm_instance.M, wdm_instance.m_H,
                       signal, wdm_instance.filter, -1)
    return tf
```

### 6.3 `static_argnames` in `_t2w_jax_impl`

The internal `_t2w_jax_impl` kernel uses `@partial(jax.jit, static_argnames=(...))` with the following static arguments:

| Static arg | Reason |
|---|---|
| `M` | Controls array shape (frequency layers = M+1, FFT size = 2*M) |
| `n_filter_taps` | Controls loop bounds and index shapes |
| `mm_eff` | Controls time-bin count |
| `return_quadrature` | Determines which branch to compile |
| `n_time_bins` | Controls vmap output shape |

Changing any of these values triggers recompilation.  Running many `(M, m_H)` combinations at once will cause cache thrashing — prefer fixing these at program start.

### 6.4 `jax.vmap` over a batch of signals

```python
# Batch forward transform over a leading batch dimension
batched_t2w = jax.vmap(
    lambda sig: t2w_jax(M, m_H, sig, filter_taps, -1)[2],
    in_axes=0,
)
batch_tf = batched_t2w(signals)  # signals: shape (batch, N)
# batch_tf: shape (batch, 2, n_time, M+1)
```

---

## 7. Embedding in Numba JIT contexts

The core JAX functions **cannot** be called from `@numba.njit` contexts because Numba cannot trace through JAX operations.

The recommended pattern is to **exit Numba**, call the JAX function, then re-enter Numba:

```python
import numba
import numpy as np
from wdm_wavelet.core.t2w import t2w_jax

# JAX call outside Numba — get a NumPy result
_, _, tf_jax = t2w_jax(M, m_H, signal, filter_taps, MM)
tf_np = np.asarray(tf_jax)  # move to CPU NumPy

@numba.njit(cache=True)
def process_tf(tf):
    # Work on the NumPy TF map inside Numba
    result = np.zeros_like(tf[0])
    for t in range(tf.shape[1]):
        for f in range(tf.shape[2]):
            result[t, f] = tf[0, t, f] ** 2  # example
    return result

output = process_tf(tf_np)
```

### What IS compiled with Numba in this package

Only TD filter construction uses Numba JIT:

```python
# These are @numba.njit(cache=True) — called automatically by set_td_filter:
_fill_td_filter1_loop(out, fourier, n_fourier, K, K2, L, size, A, B, M_int, g_norm)
_fill_td_filter2_loop(out, fourier, n_fourier, K, K2, L, size, B, g_norm)
```

When Numba is absent, these fall back to plain Python loops transparently.

---

## 8. Data types and precision

| Location | dtype | Notes |
|---|---|---|
| `filter_taps` internal | `float64` | Always; JAX x64 mode is enabled at import |
| `signal` input | converted to `float64` | Cast via `jnp.asarray(signal, dtype=jnp.float64)` |
| `tf_map` raw output | `float64` | Real and imaginary parts separately |
| `TimeFrequencyMap.data` | `complex128` | Interleaved after class wrapping |
| TD filter tables `T0`, `Tx` | `float64` | NumPy arrays |

**Important**: `jax.config.update("jax_enable_x64", True)` is called at import time in all core modules.  This must remain enabled for filter computations to be numerically accurate.

---

## 9. Quick-reference table

| Symbol | Meaning | Typical value |
|---|---|---|
| `M` | Frequency bands; Nyquist layer index | 32, 64, 128 … |
| `K` | Meyer window sharpness | `2*M` (e.g. 64 for M=32) |
| `beta_order` | Beta polynomial order | 6 |
| `precision` | Reconstruction decimal digits | 10 |
| `m_H` | Filter half-length | Auto-computed; `~K * (precision + a few)` blocks |
| `MM` | Decimation stride | -1 (quadrature), M (power-only, natural) |
| `m_layer` | Frequency layer count = M+1 | 33 for M=32 |
| `n_time` | Time bins = `aligned_length // MM` | depends on signal length |
| `L` | TD delay upsampling factor | 1 (integer delays), or higher |
| `n_coeffs` | TD filter half-width | 4–8 |
| `dT` | Delay index | integer in `[-M*L, M*L]` |

---

## Example: complete round-trip

```python
import numpy as np
from wdm_wavelet import WDM

# 1. Build instance (backend locked to JAX)
wdm = WDM(M=32, K=64, beta_order=6, precision=10)

# 2. Generate test signal
rng = np.random.default_rng(0)
signal = rng.standard_normal(4096 * 4)  # 4 s at 4096 Hz

# 3. Forward transform with full quadrature
tf_map = wdm.t2w(signal, sample_rate=4096.0)

# 4. Inverse transform
reconstructed = wdm.w2t(tf_map)

# 5. Check reconstruction quality
residual = np.max(np.abs(signal - np.asarray(reconstructed.value[:len(signal)])))
print(f"Max reconstruction error: {residual:.2e}")  # expect < 1e-10

# 6. Time-delay amplitudes
wdm.set_td_filter(n_coeffs=4, L=1)
amp_curve = wdm.get_td_vec(tf_map, pixel_index=100, K=8, mode="p")
print(f"TD power vector shape: {amp_curve.shape}")  # (17,) = 2*8+1
```

---

## Example: pure-function pipeline inside `@jax.jit`

```python
import jax
import jax.numpy as jnp
from wdm_wavelet import WDM
from wdm_wavelet.core.t2w import t2w_jax
from wdm_wavelet.core.w2t import w2t_jax

wdm = WDM(M=32, K=64, beta_order=6, precision=10)
M, m_H, filt = wdm.M, wdm.m_H, wdm.filter
m_layer = M + 1

@jax.jit
def matched_filter_pipeline(signal, template):
    _, _, tf_signal   = t2w_jax(M, m_H, signal,   filt, -1)
    _, _, tf_template = t2w_jax(M, m_H, template, filt, -1)
    
    # Whiten / correlate in TF domain
    snr_map = tf_signal[0] * tf_template[0] + tf_signal[1] * tf_template[1]
    
    return snr_map  # jax.Array, shape (n_time, M+1)

snr = matched_filter_pipeline(signal_jax, template_jax)
```
