diff --git a/src/noisepy/seis/datatypes.py b/src/noisepy/seis/datatypes.py index 9e14da7d..30b1c525 100644 --- a/src/noisepy/seis/datatypes.py +++ b/src/noisepy/seis/datatypes.py @@ -110,6 +110,12 @@ class StackMethod(Enum): ALL = "all" +class FreqNorm(Enum): + RMA = "rma" + NO = "no" + PHASE_ONLY = "phase_only" + + class ConfigParameters(BaseModel): model_config = ConfigDict(validate_default=True) @@ -139,7 +145,7 @@ class ConfigParameters(BaseModel): freqmin: float = Field(default=0.05) freqmax: float = Field(default=2.0) freq_norm: str = Field( - default="rma", description="choose between 'rma' for a soft whitenning or 'no' for no whitening" + default=FreqNorm.RMA.value, description="choose between 'rma' for a soft whitenning or 'no' for no whitening" ) # TODO: change "no"for "None", and add "one_bit"as an option # TODO: change time_norm option from "no"to "None" diff --git a/src/noisepy/seis/noise_module.py b/src/noisepy/seis/noise_module.py index eb38be71..7fb0a252 100644 --- a/src/noisepy/seis/noise_module.py +++ b/src/noisepy/seis/noise_module.py @@ -25,7 +25,7 @@ from scipy.fftpack import next_fast_len from scipy.signal import hilbert -from .datatypes import ChannelData, ConfigParameters, StackMethod +from .datatypes import ChannelData, ConfigParameters, StackMethod, FreqNorm logger = logging.getLogger(__name__) """ @@ -559,7 +559,7 @@ def noise_processing(fft_para: ConfigParameters, dataS): white = dataS # -----to whiten or not------ - if fft_para.freq_norm != "no": + if fft_para.freq_norm != FreqNorm.NO: source_white = whiten(white, fft_para) # whiten and return FFT else: Nfft = int(next_fast_len(int(dataS.shape[1]))) diff --git a/tests/test_whiten.py b/tests/test_whiten.py index 622e29e6..5ffa429d 100644 --- a/tests/test_whiten.py +++ b/tests/test_whiten.py @@ -5,6 +5,8 @@ from noisepy.seis.correlate import ConfigParameters from noisepy.seis.noise_module import moving_ave, whiten +from noisepy.seis.datatypes import FreqNorm + def whiten_original(data, fft_para: ConfigParameters): @@ -56,9 +58,9 @@ def whiten_original(data, fft_para: ConfigParameters): 1j * np.angle(FFTRawSign[:, low:left]) ) # Pass band: - if fft_para.freq_norm == "phase_only": + if fft_para.freq_norm == FreqNorm.PHASE_ONLY: FFTRawSign[:, left:right] = np.exp(1j * np.angle(FFTRawSign[:, left:right])) - elif fft_para.freq_norm == "rma": + elif fft_para.freq_norm == FreqNorm.RMA: for ii in range(data.shape[0]): tave = moving_ave(np.abs(FFTRawSign[ii, left:right]), fft_para.smooth_N) FFTRawSign[ii, left:right] = FFTRawSign[ii, left:right] / tave @@ -76,9 +78,9 @@ def whiten_original(data, fft_para: ConfigParameters): 1j * np.angle(FFTRawSign[low:left]) ) # Pass band: - if fft_para.freq_norm == "phase_only": + if fft_para.freq_norm == FreqNorm.PHASE_ONLY: FFTRawSign[left:right] = np.exp(1j * np.angle(FFTRawSign[left:right])) - elif fft_para.freq_norm == "rma": + elif fft_para.freq_norm == FreqNorm.RMA: tave = moving_ave(np.abs(FFTRawSign[left:right]), fft_para.smooth_N) FFTRawSign[left:right] = FFTRawSign[left:right] / tave # Right tapering: @@ -103,7 +105,7 @@ def whiten_original(data, fft_para: ConfigParameters): fft_para.freqmin = 0.01 fft_para.freqmax = 0.2 fft_para.smooth_N = 1 -fft_para.freq_norm = "phase_only" +fft_para.freq_norm = FreqNorm.PHASE_ONLY def whiten1d(): diff --git a/tutorials/get_started.ipynb b/tutorials/get_started.ipynb index 0357a214..609b5945 100644 --- a/tutorials/get_started.ipynb +++ b/tutorials/get_started.ipynb @@ -60,7 +60,7 @@ "source": [ "from noisepy.seis import download, cross_correlate, stack, plotting_modules, __version__\n", "from noisepy.seis.asdfstore import ASDFRawDataStore, ASDFCCStore, ASDFStackStore\n", - "from noisepy.seis.datatypes import ConfigParameters\n", + "from noisepy.seis.datatypes import ConfigParameters, FreqNorm\n", "from dateutil.parser import isoparse\n", "import os\n", "print(f\"Using NoisePy version {__version__}\")\n", @@ -119,7 +119,7 @@ "config.max_over_std = 10 # threshold to remove window of bad signals: set it to 10*9 if prefer not to remove them\n", "\n", "# TEMPORAL and SPECTRAL NORMALISATION\n", - "config.freq_norm= \"rma\" # choose between \"rma\" for a soft whitenning or \"no\" for no whitening. Pure whitening is not implemented correctly at this point.\n", + "config.freq_norm= FreqNorm.RMA # choose between \"rma\" for a soft whitenning or \"no\" for no whitening. Pure whitening is not implemented correctly at this point.\n", "config.smoothspect_N = 10 # moving window length to smooth spectrum amplitude (points)\n", " # here, choose smoothspect_N for the case of a strict whitening (e.g., phase_only)\n", "\n", @@ -248,7 +248,7 @@ }, "outputs": [], "source": [ - "config.freq_norm = \"rma\"\n", + "config.freq_norm = FreqNorm.RMA\n", "raw_store = ASDFRawDataStore(raw_data_path) # Store for reading raw data\n", "cc_store = ASDFCCStore(cc_data_path) # Store for writing CC data\n", "\n", diff --git a/tutorials/noisepy_scedc_tutorial.ipynb b/tutorials/noisepy_scedc_tutorial.ipynb index e1a93080..7b2be15b 100644 --- a/tutorials/noisepy_scedc_tutorial.ipynb +++ b/tutorials/noisepy_scedc_tutorial.ipynb @@ -67,7 +67,7 @@ "from noisepy.seis import cross_correlate, stack, plotting_modules, __version__ # noisepy core functions\n", "from noisepy.seis.asdfstore import ASDFCCStore, ASDFStackStore # Object to store ASDF data within noisepy\n", "from noisepy.seis.scedc_s3store import SCEDCS3DataStore, channel_filter # Object to query SCEDC data from on S3\n", - "from noisepy.seis.datatypes import ConfigParameters, StackMethod # Main configuration object\n", + "from noisepy.seis.datatypes import ConfigParameters, FreqNorm, StackMethod # Main configuration object\n", "from noisepy.seis.channelcatalog import XMLStationChannelCatalog # Required stationXML handling object\n", "import os\n", "from datetime import datetime\n", @@ -201,7 +201,7 @@ "config.max_over_std = 10 # threshold to remove window of bad signals: set it to 10*9 if prefer not to remove them\n", "\n", "################### SPECTRAL NORMALIZATION ############\n", - "config.freq_norm= \"rma\" # choose between \"rma\" for a soft whitening or \"no\" for no whitening. Pure whitening is not implemented correctly at this point.\n", + "config.freq_norm= FreqNorm.RMA # choose between \"rma\" for a soft whitening or \"no\" for no whitening. Pure whitening is not implemented correctly at this point.\n", "config.smoothspect_N = 10 # moving window length to smooth spectrum amplitude (points)\n", " # here, choose smoothspect_N for the case of a strict whitening (e.g., phase_only)\n", "\n",