├── MANIFEST.in ├── corner.png ├── .gitignore ├── demo.py ├── setup.py ├── LICENSE ├── README.rst ├── tests.py └── corner.py /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst LICENSE 2 | -------------------------------------------------------------------------------- /corner.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jakevdp/corner.py/master/corner.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | *.pyc 3 | *.png 4 | *.egg-info 5 | *.pdf 6 | dist 7 | build 8 | test_figures 9 | -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import numpy as np 4 | import corner 5 | 6 | # Set up the parameters of the problem. 7 | ndim, nsamples = 3, 50000 8 | 9 | # Generate some fake data. 10 | data1 = np.random.randn(ndim * 4 * nsamples / 5).reshape([4 * nsamples / 5, 11 | ndim]) 12 | data2 = (5 * np.random.rand(ndim)[None, :] 13 | + np.random.randn(ndim * nsamples / 5).reshape([nsamples / 5, ndim])) 14 | data = np.vstack([data1, data2]) 15 | 16 | # Plot it. 17 | figure = corner.corner(data, labels=[r"$x$", r"$y$", r"$\log \alpha$", 18 | r"$\Gamma \, [\mathrm{parsec}]$"], 19 | truths=[0.0, 0.0, 0.0], 20 | quantiles=[0.16, 0.5, 0.84], 21 | show_titles=True, title_args={"fontsize": 12}) 22 | figure.gca().annotate("A Title", xy=(0.5, 1.0), xycoords="figure fraction", 23 | xytext=(0, -5), textcoords="offset points", 24 | ha="center", va="top") 25 | figure.savefig("demo.png") 26 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import re 4 | import os 5 | import sys 6 | from setuptools import setup 7 | 8 | if sys.argv[-1] == "publish": 9 | os.system("python setup.py sdist upload") 10 | sys.exit() 11 | 12 | # Hackishly synchronize the version. 13 | version = re.findall(r"__version__ = \"(.*?)\"", open("corner.py").read())[0] 14 | 15 | 16 | setup( 17 | name="corner", 18 | version=version, 19 | author="Daniel Foreman-Mackey", 20 | author_email="danfm@nyu.edu", 21 | url="https://github.com/dfm/corner.py", 22 | py_modules=["corner"], 23 | description="Make some beautiful corner plots of samples.", 24 | long_description=open("README.rst").read(), 25 | package_data={"": ["LICENSE"]}, 26 | include_package_data=True, 27 | classifiers=[ 28 | "Development Status :: 5 - Production/Stable", 29 | "License :: OSI Approved :: BSD License", 30 | "Intended Audience :: Developers", 31 | "Intended Audience :: Science/Research", 32 | "Operating System :: OS Independent", 33 | "Programming Language :: Python", 34 | ], 35 | ) 36 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2013-2015 Daniel Foreman-Mackey 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | 7 | 1. Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 2. Redistributions in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 14 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 15 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 16 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR 17 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 18 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 19 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 20 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 21 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 22 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 23 | 24 | The views and conclusions contained in the software and documentation are those 25 | of the authors and should not be interpreted as representing official policies, 26 | either expressed or implied, of the FreeBSD Project. 27 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | corner.py 2 | ========= 3 | 4 | Make some beautiful corner plots. 5 | 6 | Corner plot /ˈkôrnər plät/ (noun): 7 | An illustrative representation of different projections of samples in 8 | high dimensional spaces. It is awesome. I promise. 9 | 10 | Built by `Dan Foreman-Mackey `_ and collaborators (see 11 | ``corner.__contributors__`` for the most up to date list). Licensed under 12 | the 2-clause BSD license (see ``LICENSE``). 13 | 14 | 15 | Installation 16 | ------------ 17 | 18 | Just run 19 | 20 | :: 21 | 22 | pip install corner 23 | 24 | to get the most recent stable version. 25 | 26 | 27 | Usage 28 | ----- 29 | 30 | The main entry point is the ``corner.corner`` function. You'll just use it 31 | like this: 32 | 33 | :: 34 | 35 | import numpy as np 36 | import corner 37 | 38 | ndim, nsamples = 5, 10000 39 | samples = np.random.randn(ndim * nsamples).reshape([nsamples, ndim]) 40 | figure = corner.corner(samples) 41 | figure.savefig("corner.png") 42 | 43 | With some other tweaks (see `demo.py 44 | `_) you can get 45 | something that looks awesome like: 46 | 47 | .. image:: https://raw.github.com/dfm/corner.py/master/corner.png 48 | 49 | By default, data points are shown as grayscale points with contours. 50 | Contours are shown at 0.5, 1, 1.5, and 2 sigma. 51 | 52 | For more usage examples, take a look at `tests.py 53 | `_. 54 | 55 | 56 | Documentation 57 | ------------- 58 | 59 | All the options are documented in the docstrings for the ``corner`` and 60 | ``hist2d`` functions. These can be viewed in a Python shell using: 61 | 62 | :: 63 | 64 | import corner 65 | print(corner.corner.__doc__) 66 | 67 | or, in IPython using: 68 | 69 | :: 70 | 71 | import corner 72 | corner.corner? 73 | 74 | 75 | A note about "sigmas" 76 | +++++++++++++++++++++ 77 | 78 | We are regularly asked about the "sigma" levels in the 2D histograms. These 79 | are not the 68%, *etc.* values that we're used to for 1D distributions. In two 80 | dimensions, a Gaussian density is given by: 81 | 82 | :: 83 | 84 | pdf(r) = exp(-(r/s)^2/2) / (2*pi*s^2) 85 | 86 | The integral under this density is: 87 | 88 | :: 89 | 90 | cdf(x) = Integral(r * exp(-(r/s)^2/2) / s^2, {r, 0, x}) 91 | = 1 - exp(-(x/s)^2/2) 92 | 93 | This means that within "1-sigma", the Gaussian contains ``1-exp(-0.5) ~ 0.393`` 94 | or 39.3% of the volume. Therefore the relevant 1-sigma levels for a 2D 95 | histogram of samples is 39% not 68%. If you must use 68% of the mass, use the 96 | ``levels`` keyword argument. 97 | 98 | The `"sigma-demo" notebook 99 | `_ visually 100 | demonstrates the difference between these choices of levels. 101 | 102 | 103 | Attribution 104 | ----------- 105 | 106 | .. image:: https://zenodo.org/badge/doi/10.5281/zenodo.11020.png 107 | :target: http://dx.doi.org/10.5281/zenodo.11020 108 | 109 | If you make use of this code, please `cite it 110 | `_. 111 | 112 | 113 | License 114 | ------- 115 | 116 | Copyright 2013, 2014 Dan Foreman-Mackey 117 | 118 | corner.py is free software made available under the BSD License. 119 | For details see the LICENSE file. 120 | -------------------------------------------------------------------------------- /tests.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import division, print_function 4 | 5 | __all__ = ["test_hist2d", "test_corner"] 6 | 7 | import os 8 | import numpy as np 9 | import pandas as pd 10 | import matplotlib.pyplot as pl 11 | 12 | import corner 13 | 14 | FIGURE_PATH = "test_figures" 15 | 16 | 17 | def _run_hist2d(nm, N=50000, seed=1234, **kwargs): 18 | print(" .. {0}".format(nm)) 19 | 20 | if not os.path.exists(FIGURE_PATH): 21 | os.makedirs(FIGURE_PATH) 22 | 23 | # Generate some fake data. 24 | np.random.seed(seed) 25 | x = np.random.randn(N) 26 | y = np.random.randn(N) 27 | 28 | fig, ax = pl.subplots(1, 1, figsize=(8, 8)) 29 | corner.hist2d(x, y, ax=ax, **kwargs) 30 | fig.savefig(os.path.join(FIGURE_PATH, "hist2d_{0}.png".format(nm))) 31 | pl.close(fig) 32 | 33 | 34 | def test_hist2d(): 35 | _run_hist2d("cutoff", range=[(0, 4), (0, 2.5)]) 36 | _run_hist2d("cutoff2", range=[(-4, 4), (-0.1, 0.1)], N=100000, 37 | fill_contours=True, smooth=1) 38 | _run_hist2d("basic") 39 | _run_hist2d("color", color="g") 40 | _run_hist2d("levels1", levels=[0.68, 0.95]) 41 | _run_hist2d("levels2", levels=[0.5, 0.75]) 42 | _run_hist2d("filled", fill_contours=True) 43 | _run_hist2d("smooth1", bins=50) 44 | _run_hist2d("smooth2", bins=50, smooth=(1.0, 1.5)) 45 | _run_hist2d("philsplot", plot_datapoints=False, fill_contours=True, 46 | levels=[0.68, 0.95], color="g", bins=50, smooth=1.) 47 | 48 | 49 | def _run_corner(nm, pandas=False, N=10000, seed=1234, ndim=3, ret=False, 50 | factor=None, **kwargs): 51 | print(" .. {0}".format(nm)) 52 | 53 | if not os.path.exists(FIGURE_PATH): 54 | os.makedirs(FIGURE_PATH) 55 | 56 | np.random.seed(seed) 57 | data1 = np.random.randn(ndim*4*N/5.).reshape([4*N/5., ndim]) 58 | data2 = (5 * np.random.rand(ndim)[None, :] 59 | + np.random.randn(ndim*N/5.).reshape([N/5., ndim])) 60 | data = np.vstack([data1, data2]) 61 | if factor is not None: 62 | data[:, 0] *= factor 63 | data[:, 1] /= factor 64 | if pandas: 65 | data = pd.DataFrame.from_items(zip(map("d{0}".format, range(ndim)), 66 | data.T)) 67 | 68 | fig = corner.corner(data, **kwargs) 69 | fig.savefig(os.path.join(FIGURE_PATH, "corner_{0}.png".format(nm))) 70 | if ret: 71 | return fig 72 | else: 73 | pl.close(fig) 74 | 75 | 76 | def test_corner(): 77 | _run_corner("basic") 78 | _run_corner("labels", labels=["a", "b", "c"]) 79 | _run_corner("quantiles", quantiles=[0.16, 0.5, 0.84]) 80 | _run_corner("color", color="g") 81 | fig = _run_corner("color-filled", color="g", fill_contours=True, 82 | ret=True) 83 | _run_corner("overplot", seed=15, color="b", fig=fig, fill_contours=True) 84 | _run_corner("smooth1", bins=50) 85 | _run_corner("smooth2", bins=50, smooth=1.0) 86 | _run_corner("smooth1d", bins=50, smooth=1.0, smooth1d=1.0) 87 | _run_corner("titles1", show_titles=True) 88 | _run_corner("titles2", show_titles=True, title_fmt=None, 89 | labels=["a", "b", "c"]) 90 | _run_corner("top-ticks", top_ticks=True) 91 | _run_corner("pandas", pandas=True) 92 | _run_corner("truths", truths=[0.0, None, 0.15]) 93 | _run_corner("no-fill-contours", no_fill_contours=True) 94 | 95 | # _run_corner("mathtext", factor=1e8, use_math_text=True) 96 | 97 | fig = _run_corner("tight", ret=True) 98 | pl.tight_layout() 99 | fig.savefig(os.path.join(FIGURE_PATH, "corner_tight.png")) 100 | pl.close(fig) 101 | 102 | 103 | if __name__ == "__main__": 104 | print("Testing 'hist2d'") 105 | test_hist2d() 106 | 107 | print("Testing 'corner'") 108 | test_corner() 109 | -------------------------------------------------------------------------------- /corner.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import print_function, absolute_import, unicode_literals 4 | 5 | __all__ = ["corner", "hist2d"] 6 | __version__ = "1.0.0" 7 | __author__ = "Dan Foreman-Mackey (danfm@nyu.edu)" 8 | __copyright__ = "Copyright 2013-2015 Daniel Foreman-Mackey" 9 | __contributors__ = [ 10 | # Alphabetical by first name. 11 | "Adrian Price-Whelan @adrn", 12 | "Brendon Brewer @eggplantbren", 13 | "Ekta Patel @ekta1224", 14 | "Emily Rice @emilurice", 15 | "Geoff Ryan @geoffryan", 16 | "Guillaume @ceyzeriat", 17 | "Gregory Ashton @ga7g08", 18 | "Kelle Cruz @kelle", 19 | "Kyle Barbary @kbarbary", 20 | "Marco Tazzari @mtazzari", 21 | "Matt Pitkin @mattpitkin", 22 | "Phil Marshall @drphilmarshall", 23 | "Pierre Gratier @pirg", 24 | "Stephan Hoyer @shoyer", 25 | "Will Vousden @willvousden", 26 | "Wolfgang Kerzendorf @wkerzendorf", 27 | ] 28 | 29 | import logging 30 | import numpy as np 31 | import matplotlib.pyplot as pl 32 | from matplotlib.ticker import MaxNLocator 33 | from matplotlib.colors import LinearSegmentedColormap, colorConverter 34 | from matplotlib.ticker import ScalarFormatter 35 | 36 | try: 37 | from scipy.ndimage import gaussian_filter 38 | except ImportError: 39 | gaussian_filter = None 40 | 41 | 42 | def corner(xs, bins=20, range=None, weights=None, color="k", 43 | smooth=None, smooth1d=None, 44 | labels=None, label_kwargs=None, 45 | show_titles=False, title_fmt=".2f", title_kwargs=None, 46 | truths=None, truth_color="#4682b4", 47 | scale_hist=False, quantiles=None, verbose=False, fig=None, 48 | max_n_ticks=5, top_ticks=False, use_math_text=False, 49 | hist_kwargs=None, **hist2d_kwargs): 50 | """ 51 | Make a *sick* corner plot showing the projections of a data set in a 52 | multi-dimensional space. kwargs are passed to hist2d() or used for 53 | `matplotlib` styling. 54 | 55 | Parameters 56 | ---------- 57 | xs : array_like (nsamples, ndim) 58 | The samples. This should be a 1- or 2-dimensional array. For a 1-D 59 | array this results in a simple histogram. For a 2-D array, the zeroth 60 | axis is the list of samples and the next axis are the dimensions of 61 | the space. 62 | 63 | bins : int or array_like (ndim,) (optional) 64 | The number of bins to use in histograms, either as a fixed value for 65 | all dimensions, or as a list of integers for each dimension. 66 | 67 | weights : array_like (nsamples,) 68 | The weight of each sample. If `None` (default), samples are given 69 | equal weight. 70 | 71 | color : str (optional) 72 | A ``matplotlib`` style color for all histograms. 73 | 74 | smooth, smooth1d : float (optional) 75 | The standard deviation for Gaussian kernel passed to 76 | `scipy.ndimage.gaussian_filter` to smooth the 2-D and 1-D histograms 77 | respectively. If `None` (default), no smoothing is applied. 78 | 79 | labels : iterable (ndim,) (optional) 80 | A list of names for the dimensions. If a ``xs`` is a 81 | ``pandas.DataFrame``, labels will default to column names. 82 | 83 | label_kwargs : dict (optional) 84 | Any extra keyword arguments to send to the `set_xlabel` and 85 | `set_ylabel` methods. 86 | 87 | show_titles : bool (optional) 88 | Displays a title above each 1-D histogram showing the 0.5 quantile 89 | with the upper and lower errors supplied by the quantiles argument. 90 | 91 | title_fmt : string (optional) 92 | The format string for the quantiles given in titles. If you explicitly 93 | set ``show_titles=True`` and ``title_fmt=None``, the labels will be 94 | shown as the titles. (default: ``.2f``) 95 | 96 | title_kwargs : dict (optional) 97 | Any extra keyword arguments to send to the `set_title` command. 98 | 99 | range : iterable (ndim,) (optional) 100 | A list where each element is either a length 2 tuple containing 101 | lower and upper bounds or a float in range (0., 1.) 102 | giving the fraction of samples to include in bounds, e.g., 103 | [(0.,10.), (1.,5), 0.999, etc.]. 104 | If a fraction, the bounds are chosen to be equal-tailed. 105 | 106 | truths : iterable (ndim,) (optional) 107 | A list of reference values to indicate on the plots. Individual 108 | values can be omitted by using ``None``. 109 | 110 | truth_color : str (optional) 111 | A ``matplotlib`` style color for the ``truths`` makers. 112 | 113 | scale_hist : bool (optional) 114 | Should the 1-D histograms be scaled in such a way that the zero line 115 | is visible? 116 | 117 | quantiles : iterable (optional) 118 | A list of fractional quantiles to show on the 1-D histograms as 119 | vertical dashed lines. 120 | 121 | verbose : bool (optional) 122 | If true, print the values of the computed quantiles. 123 | 124 | plot_contours : bool (optional) 125 | Draw contours for dense regions of the plot. 126 | 127 | use_math_text : bool (optional) 128 | If true, then axis tick labels for very large or small exponents will 129 | be displayed as powers of 10 rather than using `e`. 130 | 131 | max_n_ticks: int (optional) 132 | Maximum number of ticks to try to use 133 | 134 | top_ticks : bool (optional) 135 | If true, label the top ticks of each axis 136 | 137 | fig : matplotlib.Figure (optional) 138 | Overplot onto the provided figure object. 139 | 140 | hist_kwargs : dict (optional) 141 | Any extra keyword arguments to send to the 1-D histogram plots. 142 | 143 | **hist2d_kwargs : (optional) 144 | Any remaining keyword arguments are sent to `corner.hist2d` to generate 145 | the 2-D histogram plots. 146 | """ 147 | if quantiles is None: 148 | quantiles = [] 149 | if title_kwargs is None: 150 | title_kwargs = dict() 151 | if label_kwargs is None: 152 | label_kwargs = dict() 153 | 154 | # Try filling in labels from pandas.DataFrame columns. 155 | if labels is None: 156 | try: 157 | labels = xs.columns 158 | except AttributeError: 159 | pass 160 | 161 | # Deal with 1D sample lists. 162 | xs = np.atleast_1d(xs) 163 | if len(xs.shape) == 1: 164 | xs = np.atleast_2d(xs) 165 | else: 166 | assert len(xs.shape) == 2, "The input sample array must be 1- or 2-D." 167 | xs = xs.T 168 | assert xs.shape[0] <= xs.shape[1], "I don't believe that you want more " \ 169 | "dimensions than samples!" 170 | 171 | # Parse the weight array. 172 | if weights is not None: 173 | weights = np.asarray(weights) 174 | if weights.ndim != 1: 175 | raise ValueError("Weights must be 1-D") 176 | if xs.shape[1] != weights.shape[0]: 177 | raise ValueError("Lengths of weights must match number of samples") 178 | 179 | # Parse the parameter ranges. 180 | if range is None: 181 | if "extents" in hist2d_kwargs: 182 | logging.warn("Deprecated keyword argument 'extents'. " 183 | "Use 'range' instead.") 184 | range = hist2d_kwargs.pop("extents") 185 | else: 186 | range = [[x.min(), x.max()] for x in xs] 187 | # Check for parameters that never change. 188 | m = np.array([e[0] == e[1] for e in range], dtype=bool) 189 | if np.any(m): 190 | raise ValueError(("It looks like the parameter(s) in " 191 | "column(s) {0} have no dynamic range. " 192 | "Please provide a `range` argument.") 193 | .format(", ".join(map( 194 | "{0}".format, np.arange(len(m))[m])))) 195 | 196 | else: 197 | # If any of the extents are percentiles, convert them to ranges. 198 | # Also make sure it's a normal list. 199 | range = list(range) 200 | for i, _ in enumerate(range): 201 | try: 202 | emin, emax = range[i] 203 | except TypeError: 204 | q = [0.5 - 0.5*range[i], 0.5 + 0.5*range[i]] 205 | range[i] = quantile(xs[i], q, weights=weights) 206 | 207 | if len(range) != xs.shape[0]: 208 | raise ValueError("Dimension mismatch between samples and range") 209 | 210 | # Parse the bin specifications. 211 | try: 212 | bins = [float(bins) for _ in range] 213 | except TypeError: 214 | if len(bins) != len(range): 215 | raise ValueError("Dimension mismatch between bins and range") 216 | 217 | # Some magic numbers for pretty axis layout. 218 | K = len(xs) 219 | factor = 2.0 # size of one side of one panel 220 | lbdim = 0.5 * factor # size of left/bottom margin 221 | trdim = 0.2 * factor # size of top/right margin 222 | whspace = 0.05 # w/hspace size 223 | plotdim = factor * K + factor * (K - 1.) * whspace 224 | dim = lbdim + plotdim + trdim 225 | 226 | # Create a new figure if one wasn't provided. 227 | if fig is None: 228 | fig, axes = pl.subplots(K, K, figsize=(dim, dim)) 229 | else: 230 | try: 231 | axes = np.array(fig.axes).reshape((K, K)) 232 | except: 233 | raise ValueError("Provided figure has {0} axes, but data has " 234 | "dimensions K={1}".format(len(fig.axes), K)) 235 | 236 | # Format the figure. 237 | lb = lbdim / dim 238 | tr = (lbdim + plotdim) / dim 239 | fig.subplots_adjust(left=lb, bottom=lb, right=tr, top=tr, 240 | wspace=whspace, hspace=whspace) 241 | 242 | # Set up the default histogram keywords. 243 | if hist_kwargs is None: 244 | hist_kwargs = dict() 245 | hist_kwargs["color"] = hist_kwargs.get("color", color) 246 | if smooth1d is None: 247 | hist_kwargs["histtype"] = hist_kwargs.get("histtype", "step") 248 | 249 | for i, x in enumerate(xs): 250 | # Deal with masked arrays. 251 | if hasattr(x, "compressed"): 252 | x = x.compressed() 253 | 254 | if np.shape(xs)[0] == 1: 255 | ax = axes 256 | else: 257 | ax = axes[i, i] 258 | # Plot the histograms. 259 | if smooth1d is None: 260 | n, _, _ = ax.hist(x, bins=bins[i], weights=weights, 261 | range=range[i], **hist_kwargs) 262 | else: 263 | if gaussian_filter is None: 264 | raise ImportError("Please install scipy for smoothing") 265 | n, b = np.histogram(x, bins=bins[i], weights=weights, 266 | range=range[i]) 267 | n = gaussian_filter(n, smooth1d) 268 | x0 = np.array(list(zip(b[:-1], b[1:]))).flatten() 269 | y0 = np.array(list(zip(n, n))).flatten() 270 | ax.plot(x0, y0, **hist_kwargs) 271 | 272 | if truths is not None and truths[i] is not None: 273 | ax.axvline(truths[i], color=truth_color) 274 | 275 | # Plot quantiles if wanted. 276 | if len(quantiles) > 0: 277 | qvalues = quantile(x, quantiles, weights=weights) 278 | for q in qvalues: 279 | ax.axvline(q, ls="dashed", color=color) 280 | 281 | if verbose: 282 | print("Quantiles:") 283 | print([item for item in zip(quantiles, qvalues)]) 284 | 285 | if show_titles: 286 | title = None 287 | if title_fmt is not None: 288 | # Compute the quantiles for the title. This might redo 289 | # unneeded computation but who cares. 290 | q_16, q_50, q_84 = quantile(x, [0.16, 0.5, 0.84], 291 | weights=weights) 292 | q_m, q_p = q_50-q_16, q_84-q_50 293 | 294 | # Format the quantile display. 295 | fmt = "{{0:{0}}}".format(title_fmt).format 296 | title = r"${{{0}}}_{{-{1}}}^{{+{2}}}$" 297 | title = title.format(fmt(q_50), fmt(q_m), fmt(q_p)) 298 | 299 | # Add in the column name if it's given. 300 | if labels is not None: 301 | title = "{0} = {1}".format(labels[i], title) 302 | 303 | elif labels is not None: 304 | title = "{0}".format(labels[i]) 305 | 306 | if title is not None: 307 | ax.set_title(title, **title_kwargs) 308 | 309 | # Set up the axes. 310 | ax.set_xlim(range[i]) 311 | if scale_hist: 312 | maxn = np.max(n) 313 | ax.set_ylim(-0.1 * maxn, 1.1 * maxn) 314 | else: 315 | ax.set_ylim(0, 1.1 * np.max(n)) 316 | ax.set_yticklabels([]) 317 | ax.xaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower")) 318 | 319 | if i < K - 1: 320 | if top_ticks: 321 | ax.xaxis.set_ticks_position("top") 322 | [l.set_rotation(45) for l in ax.get_xticklabels()] 323 | else: 324 | ax.set_xticklabels([]) 325 | else: 326 | [l.set_rotation(45) for l in ax.get_xticklabels()] 327 | if labels is not None: 328 | ax.set_xlabel(labels[i], **label_kwargs) 329 | ax.xaxis.set_label_coords(0.5, -0.3) 330 | 331 | # use MathText for axes ticks 332 | ax.xaxis.set_major_formatter( 333 | ScalarFormatter(useMathText=use_math_text)) 334 | 335 | for j, y in enumerate(xs): 336 | if np.shape(xs)[0] == 1: 337 | ax = axes 338 | else: 339 | ax = axes[i, j] 340 | if j > i: 341 | ax.set_frame_on(False) 342 | ax.set_xticks([]) 343 | ax.set_yticks([]) 344 | continue 345 | elif j == i: 346 | continue 347 | 348 | # Deal with masked arrays. 349 | if hasattr(y, "compressed"): 350 | y = y.compressed() 351 | 352 | hist2d(y, x, ax=ax, range=[range[j], range[i]], weights=weights, 353 | color=color, smooth=smooth, bins=[bins[j], bins[i]], 354 | **hist2d_kwargs) 355 | 356 | if truths is not None: 357 | if truths[i] is not None and truths[j] is not None: 358 | ax.plot(truths[j], truths[i], "s", color=truth_color) 359 | if truths[j] is not None: 360 | ax.axvline(truths[j], color=truth_color) 361 | if truths[i] is not None: 362 | ax.axhline(truths[i], color=truth_color) 363 | 364 | ax.xaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower")) 365 | ax.yaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower")) 366 | 367 | if i < K - 1: 368 | ax.set_xticklabels([]) 369 | else: 370 | [l.set_rotation(45) for l in ax.get_xticklabels()] 371 | if labels is not None: 372 | ax.set_xlabel(labels[j], **label_kwargs) 373 | ax.xaxis.set_label_coords(0.5, -0.3) 374 | 375 | # use MathText for axes ticks 376 | ax.xaxis.set_major_formatter( 377 | ScalarFormatter(useMathText=use_math_text)) 378 | 379 | if j > 0: 380 | ax.set_yticklabels([]) 381 | else: 382 | [l.set_rotation(45) for l in ax.get_yticklabels()] 383 | if labels is not None: 384 | ax.set_ylabel(labels[i], **label_kwargs) 385 | ax.yaxis.set_label_coords(-0.3, 0.5) 386 | 387 | # use MathText for axes ticks 388 | ax.yaxis.set_major_formatter( 389 | ScalarFormatter(useMathText=use_math_text)) 390 | 391 | return fig 392 | 393 | 394 | def quantile(x, q, weights=None): 395 | """ 396 | Like numpy.percentile, but: 397 | 398 | * Values of q are quantiles [0., 1.] rather than percentiles [0., 100.] 399 | * scalar q not supported (q must be iterable) 400 | * optional weights on x 401 | 402 | """ 403 | if weights is None: 404 | return np.percentile(x, [100. * qi for qi in q]) 405 | else: 406 | idx = np.argsort(x) 407 | xsorted = x[idx] 408 | cdf = np.add.accumulate(weights[idx]) 409 | cdf /= cdf[-1] 410 | return np.interp(q, cdf, xsorted).tolist() 411 | 412 | 413 | def hist2d(x, y, bins=20, range=None, weights=None, levels=None, smooth=None, 414 | ax=None, color=None, plot_datapoints=True, plot_density=True, 415 | plot_contours=True, no_fill_contours=False, fill_contours=False, 416 | contour_kwargs=None, contourf_kwargs=None, data_kwargs=None, 417 | **kwargs): 418 | """ 419 | Plot a 2-D histogram of samples. 420 | 421 | Parameters 422 | ---------- 423 | x, y : array_like (nsamples,) 424 | The samples. 425 | 426 | levels : array_like 427 | The contour levels to draw. 428 | 429 | ax : matplotlib.Axes (optional) 430 | A axes instance on which to add the 2-D histogram. 431 | 432 | plot_datapoints : bool (optional) 433 | Draw the individual data points. 434 | 435 | plot_density : bool (optional) 436 | Draw the density colormap. 437 | 438 | plot_contours : bool (optional) 439 | Draw the contours. 440 | 441 | no_fill_contours : bool (optional) 442 | Add no filling at all to the contours (unlike setting 443 | ``fill_contours=False``, which still adds a white fill at the densest 444 | points). 445 | 446 | fill_contours : bool (optional) 447 | Fill the contours. 448 | 449 | contour_kwargs : dict (optional) 450 | Any additional keyword arguments to pass to the `contour` method. 451 | 452 | contourf_kwargs : dict (optional) 453 | Any additional keyword arguments to pass to the `contourf` method. 454 | 455 | data_kwargs : dict (optional) 456 | Any additional keyword arguments to pass to the `plot` method when 457 | adding the individual data points. 458 | """ 459 | if ax is None: 460 | ax = pl.gca() 461 | 462 | # Set the default range based on the data range if not provided. 463 | if range is None: 464 | if "extent" in kwargs: 465 | logging.warn("Deprecated keyword argument 'extent'. " 466 | "Use 'range' instead.") 467 | range = kwargs["extent"] 468 | else: 469 | range = [[x.min(), x.max()], [y.min(), y.max()]] 470 | 471 | # Set up the default plotting arguments. 472 | if color is None: 473 | color = "k" 474 | 475 | # Choose the default "sigma" contour levels. 476 | if levels is None: 477 | levels = 1.0 - np.exp(-0.5 * np.arange(0.5, 2.1, 0.5) ** 2) 478 | 479 | # This is the color map for the density plot, over-plotted to indicate the 480 | # density of the points near the center. 481 | density_cmap = LinearSegmentedColormap.from_list( 482 | "density_cmap", [color, (1, 1, 1, 0)]) 483 | 484 | # This color map is used to hide the points at the high density areas. 485 | white_cmap = LinearSegmentedColormap.from_list( 486 | "white_cmap", [(1, 1, 1), (1, 1, 1)], N=2) 487 | 488 | # This "color map" is the list of colors for the contour levels if the 489 | # contours are filled. 490 | rgba_color = colorConverter.to_rgba(color) 491 | contour_cmap = [rgba_color] + [list(rgba_color) for l in levels] 492 | for i, l in enumerate(levels): 493 | contour_cmap[i+1][-1] *= float(len(levels) - i) / (len(levels)+1) 494 | 495 | # We'll make the 2D histogram to directly estimate the density. 496 | try: 497 | H, X, Y = np.histogram2d(x.flatten(), y.flatten(), bins=bins, 498 | range=range, weights=weights) 499 | except ValueError: 500 | raise ValueError("It looks like at least one of your sample columns " 501 | "have no dynamic range. You could try using the " 502 | "'range' argument.") 503 | 504 | if smooth is not None: 505 | if gaussian_filter is None: 506 | raise ImportError("Please install scipy for smoothing") 507 | H = gaussian_filter(H, smooth) 508 | 509 | # Compute the density levels. 510 | Hflat = H.flatten() 511 | inds = np.argsort(Hflat)[::-1] 512 | Hflat = Hflat[inds] 513 | sm = np.cumsum(Hflat) 514 | sm /= sm[-1] 515 | V = np.empty(len(levels)) 516 | for i, v0 in enumerate(levels): 517 | try: 518 | V[i] = Hflat[sm <= v0][-1] 519 | except: 520 | V[i] = Hflat[0] 521 | 522 | # Compute the bin centers. 523 | X1, Y1 = 0.5 * (X[1:] + X[:-1]), 0.5 * (Y[1:] + Y[:-1]) 524 | 525 | # Extend the array for the sake of the contours at the plot edges. 526 | H2 = H.min() + np.zeros((H.shape[0] + 4, H.shape[1] + 4)) 527 | H2[2:-2, 2:-2] = H 528 | H2[2:-2, 1] = H[:, 0] 529 | H2[2:-2, -2] = H[:, -1] 530 | H2[1, 2:-2] = H[0] 531 | H2[-2, 2:-2] = H[-1] 532 | H2[1, 1] = H[0, 0] 533 | H2[1, -2] = H[0, -1] 534 | H2[-2, 1] = H[-1, 0] 535 | H2[-2, -2] = H[-1, -1] 536 | X2 = np.concatenate([ 537 | X1[0] + np.array([-2, -1]) * np.diff(X1[:2]), 538 | X1, 539 | X1[-1] + np.array([1, 2]) * np.diff(X1[-2:]), 540 | ]) 541 | Y2 = np.concatenate([ 542 | Y1[0] + np.array([-2, -1]) * np.diff(Y1[:2]), 543 | Y1, 544 | Y1[-1] + np.array([1, 2]) * np.diff(Y1[-2:]), 545 | ]) 546 | 547 | if plot_datapoints: 548 | if data_kwargs is None: 549 | data_kwargs = dict() 550 | data_kwargs["color"] = data_kwargs.get("color", color) 551 | data_kwargs["ms"] = data_kwargs.get("ms", 2.0) 552 | data_kwargs["mec"] = data_kwargs.get("mec", "none") 553 | data_kwargs["alpha"] = data_kwargs.get("alpha", 0.1) 554 | ax.plot(x, y, "o", zorder=-1, rasterized=True, **data_kwargs) 555 | 556 | # Plot the base fill to hide the densest data points. 557 | if (plot_contours or plot_density) and not no_fill_contours: 558 | ax.contourf(X2, Y2, H2.T, [V[-1], H.max()], 559 | cmap=white_cmap, antialiased=False) 560 | 561 | if plot_contours and fill_contours: 562 | if contourf_kwargs is None: 563 | contourf_kwargs = dict() 564 | contourf_kwargs["colors"] = contourf_kwargs.get("colors", contour_cmap) 565 | contourf_kwargs["antialiased"] = contourf_kwargs.get("antialiased", 566 | False) 567 | ax.contourf(X2, Y2, H2.T, np.concatenate([[H.max()], V, [0]]), 568 | **contourf_kwargs) 569 | 570 | # Plot the density map. This can't be plotted at the same time as the 571 | # contour fills. 572 | elif plot_density: 573 | ax.pcolor(X, Y, H.max() - H.T, cmap=density_cmap) 574 | 575 | # Plot the contour edge colors. 576 | if plot_contours: 577 | if contour_kwargs is None: 578 | contour_kwargs = dict() 579 | contour_kwargs["colors"] = contour_kwargs.get("colors", color) 580 | ax.contour(X2, Y2, H2.T, V, **contour_kwargs) 581 | 582 | ax.set_xlim(range[0]) 583 | ax.set_ylim(range[1]) 584 | --------------------------------------------------------------------------------