Add plot_density option to plot_corner_multi

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,
             )

dingo/core/utils/plotting.py

@@ -8,7 +8,12 @@
 
 
 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.
@@ -20,10 +25,12 @@ 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.
+    plot_density : bool
+        Whether to normalize the posteriors. Defaults to False.
     **kwargs :
         Forwarded to corner.corner.
     """
@@ -61,6 +68,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)):
@@ -72,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)