Coherence Sorting Prototype

Wrote and tested a simple coherence sorting.  Mostly tested on IAK34
dataset but did not implement -- i.e. commented it out in
transfer_function_helpers.py for this commit.

If this commit passes I will test moving mth5, mt_metadata to master branches

[Issue: #119]

aurora/pipelines/transfer_function_helpers.py

@@ -4,6 +4,7 @@
 
 
 """
+import numpy as np
 from aurora.time_series.frequency_band_helpers import extract_band
 
 # from aurora.time_series.xarray_helpers import cast_3d_stft_to_2d_observations
@@ -195,6 +196,7 @@ def process_transfer_functions(
     -------
 
     """
+    # PUT COHERENCE SORTING HERE IF WIDE BAND?
     regression_class = get_regression_class(config)
     for band in transfer_function_obj.frequency_bands.bands():
         iter_control = set_up_iter_control(config)
@@ -210,17 +212,36 @@ def process_transfer_functions(
             RR = RR.stack(observation=("frequency", "time"))
 
         W = effective_degrees_of_freedom_weights(X, RR, edf_obj=None)
+        W[W == 0] = np.nan  # use this to drop values in the handle_nan
         # apply weights
         X *= W
         Y *= W
         if RR is not None:
             RR *= W
+        X = X.dropna(dim="observation")
+        Y = Y.dropna(dim="observation")
+        if RR is not None:
+            RR = RR.dropna(dim="observation")
+
+        # < INSERT COHERENCE SORTING HERE>
+        # coh_type = "local"
+        # if config.decimation_level_id == 0:
+        #     from aurora.transfer_function.weights.coherence_weights import compute_coherence_weights
+        #     X, Y, RR = compute_coherence_weights(X,Y,RR, coh_type=coh_type)
+        # </ INSERT COHERENCE SORTING HERE>
 
         if config.estimate_per_channel:
             for ch in config.output_channels:
                 Y_ch = Y[ch].to_dataset()  # keep as a dataset, maybe not needed
 
                 X_, Y_, RR_ = handle_nan(X, Y_ch, RR, drop_dim="observation")
+
+                # W = effective_degrees_of_freedom_weights(X_, RR_, edf_obj=None)
+                # X_ *= W
+                # Y_ *= W
+                # if RR is not None:
+                #     RR_ *= W
+
                 regression_estimator = regression_class(
                     X=X_, Y=Y_, Z=RR_, iter_control=iter_control
                 )

aurora/test_utils/dataset_definitions.py

@@ -95,6 +95,74 @@ def make_iak34_test_00_config():
     return test_data_set
 
 
+def make_iak34_test_01_config():
+    test_data_set = FDSNDatasetConfig()
+    test_data_set.dataset_id = "iak34_test_01_long_ss"
+    test_data_set.network = "EM"
+    test_data_set.station = "IAK34"
+    # <ORIGINAL>
+    # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z")
+    # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z")
+    # </ORIGINAL>
+    test_data_set.starttime = UTCDateTime("2013-04-26T00:00:00.000000Z")
+    test_data_set.endtime = UTCDateTime("2013-05-11T00:00:00.000000Z")
+    test_data_set.channel_codes = None
+    test_data_set.description = "earthscope example dataset IAK34"
+    test_data_set.components_list = ["hx", "hy", "ex", "ey"]
+    return test_data_set
+
+
+def make_iak34_test_02_config():
+    test_data_set = FDSNDatasetConfig()
+    test_data_set.dataset_id = "iak34_test_02_long_rr"
+    test_data_set.network = "EM"
+    test_data_set.station = "IAK34,NEK33"
+    # <ORIGINAL>
+    # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z")
+    # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z")
+    # </ORIGINAL>
+    test_data_set.starttime = UTCDateTime("2013-04-26T00:00:00.000000Z")
+    test_data_set.endtime = UTCDateTime("2013-05-10T00:00:00.000000Z")
+    test_data_set.channel_codes = None
+    test_data_set.description = "earthscope example dataset IAK34"
+    test_data_set.components_list = ["hx", "hy", "ex", "ey"]
+    return test_data_set
+
+
+def make_iak34_test_03_config():
+    test_data_set = FDSNDatasetConfig()
+    test_data_set.dataset_id = "iak34_test_03_long_rr"
+    test_data_set.network = "EM"
+    test_data_set.station = "IAK34,NEK33"
+    # <ORIGINAL>
+    # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z")
+    # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z")
+    # </ORIGINAL>
+    test_data_set.starttime = UTCDateTime("2013-05-15T00:00:00.000000Z")
+    test_data_set.endtime = UTCDateTime("2013-05-26T00:00:00.000000Z")
+    test_data_set.channel_codes = None
+    test_data_set.description = "earthscope example dataset IAK34"
+    test_data_set.components_list = ["hx", "hy", "ex", "ey"]
+    return test_data_set
+
+
+def make_iak34_test_04_config():
+    test_data_set = FDSNDatasetConfig()
+    test_data_set.dataset_id = "iak34_test_04_rr"
+    test_data_set.network = "EM"
+    test_data_set.station = "IAK34,NEN34"  # NEK33
+    # <ORIGINAL>
+    # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z")
+    # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z")
+    # </ORIGINAL>
+    test_data_set.starttime = UTCDateTime("2013-04-28T00:00:00.000000Z")
+    test_data_set.endtime = UTCDateTime("2013-04-29T00:00:00.000000Z")
+    test_data_set.channel_codes = None
+    test_data_set.description = "earthscope example dataset IAK34"
+    test_data_set.components_list = ["hx", "hy", "ex", "ey"]
+    return test_data_set
+
+
 # def make_iak34_nen34_test_00_config():
 #     test_data_set = FDSNDatasetConfig()
 #     test_data_set.dataset_id = "iak34_nen34_test_00"
@@ -143,6 +211,14 @@ def make_test_configs():
     # <IAK34SS>
     test_data_set = make_iak34_test_00_config()
     test_data_set_configs[test_data_set.dataset_id] = test_data_set
+    test_data_set = make_iak34_test_01_config()
+    test_data_set_configs[test_data_set.dataset_id] = test_data_set
+    test_data_set = make_iak34_test_02_config()
+    test_data_set_configs[test_data_set.dataset_id] = test_data_set
+    test_data_set = make_iak34_test_03_config()
+    test_data_set_configs[test_data_set.dataset_id] = test_data_set
+    test_data_set = make_iak34_test_04_config()
+    test_data_set_configs[test_data_set.dataset_id] = test_data_set
     # </IAK34SS>
 
     return test_data_set_configs

aurora/transfer_function/weights/coherence_weights.py

@@ -43,7 +43,7 @@ def coherence_from_fc_series(c_xy, c_xx, c_yy):
     return coh
 
 
-def coherence_weights_v00(x, y, threshold=0.92):  # 975):#0.98
+def coherence_weights_v00(x, y, threshold=0.95):  # 975):#0.98
     """
 
     Parameters
@@ -83,3 +83,70 @@ def coherence_weights_v00(x, y, threshold=0.92):  # 975):#0.98
     keepers = worst_to_best[clip_point:]
     W[keepers] = 1
     return W
+
+
+def compute_coherence_weights(X, Y, RR, coh_type="local"):
+    """
+
+    Parameters
+    ----------
+    X
+    Y
+    RR
+    coh_type: "local" or "remote"
+
+    Returns
+    -------
+
+    """
+    remote_threshold = 0.8
+    local_threshold = 0.95
+
+    X = X.dropna(dim="observation")
+    Y = Y.dropna(dim="observation")
+    if RR is not None:
+        RR = RR.dropna(dim="observation")
+
+    # < INSERT COHERENCE SORTING HERE>    y_type = "remote"
+    null_indices = X["hx"].isnull()  # not robust -- hail mary
+    finite_indices = ~null_indices
+    W = np.zeros(len(X.observation))
+
+    x = X["hx"]  # .dropna(dim="observation").data
+
+    if coh_type == "local":
+        y = Y["ey"]  # .dropna(dim="observation").data
+        threshold = local_threshold
+    elif coh_type == "remote":
+        y = RR["hx"]  # .dropna(dim="observation").data
+        threshold = remote_threshold
+
+    W1 = coherence_weights_v00(x, y, threshold=threshold)
+
+    W[finite_indices] = W1
+
+    W[W == 0] = np.nan
+    X["hx"].data *= W
+    Y["ey"].data *= W
+
+    # x = X["hy"].data
+    # y = Y["ex"].data
+    W = np.zeros(len(finite_indices))
+    x = X["hy"]  # .dropna(dim="observation").data
+    if coh_type == "local":
+        y = Y["ex"]  # .dropna(dim="observation").data
+        threshold = local_threshold
+    elif coh_type == "remote":
+        y = RR["hy"]  # .dropna(dim="observation").data
+        threshold = remote_threshold
+    W2 = coherence_weights_v00(x, y)
+    W[finite_indices] = W2
+    W[W == 0] = np.nan
+    X["hy"].data *= W
+    Y["ex"].data *= W
+    # W = W*W2
+    # X *= W
+    # Y *= W
+    if RR is not None:
+        RR *= W
+    return X, Y, RR

tests/synthetic/test_synthetic_driver.py

@@ -70,15 +70,15 @@ def process_synthetic_rr12():
     test_config = CONFIG_PATH.joinpath("test12rr-RR_test2_run_config.json")
     # test_config = Path("config", "test12rr_run_config.json")
     run_id = STATION_01_CFG["run_id"]
-    process_mth5_run(test_config, run_id, units="MT")
+    process_mth5_run(test_config, run_id, units="MT", show_plot=False)
 
 
 def test_process_mth5():
     # process_synthetic_1_underdetermined()
     # process_synthetic_1_with_nans()
     # create_run_config_for_test_case("test1")
-    # process_synthetic_1()
-    # process_synthetic_2()
+    process_synthetic_1()
+    process_synthetic_2()
     process_synthetic_rr12()