Let plot_probe accept xlims, ylims and zlims independently

src/probeinterface/plotting.py

@@ -243,7 +243,12 @@ def on_press(event):
             ax.text(x, y, txt, ha="center", va="center", clip_on=True)
 
     if xlims is None or ylims is None or (zlims is None and probe.ndim == 3):
-        xlims, ylims, zlims = get_auto_lims(probe)
+        # Only fill in the limits the caller left unset. Rebinding all three would
+        # discard a caller-supplied xlims whenever ylims was omitted, and vice versa.
+        auto_xlims, auto_ylims, auto_zlims = get_auto_lims(probe)
+        xlims = auto_xlims if xlims is None else xlims
+        ylims = auto_ylims if ylims is None else ylims
+        zlims = auto_zlims if zlims is None else zlims
 
     ax.set_xlim(*xlims)
     ax.set_ylim(*ylims)

tests/test_plotting.py

@@ -1,6 +1,7 @@
 from probeinterface import Probe, ProbeGroup
 from probeinterface import generate_dummy_probe, generate_dummy_probe_group
 from probeinterface.plotting import plot_probe, plot_probegroup
+from probeinterface.utils import get_auto_lims
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -43,6 +44,57 @@ def test_plot_probegroup():
     plot_probegroup(probegroup_3d, same_axes=True)
 
 
+def test_plot_probe_partial_lims():
+    """Passing only one of xlims/ylims/zlims must not discard the ones that were given."""
+    probe = generate_dummy_probe()
+    auto_xlims, auto_ylims, _ = get_auto_lims(probe)
+
+    # x only: xlims is honoured, ylims falls back to auto
+    _, ax = plt.subplots()
+    plot_probe(probe, ax=ax, xlims=(-11, 22))
+    assert ax.get_xlim() == (-11, 22)
+    assert ax.get_ylim() == pytest.approx(auto_ylims)
+
+    # y only: ylims is honoured, xlims falls back to auto
+    _, ax = plt.subplots()
+    plot_probe(probe, ax=ax, ylims=(-33, 44))
+    assert ax.get_ylim() == (-33, 44)
+    assert ax.get_xlim() == pytest.approx(auto_xlims)
+
+    # both given: neither is touched
+    _, ax = plt.subplots()
+    plot_probe(probe, ax=ax, xlims=(-11, 22), ylims=(-33, 44))
+    assert ax.get_xlim() == (-11, 22)
+    assert ax.get_ylim() == (-33, 44)
+
+    # neither given: both are auto
+    _, ax = plt.subplots()
+    plot_probe(probe, ax=ax)
+    assert ax.get_xlim() == pytest.approx(auto_xlims)
+    assert ax.get_ylim() == pytest.approx(auto_ylims)
+
+
+def test_plot_probe_partial_lims_3d():
+    """In 3D, omitting zlims must not discard a supplied xlims or ylims."""
+    probe_3d = generate_dummy_probe().to_3d(axes="xz")
+    auto_xlims, auto_ylims, auto_zlims = get_auto_lims(probe_3d)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection="3d")
+    plot_probe(probe_3d, ax=ax, xlims=(-11, 22), ylims=(-33, 44))
+    assert ax.get_xlim() == (-11, 22)
+    assert ax.get_ylim() == (-33, 44)
+    assert ax.get_zlim() == pytest.approx(auto_zlims)
+
+    # zlims only: x and y fall back to auto
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection="3d")
+    plot_probe(probe_3d, ax=ax, zlims=(-55, 66))
+    assert ax.get_zlim() == (-55, 66)
+    assert ax.get_xlim() == pytest.approx(auto_xlims)
+    assert ax.get_ylim() == pytest.approx(auto_ylims)
+
+
 def test_plot_probe_two_side():
     probe = Probe()
     probe.set_contacts(