merge from widescale

simpeg · Sep 15, 2023 · 60bb3b9 · 60bb3b9
1 parent 0a4e991
commit 60bb3b9
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/aurora/pipelines/process_mth5.py b/aurora/pipelines/process_mth5.py
@@ -288,10 +288,12 @@ def station_obj_from_row(row):
                 continue
 
             run_xrds = row["run_dataarray"].to_dataset("channel")
+            print(f"DEBUG Issue 289: TS {row.station_id} {row.run_id} {run_xrds.time.shape} {row.start} {row.end}")
             run_obj = row.mth5_obj.from_reference(row.run_reference)
             stft_obj = make_stft_objects(
                 tfk.config, i_dec_level, run_obj, run_xrds, units, row.station_id
             )
+            print(f"DEBUG Issue 289: FC {row.station_id} {row.run_id} {stft_obj.time.shape} {stft_obj.time.min()} {stft_obj.time.max()}")
             # Pack FCs into h5
             if dec_level_config.save_fcs:
                 if dec_level_config.save_fcs_type == "csv":
@@ -320,6 +322,26 @@ def station_obj_from_row(row):
                 remote_stfts.append(stft_obj)
 
         # Merge STFTs
+
+        # Timing Error Workaround See Aurora Issue #289
+        if tfk.config.stations.remote:
+            n_chunks = len(local_stfts)
+            for i_chunk in range(n_chunks):
+                ok = local_stfts[i_chunk].time.shape == remote_stfts[i_chunk].time.shape
+                if not ok:
+                    print(f"Mismatch in FC array lengths detected -- Issue #289")
+                    glb = max(local_stfts[i_chunk].time.min(), remote_stfts[i_chunk].time.min())
+                    lub = min(local_stfts[i_chunk].time.max(), remote_stfts[i_chunk].time.max())
+                    cond1 = local_stfts[i_chunk].time >= glb
+                    cond2 = local_stfts[i_chunk].time <= lub
+                    local_stfts[i_chunk] = local_stfts[i_chunk].where(cond1 & cond2, drop=True)
+                    cond1 = remote_stfts[i_chunk].time >= glb
+                    cond2 = remote_stfts[i_chunk].time <= lub
+                    remote_stfts[i_chunk] = remote_stfts[i_chunk].where(cond1 & cond2, drop=True)
+                    assert (local_stfts[i_chunk].time.shape==remote_stfts[i_chunk].time.shape)
+
+
+
         local_merged_stft_obj = xr.concat(local_stfts, "time")
 
         if tfk.config.stations.remote:

diff --git a/aurora/transfer_function/kernel_dataset.py b/aurora/transfer_function/kernel_dataset.py
@@ -374,8 +374,7 @@ def initialize_dataframe_for_processing(self, mth5_objs):
             #    continue
             # the line below is not lazy, See Note #2
             run_ts = run_obj.to_runts(start=row.start, end=row.end)
-            xr_ds = run_ts.dataset
-            self.df["run_dataarray"].at[i] = xr_ds.to_array("channel")
+            self.df["run_dataarray"].at[i] = run_ts.dataset.to_array("channel")
 
             if i == 0:
                 if run_ts.survey_metadata.id in self.survey_metadata.keys():