Repository: scipp/plopp

Issue:
ax arg in 2d plots is broken
```python
fig, ax = plt.subplots(1, 1)
pp.plot(sc.DataArray(sc.zeros(sizes={'x': 10, 'y': 10})), ax=ax)
```
raises
```

File ~/Software/mambaforge/envs/uncertain/lib/python3.8/site-packages/plopp/graphics/mesh.py:100, in Mesh.__init__(self, canvas, data, shading, rasterized, **kwargs)
     94 self._mesh = None
     96 x, y, z = _from_data_array_to_pcolormesh(data=self._data.data,
     97                                          coords=bin_edge_coords,
     98                                          dim_1d=self._dim_1d,
     99                                          dim_2d=self._dim_2d)
--> 100 self._mesh = self._ax.pcolormesh(x.values,
    101                                  y.values,
    102                                  z.values,
    103                                  shading=shading,
    104                                  rasterized=rasterized,
    105                                  **kwargs)
    106 self._mesh.set_array(None)
    108 for xy, var in string_labels.items():

TypeError: inner() got multiple values for argument 'ax'
```

Discussion:

Patch:
diff --git a/src/plopp/graphics/canvas.py b/src/plopp/graphics/canvas.py
index c106cb6..c6ee56d 100644
--- a/src/plopp/graphics/canvas.py
+++ b/src/plopp/graphics/canvas.py
@@ -7,6 +7,7 @@ from .utils import fig_to_bytes, silent_mpl_figure
 
 import matplotlib.pyplot as plt
 from matplotlib.collections import QuadMesh
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 import numpy as np
 import scipp as sc
 from typing import Dict, Literal, Tuple
@@ -27,6 +28,9 @@ class Canvas:
     ax:
         If supplied, use these axes to create the figure. If none are supplied, the
         canvas will create its own axes.
+    cax:
+        If supplied, use these axes for the colorbar. If none are supplied, and a
+        colorbar is required, the canvas will create its own axes.
     figsize:
         The width and height of the figure, in inches.
     title:
@@ -48,6 +52,7 @@ class Canvas:
 
     def __init__(self,
                  ax: plt.Axes = None,
+                 cax: plt.Axes = None,
                  figsize: Tuple[float, float] = None,
                  title: str = None,
                  grid: bool = False,
@@ -59,37 +64,28 @@ class Canvas:
 
         self.fig = None
         self.ax = ax
+        self.cax = cax
         self._user_vmin = vmin
         self._user_vmax = vmax
         self._scale = {} if scale is None else scale
         self.xunit = None
         self.yunit = None
+        self._own_axes = False
 
         if self.ax is None:
-            if figsize is None:
-                figsize = (6, 4)
+            self._own_axes = True
             with silent_mpl_figure():
-                self.fig = plt.figure(figsize=figsize, dpi=96)
-            left = 0.11
-            right = 0.9
-            bottom = 0.11
-            top = 0.95
-            if cbar:
-                cbar_width = 0.03
-                cbar_gap = 0.04
-                self.ax = self.fig.add_axes(
-                    [left, bottom, right - left - cbar_width - cbar_gap, top - bottom])
-                self.cax = self.fig.add_axes(
-                    [right - cbar_width, bottom, cbar_width, top - bottom])
-            else:
-                self.ax = self.fig.add_axes([left, bottom, right - left, top - bottom])
-                self.cax = None
+                self.fig, self.ax = plt.subplots(figsize=figsize)
             if hasattr(self.fig.canvas, "on_widget_constructed"):
                 self.fig.canvas.toolbar_visible = False
                 self.fig.canvas.header_visible = False
         else:
             self.fig = self.ax.get_figure()
 
+        if cbar and self.cax is None:
+            divider = make_axes_locatable(self.ax)
+            self.cax = divider.append_axes("right", "4%", pad="5%")
+
         self.ax.set_aspect(aspect)
         self.ax.set_title(title)
         self.ax.grid(grid)
@@ -271,3 +267,10 @@ class Canvas:
         self._ymin = np.inf
         self._ymax = np.NINF
         self.autoscale()
+
+    def fit_to_page(self):
+        """
+        Trim the margins around the figure.
+        """
+        if self._own_axes:
+            self.fig.tight_layout()
diff --git a/src/plopp/graphics/colormapper.py b/src/plopp/graphics/colormapper.py
index 88e927a..4b1c68d 100644
--- a/src/plopp/graphics/colormapper.py
+++ b/src/plopp/graphics/colormapper.py
@@ -109,10 +109,10 @@ class ColorMapper:
             height_inches = (figsize[1] / dpi) if figsize is not None else 6
             with silent_mpl_figure():
                 fig = plt.figure(figsize=(height_inches * 0.2, height_inches))
-            self.cax = fig.add_axes([0.05, 0.02, 0.2, 1.0])
+            self.cax = fig.add_axes([0.05, 0.02, 0.2, 0.98])
 
         self.colorbar = ColorbarBase(self.cax, cmap=self.cmap, norm=self.normalizer)
-        self.cax.yaxis.set_label_coords(-1.1, 0.5)
+        self.cax.yaxis.set_label_coords(-0.9, 0.5)
 
     def __setitem__(self, key, val):
         self.artists[key] = val
diff --git a/src/plopp/graphics/fig1d.py b/src/plopp/graphics/fig1d.py
index 01247cc..51dcafa 100644
--- a/src/plopp/graphics/fig1d.py
+++ b/src/plopp/graphics/fig1d.py
@@ -7,7 +7,7 @@ from .canvas import Canvas
 from .line import Line
 
 import scipp as sc
-from typing import Dict, Literal, Tuple
+from typing import Any, Dict, Literal, Tuple
 
 
 class Figure1d(BaseFig):
@@ -43,6 +43,9 @@ class Figure1d(BaseFig):
         The figure title.
     figsize:
         The width and height of the figure, in inches.
+    ax:
+        If supplied, use these axes to create the figure. If none are supplied, the
+        canvas will create its own axes.
     **kwargs:
         All other kwargs are forwarded to Matplotlib:
 
@@ -63,6 +66,7 @@ class Figure1d(BaseFig):
                  crop: Dict[str, Dict[str, sc.Variable]] = None,
                  title: str = None,
                  figsize: Tuple[float, float] = None,
+                 ax: Any = None,
                  **kwargs):
 
         super().__init__(*nodes)
@@ -76,6 +80,7 @@ class Figure1d(BaseFig):
                              grid=grid,
                              figsize=figsize,
                              title=title,
+                             ax=ax,
                              vmin=vmin,
                              vmax=vmax)
         self.canvas.yscale = norm
@@ -84,6 +89,7 @@ class Figure1d(BaseFig):
         self.canvas.autoscale()
         if crop is not None:
             self.crop(**crop)
+        self.canvas.fit_to_page()
 
     def update(self, new_values: sc.DataArray, key: str, draw: bool = True):
         """
diff --git a/src/plopp/graphics/fig2d.py b/src/plopp/graphics/fig2d.py
index 4aa4f22..029bf03 100644
--- a/src/plopp/graphics/fig2d.py
+++ b/src/plopp/graphics/fig2d.py
@@ -8,7 +8,7 @@ from .colormapper import ColorMapper
 from .mesh import Mesh
 
 import scipp as sc
-from typing import Dict, Literal, Tuple
+from typing import Any, Dict, Literal, Tuple
 
 
 class Figure2d(BaseFig):
@@ -49,6 +49,12 @@ class Figure2d(BaseFig):
         The figure title.
     figsize:
         The width and height of the figure, in inches.
+    ax:
+        If supplied, use these axes to create the figure. If none are supplied, the
+        canvas will create its own axes.
+    cax:
+        If supplied, use these axes for the colorbar. If none are supplied, and a
+        colorbar is required, the canvas will create its own axes.
     **kwargs:
         All other kwargs are forwarded to the Mesh artist.
     """
@@ -67,6 +73,8 @@ class Figure2d(BaseFig):
                  cbar: bool = True,
                  title: str = None,
                  figsize: Tuple[float, float] = None,
+                 ax: Any = None,
+                 cax: Any = None,
                  **kwargs):
 
         super().__init__(*nodes)
@@ -77,7 +85,9 @@ class Figure2d(BaseFig):
                              aspect=aspect,
                              grid=grid,
                              title=title,
-                             figsize=figsize)
+                             figsize=figsize,
+                             ax=ax,
+                             cax=cax)
         self.colormapper = ColorMapper(cmap=cmap,
                                        mask_cmap=mask_cmap,
                                        norm=norm,
@@ -89,6 +99,7 @@ class Figure2d(BaseFig):
         self.canvas.autoscale()
         if crop is not None:
             self.crop(**crop)
+        self.canvas.fit_to_page()
 
     def update(self, new_values: sc.DataArray, key: str, draw: bool = True):
         """


Test patch:
diff --git a/tests/graphics/fig1d_test.py b/tests/graphics/fig1d_test.py
index d584c51..2439621 100644
--- a/tests/graphics/fig1d_test.py
+++ b/tests/graphics/fig1d_test.py
@@ -5,6 +5,7 @@ from plopp.data import data_array
 from plopp.graphics.fig1d import Figure1d
 from plopp.graphics.line import Line
 from plopp import input_node
+import matplotlib.pyplot as plt
 import numpy as np
 import scipp as sc
 import pytest
@@ -161,3 +162,11 @@ def test_vmin_vmax():
                    vmin=sc.scalar(-0.5, unit='m/s'),
                    vmax=sc.scalar(0.68, unit='m/s'))
     assert np.allclose(fig.canvas.ax.get_ylim(), [-0.5, 0.68])
+
+
+def test_ax():
+    fig, ax = plt.subplots()
+    assert len(ax.lines) == 0
+    da = data_array(ndim=1)
+    _ = Figure1d(input_node(da), ax=ax)
+    assert len(ax.lines) > 0
diff --git a/tests/graphics/fig2d_test.py b/tests/graphics/fig2d_test.py
index 1c09e29..78b5f70 100644
--- a/tests/graphics/fig2d_test.py
+++ b/tests/graphics/fig2d_test.py
@@ -5,6 +5,7 @@ from plopp.data import data_array
 from plopp.graphics.fig2d import Figure2d
 from plopp.graphics.mesh import Mesh
 from plopp import input_node
+import matplotlib.pyplot as plt
 import numpy as np
 import scipp as sc
 import pytest
@@ -163,3 +164,20 @@ def test_grid():
     da = data_array(ndim=2)
     fig = Figure2d(input_node(da), grid=True)
     assert fig.canvas.ax.xaxis.get_gridlines()[0].get_visible()
+
+
+def test_ax():
+    fig, ax = plt.subplots()
+    assert len(ax.collections) == 0
+    da = data_array(ndim=2)
+    _ = Figure2d(input_node(da), ax=ax)
+    assert len(ax.collections) == 1
+
+
+def test_cax():
+    fig, ax = plt.subplots()
+    cax = fig.add_axes([0.9, 0.02, 0.05, 0.98])
+    assert len(cax.collections) == 0
+    da = data_array(ndim=2)
+    _ = Figure2d(input_node(da), ax=ax, cax=cax)
+    assert len(cax.collections) > 0