diff --git a/.gitignore b/.gitignore index 60f9f7c..0bb4783 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,6 @@ dist build test_figures .coverage -!corner/tests/baseline_images/*/*.png +!tests/baseline_images/*/*.png .pytest_cache corner_version.py diff --git a/src/corner/core.py b/src/corner/core.py index 53fca73..f8f9305 100644 --- a/src/corner/core.py +++ b/src/corner/core.py @@ -41,6 +41,7 @@ def corner_impl( truth_color="#4682b4", scale_hist=False, quantiles=None, + title_quantiles=None, verbose=False, fig=None, max_n_ticks=5, @@ -62,6 +63,20 @@ def corner_impl( if titles is None: titles = labels + # deal with title quantiles so they much quantiles unless desired otherwise + if title_quantiles is None: + if len(quantiles) > 0: + title_quantiles = quantiles + else: + # a default for when quantiles not supplied. + title_quantiles = [0.16, 0.5, 0.84] + + if show_titles and len(title_quantiles) != 3: + raise ValueError( + "'title_quantiles' must contain exactly three values; " + "pass a length-3 list or array using the 'title_quantiles' argument" + ) + # Deal with 1D sample lists. xs = _parse_input(xs) assert xs.shape[0] <= xs.shape[1], ( @@ -212,15 +227,15 @@ def corner_impl( if title_fmt is not None: # Compute the quantiles for the title. This might redo # unneeded computation but who cares. - q_16, q_50, q_84 = quantile( - x, [0.16, 0.5, 0.84], weights=weights + q_lo, q_mid, q_hi = quantile( + x, title_quantiles, weights=weights ) - q_m, q_p = q_50 - q_16, q_84 - q_50 + q_m, q_p = q_mid - q_lo, q_hi - q_mid # Format the quantile display. fmt = "{{0:{0}}}".format(title_fmt).format title = r"${{{0}}}_{{-{1}}}^{{+{2}}}$" - title = title.format(fmt(q_50), fmt(q_m), fmt(q_p)) + title = title.format(fmt(q_mid), fmt(q_m), fmt(q_p)) # Add in the column name if it's given. if titles is not None: diff --git a/src/corner/corner.py b/src/corner/corner.py index 108af85..da058a3 100644 --- a/src/corner/corner.py +++ b/src/corner/corner.py @@ -29,6 +29,7 @@ def corner( label_kwargs=None, titles=None, show_titles=False, + title_quantiles=None, title_fmt=".2f", title_kwargs=None, truths=None, @@ -135,6 +136,11 @@ def corner( Displays a title above each 1-D histogram showing the 0.5 quantile with the upper and lower errors supplied by the quantiles argument. + title_quantiles : iterable + A list of 3 fractional quantiles to show as the the upper and lower + errors. If `None` (default), inherit the values from quantiles, unless + quantiles is `None`, in which case it defaults to [0.16,0.5,0.84] + title_fmt : string The format string for the quantiles given in titles. If you explicitly set ``show_titles=True`` and ``title_fmt=None``, the labels will be @@ -241,6 +247,7 @@ def corner( label_kwargs=label_kwargs, titles=titles, show_titles=show_titles, + title_quantiles=title_quantiles, title_fmt=title_fmt, title_kwargs=title_kwargs, truths=truths, @@ -271,6 +278,7 @@ def corner( label_kwargs=label_kwargs, titles=titles, show_titles=show_titles, + title_quantiles=title_quantiles, title_fmt=title_fmt, title_kwargs=title_kwargs, truths=truths, diff --git a/tests/baseline_images/test_corner/title_quantiles.png b/tests/baseline_images/test_corner/title_quantiles.png new file mode 100644 index 0000000..3847c29 Binary files /dev/null and b/tests/baseline_images/test_corner/title_quantiles.png differ diff --git a/tests/baseline_images/test_corner/title_quantiles_default.png b/tests/baseline_images/test_corner/title_quantiles_default.png new file mode 100644 index 0000000..b74178a Binary files /dev/null and b/tests/baseline_images/test_corner/title_quantiles_default.png differ diff --git a/tests/baseline_images/test_corner/title_quantiles_raises.png b/tests/baseline_images/test_corner/title_quantiles_raises.png new file mode 100644 index 0000000..3aebbc9 Binary files /dev/null and b/tests/baseline_images/test_corner/title_quantiles_raises.png differ diff --git a/tests/test_corner.py b/tests/test_corner.py index 5037d68..323212a 100644 --- a/tests/test_corner.py +++ b/tests/test_corner.py @@ -80,6 +80,39 @@ def test_quantiles(): _run_corner(quantiles=[0.16, 0.5, 0.84]) +@image_comparison( + baseline_images=["title_quantiles"], remove_text=False, extensions=["png"] +) +def test_title_quantiles(): + _run_corner( + quantiles=[0.16, 0.5, 0.84], + title_quantiles=[0.05, 0.5, 0.95], + show_titles=True, + ) + + +@image_comparison( + baseline_images=["title_quantiles_default"], + remove_text=False, + extensions=["png"], +) +def test_title_quantiles_default(): + _run_corner(quantiles=[0.16, 0.5, 0.84], show_titles=True) + + +@image_comparison( + baseline_images=["title_quantiles_raises"], + remove_text=False, + extensions=["png"], +) +def test_title_quantiles_raises(): + with pytest.raises(ValueError): + _run_corner(quantiles=[0.05, 0.16, 0.5, 0.84, 0.95], show_titles=True) + + # This one shouldn't raise since show_titles isn't provided + _run_corner(quantiles=[0.05, 0.16, 0.5, 0.84, 0.95]) + + @image_comparison( baseline_images=["color"], remove_text=True, extensions=["png"] )