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 0ac02f40b..e89981e5e 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", **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)