I would like to implement continuous wavelet transform (CWT) using JAX. According to ChatGPT, it is in practice computed by performing a discrete convolution with a sampled wavelet function at different scales. I implemented it together with the scalogram as follows:

import jax.numpy as jnp
from jax import 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))

I used as a test signal cos(2πx) sampled at 500Hz for 10s:

import numpy as np

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

I visualized the CWT scalogram using my implementation of CWT:

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

For comparison, I visualized the CWT scalogram using ssqueezepy (recommended by PyWavelets) implementation of CWT:

import ssqueezepy

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

Even disregarding the fact that the choice of scales is different (ssqueezespy uses "log-piecewise" scale), the two images still look different - using my implementation, there appears to be a cone in the middle. Why? Is my implementation correct?