Fix non-matching dimension lengths on core and anvil

w-k-jones · May 22, 2023 · 1dbc9ae · 1dbc9ae
1 parent 8c69811
commit 1dbc9ae
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/tobac_flow/utils/filter_utils.py b/src/tobac_flow/utils/filter_utils.py
@@ -199,10 +199,10 @@ def max_t_diff(x, *args, **kwargs):
         .data,
         {"anvil": dataset.anvil},
     )
-    wh_core_no_anvil = dataset.core_anvil_index == 0
+    wh_core_has_anvil = np.isin(dataset.core_anvil_index, dataset.anvil)
     anvil_max_core_area = xr.DataArray(
-        dataset.core_max_area[np.logical_not(wh_core_no_anvil)]
-        .groupby(dataset.core_anvil_index[np.logical_not(wh_core_no_anvil)])
+        dataset.core_max_area[wh_core_has_anvil]
+        .groupby(dataset.core_anvil_index[wh_core_has_anvil])
         .max()
         .data,
         {"anvil": dataset.anvil},
@@ -219,8 +219,8 @@ def max_t_diff(x, *args, **kwargs):
         {"anvil": dataset.anvil},
     )
     anvil_core_start_t = xr.DataArray(
-        dataset.core_start_t[np.logical_not(wh_core_no_anvil)]
-        .groupby(dataset.core_anvil_index[np.logical_not(wh_core_no_anvil)])
+        dataset.core_start_t[wh_core_has_anvil]
+        .groupby(dataset.core_anvil_index[wh_core_has_anvil])
         .min()
         .data,
         {"anvil": dataset.anvil},
@@ -239,8 +239,8 @@ def max_t_diff(x, *args, **kwargs):
         {"anvil": dataset.anvil},
     )
     anvil_core_end_t = xr.DataArray(
-        dataset.core_end_t[np.logical_not(wh_core_no_anvil)]
-        .groupby(dataset.core_anvil_index[np.logical_not(wh_core_no_anvil)])
+        dataset.core_end_t[wh_core_has_anvil]
+        .groupby(dataset.core_anvil_index[wh_core_has_anvil])
         .max()
         .data,
         {"anvil": dataset.anvil},