Merge pull request #82 from Laetitia1402/PR-Laetitia

Add field option to plot_iso

Author: amorison

stagpy/field.py

@@ -259,18 +259,27 @@ def plot_scalar(step: Step, var: str, field: Optional[ndarray] = None,
     return fig, axis, surf, cbar
 
 
-def plot_iso(axis: Axes, step: Step, var: str, **extra: Any) -> None:
+def plot_iso(axis: Axes, step: Step, var: str,
+             field: Optional[ndarray] = None, **extra: Any) -> None:
     """Plot isocontours of scalar field.
 
     Args:
         axis: the :class:`matplotlib.axes.Axes` of an existing matplotlib
             figure where the isocontours should be plotted.
         step: a :class:`~stagpy._step.Step` of a StagyyData instance.
         var: the scalar field name.
+        field: if not None, it is plotted instead of step.fields[var].  This is
+            useful to plot a masked or rescaled array.  Note that if
+            conf.scaling.dimensional is True, this field will be scaled
+            accordingly.
         extra: options that will be passed on to
             :func:`matplotlib.axes.Axes.contour`.
     """
     xmesh, ymesh, fld, _ = get_meshes_fld(step, var)
+
+    if field is not None:
+        fld = field
+
     if conf.field.shift:
         fld = np.roll(fld, conf.field.shift, axis=0)
     extra_opts: Dict[str, Any] = dict(linewidths=1)