Plot multiple datasets in the same corner plot

user_1189239 0
Viewed 1800

I was wondering if it is possible to plot more than one dataset together using corner plot. Please don't suggest to use seaborn pairplot.

Here the script from the documentation, it is a simple corner plot:

import corner
import numpy as np

ndim, nsamples = 2, 10000
np.random.seed(42)
samples = np.random.randn(ndim * nsamples).reshape([nsamples, ndim])
figure = corner.corner(samples)

and the result: enter image description here

Now, I would like to plot a second dataset, maybe with a different color, over the three images above. Is there a way to do it?

matplotlib python
2 Answers

Yes, this is how:

figure = corner.corner(samples)
corner.corner(samples2,fig=figure)
user_104350270

Below is some code to overlay multiple sample sets in one corner plot, with a legend.

enter image description here

# -*- coding: utf-8 -*-
"""Demo to overlay multiple corners on top of each other"""
import corner
import matplotlib.lines as mlines
import matplotlib.pyplot as plt
import numpy as np

CORNER_KWARGS = dict(
    smooth=0.9,
    label_kwargs=dict(fontsize=16),
    title_kwargs=dict(fontsize=16),
    quantiles=[0.16, 0.84],
    levels=(1 - np.exp(-0.5), 1 - np.exp(-2), 1 - np.exp(-9 / 2.)),
    plot_density=False,
    plot_datapoints=False,
    fill_contours=True,
    show_titles=True,
    max_n_ticks=3,
)


def overlaid_corner(samples_list, sample_labels):
    """Plots multiple corners on top of each other"""
    # get some constants
    n = len(samples_list)
    _, ndim = samples_list[0].shape
    max_len = max([len(s) for s in samples_list])
    cmap = plt.cm.get_cmap('gist_rainbow', n)
    colors = [cmap(i) for i in range(n)]

    plot_range = []
    for dim in range(ndim):
        plot_range.append(
            [
                min([min(samples_list[i].T[dim]) for i in range(n)]),
                max([max(samples_list[i].T[dim]) for i in range(n)]),
            ]
        )

    CORNER_KWARGS.update(range=plot_range)

    fig = corner.corner(
        samples_list[0],
        color=colors[0],
        **CORNER_KWARGS
    )

    for idx in range(1, n):
        fig = corner.corner(
            samples_list[idx],
            fig=fig,
            weights=get_normalisation_weight(len(samples_list[idx]), max_len),
            color=colors[idx],
            **CORNER_KWARGS
        )

    plt.legend(
        handles=[
            mlines.Line2D([], [], color=colors[i], label=sample_labels[i])
            for i in range(n)
        ],
        fontsize=20, frameon=False,
        bbox_to_anchor=(1, ndim), loc="upper right"
    )
    plt.savefig("corner.png")
    plt.close()


def get_normalisation_weight(len_current_samples, len_of_longest_samples):
    return np.ones(len_current_samples) * (len_of_longest_samples / len_current_samples)


def main():
    ndim, nsamples = 3, 10000
    samples = np.random.randn(ndim * nsamples).reshape([nsamples, ndim])

    overlaid_corner(
        [samples * 3, samples * 2, samples],
        ["samples x 3", "samples x 2", "samples"]
    )


if __name__ == "__main__":
    main()
user_101777590
Related