Skip to content
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.

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

ENH: Add STFT function to Function class #620

Merged
merged 23 commits into from
Aug 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
6013462
Add STFT function to Function class
AdvaitChandorkar07 Jun 12, 2024
027e7f0
Added feature: Short-Time Fourier Transform function
AdvaitChandorkar07 Jun 12, 2024
98fbb6d
Added feature: Short-Time Fourier Transform function
AdvaitChandorkar07 Jun 12, 2024
f61f681
Merge branch 'develop' into feature/stft-function
AdvaitChandorkar07 Jul 12, 2024
4f682bc
"Variable name changes in stft"
AdvaitChandorkar07 Jul 13, 2024
a319a54
"Variable and function name formatting"
AdvaitChandorkar07 Jul 13, 2024
e47ec12
Merge branch 'develop' into feature/stft-function
AdvaitChandorkar07 Jul 13, 2024
257238c
Merge branch 'feature/stft-function' of github.com:AdvaitChandorkar07…
AdvaitChandorkar07 Jul 13, 2024
cd35030
"Better Example"
AdvaitChandorkar07 Jul 14, 2024
511f000
Add STFT function to Function class
AdvaitChandorkar07 Jun 12, 2024
a2d373f
Added feature: Short-Time Fourier Transform function
AdvaitChandorkar07 Jun 12, 2024
5f60c54
Added feature: Short-Time Fourier Transform function
AdvaitChandorkar07 Jun 12, 2024
236552a
"Variable name changes in stft"
AdvaitChandorkar07 Jul 13, 2024
f141fe2
"Variable and function name formatting"
AdvaitChandorkar07 Jul 13, 2024
efc2583
"Better Example"
AdvaitChandorkar07 Jul 14, 2024
c4b7f62
Merge branch 'feature/stft-function' of github.com:AdvaitChandorkar07…
AdvaitChandorkar07 Jul 18, 2024
03faf61
Fixed the doctest
AdvaitChandorkar07 Jul 18, 2024
3c7f548
"Spectrogram example"
AdvaitChandorkar07 Jul 25, 2024
c5d2b96
Merge branch 'develop' into feature/stft-function
AdvaitChandorkar07 Jul 25, 2024
fee5e71
Merge branch 'develop' into feature/stft-function
Gui-FernandesBR Aug 4, 2024
403b5f6
Merge branch 'develop' into feature/stft-function
Gui-FernandesBR Aug 13, 2024
5f23551
small fixes to STFT function
Gui-FernandesBR Aug 18, 2024
526ab11
Merge branch 'develop' into feature/stft-function
Gui-FernandesBR Aug 18, 2024
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
3 changes: 3 additions & 0 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@
"copybutton",
"cstride",
"csys",
"cumsum",
"datapoints",
"datetime",
"dcsys",
Expand Down Expand Up @@ -149,6 +150,7 @@
"IGRA",
"imageio",
"imread",
"imshow",
"intc",
"interp",
"Interquartile",
Expand Down Expand Up @@ -258,6 +260,7 @@
"SRTM",
"SRTMGL",
"Stano",
"STFT",
"subintervals",
"suptitle",
"ticklabel",
Expand Down
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,8 @@ Attention: The newest changes should be on top -->

### Added

- ENH: Rocket Axis Definition [#635](https://github.com/RocketPy-Team/RocketPy/pull/635)
- ENH: Add STFT function to Function class [#620](https://github.com/RocketPy-Team/RocketPy/pull/620)
- ENH: Rocket Axis Definition [#635](https://github.com/RocketPy-Team/RocketPy/pull/635)

### Changed

Expand Down
136 changes: 136 additions & 0 deletions rocketpy/mathutils/function.py
Original file line number Diff line number Diff line change
Expand Up @@ -1007,6 +1007,142 @@ def to_frequency_domain(self, lower, upper, sampling_frequency, remove_dc=True):
extrapolation="zero",
)

def short_time_fft(
self,
lower,
upper,
sampling_frequency,
window_size,
step_size,
remove_dc=True,
only_positive=True,
):
r"""
Performs the Short-Time Fourier Transform (STFT) of the Function and
returns the result. The STFT is computed by applying the Fourier
transform to overlapping windows of the Function.

Parameters
----------
lower : float
Lower bound of the time range.
upper : float
Upper bound of the time range.
sampling_frequency : float
Sampling frequency at which to perform the Fourier transform.
window_size : float
Size of the window for the STFT, in seconds.
step_size : float
Step size for the window, in seconds.
remove_dc : bool, optional
If True, the DC component is removed from each window before
computing the Fourier transform.
only_positive: bool, optional
If True, only the positive frequencies are returned.

Returns
-------
list[Function]
A list of Functions, each representing the STFT of a window.

Examples
--------
Gui-FernandesBR marked this conversation as resolved.
Show resolved Hide resolved

>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> from rocketpy import Function

Generate a signal with varying frequency:

>>> T_x, N = 1 / 20 , 1000 # 20 Hz sampling rate for 50 s signal
>>> t_x = np.arange(N) * T_x # time indexes for signal
>>> f_i = 1 * np.arctan((t_x - t_x[N // 2]) / 2) + 5 # varying frequency
>>> signal = np.sin(2 * np.pi * np.cumsum(f_i) * T_x) # the signal

Create the Function object and perform the STFT:

>>> time_domain = Function(np.array([t_x, signal]).T)
>>> stft_result = time_domain.short_time_fft(
... lower=0,
... upper=50,
... sampling_frequency=95,
... window_size=2,
... step_size=0.5,
... )

Plot the spectrogram:

>>> Sx = np.abs([window[:, 1] for window in stft_result])
>>> t_lo, t_hi = t_x[0], t_x[-1]
>>> fig1, ax1 = plt.subplots(figsize=(10, 6))
>>> im1 = ax1.imshow(
... Sx.T,
... origin='lower',
... aspect='auto',
... extent=[t_lo, t_hi, 0, 50],
... cmap='viridis'
... )
>>> _ = ax1.set_title(rf"STFT (2$\,s$ Gaussian window, $\sigma_t=0.4\,$s)")
>>> _ = ax1.set(
... xlabel=f"Time $t$ in seconds",
... ylabel=f"Freq. $f$ in Hz)",
... xlim=(t_lo, t_hi)
... )
>>> _ = ax1.plot(t_x, f_i, 'r--', alpha=.5, label='$f_i(t)$')
>>> _ = fig1.colorbar(im1, label="Magnitude $|S_x(t, f)|$")
>>> # Shade areas where window slices stick out to the side
>>> for t0_, t1_ in [(t_lo, 1), (49, t_hi)]:
... _ = ax1.axvspan(t0_, t1_, color='w', linewidth=0, alpha=.2)
>>> # Mark signal borders with vertical line
>>> for t_ in [t_lo, t_hi]:
... _ = ax1.axvline(t_, color='y', linestyle='--', alpha=0.5)
>>> # Add legend and finalize plot
>>> _ = ax1.legend()
>>> fig1.tight_layout()
>>> # plt.show() # uncomment to show the plot

References
----------
Example adapted from the SciPy documentation:
https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.ShortTimeFFT.html
"""
# Get the time domain data
sampling_time_step = 1.0 / sampling_frequency
sampling_range = np.arange(lower, upper, sampling_time_step)
sampled_points = self(sampling_range)
samples_per_window = int(window_size * sampling_frequency)
samples_skipped_per_step = int(step_size * sampling_frequency)
stft_results = []

max_start = len(sampled_points) - samples_per_window + 1

for start in range(0, max_start, samples_skipped_per_step):
windowed_samples = sampled_points[start : start + samples_per_window]
if remove_dc:
windowed_samples -= np.mean(windowed_samples)
Gui-FernandesBR marked this conversation as resolved.
Show resolved Hide resolved
fourier_amplitude = np.abs(
np.fft.fft(windowed_samples) / (samples_per_window / 2)
)
fourier_frequencies = np.fft.fftfreq(samples_per_window, sampling_time_step)

# Filter to keep only positive frequencies if specified
if only_positive:
positive_indices = fourier_frequencies > 0
fourier_frequencies = fourier_frequencies[positive_indices]
fourier_amplitude = fourier_amplitude[positive_indices]

stft_results.append(
Function(
source=np.array([fourier_frequencies, fourier_amplitude]).T,
inputs="Frequency (Hz)",
outputs="Amplitude",
interpolation="linear",
extrapolation="zero",
)
)

return stft_results

def low_pass_filter(self, alpha, file_path=None):
"""Implements a low pass filter with a moving average filter. This does
not mutate the original Function object, but returns a new one with the
Expand Down
Loading