Skip to content

Fix mult_f boundary conventions, window grid, and mg time-domain path #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mcdigman wants to merge 1 commit into
base: working
Choose a base branch
from
Open

Fix mult_f boundary conventions, window grid, and mg time-domain path #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 42 additions & 129 deletions WDMWaveletTransforms/inverse_wavelet_freq_funcs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,72 +6,6 @@

import WDMWaveletTransforms.fft_funcs as fft

# @njit()
# def unpack_wave_inverse(
# m: int,
# Nt: int,
# Nf: int,
# mult_f: int,
# phif: NDArray[np.floating],
# fft_prefactor2s: NDArray[np.complexfloating],
# res: NDArray[np.complexfloating],
# ) -> None:
# """Helper for unpacking results of frequency domain inverse transform"""
# ND = Nf * Nt
# K = mult_f * Nt
#
# assert 0 <= m <= Nf
# assert Nt % 2 == 0
# assert Nf % 2 == 0
#
# assert Nf > 0
# assert Nt > 0
# assert mult_f > 0
#
# f_size = int(ND//2 + 1)
# half_K = int(K //2)
# half_Nt = int(Nt // 2)
#
# assert res.shape == (f_size,)
# assert fft_prefactor2s.shape == (K,)
# assert phif.shape == (half_K+1,)
#
# i_midpoint = m * half_Nt
#
# if m in (0, Nf):
# for i_ind in range(half_K):
# i = abs(i_midpoint - i_ind)
# ind3 = (2 * i) % K
# #for r in range(mult_f):
# res[i] += fft_prefactor2s[ind3] * phif[i_ind]
# if m == Nf:
# i_ind = half_K
# i = abs(i_midpoint - i_ind)
# ind3 = 0 * mult_f
# #for r in range(mult_f):
# res[i] += fft_prefactor2s[ind3] * phif[i_ind]
# else:
# ind31 = i_midpoint % K
# ind32 = i_midpoint % K
# for i_ind in range(half_K):
# i1 = i_midpoint - i_ind
# i2 = i_midpoint + i_ind
#
# if i1 >= 0:
# #for r in range(mult_f):
# res[i1] += fft_prefactor2s[ind31] * phif[i_ind]
# if i2 < f_size:
# #for r in range(mult_f):
# res[i2] += fft_prefactor2s[ind32] * phif[i_ind]
# ind31 -= 1
# ind32 += 1
# if ind31 < 0:
# ind31 = K - 1
# if ind32 == K:
# ind32 = 0
# #for r in range(mult_f):
# res[i_midpoint] = fft_prefactor2s[(mult_f * i_midpoint) % K] * phif[0] #* np.sqrt(mult_f)


@njit()
def unpack_wave_inverse(
Expand All @@ -83,7 +17,18 @@ def unpack_wave_inverse(
fft_prefactor2s: NDArray[np.complexfloating],
res: NDArray[np.complexfloating],
) -> None:
"""Helper for unpacking results of frequency domain inverse transform"""
"""Helper for unpacking results of frequency domain inverse transform.

fft_prefactor2s is the length-Nt fft of the packed pixel coefficients; window
tap delta contributes fft_prefactor2s[jj mod Nt] * phif[|delta|] at spectrum
bin jj = m*Nt/2 + delta, over the open support delta in (-K/2, K/2).

For interior bands, contributions at bins outside the rfft range [0, ND/2]
are the mirror lobe of the real wavelet and fold back conjugated; the
self-conjugate bins 0 and ND/2 receive both lobes at once (2 Re). The
self-conjugate bands m = 0 and m = Nf scatter only their own half-plane,
which is already the complete synthesis for reflection-symmetric atoms.
"""
ND = Nf * Nt
K = mult_f * Nt

Expand All @@ -93,93 +38,67 @@ def unpack_wave_inverse(

assert Nf > 0
assert Nt > 0
assert mult_f > 0
assert 0 < mult_f <= Nf

f_size = int(ND // 2 + 1)
half_K = int(K // 2)
half_Nt = int(Nt // 2)
half_ND = int(ND // 2)

assert res.shape == (f_size,)
assert fft_prefactor2s.shape == (K,)
assert fft_prefactor2s.shape == (Nt,)
assert phif.shape == (half_K + 1,)

i_midpoint = m * half_Nt

if m in (0, Nf):
# taps at delta = 0..K/2-1 for m = 0, delta = 1-K/2..0 for m = Nf; both
# stay inside [0, ND/2] since mult_f <= Nf
for i_ind in range(half_K):
i = abs(i_midpoint - i_ind)
ind3 = (2 * i) % K
# for r in range(mult_f):
res[i] += fft_prefactor2s[ind3] * phif[i_ind]
if m == Nf:
i_ind = half_K
i = abs(i_midpoint - i_ind)
ind3 = 0 * mult_f
# for r in range(mult_f):
ind3 = (2 * i) % Nt
res[i] += fft_prefactor2s[ind3] * phif[i_ind]
else:
ind31 = (i_midpoint - 1) % K
for i_ind in range(1, half_K):
i1 = i_midpoint - i_ind
val = fft_prefactor2s[i1 % Nt] * phif[i_ind]
if i1 > 0:
res[i1] += val
elif i1 == 0:
# self-conjugate bin: both lobes land here
res[0] += 2.0 * val.real
else:
# mirror lobe of the real atom folds back conjugated
res[-i1] += np.conj(val)

if i1 >= 0: # and i1%half_Nt >0:
# for r in range(mult_f):
res[i1] += fft_prefactor2s[ind31] * phif[i_ind]
ind31 -= 1
if ind31 < 0:
ind31 = K - 1

ind32 = (i_midpoint + 1) % K
for i_ind in range(1, half_K):
i2 = i_midpoint + i_ind
if i2 < f_size: # and i2%half_Nt >0:
# for r in range(mult_f):
res[i2] += fft_prefactor2s[ind32] * phif[i_ind]
ind32 += 1
if ind32 == K:
ind32 = 0
# for r in range(mult_f):
# TODO why is only this one conjugated? maybe some other indices should be but the effect is smaller?
if mult_f % 2 or m % 2 == 0:
res[i_midpoint] += fft_prefactor2s[(mult_f * i_midpoint) % K] * phif[0] # * np.sqrt(mult_f)
else:
res[i_midpoint] += -np.conjugate(fft_prefactor2s[(mult_f * i_midpoint) % K]) * phif[0] # * np.sqrt(mult_f)


# @njit()
# def unpack_wave_inverse(m, Nt, Nf, phif, fft_prefactor2s, res):
# """helper for unpacking results of frequency domain inverse transform"""
# ND = Nt*Nf
# i_min2 = min(max(Nt//2*(m-1), 0), ND//2+1)
# i_max2 = min(max(Nt//2*(m+1), 0), ND//2+1)
# for i in range(i_min2, i_max2):
# i_ind = np.abs(i-Nt//2*m)
# if i_ind > Nt//2:
# continue
# if m == 0:
# res[i] += fft_prefactor2s[(2*i) % Nt]*phif[i_ind]
# elif m == Nf:
# res[i] += fft_prefactor2s[(2*i) % Nt]*phif[i_ind]
# else:
# res[i] += fft_prefactor2s[i % Nt]*phif[i_ind]
val = fft_prefactor2s[i2 % Nt] * phif[i_ind]
if i2 < half_ND:
res[i2] += val
elif i2 == half_ND:
res[half_ND] += 2.0 * val.real
else:
res[ND - i2] += np.conj(val)

# center tap; i_midpoint is strictly inside (0, ND/2) for interior m
res[i_midpoint] += fft_prefactor2s[i_midpoint % Nt] * phif[0]


@njit()
def pack_wave_inverse(
m: int,
Nt: int,
Nf: int,
mult_f: int,
prefactor2s: NDArray[np.complexfloating],
wave_in: NDArray[np.floating],
) -> None:
"""Helper for fast frequency domain inverse transform to prepare for fourier transform"""
if m == 0:
for n in range(Nt):
prefactor2s[n * mult_f] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt, 0]
prefactor2s[n] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt, 0]
elif m == Nf:
for n in range(Nt):
prefactor2s[n * mult_f] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt + 1, 0]
prefactor2s[n] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt + 1, 0]
else:
for n in range(Nt):
val = float(wave_in[n, m])
Expand All @@ -188,10 +107,9 @@ def pack_wave_inverse(
else:
mult2 = 1

prefactor2s[n * mult_f] = mult2 * val
prefactor2s[n] = mult2 * val


# @njit()
def inverse_wavelet_freq_helper_fast(
wave_in: NDArray[np.floating],
phif: NDArray[np.floating],
Expand All @@ -202,16 +120,11 @@ def inverse_wavelet_freq_helper_fast(
"""Jit compatible loop for inverse_wavelet_freq"""
ND = Nf * Nt

prefactor2s = np.zeros(Nt * mult_f, dtype=complex)
prefactor2s = np.zeros(Nt, dtype=complex)
res = np.zeros(ND // 2 + 1, dtype=complex)

phif_alt = phif[: Nt // 2 + 1]
prefactor2s_alt = np.zeros(Nt, dtype=complex)
res_alt = np.zeros(ND // 2 + 1, dtype=complex)

for m in range(Nf + 1):
prefactor2s[:] = 0.0
pack_wave_inverse(m, Nt, Nf, mult_f, prefactor2s, wave_in)
pack_wave_inverse(m, Nt, Nf, prefactor2s, wave_in)
fft_prefactor2s = fft.fft(prefactor2s)
unpack_wave_inverse(m, Nt, Nf, mult_f, phif, fft_prefactor2s, res)

Expand Down
Loading