Proposed fix for title errorbars/quantiles bug (#193)

* initial attempt to fix hardwired title quantiles

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* updating tests

* adding baseline images

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Dan F-M <foreman.mackey@gmail.com>

.gitignore

@@ -7,6 +7,6 @@ dist
 build
 test_figures
 .coverage
-!corner/tests/baseline_images/*/*.png
+!tests/baseline_images/*/*.png
 .pytest_cache
 corner_version.py

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:

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,

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"]
 )