Generalise out dimensions detection

boutproject · Jan 8, 2024 · ad46e9a · ad46e9a
1 parent 8bf6e89
commit ad46e9a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/xbout/fci/evaluate_at.py b/xbout/fci/evaluate_at.py
@@ -175,10 +175,12 @@ def evaluate_at_keys(ds, keys, key, fill_value=np.nan, slow=False):
         except AttributeError:
             pass
         outd = []
+        havenew = False
         for d in ds[k].dims:
-            if d == "x":
-                outd += weights.dims[:-3]
             if d in "xyz":
+                if not havenew:
+                    havenew = True
+                    outd += weights.dims[:-3]
                 continue
             outd.append(d)
         tmp = np.einsum("...ijk,...ijk->...", theisel, weights)