From 9d7ca1a3f881b4cb4605f614f882ef841264776f Mon Sep 17 00:00:00 2001
From: Nicolas Tessore <n.tessore@ucl.ac.uk>
Date: Mon, 9 Oct 2023 11:11:35 +0100
Subject: [PATCH] ENH(plot): plot error bands in postage stamps (#50)

Add the possibility to show error bands in `postage_stamps()` when
individual cls contain an `ERR` column.

Closes: #46
---
 heracles/plot.py | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/heracles/plot.py b/heracles/plot.py
index 3afb14b..a6eae46 100644
--- a/heracles/plot.py
+++ b/heracles/plot.py
@@ -167,8 +167,11 @@ def postage_stamps(
 
             if cl.dtype.names is None:
                 ell = np.arange(len(cl), dtype=float)
+                err = None
             else:
-                ell, cl = cl["L"].astype(float), cl["CL"]
+                ell = cl["L"].astype(float)
+                err = cl["ERR"] if "ERR" in cl.dtype.names else None
+                cl = cl["CL"]
 
             if scale is None:
                 pass
@@ -180,13 +183,21 @@ def postage_stamps(
             xmin = np.nanmin(ell[ell != 0], initial=xmin)
             xmax = np.nanmax(ell[ell != 0], initial=xmax)
             if n < len(keys):
-                ymin = np.nanmin(cl, initial=ymin)
-                ymax = np.nanmax(cl, initial=ymax)
+                if err is None:
+                    ymin = np.nanmin(cl, initial=ymin)
+                    ymax = np.nanmax(cl, initial=ymax)
+                else:
+                    ymin = np.nanmin(cl - err, initial=ymin)
+                    ymax = np.nanmax(cl + err, initial=ymax)
             else:
-                trymin = np.nanmin(cl, initial=trymin)
-                trymax = np.nanmax(cl, initial=trymax)
+                if err is None:
+                    trymin = np.nanmin(cl, initial=trymin)
+                    trymax = np.nanmax(cl, initial=trymax)
+                else:
+                    trymin = np.nanmin(cl - err, initial=trymin)
+                    trymax = np.nanmax(cl + err, initial=trymax)
 
-            ax.plot(
+            lines = ax.plot(
                 ell,
                 cl,
                 label=label,
@@ -194,6 +205,16 @@ def postage_stamps(
                 **{**oprop, **iprop},
             )
 
+            if err is not None:
+                ax.fill_between(
+                    ell,
+                    cl - err,
+                    cl + err,
+                    ec="none",
+                    fc=lines[-1].get_color(),
+                    alpha=0.2,
+                )
+
             # prevent multiple labels with same colour
             label = None