wavelet transform fixes

[selipot]

This PR addresses issue #256 by reordering the shape/axes of output of [morse_]wavelet_transform as ((additional input axes),time,freq_axis,order_axis) which allows this function to be used with apply_ragged. The example below works. I am not 100% happy with my proposed fix because the wavelet argument of wavelet_transform should be by default (order_axis, freq_axis, time) so that the output is kind of inconsistent with that. An alternative would be a bigger rewrite morse_wavelet and wavelet_transform.

length = 1000
m = 10
x = np.random.random((length))

gamma = 3
beta = 10
radian_frequency = 2 * np.pi * np.array([0.1, 0.2, 0.3])
order = 2
rowsize = np.array([300, 400, 140, 160])
wtx = apply_ragged(morse_wavelet_transform,x,rowsize,gamma,beta,
                    radian_frequency, order = order)

print(np.shape(x))
print(np.shape(wtx))

[selipot]

Note that this requires x to be a unidimensional ragged array. Should apply_ragged include an option to specify the ragged dimension as suggested by @philippemiron ? Is that possible?

[selipot]

In fact, with this modification of wavelet, apply_ragged works with [morse_]wavelet_transform even if the input x is multidimensional. The only condition is for the ragged dimension (i.e. time dimension) to be in position 0. Can this be generalized by modifying apply_ragged to allow for specifying the ragged dimension as suggested by @philippemiron ?

This works:

length = 1000
m = 9
x = np.random.random((length,2*m,m))
print(np.shape(x))

gamma = 3
beta = 10
radian_frequency = 2 * np.pi * np.array([0.1, 0.2, 0.3])
order = 2
time_axis = 0
rowsize = np.array([300, 400, 140, 160])
wtx = morse_wavelet_transform(x, gamma, beta,radian_frequency, 
                                order = order,time_axis=time_axis)
print(np.shape(wtx))

wtx2 = apply_ragged(morse_wavelet_transform,x,rowsize,gamma,beta,
                    radian_frequency, order = order,time_axis=time_axis)
print(np.shape(wtx2))

[philippemiron]

Yes, we could modify apply_ragged to have the axis of the concatenation as an argument.

I don't really know enough about the wavelets, so I can't recommend which order would be better. But, I would look at other code and see what they do.

[philippemiron]

I'm looking at morse_wavelet and don't really understand the dimensions.

In the docstring it says:

Returns
-------

wavelet : np.ndarray
        Time-domain wavelets. ``wavelet`` will be of shape (length,np.size(radian_frequency),order).

but then you have those examples:

>>> wavelet, wavelet_fft = morse_wavelet(1024, 3, 4, np.array([2*np.pi*0.2]))
>>> np.shape(wavelet)
(1, 1, 1024)

which reading the docstring I would have expected the size to be (1024, 1, 1).

[philippemiron]

It is actually defined:

wavelet = np.zeros((length, order, len(radian_frequency)), dtype=np.cdouble)

but then gets reorder before outputting.

wavelet = np.moveaxis(wavelet, [0, 1, 2], [2, 0, 1])

I would suggest we remove the moveaxis and rewrite everything to be consistent.

[selipot]

I'm looking at morse_wavelet and don't really understand the dimensions.

In the docstring it says:

Returns
-------

wavelet : np.ndarray
        Time-domain wavelets. ``wavelet`` will be of shape (length,np.size(radian_frequency),order).

but then you have those examples:

>>> wavelet, wavelet_fft = morse_wavelet(1024, 3, 4, np.array([2*np.pi*0.2]))
>>> np.shape(wavelet)
(1, 1, 1024)

which reading the docstring I would have expected the size to be (1024, 1, 1).

Yes, that's a mistake.

[selipot]

It is actually defined:

wavelet = np.zeros((length, order, len(radian_frequency)), dtype=np.cdouble)

but then gets reorder before outputting.

wavelet = np.moveaxis(wavelet, [0, 1, 2], [2, 0, 1])

I would suggest we remove the moveaxis and rewrite everything to be consistent.

Indeed, that was my quick fix to make it work with apply_ragged but I agree it would be better to start with these shape/axes.

[milancurcic]

I'm not familiar with wavelets but I reviewed the changes and ran tests locally. Let's release v0.22.0 once this is in.

[selipot]

It is actually defined:

wavelet = np.zeros((length, order, len(radian_frequency)), dtype=np.cdouble)

but then gets reorder before outputting.

wavelet = np.moveaxis(wavelet, [0, 1, 2], [2, 0, 1])

I would suggest we remove the moveaxis and rewrite everything to be consistent.

It's looking harder than I thought to rewrite the whole function. I am wondering if @philippemiron you'd be happy if i kept the axis reordering at the end after making sure the docstrings are correct? It is working as it is. The only thing is that the time axis need to come first to be used with apply_ragged.

[milancurcic]

We can always improve it further in a future PR.

[selipot]

We are abandoning the idea of reorganizing the outputs of wavelet yet I am continuing this PR with cosmetic/consistency changes to the docstring. If that is acceptable to you we can merge.