diff --git a/src/corner/arviz_corner.py b/src/corner/arviz_corner.py index 9f972af..55e70c2 100644 --- a/src/corner/arviz_corner.py +++ b/src/corner/arviz_corner.py @@ -125,7 +125,7 @@ def arviz_corner( if isinstance(data, np.ndarray): is_np = True if data.ndim == 1: - data = data[None, :, :] + data = data[None, :, None] elif data.ndim == 2: data = data[None, :, :] elif data.ndim != 3: diff --git a/src/corner/core.py b/src/corner/core.py index 0b160d2..f8ed1f0 100644 --- a/src/corner/core.py +++ b/src/corner/core.py @@ -845,6 +845,7 @@ def overplot_lines(fig, xs, reverse=False, **kwargs): """ K = len(xs) axes, _ = _get_fig_axes(fig, K) + axes = np.atleast_2d(axes) if reverse: for k1 in range(K): if xs[k1] is not None: diff --git a/tests/test_corner.py b/tests/test_corner.py index e44f1bb..549cdcc 100644 --- a/tests/test_corner.py +++ b/tests/test_corner.py @@ -461,3 +461,50 @@ def test_range_fig_arg(): def test_1d_fig_argument(): fig = _run_corner(ndim=1, seed=0) _run_corner(ndim=1, seed=1, fig=fig) + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("value", [None, 0.25]) +@pytest.mark.parametrize("input_1d", [False, True]) +def test_overplot_lines_1d(reverse, value, input_1d): + data = np.linspace(-1, 1, 101) + if not input_1d: + data = data[:, None] + fig = corner.corner(data, reverse=reverse) + ax = fig.axes[0] + initial_lines = len(ax.lines) + + corner.overplot_lines( + fig, [value], reverse=reverse, color="red", linewidth=2.5 + ) + + if value is None: + assert len(ax.lines) == initial_lines + else: + assert len(ax.lines) == initial_lines + 1 + line = ax.lines[-1] + np.testing.assert_array_equal(line.get_xdata(), [value, value]) + np.testing.assert_array_equal(line.get_ydata(), [0, 1]) + assert line.get_color() == "red" + assert line.get_linewidth() == 2.5 + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("use_fig", [False, True]) +@pytest.mark.parametrize("input_1d", [False, True]) +def test_truths_1d(reverse, use_fig, input_1d): + data = np.linspace(-1, 1, 101) + if not input_1d: + data = data[:, None] + fig = pl.figure() if use_fig else None + fig = corner.corner( + data, + truths=[0.25], + truth_color="red", + reverse=reverse, + fig=fig, + ) + assert len(fig.axes) == 1 + line = fig.axes[0].lines[-1] + np.testing.assert_array_equal(line.get_xdata(), [0.25, 0.25]) + assert line.get_color() == "red"