From 8ed5c31e0a43c6e5b7e3fb4cb490922f9fadeae3 Mon Sep 17 00:00:00 2001 From: Sebastian Heimann Date: Wed, 6 Apr 2022 10:38:26 +0200 Subject: [PATCH] get test running --- setup.py | 9 +- src/__init__.py | 2 + test/noisecut.py | 195 ------------------------------------------ test/test_noisecut.py | 14 +-- 4 files changed, 14 insertions(+), 206 deletions(-) delete mode 100644 test/noisecut.py diff --git a/setup.py b/setup.py index d59e635..784f806 100644 --- a/setup.py +++ b/setup.py @@ -11,12 +11,9 @@ setup( packages=[ 'noisecut'], install_requires=[ - 'numpy', + 'numpy<=2.21', + 'matplotlib', 'obspy', 'librosa', ], - extras_require={ - 'obspy_compatibility': ['obspy'], - 'pyrocko_compatibility': ['pyrocko'], - }, -) \ No newline at end of file +) diff --git a/src/__init__.py b/src/__init__.py index 1b1569a..b1c013f 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -6,3 +6,5 @@ # # The NoiseCut Developers, 21st century # ----------------------------------------------------------------------------- + +from .noisecut import * # noqa diff --git a/test/noisecut.py b/test/noisecut.py deleted file mode 100644 index ec50ee9..0000000 --- a/test/noisecut.py +++ /dev/null @@ -1,195 +0,0 @@ -# ----------------------------------------------------------------------------- -# NoiseCut -# -# This file is part of the NoiseCut library. For licensing information see the -# accompanying file `LICENSE`. -# ----------------------------------------------------------------------------- - -import math -import numpy as np -import librosa -from obspy import Trace -import librosa.display -import matplotlib.pyplot as plt - - - -def _next_pow2(n): - return int(round(2**np.ceil(np.log2(n)))) - - -def _valid_win_length_samples(win_length_samples, win_length, sampling_rate): - if win_length_samples is None and win_length is None: - # fully automatic window length - win_length_samples = _next_pow2(120*sampling_rate) - - elif win_length_samples is None and win_length is not None: - win_length_samples = _next_pow2(win_length*sampling_rate) - - elif win_length_samples is not None and win_length_samples is not None: - raise ValueError( - 'Parameters win_length and win_length_samples are mutually ' - 'exclusive.') - elif win_length_samples is not None and win_length is None: - # check win_length_samples is a power of 2 - win_length_samples = int(win_length_samples) - if win_length_samples != _next_pow2(win_length_samples): - raise ValueError( - 'Parameter win_length_samples must be a power of 2.') - - return win_length_samples - - -def noisecut( - trace, - ret_spectrograms=False, - win_length_samples=None, - win_length=None): - ''' - Reduce noise from all the components of the OBS data using HPS noise - reduction algorithms. - - :param win_length_samples: - Window length in samples. Must be a power of 2. Alternatively it can be - set with `win_length`. - :type win_length_samples: - int - - :param win_length: - Window length [s]. Alternatively it can be set with - `win_length_samples`. - :type win_length: - float - - :returns: - The HPS trace and the spectrograms of the original, noise, and hps - trace as well as an array with the frequencies. - :return_type: - tuple ``(hps_trace, (s_original, s_noise, s_hps, frequencies))`` - ''' - - x = trace.data.astype(float) - - win_length_samples = _valid_win_length_samples( - win_length_samples, win_length, trace.stats.sampling_rate) - - hop_length = win_length_samples // 4 - n_fft = win_length_samples - - # Compute the spectrogram amplitude and phase - S_full, phase = librosa.magphase(librosa.stft( - x, - n_fft=n_fft, - hop_length=hop_length, - win_length=win_length_samples)) - - l1 = math.floor((0.1 * win_length_samples) / trace.stats.sampling_rate) - l2 = math.ceil((1 * win_length_samples) / trace.stats.sampling_rate) - - # We consider the frequency range out of the [0.1-1] Hz for the first step - S_full2 = np.zeros((S_full.shape[0], S_full.shape[1])) - S_full2[l1:l2, :] = S_full[l1:l2, :] - - # We consider the frequency range of [0.1-1] Hz for the second step - S_full1 = np.zeros((S_full.shape[0], S_full.shape[1])) - S_full1[:l1, :] = S_full[:l1, :] - S_full1[l2:, :] = S_full[l2:, :] - - # We'll compare frames using cosine similarity, and aggregate similar - # frames by taking their (per-frequency) median value. - S_filter = librosa.decompose.nn_filter( - S_full1, - aggregate=np.median, - metric='cosine', width=200) - - # The output of the filter shouldn't be greater than the input - S_filter = np.minimum(np.abs(S_full1), np.abs(S_filter)) - margin_i = 1 - power = 2 - - # Once we have the masks, simply multiply them with the input spectrogram - # to separate the components. - mask_i = librosa.util.softmask( - S_filter, - margin_i * (S_full1 - S_filter), - power=power) - - S_background = mask_i * S_full1 - - # In the second step we apply a median filter - D_harmonic, D_percussive = librosa.decompose.hpss( - S_full2, - kernel_size=80, - margin=5) - - S_background = S_background + D_harmonic - - f = S_background * phase - L = x.shape[0] - new = librosa.istft( - f, - hop_length=hop_length, - win_length=win_length_samples, - window='hann', - length=L) - - z = x - new - stats = trace.stats - stats.location = 'NC' - - hps_trace = Trace(data=z, header=stats) - hps_trace.write( 'NoiseCut.mseed', format='MSEED', encoding=5, reclen=4096) - - if ret_spectrograms: - - S_hps = S_full - S_background - - df = trace.stats.sampling_rate/win_length_samples - frequencies = np.arange(S_hps.shape[0]) * df - times = np.arange(S_hps.shape[1]) * hop_length - - return hps_trace, (S_full, S_background, S_hps, frequencies, times) - else: - return hps_trace - - - -def plot_noisecut_spectrograms( - S_full, S_background, S_hps, frequencies, times): - - fig= plt.figure(figsize=(15, 9)) - - axs=fig.add_subplot(311) - pcm=axs.pcolormesh(times, frequencies, librosa.power_to_db(np.abs(S_full)), cmap = 'magma', shading= 'auto') - plt.ylim(0,1) - plt.title('Full spectrogram', fontsize=14) - plt.ylabel('Frequency (Hz)', fontsize=14) - plt.yticks (fontsize= 14) - axs.set_xticks([]) - cbar=fig.colorbar(pcm, ax=axs, pad= 0.01) - cbar.ax.tick_params(labelsize=14) - - axs=fig.add_subplot(312) - pcm=axs.pcolormesh(times, frequencies, librosa.power_to_db(np.abs(S_background)), cmap = 'magma', shading= 'auto') - plt.ylim(0,1) - plt.title('Noise spectrogram', fontsize=14) - plt.ylabel('Frequency (Hz)', fontsize=14) - plt.yticks (fontsize= 14) - axs.set_xticks([]) - cbar=fig.colorbar(pcm, ax=axs, pad= 0.01) - cbar.ax.tick_params(labelsize=14) - - axs=fig.add_subplot(313) - pcm=axs.pcolormesh(times, frequencies, librosa.power_to_db(np.abs(S_hps)), cmap = 'magma', shading= 'auto') - plt.ylim(0,1) - plt.title('Denoised spectrogram', fontsize=14) - plt.ylabel('Frequency (Hz)', fontsize=14) - plt.yticks (fontsize= 14) - cbar=fig.colorbar(pcm, ax=axs, pad= 0.01) - cbar.ax.tick_params(labelsize=14) - labelsx = [0,4,8,12,16,20,24] - plt.xticks(np.arange(0,times[-1],(times[-1]/6)-1), labelsx, fontsize= 14) - - fig.savefig ('NoiseCut spectrograms.png', dpi=100) - plt.show() - plt.close(fig) diff --git a/test/test_noisecut.py b/test/test_noisecut.py index 8f4a978..ac4b4be 100644 --- a/test/test_noisecut.py +++ b/test/test_noisecut.py @@ -1,12 +1,16 @@ +import os import noisecut import obspy -st= obspy.read('D10.DO.HH4..D.2012.080.000000') -hps_trace, spectrograms = noisecut.noisecut (st[0], ret_spectrograms= True) -noisecut.plot_noisecut_spectrograms (*spectrograms) - - +def test_noisecut(): + print(__file__) + fn = os.path.join( + os.path.dirname(__file__), + 'D10.DO.HH4..D.2012.080.000000') + st = obspy.read(fn) + hps_trace, spectrograms = noisecut.noisecut(st[0], ret_spectrograms=True) + noisecut.plot_noisecut_spectrograms(*spectrograms)