Some plotting improvements

pastas/plotting/modelcompare.py

@@ -3,7 +3,7 @@
 
 from itertools import combinations
 from logging import getLogger
-from typing import List, Optional, Tuple
+from typing import List, Optional
 from warnings import warn
 
 import matplotlib.pyplot as plt
@@ -92,9 +92,9 @@ def __init__(self, models: List[Model], names: Optional[List[str]] = None) -> No
     def initialize_figure(
         self,
         mosaic: Optional[List[List[str]]] = None,
-        figsize: Tuple[int, int] = (10, 8),
         cmap: str = "tab10",
         return_ax: bool = False,
+        **fig_kwargs,
     ) -> None:
         """initialize a custom figure based on a mosaic.
 
@@ -112,7 +112,7 @@ def initialize_figure(
 
         self.cmap = plt.get_cmap(cmap)
 
-        figure, axes = plt.subplot_mosaic(mosaic, figsize=figsize)
+        figure, axes = plt.subplot_mosaic(mosaic, **fig_kwargs)
         if return_ax:
             return axes
 
@@ -123,9 +123,9 @@ def initialize_figure(
     def initialize_adjust_height_figure(
         self,
         mosaic: Optional[List[List[str]]] = None,
-        figsize: Tuple[int] = (10, 8),
         cmap: str = "tab10",
         smdict: Optional[dict] = None,
+        **fig_kwargs,
     ) -> None:
         """initialize subplots based on a mosaic with equal vertical scales.
 
@@ -226,8 +226,8 @@ def initialize_adjust_height_figure(
         self.mosaic = mosaic
         fig, axes = plt.subplot_mosaic(
             self.mosaic,
-            figsize=figsize,
             gridspec_kw=dict(height_ratios=heights_list),
+            **fig_kwargs,
         )
 
         self.figure = fig
@@ -577,10 +577,14 @@ def plot_response(
                         continue
                     if response == "step":
                         kwargs = {}
+                        p = None
                         if ml.stressmodels[smn].rfunc is not None:
                             if isinstance(ml.stressmodels[smn].rfunc, HantushWellModel):
                                 kwargs = {"warn": False}
-                        step = ml.get_step_response(smn, add_0=True, **kwargs)
+                                p = ml.stressmodels[smn].get_parameters(
+                                    model=ml, istress=0
+                                )
+                        step = ml.get_step_response(smn, p=p, add_0=True, **kwargs)
                         if step is None:
                             continue
                         if self.axes is None:
@@ -598,10 +602,14 @@ def plot_response(
                             )
                     elif response == "block":
                         kwargs = {}
+                        p = None
                         if ml.stressmodels[smn].rfunc is not None:
                             if isinstance(ml.stressmodels[smn].rfunc, HantushWellModel):
                                 kwargs = {"warn": False}
-                        block = ml.get_block_response(smn, **kwargs)
+                                p = ml.stressmodels[smn].get_parameters(
+                                    model=ml, istress=0
+                                )
+                        block = ml.get_block_response(smn, p=p, add_0=True, **kwargs)
                         if block is None:
                             continue
                         if self.axes is None:
@@ -897,11 +905,11 @@ def plot(
         smdict: Optional[dict] = None,
         normalized: bool = False,
         param_selection: Optional[list] = None,
-        figsize: Optional[tuple] = (10, 8),
         grid: bool = True,
         legend: bool = True,
         adjust_height: bool = False,
         legend_kwargs: Optional[dict] = None,
+        **fig_kwargs,
     ) -> None:
         """plot the models in a comparison plot.
 
@@ -921,8 +929,6 @@ def plot(
             zero, by default False.
         param_selection : list, optional
             list of (sub)strings of which parameters to show in table, by default None.
-        figsize : tuple, optional
-            figure size, by default (10, 8).
         grid : bool, optional
             grid in each subplot, by default True.
         legend : bool, optional
@@ -935,10 +941,12 @@ def plot(
             pass legend keyword arguments to plots.
         """
         self.adjust_height = adjust_height
+        if "figsize" not in fig_kwargs:
+            fig_kwargs["figsize"] = (10, 8)
         if self.axes is None and not self.adjust_height:
-            self.initialize_figure(figsize=figsize)
+            self.initialize_figure(**fig_kwargs)
         if self.axes is None and self.adjust_height:
-            self.initialize_adjust_height_figure(smdict=smdict, figsize=figsize)
+            self.initialize_adjust_height_figure(smdict=smdict, **fig_kwargs)
 
         # sim
         _ = self.plot_oseries()

pastas/plotting/modelplots.py

@@ -1,4 +1,4 @@
-"""This module contains all the plotting methods for Pastas Models."""
+"""This module contains plotting methods for Pastas Models."""
 
 import logging
 
@@ -267,11 +267,15 @@ def results(
 
             # plot the step response
             rkwargs = {}
+            p = None
             if self.ml.stressmodels[sm_name].rfunc is not None:
                 if isinstance(self.ml.stressmodels[sm_name].rfunc, HantushWellModel):
                     rkwargs = {"warn": False}
+                    p = self.ml.stressmodels[sm_name].get_parameters(
+                        model=self.ml, istress=0
+                    )
             response = self.ml._get_response(
-                block_or_step=block_or_step, name=sm_name, add_0=True, **rkwargs
+                block_or_step=block_or_step, name=sm_name, p=p, add_0=True, **rkwargs
             )
 
             if response is not None:
@@ -907,7 +911,7 @@ def custom_sort(t):
                         contributions.append((name, h))
 
                         # plot step responses for each well, scaled with distance
-                        p = sml.get_parameters(istress=istress)
+                        p = sml.get_parameters(model=self.ml, istress=istress)
                         step = self.ml.get_step_response(sm, p=p)
                         ax_step.plot(step.index, step, c=stackcolors[name], label=name)
                         # recalculate y-limits step response axes

pastas/plotting/plotly.py

@@ -1,3 +1,6 @@
+"""This module contains interactive plots for Pastas models.
+"""
+
 import numpy as np
 import pandas as pd
 import plotly.graph_objs as go
@@ -10,6 +13,7 @@
     _table_formatter_params,
     _table_formatter_stderr,
 )
+from pastas.rfunc import HantushWellModel
 from pastas.stats import acf
 
 
@@ -246,8 +250,16 @@ def results(self, tmin=None, tmax=None):
             traces.append(trace_c)
 
             # response
+            rkwargs = {}
+            p = None
+            if self._model.stressmodels[c.name].rfunc is not None:
+                if isinstance(self._model.stressmodels[c.name].rfunc, HantushWellModel):
+                    rkwargs = {"warn": False}
+                    p = self._model.stressmodels[c.name].get_parameters(
+                        model=self._model, istress=0
+                    )
             response = self._model._get_response(
-                block_or_step="step", name=c.name, add_0=True
+                block_or_step="step", name=c.name, p=p, add_0=True, **rkwargs
             )
             trace_r = go.Scatter(
                 x=response.index,

pastas/plotting/plots.py

@@ -1,4 +1,4 @@
-"""This module contains all the plotting methods in Pastas."""
+"""This module contains plotting methods for Pastas."""
 
 import logging
 from typing import Dict, List, Optional, Tuple, Union

pastas/plotting/plotutil.py

@@ -1,3 +1,6 @@
+"""This module contains utility functions for plotting.
+"""
+
 from typing import List, Union
 
 import numpy as np

pastas/timeseries_utils.py

@@ -267,7 +267,7 @@ def timestep_weighted_resample(s: Series, index: Index, fast: bool = False) -> S
         # calculate the cumulative sum
         s_new = s_new.cumsum()
 
-        # add NaNs at none-existing values in series at index
+        # add NaNs at non-existing values in series at index
         s_new = s_new.combine_first(Series(np.NaN, index))
 
         # interpolate these NaN's, only keep values at index
@@ -276,7 +276,7 @@ def timestep_weighted_resample(s: Series, index: Index, fast: bool = False) -> S
         # calculate the diff again (inverse of cumsum)
         s_new = s_new.diff()
 
-        # devide by the timestep again
+        # divide by the timestep again
         s_new = s_new / _get_dt_array(s_new.index)
 
         # set values after the end of the original series to NaN