From e31e3c03976e3ef2886affb6bee264ef53bed0f4 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Thu, 17 Jun 2021 20:19:52 +0200
Subject: [PATCH] [postprocesor] weighted and ranged mean by bins

---
 postprocessor.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/postprocessor.py b/postprocessor.py
index d796b48..5a3a049 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -96,6 +96,8 @@ def mean_by_bins(
     y,
     bins=100,
     logbins=False,
+    weights=None,
+    range=None,
 ):
     """
     Compute the mean of y in bins of x
@@ -105,16 +107,21 @@ def mean_by_bins(
     x, y : np.array of same dimensions
     bins : int, number of bins
     logbins : bool, if true, the bins will be logaritmically distributed
+    weights : np.array, same size as y, default None.
+    range : tuple of float, default None.
     """
     mask = np.isfinite(x) & np.isfinite(y)
     x = x[mask].flatten()
     y = y[mask].flatten()
 
     if logbins:
-        minvalue = np.min(x[x > 0])
-        x_bins = np.logspace(np.log10(minvalue), np.log10(np.max(x)), bins, base=10)
+        if range is None:
+            range = (np.min(x[x > 0]), np.max(x))
+        x_bins = np.logspace(np.log10(range[0]), np.log10(range[1]), bins, base=10)
     else:
-        x_bins = np.linspace(np.min(x), np.max(x), bins)
+        if range is None:
+            range = (np.min(x), np.max(x))
+        x_bins = np.linspace(range[0], range[1], bins)
 
     # For each cell, bin_number contains the number of the bins it belongs to
     bin_number = np.zeros(len(y))
@@ -125,8 +132,14 @@ def mean_by_bins(
 
     # Compute the mean in each bin
     y_mean = np.zeros(len(x_bins) - 1)
-    for i in range(len(y_mean)):
-        y_mean[i] = np.mean(y[bin_number == i])
+    for i in np.arange(len(y_mean)):
+        mask_bin = bin_number == i
+        if weights is None:
+            y_mean[i] = np.mean(y[mask_bin])
+        else:
+            y_mean[i] = np.sum(y[mask_bin] * weights[mask_bin]) / np.sum(
+                weights[mask_bin]
+            )
 
     # Get the center of each bin
     if logbins: