Правильна ли моя реализация JAX непрерывного вейвлет-преобразования? ⇐ Python

[Anonymous]

Я хотел бы реализовать непрерывное вейвлет-преобразование (CWT) с помощью JAX. Согласно ChatGPT, на практике он вычисляется путем выполнения дискретной свертки с выборочной вейвлет-функцией в разных масштабах. Я реализовал это вместе со скалограммой следующим образом:

Код: Выделить всё

import jax.numpy as jnp
from jax import image, jit, vmap

HEIGHT = 224
WIDTH = 224

@jit
def morlet(t, sigma):
"""Compute the Morlet wavelet."""
return (
(1 + jnp.exp(-(sigma**2)) - 2 * jnp.exp(-0.75 * sigma**2)) ** -0.5
* jnp.pi**-0.25
* jnp.exp(-0.5 * t**2)
* (jnp.exp(1j * sigma * t) - jnp.exp(-0.5 * sigma**2))
)

@jit
def cwt(signal):
"""Perform continuous wavelet transform."""
n = signal.shape[0]
t = jnp.arange(-n // 2, n // 2)
signal_fft = jnp.fft.fft(signal)

def transform(scale):
wavelet = morlet(t / scale, 16.0) / jnp.sqrt(scale)
wavelet_fft = jnp.fft.fft(wavelet)
return jnp.fft.ifft(signal_fft * wavelet_fft)

return vmap(transform)(jnp.logspace(1, 8, num=HEIGHT, base=jnp.e))

def lsa_resize_rescale(transformed):
"""Compute the log-squared-absolute values, resize and rescale."""
lsa = jnp.log(jnp.abs(transformed) ** 2 + jnp.finfo(jnp.float32).eps)
resized = image.resize(lsa, (HEIGHT, WIDTH), "lanczos5")
vmin, vmax = resized.min(), resized.max()
return (resized - vmin) / (vmax - vmin)

@jit
def make_cwt_scalogram(signal):
"""Make a CWT scalogram."""
return lsa_resize_rescale(cwt(signal))

В качестве тестового сигнала я использовал cos(2πx), дискретизированный на частоте 500 Гц в течение 10 секунд:

Код: Выделить всё

import numpy as np

grid = np.linspace(0, 10, num=5000)
signal = np.cos(2 * np.pi * grid)

Я визуализировал скалограмму CWT, используя свою реализацию CWT:

Код: Выделить всё

fig, ax = plt.subplots(figsize=(4, 4))
ax.imshow(make_cwt_scalogram(signal), cmap='gray')
ax.set_axis_off();

Для сравнения я визуализировал скалограмму CWT с помощью ssqueezepy (рекомендуется PyWavelets) реализации CWT:

Код: Выделить всё

import ssqueezepy

fig, ax = plt.subplots(figsize=(4, 4))
ax.imshow(
lsa_resize_rescale(
ssqueezepy.cwt(signal, wavelet='morlet', padtype=None)[0]
),
cmap='gray'
)
ax.set_axis_off();

Даже не учитывая тот факт, что выбор масштабов разный (

Код: Выделить всё

ssqueezespy

использует «логарифмический» масштаб), два изображения по-прежнему выглядят по-разному — в моей реализации в середине кажется конус. Почему? Правильна ли моя реализация?
Обновление. По предложению @jakevdp я попробовал использовать t = jnp.fft.ifftshift(jnp.arange(-n // 2, n // 2)) вместо t = jnp.arange(-n // 2, n // 2). Это удалило конус посередине, но создало эффект расхождения по бокам.



Подробнее здесь: https://stackoverflow.com/questions/798 ... rm-correct