From 9839d1d25fec8f90a1e29d8134aedd9d6dbb703e Mon Sep 17 00:00:00 2001
From: Javier Vegas-Regidor <javier.vegas@bsc.es>
Date: Tue, 8 Jun 2021 12:04:12 +0200
Subject: [PATCH] Avoid doing extra work in regmean

---
 diagonals/regmean.py | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/diagonals/regmean.py b/diagonals/regmean.py
index 0523d0b..03a7802 100644
--- a/diagonals/regmean.py
+++ b/diagonals/regmean.py
@@ -175,19 +175,14 @@ def _compute_regmean_levels_cpu(var, basins, volume):
     regmean_total: float32
         List containing regional mean at every depth level for variable var.
     """
-    times = var.shape[0]
-    levs = var.shape[1]
     regmean_total = {}
     for basin, mask in basins.items():
-        regmean = np.empty((times, levs))
         weights = _compute_weights_3d(mask, volume)
-        for time in range(times):
-            for lev in range(levs):
-                regmean = da.ma.average(
-                    var,
-                    axis=(2, 3),
-                    weights=da.broadcast_to(weights, var.shape)
-                )
+        regmean = da.ma.average(
+            var,
+            axis=(2, 3),
+            weights=da.broadcast_to(weights, var.shape)
+        )
         regmean_total[basin] = regmean
     return regmean_total
 
-- 
GitLab