Merge branch 'master' of https://github.com/JoranAngevaare/optim_esm_…

…tools
JoranAngevaare · Feb 24, 2025 · d59e655 · d59e655
2 parents 09842b2 + 9aa191b
commit d59e655
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/optim_esm_tools/analyze/clustering.py b/optim_esm_tools/analyze/clustering.py
@@ -12,7 +12,7 @@
 
 from optim_esm_tools.config import config
 from optim_esm_tools.config import get_logger
-from optim_esm_tools.utils import timed
+from optim_esm_tools.utils import timed, deprecated
 from optim_esm_tools.utils import tqdm
 
 
@@ -232,11 +232,17 @@ def _check_input(
     return lat, lon
 
 
-def _split_to_continous(
+@deprecated
+def _split_to_continous(*a, **kw):
+    return _split_to_continuous(*a, **kw)
+
+
+def _split_to_continuous(
     masks: ty.List,
+    **kw,
 ) -> ty.List[np.ndarray]:
     no_group = -1
-    mask_groups = masks_array_to_coninuous_sets(masks, no_group_value=no_group)
+    mask_groups = masks_array_to_coninuous_sets(masks, no_group_value=no_group, **kw)
     continous_masks = []
     for grouped_members in mask_groups:
         for group_id in np.unique(grouped_members):
@@ -296,7 +302,7 @@ def _build_cluster_with_kw(
         masks.append(np.array(full_2d_mask))
 
     if force_continuity:
-        masks = _split_to_continous(masks=masks)
+        masks = _split_to_continuous(masks=masks)
 
         clusters = [_find_lat_lon_values(m, lats=lat, lons=lon) for m in masks]
 
@@ -493,18 +499,21 @@ def masks_array_to_coninuous_sets(
 
     result_groups = np.ones_like(masks[0], dtype=np.int64) * no_group_value
     check_buffer = np.zeros_like(masks[0], dtype=np.bool_)
-
+    kw_cont_sets = dict(
+        len_x=len_x,
+        len_y=len_y,
+        add_diagonal=add_diagonal,
+    )
+    kw_cont_sets.update(kw)
     # Warning, do notice that the result_buffer and check_buffer are modified in place! However, _group_mask_in_continous_sets does reset the buffer each time
     # Therefore, we have to copy the result each time! Otherwise that result will be overwritten in the next iteration
     return [
         _group_mask_in_continous_sets(
             mask=mask,
             no_group_value=no_group_value,
-            add_diagonal=add_diagonal,
-            len_x=len_x,
-            len_y=len_y,
             result_buffer=result_groups,
             check_buffer=check_buffer,
+            **kw_cont_sets,
         ).copy()
         for mask in masks
     ]

diff --git a/optim_esm_tools/analyze/merge_candidate_regions.py b/optim_esm_tools/analyze/merge_candidate_regions.py
@@ -212,6 +212,7 @@ def merge_datasets(self):
             pbar.display()
             self.log.info(pbar)
             doc = self._group_to_first(candidates)
+
             if not self.pass_criteria(**doc['stats']):
                 self.log.info(
                     f'Discarding group {doc["merged"]} because {doc["stats"]} does not pass',
@@ -230,6 +231,10 @@ def merge_datasets(self):
                 )
                 groups.append(doc)
             candidates = [c for i, c in enumerate(candidates) if i not in doc['merged']]
+
+            if doc.get('force_break', False):
+                self.log.warning('Breaking forcefully')
+                candidates = []
         pbar.n = pbar.total
         pbar.close()
         pbar.display()
@@ -268,7 +273,8 @@ def _group_to_first(
             self.log.info(
                 f"Exhausted passing regions, so next candidates are ignored ({len(candidates)-1} remaining)",
             )
-            return dict(stats=first_doc, ds=candidates[0], merged=[0])
+            # We are going to pass one additional argument that allows us to break the overencompasing loop
+            return dict(stats=first_doc, ds=candidates[0], merged=[0], force_break=True)
         while something_merged:
             something_merged = False
             for i, ds_alt in global_masks.items():

diff --git a/optim_esm_tools/region_finding/keep_all.py b/optim_esm_tools/region_finding/keep_all.py
@@ -39,7 +39,7 @@ def get_masks(
                 coords.append(np.array(this_coords))
                 masks.append(this_mask)
         if force_continuity:
-            masks = oet.analyze.clustering._split_to_continous(masks=masks)
+            masks = oet.analyze.clustering._split_to_continuous(masks=masks)
             lat, lon = np.meshgrid(lats, lons)
             coords = [
                 oet.analyze.clustering._find_lat_lon_values(m, lats=lat.T, lons=lon.T)