diff --git a/aurora/pipelines/transfer_function_helpers.py b/aurora/pipelines/transfer_function_helpers.py index 5d4d37ee..b870fbc3 100644 --- a/aurora/pipelines/transfer_function_helpers.py +++ b/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) + # 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 ) diff --git a/aurora/test_utils/dataset_definitions.py b/aurora/test_utils/dataset_definitions.py index 6fd761c7..cea33e8f 100644 --- a/aurora/test_utils/dataset_definitions.py +++ b/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" + # + # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z") + # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z") + # + 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" + # + # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z") + # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z") + # + 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" + # + # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z") + # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z") + # + 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 + # + # test_data_set.starttime = UTCDateTime("2013-04-25T20:10:08.000000Z") + # test_data_set.endtime = UTCDateTime("2013-05-13T21:18:53.000000Z") + # + 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(): # 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 # return test_data_set_configs diff --git a/aurora/transfer_function/weights/coherence_weights.py b/aurora/transfer_function/weights/coherence_weights.py index 24681640..5037276c 100644 --- a/aurora/transfer_function/weights/coherence_weights.py +++ b/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 diff --git a/tests/synthetic/test_synthetic_driver.py b/tests/synthetic/test_synthetic_driver.py index e7b608cf..baa24f9f 100644 --- a/tests/synthetic/test_synthetic_driver.py +++ b/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()