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)
+ # 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
)
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()