xarray can squeeze groupby values, making them non-iterable. Enforce.

deeplycloudy · Feb 4, 2023 · f0170c3 · f0170c3
1 parent 02a4e40
commit f0170c3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pyxlma/lmalib/traversal.py b/pyxlma/lmalib/traversal.py
@@ -219,7 +219,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids):
         for e_var, p_var in self._descend():
             if prune == True:
                 p_group = self.parent_groups[p_var].groups
-                e_iter = (p_group[eid] for eid in last_entity_ids
+                e_iter = (np.atleast_1d(p_group[eid]) for eid in last_entity_ids
                           if eid in p_group)
                 e_idx = list(itertools.chain.from_iterable(e_iter))
                 if len(e_idx) == 0:
@@ -239,7 +239,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids):
                 # through.
                 prune = True
                 e_group = self.entity_groups[e_var].groups
-                e_iter = (e_group[eid] for eid in entity_ids
+                e_iter = (np.atleast_1d(e_group[eid]) for eid in entity_ids
                           if eid in e_group)
                 e_idx = list(itertools.chain.from_iterable(e_iter))
                 last_entity_ids = entity_ids # == dataset[e_var].data
@@ -263,7 +263,7 @@ def reduce_to_entities(self, entity_id_var, entity_ids):
         for e_var, p_var in self._ascend():
             if (prune == True):
                 e_group = self.entity_groups[e_var].groups
-                e_iter = (e_group[eid] for eid in last_entity_ids
+                e_iter = (np.atleast_1d(e_group[eid]) for eid in last_entity_ids
                           if eid in e_group)
                 e_idx = list(itertools.chain.from_iterable(e_iter))
                 if len(e_idx) == 0: