Update README.md
Signed-off-by: David Rotermund <54365609+davrot@users.noreply.github.com>
This commit is contained in:
parent
400bb20720
commit
d1eb2167b5
1 changed files with 130 additions and 7 deletions
|
@ -145,6 +145,21 @@ import numpy as np
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
import pywt
|
import pywt
|
||||||
|
|
||||||
|
# Calculate the wavelet scales we requested
|
||||||
|
def calculate_wavelet_scale(
|
||||||
|
number_of_frequences: int,
|
||||||
|
frequency_range_min: float,
|
||||||
|
frequency_range_max: float,
|
||||||
|
dt: float,
|
||||||
|
) -> np.ndarray:
|
||||||
|
s_spacing: np.ndarray = (1.0 / (number_of_frequences - 1)) * np.log2(
|
||||||
|
frequency_range_max / frequency_range_min
|
||||||
|
)
|
||||||
|
scale: np.ndarray = np.power(2, np.arange(0, number_of_frequences) * s_spacing)
|
||||||
|
frequency_axis_request: np.ndarray = frequency_range_min * np.flip(scale)
|
||||||
|
return 1.0 / (frequency_axis_request * dt)
|
||||||
|
|
||||||
|
|
||||||
f_test: float = 50 # Hz
|
f_test: float = 50 # Hz
|
||||||
number_of_test_samples: int = 1000
|
number_of_test_samples: int = 1000
|
||||||
|
|
||||||
|
@ -160,15 +175,16 @@ dt: float = 1.0 / 1000 # sec
|
||||||
t_test: np.ndarray = np.arange(0, number_of_test_samples) * dt
|
t_test: np.ndarray = np.arange(0, number_of_test_samples) * dt
|
||||||
test_data: np.ndarray = np.sin(2 * np.pi * f_test * t_test)
|
test_data: np.ndarray = np.sin(2 * np.pi * f_test * t_test)
|
||||||
|
|
||||||
# Calculate the wavelet scales we requested
|
wave_scales = calculate_wavelet_scale(
|
||||||
s_spacing: np.ndarray = (1.0 / (number_of_frequences - 1)) * np.log2(
|
number_of_frequences=number_of_frequences,
|
||||||
frequency_range_max / frequency_range_min
|
frequency_range_min=frequency_range_min,
|
||||||
|
frequency_range_max=frequency_range_max,
|
||||||
|
dt=dt,
|
||||||
)
|
)
|
||||||
scale: np.ndarray = np.power(2, np.arange(0, number_of_frequences) * s_spacing)
|
|
||||||
frequency_axis_request: np.ndarray = frequency_range_min * np.flip(scale)
|
|
||||||
wave_scales: np.ndarray = 1.0 / (frequency_axis_request * dt)
|
|
||||||
|
|
||||||
complex_spectrum, frequency_axis = pywt.cwt(test_data, wave_scales, mother, dt)
|
complex_spectrum, frequency_axis = pywt.cwt(
|
||||||
|
data=test_data, scales=wave_scales, wavelet=mother, sampling_period=dt
|
||||||
|
)
|
||||||
|
|
||||||
plt.imshow(abs(complex_spectrum) ** 2, cmap="hot", aspect="auto")
|
plt.imshow(abs(complex_spectrum) ** 2, cmap="hot", aspect="auto")
|
||||||
plt.colorbar()
|
plt.colorbar()
|
||||||
|
@ -182,4 +198,111 @@ plt.show()
|
||||||
```
|
```
|
||||||
![figure 5](image5.png)
|
![figure 5](image5.png)
|
||||||
|
|
||||||
|
**Done** ?!?!
|
||||||
|
|
||||||
|
### Fixing the problems -- the axis of the plot
|
||||||
|
|
||||||
|
The axis look horrible! Let us fix that.
|
||||||
|
|
||||||
|
```python
|
||||||
|
import numpy as np
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import pywt
|
||||||
|
|
||||||
|
|
||||||
|
# Calculate the wavelet scales we requested
|
||||||
|
def calculate_wavelet_scale(
|
||||||
|
number_of_frequences: int,
|
||||||
|
frequency_range_min: float,
|
||||||
|
frequency_range_max: float,
|
||||||
|
dt: float,
|
||||||
|
) -> np.ndarray:
|
||||||
|
s_spacing: np.ndarray = (1.0 / (number_of_frequences - 1)) * np.log2(
|
||||||
|
frequency_range_max / frequency_range_min
|
||||||
|
)
|
||||||
|
scale: np.ndarray = np.power(2, np.arange(0, number_of_frequences) * s_spacing)
|
||||||
|
frequency_axis_request: np.ndarray = frequency_range_min * np.flip(scale)
|
||||||
|
return 1.0 / (frequency_axis_request * dt)
|
||||||
|
|
||||||
|
|
||||||
|
def get_y_ticks(
|
||||||
|
reduction_to_ticks: int, frequency_axis: np.ndarray, round: int
|
||||||
|
) -> tuple[np.ndarray, np.ndarray]:
|
||||||
|
output_ticks = np.arange(
|
||||||
|
0,
|
||||||
|
frequency_axis.shape[0],
|
||||||
|
int(np.floor(frequency_axis.shape[0] / reduction_to_ticks)),
|
||||||
|
)
|
||||||
|
if round < 0:
|
||||||
|
output_freq = frequency_axis[output_ticks]
|
||||||
|
else:
|
||||||
|
output_freq = np.round(frequency_axis[output_ticks], round)
|
||||||
|
return output_ticks, output_freq
|
||||||
|
|
||||||
|
|
||||||
|
def get_x_ticks(
|
||||||
|
reduction_to_ticks: int, dt: float, number_of_timesteps: int, round: int
|
||||||
|
) -> tuple[np.ndarray, np.ndarray]:
|
||||||
|
time_axis = dt * np.arange(0, number_of_timesteps)
|
||||||
|
output_ticks = np.arange(
|
||||||
|
0, time_axis.shape[0], int(np.floor(time_axis.shape[0] / reduction_to_ticks))
|
||||||
|
)
|
||||||
|
if round < 0:
|
||||||
|
output_time_axis = time_axis[output_ticks]
|
||||||
|
else:
|
||||||
|
output_time_axis = np.round(time_axis[output_ticks], round)
|
||||||
|
return output_ticks, output_time_axis
|
||||||
|
|
||||||
|
|
||||||
|
f_test: float = 50 # Hz
|
||||||
|
number_of_test_samples: int = 1000
|
||||||
|
|
||||||
|
# The wavelet we want to use
|
||||||
|
mother = pywt.ContinuousWavelet("cmor1.5-1.0")
|
||||||
|
|
||||||
|
# Parameters for the wavelet transform
|
||||||
|
number_of_frequences: int = 25 # frequency bands
|
||||||
|
frequency_range_min: float = 15 # Hz
|
||||||
|
frequency_range_max: float = 200 # Hz
|
||||||
|
dt: float = 1.0 / 1000 # sec
|
||||||
|
|
||||||
|
t_test: np.ndarray = np.arange(0, number_of_test_samples) * dt
|
||||||
|
test_data: np.ndarray = np.sin(2 * np.pi * f_test * t_test)
|
||||||
|
|
||||||
|
wave_scales = calculate_wavelet_scale(
|
||||||
|
number_of_frequences=number_of_frequences,
|
||||||
|
frequency_range_min=frequency_range_min,
|
||||||
|
frequency_range_max=frequency_range_max,
|
||||||
|
dt=dt,
|
||||||
|
)
|
||||||
|
|
||||||
|
complex_spectrum, frequency_axis = pywt.cwt(
|
||||||
|
data=test_data, scales=wave_scales, wavelet=mother, sampling_period=dt
|
||||||
|
)
|
||||||
|
|
||||||
|
plt.imshow(abs(complex_spectrum) ** 2, cmap="hot", aspect="auto")
|
||||||
|
plt.colorbar()
|
||||||
|
|
||||||
|
y_ticks, y_labels = get_y_ticks(
|
||||||
|
reduction_to_ticks=10, frequency_axis=frequency_axis, round=1
|
||||||
|
)
|
||||||
|
|
||||||
|
x_ticks, x_labels = get_x_ticks(
|
||||||
|
reduction_to_ticks=10, dt=dt, number_of_timesteps=complex_spectrum.shape[1], round=2
|
||||||
|
)
|
||||||
|
|
||||||
|
plt.yticks(y_ticks, y_labels)
|
||||||
|
plt.xticks(x_ticks, x_labels)
|
||||||
|
|
||||||
|
plt.xlabel("Time [sec]")
|
||||||
|
plt.ylabel("Frequency [Hz]")
|
||||||
|
plt.show()
|
||||||
|
```
|
||||||
|
![figure 6](image6.png)
|
||||||
|
|
||||||
|
This looks already better...
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue