refactor: change freq_norm from string to enum

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"

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])))

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():

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",

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",