From 910563400af9400bbb3bb5dc3854e80d0ec01e88 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 25 Sep 2023 13:20:39 +0300
Subject: [PATCH 1/4] Add plot_density option to plot_corner_multi. Defaults to
 False.

---
 dingo/core/utils/plotting.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/dingo/core/utils/plotting.py b/dingo/core/utils/plotting.py
index 0ac02f40b..ab9b54a30 100644
--- a/dingo/core/utils/plotting.py
+++ b/dingo/core/utils/plotting.py
@@ -8,7 +8,7 @@
 
 
 def plot_corner_multi(
-    samples, weights=None, labels=None, filename="corner.pdf", **kwargs
+    samples, weights=None, labels=None, filename="corner.pdf", plot_density=False, **kwargs
 ):
     """
     Generate a corner plot for multiple posteriors.
@@ -24,6 +24,8 @@ def plot_corner_multi(
         Labels for the posteriors.
     filename : str
         Where to save samples.
+    plot_density : bool
+        Whether to normalize the posteriors.
     **kwargs :
         Forwarded to corner.corner.
     """
@@ -61,6 +63,15 @@ def plot_corner_multi(
         if p in set.intersection(*(set(s.columns) for s in samples))
     ]
 
+    if plot_density:
+        # Redefine all sets of weights to have unit sum.
+        for i, w in enumerate(weights):
+            if w is not None:
+                weights[i] = w / np.sum(w)
+            else:
+                n = len(samples[i])
+                weights[i] = np.ones(n) / n
+
     fig = None
     handles = []
     for i, (s, w, l) in enumerate(zip_longest(samples, weights, labels)):

From e33edd711cf773eb6ec607740fde5f8f582cd322 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 25 Sep 2023 13:23:21 +0300
Subject: [PATCH 2/4] Update docstring to indicate default of False.

---
 dingo/core/utils/plotting.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dingo/core/utils/plotting.py b/dingo/core/utils/plotting.py
index ab9b54a30..d4c8e902a 100644
--- a/dingo/core/utils/plotting.py
+++ b/dingo/core/utils/plotting.py
@@ -25,7 +25,7 @@ def plot_corner_multi(
     filename : str
         Where to save samples.
     plot_density : bool
-        Whether to normalize the posteriors.
+        Whether to normalize the posteriors. Defaults to False.
     **kwargs :
         Forwarded to corner.corner.
     """

From f9760b4c78f5029de9ac6e765e1ef8a3df6bd0dc Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 25 Sep 2023 13:31:08 +0300
Subject: [PATCH 3/4] Run black

---
 dingo/core/utils/plotting.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/dingo/core/utils/plotting.py b/dingo/core/utils/plotting.py
index d4c8e902a..f2d32ff6e 100644
--- a/dingo/core/utils/plotting.py
+++ b/dingo/core/utils/plotting.py
@@ -8,7 +8,12 @@
 
 
 def plot_corner_multi(
-    samples, weights=None, labels=None, filename="corner.pdf", plot_density=False, **kwargs
+    samples,
+    weights=None,
+    labels=None,
+    filename="corner.pdf",
+    plot_density=False,
+    **kwargs,
 ):
     """
     Generate a corner plot for multiple posteriors.
@@ -83,7 +88,7 @@ def plot_corner_multi(
             color=color,
             no_fill_contours=True,
             fig=fig,
-            **corner_params
+            **corner_params,
         )
         handles.append(
             plt.Line2D([], [], color=color, label=l, linewidth=5, markersize=20)

From 36c7fd3b46e8a80d4aa3805a7930b6762a502463 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Nov 2023 11:09:25 +0000
Subject: [PATCH 4/4] Fix bug in plotting with weights

---
 dingo/core/result.py         | 6 ++++--
 dingo/core/utils/plotting.py | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/dingo/core/result.py b/dingo/core/result.py
index 2974a0ab7..cd86439d2 100644
--- a/dingo/core/result.py
+++ b/dingo/core/result.py
@@ -600,14 +600,16 @@ def plot_corner(self, parameters=None, filename="corner.pdf"):
         # delta_log_prob_target is not interesting so never plot it.
         theta = theta.drop(columns="delta_log_prob_target", errors="ignore")
 
+        weights = theta.get("weights")
+
         # User option to plot specific parameters.
         if parameters:
             theta = theta[parameters]
 
-        if "weights" in theta:
+        if weights:
             plot_corner_multi(
                 [theta, theta],
-                weights=[None, theta["weights"].to_numpy()],
+                weights=[None, weights.to_numpy()],
                 labels=["Dingo", "Dingo-IS"],
                 filename=filename,
             )
diff --git a/dingo/core/utils/plotting.py b/dingo/core/utils/plotting.py
index f2d32ff6e..e89981e5e 100644
--- a/dingo/core/utils/plotting.py
+++ b/dingo/core/utils/plotting.py
@@ -25,7 +25,7 @@ def plot_corner_multi(
     weights : list[np.ndarray or None] or None
         List of weights sets. The length of each array should be the same as the length of
         the corresponding samples.
-    labels : list[str or None] or None
+    labels : list[str or None] or None or str
         Labels for the posteriors.
     filename : str
         Where to save samples.