make more flexible options settable to layer plot

src/swe2hs/visualization.py

@@ -37,7 +37,13 @@ def layer_plot(
     color_variable='layer_densities',
     vmin=50.,
     vmax=550.,
-    cbar_label='Density [kg/m$^{3}$]'):
+    cmap='cool',
+    cbar_label='Density [kg/m$^{3}$]',
+    top_line_kwargs=None,
+    layer_line_kwargs=None,
+    cax_kwargs=None,
+):
+
     """
     Plot the layers in the modeled snowpack with optional coloring relating to
     one of the state variables 'layer_heights', 'layer_densities', or 
@@ -57,9 +63,19 @@ def layer_plot(
         Minimum for the colormap. The default is 50.
     vmax : float, optional
         Maximum for the colormap. The default is 550.
+    cmap : str or :class:`matplotlib.colors.Colormap`
+        Colormap for the layer coloring.
     cbar_label : str, optional
         Label of the colorbar. The default is 'Density [kg/m$^{3}$]'.
-
+    top_line_kwargs : dict or None, optional
+        Keyword arguments for the line at top of the snowpack. The default is
+        None which sets reasonable defaults.
+    layer_line_kwargs : dict or None, optional
+        Keyword arguments for the lines between the snow layers. The default is
+        None which sets reasonable defaults.
+    cax_kwargs : dict or None, optional
+        Keyword arguments passed to :func:`mpl.colorbar.make_axes`. The default
+        is None.
     Returns
     -------
     None.
@@ -77,7 +93,7 @@ def layer_plot(
         d_small = d_large.where(d_large['layer_heights'] != 0, drop=True)
         d_small['layer_tops'] = d_small['layer_heights'].cumsum(dim='layers', skipna=True)
         color_norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
-        color_mapper = mpl.cm.ScalarMappable(norm=color_norm, cmap=mpl.cm.cool)
+        color_mapper = mpl.cm.ScalarMappable(norm=color_norm, cmap=cmap)
         for l in d_small['layers'].to_index():
             for t in d_small['time'].to_index():
                 if not np.isnan(d_small['layer_heights'].sel(layers=l, time=t)):
@@ -99,14 +115,22 @@ def layer_plot(
     # there would be nans and and the layertops would be connected from arbitrary
     # locations
     d_large['layer_tops'] = d_large['layer_heights'].cumsum(dim='layers', skipna=True)
-    # thin layerborders in black
-    ax.plot(d_large['time'].to_index(), d_large['layer_tops'].to_numpy().T, lw=0.5, c='k')
-    # thicker top of snowcover in black
-    ax.plot(d_large['time'].to_index(), d_large['hs'].to_pandas(), c='k', lw=2, label='HS modeled')
+    # draw layerborders
+    l_kwargs = {'lw': 0.5, 'c': 'k'}
+    if layer_line_kwargs is not None:
+        l_kwargs.update(layer_line_kwargs)
+    ax.plot(d_large['time'].to_index(), d_large['layer_tops'].to_numpy().T, **l_kwargs)
+    # draw line at top of the snowcover
+    t_kwargs = {'lw': 2, 'c': 'k', 'label': 'HS modeled'}
+    if top_line_kwargs is not None:
+        t_kwargs.update(top_line_kwargs)
+    ax.plot(d_large['time'].to_index(), d_large['hs'].to_pandas(), **t_kwargs)
 
     if color_variable is not None:
-        # colorbar for density:
-        cax, _ = mpl.colorbar.make_axes(ax)
+        c_kwargs = {'pad': 0.01}
+        if cax_kwargs is not None:
+            c_kwargs.update(cax_kwargs)
+        cax, _ = mpl.colorbar.make_axes(ax, **c_kwargs)
         cbar = plt.colorbar(color_mapper, cax, ax)
         cbar.set_label(cbar_label)
     return None