# WDM Package: Equations and Code Map

This document explains how the `wdm_wavelet` package maps the WDM equations in:

- V. Necula et al., *J. Phys.: Conf. Ser.* **363** (2012) 012032
- `docs/ref/WDM.cc` (reference C++ implementation)

It focuses on:

1. Filter construction (`get_filter`)
2. Forward transform (`t2w`)
3. Inverse transform (`w2t`)
4. Time-delay filters and delayed amplitudes (`time_delay`)

---

## 1) Package Structure

Main entrypoint:

- `wdm_wavelet/wdm.py`: class `WDM`

Core modules:

- `wdm_wavelet/core/get_filter.py`
- `wdm_wavelet/core/t2w.py`
- `wdm_wavelet/core/w2t.py`
- `wdm_wavelet/core/time_delay.py`

Data container:

- `wdm_wavelet/types/time_frequency_map.py`: `TimeFrequencyMap`

---

## 2) Notation and Parameters

Package parameters:

- \( M \): maximum frequency layer index (`WDM.M`)
- \( K \): Meyer-like transition parameter (`WDM.K`)
- \( \beta \): beta order (`WDM.beta_order`)
- \( L \): delay upsampling factor in TD filters (`set_td_filter(..., L=...)`)

Useful constants used in code:

$$
B = \frac{\pi}{K}, \qquad
A = \frac{(K-M)\pi}{2KM}
$$

These correspond to the generalized Meyer window parameterization (paper Eq. (13)–(16), Eq. (33) setting).

---

## 3) Filter Construction (`get_filter.py`)

### 3.1 Frequency-domain window model

The implementation follows the paper’s WDM window construction:

$$
\phi(\omega) =
\begin{cases}
\sqrt{\Delta\Omega}, & |\omega| < A, \\
\sqrt{\Delta\Omega}\cos\!\left[\nu_n\!\left(\frac{|\omega|-A}{B}\right)\right], & A \le |\omega| < A+B, \\
0, & \text{otherwise},
\end{cases}
$$

with:

$$
2A + B = \Delta\Omega, \qquad B = \frac{\Omega}{K}.
$$

The smooth transition uses the beta-function profile (paper Eq. (16)).

### 3.2 Code mapping

`get_filter_numba` evaluates the inverse-Fourier closed form for discrete taps \(\phi[k]\).

```python
# wdm_wavelet/core/get_filter.py
transition_width_B = math.pi / K
flat_half_width_A = (K - M) * math.pi / (2.0 * K * M)
global_norm = math.sqrt(2 * M) / math.pi

phi_taps[0] = (flat_half_width_A + coeffs[0] * transition_width_B) * global_norm
...
phi_taps[sample_idx] = global_norm * (
    math.sin(sample_idx * flat_half_width_A) / sample_float
    + math.sqrt(2) * transition_integral
)
```

JAX path (`get_filter_jax`) is numerically equivalent and selected by default when JAX is available.

---

## 4) Forward Transform (`t2w.py`)

### 4.1 Paper equations used

The implementation is based on the FFT reformulation:

$$
\tilde{X}_n[j] = \sum_k e^{-i\frac{2\pi}{2M}jk}\, X_n[k] \quad \text{(paper Eq. (22))}
$$

with grouped samples:

$$
X_n[j] = \sum_k x[nM + 2kM + j]\,\phi[2kM + j] \quad \text{(paper Eq. (23))}
$$

Then \(w_{nm}\) and \(\hat{w}_{nm}\) are formed from parity/phase combinations (paper Eq. (19)–(21), Eq. (24)–(26)).

### 4.2 Code mapping

Eq. (23) vector build:

```python
# wdm_wavelet/core/t2w.py
def _build_eq23_vector(dst, extended_signal, center_index, filter_taps, fft_size):
    ...
    dst[freq_bin] += (
        extended_signal[forward_start + freq_bin] * filter_taps[tap_block_start + freq_bin]
        + extended_signal[mirrored_start + freq_bin] * filter_taps[last_filter_index - freq_bin]
    )
```

Eq. (22) FFT + coefficient assembly:

```python
# wdm_wavelet/core/t2w.py
_build_eq23_vector(fft_input, extended_signal, center_index, filter_taps, fft_size)
fft_forward()

fft_real[0] = fft_imag[0] = fft_real[0] / sqrt2
fft_real[M] = fft_imag[M] = fft_real[M] / sqrt2
```

Parity/phase assignment for \(w_{nm}\), \(\hat{w}_{nm}\):

```python
if (freq_idx + parity) & 1:
    tf_map[0, time_idx, freq_idx] = sqrt2 * fft_imag[freq_idx]
    tf_map[1, time_idx, freq_idx] = sqrt2 * fft_real[freq_idx]
else:
    tf_map[0, time_idx, freq_idx] = sqrt2 * fft_real[freq_idx]
    tf_map[1, time_idx, freq_idx] = -sqrt2 * fft_imag[freq_idx]
```

Power-only mode (`MM == 0`) stores:

$$
p_{nm} = w_{nm}^2 + \hat{w}_{nm}^2.
$$

---

## 5) Inverse Transform (`w2t.py`)

### 5.1 Mathematical idea

`w2t` inverts:

1. parity/phase mapping for \((w,\hat{w}) \leftrightarrow\) Fourier bins
2. FFT step (inverse of Eq. (22))
3. overlap-add reconstruction (inverse of Eq. (23) grouping)

### 5.2 Code mapping

Rebuild Fourier bins:

```python
# wdm_wavelet/core/w2t.py
def _fill_inverse_spectrum_bins(real_bins, imag_bins, coeff_map, time_idx, M, sqrt2):
    ...
```

Inverse FFT + overlap-add:

```python
# wdm_wavelet/core/w2t.py
fft_inverse()
eq23_time_vector = fft_output.real * (2 * fft_size)
_overlap_add_inverse(reconstructed, eq23_time_vector, filter_taps, n_filter_taps, time_idx, M)
```

JAX inverse (`w2t_jax`) mirrors the same operations.

---

## 6) Time-Delay Filters (`time_delay.py`)

This module implements paper Sec. 4, Eq. (27)–(32).

### 6.1 Delayed TF coefficient (Eq. (27))

$$
w_{nm}(\delta t) = \sum_{l,k} w_{lk}\,\langle g_{nm}(t+\delta t), g_{lk}(t)\rangle.
$$

### 6.2 Same-band overlap (Eq. (29), Eq. (30))

For \(k=m\), the overlap uses \(T_l\):

$$
T_l = \int_{-\Omega}^{\Omega}
e^{i\omega(Ml\tau+\delta t)}
\left|\phi(\omega)\right|^2\,d\omega.
$$

Code:

```python
# wdm_wavelet/core/time_delay.py
def get_td_filter1(...):
    # Computes T_l (Eq. 30)
```

### 6.3 Neighbor-band overlap (Eq. (31), Eq. (32))

For \(k=m\pm1\), the overlap uses \(T'_l\):

$$
T'_l = \int_{-\Omega}^{\Omega}
e^{i\omega(Ml\tau+\delta t)}
\phi^*(\omega-\Delta\Omega/2)\phi(\omega+\Delta\Omega/2)\,d\omega.
$$

Code:

```python
# wdm_wavelet/core/time_delay.py
def get_td_filter2(...):
    # Computes T'_l (Eq. 32)
```

### 6.4 Filter-bank construction

`set_td_filter` stores symmetric tables \(T_0\) and \(T_x\) for delays \(dT\in[-ML,ML]\):

$$
T_0[dT,0] = T_{|dT|},\quad
T_0[dT,\pm j] = T_{dT \pm J\,j}
$$

$$
T_x[dT,0] = \frac{1}{2}T'_{|dT|},\quad
T_x[dT,\pm j] = \frac{1}{2}T'_{dT \pm J\,j}\times(\text{C++ sign convention})
$$

(Here code uses \(J=M L\); see `time_delay.py` for exact indexing.)

### 6.5 Delayed pixel amplitude

`get_pixel_amplitude` evaluates Eq. (29)/(31) terms with:

- even/odd \(l\)-parity sums,
- phase factors \( \cos(m\pi dt/M), \sin(m\pi dt/M)\),
- neighbor phases \( \sin((2m\pm1)\pi dt/(2M))\),
- endpoint factor \(\sqrt{2}\) (paper Eq. (31) note).

Code:

```python
# wdm_wavelet/core/time_delay.py
same_even, same_odd = _sum_even_odd_for_layer(...)
...
low_term  = (...) * sin((2*m-1)*pi*dt/(2*M))
high_term = (...) * sin((2*m+1)*pi*dt/(2*M))
```

### 6.6 Scalar and vector TD APIs

Global delay index decomposition:

$$
k = \text{wdm\_shift}\cdot J + \text{local\_delay}, \quad J = M L.
$$

Power mode:

$$
p = \frac{a_{00}^2 + a_{90}^2}{2}.
$$

Code:

```python
# wdm_wavelet/core/time_delay.py
def get_td_amp(...):
    ...
    return 0.5 * (a00 * a00 + a90 * a90)
```

---

## 7) High-Level API (`wdm.py`)

`WDM` ties all pieces together:

- builds filter taps in constructor
- `t2w(...)` returns `TimeFrequencyMap`
- `w2t(...)` reconstructs time series
- TD functions:
  - `set_td_filter(...)`
  - `get_pixel_amplitude(...)`
  - `get_td_amp(...)`
  - `get_td_vec(...)`

Example:

```python
from wdm_wavelet.wdm import WDM

wdm = WDM(M=16, K=32, beta_order=6, precision=8)
tf_map = wdm.t2w(signal, sample_rate=4096.0, t0=0.0, MM=-1)

wdm.set_td_filter(n_coeffs=8, L=2)
pixel = (tf_map.n_time // 2) * (wdm.M + 1) + (wdm.M // 2)

a00 = wdm.get_td_amp(tf_map, pixel, delay_index=3, mode="a")
a90 = wdm.get_td_amp(tf_map, pixel, delay_index=3, mode="A")
pwr = wdm.get_td_amp(tf_map, pixel, delay_index=3, mode="p")
```

---

## 8) Important Implementation Notes

- TD requires quadrature map (00/90), not power-only maps.
- TD for odd \(M\) is explicitly rejected (matching C++ limitation note).
- NumPy/FFTW and JAX implementations are both available for forward/inverse/filter paths.
