minires.plotting

  1"""Convenient plot functions for reservoir model."""
  2
  3from typing import TYPE_CHECKING, Any, Optional, cast
  4
  5import matplotlib as mpl
  6import matplotlib.pyplot as plt
  7import numpy as np
  8from matplotlib import animation
  9from matplotlib.colors import BoundaryNorm
 10from matplotlib.ticker import MaxNLocator, MultipleLocator
 11
 12if TYPE_CHECKING:
 13    from minires import ResSim
 14
 15
 16class Plot2D:
 17    """Plots specialized for 2D fields.
 18
 19    This mixin is not standalone but reads grid and well attributes of the
 20    `ResSim` it gets composed into, rather than re-declaring (avoids stale).
 21    """
 22
 23    def plt_field(
 24        self: "ResSim",
 25        ax: Any,
 26        Z: np.ndarray,
 27        style: str = "default",
 28        wells: Any = True,
 29        argmax: bool = False,
 30        colorbar: Any = True,
 31        labels: bool = True,
 32        grid: bool = False,
 33        finalize: bool = True,
 34        cellwise: bool = False,
 35        **kwargs,
 36    ) -> Any:
 37        """Contour-plot of the (flat) unravelled field `Z`.
 38
 39        `kwargs` falls back to `styles[style]`, which falls back to `styles['defaults']`.
 40
 41        Inactive cells (`ResSim.active`) are masked out (left blank). Note that
 42        `contourf` interpolates between cell *centres*, so it also leaves blank the
 43        half-cell margins around them (as around the domain). `cellwise=True`
 44        instead paints each cell flat (`pcolormesh`), which is exact about the
 45        cells -- the shape of a mask, a fault, the resolution -- at the cost of the
 46        smoothness; the colour levels (hence the colorbar) are the same either way.
 47
 48        `wells` marks the completions (ref `well_scatter`): `True`, `"color"`
 49        (the producers coloured as in `plt_production`), or a `dict` of options
 50        for `well_scatter` -- where `exclude=[names]` hides the wells so named
 51        (an aquifer's ring of contacts, say; ref `minires.wells.aquifer_WI`),
 52        and `color` may also be given *per well*, as `{name: color}` (a list of
 53        colors, one per completion, for a multi-completion well; the wells left
 54        out keep the default of their sign). Both take glob patterns
 55        (`"Prd*"`), being matched by `minires.wells.Wells.which`.
 56        """
 57        # Populate kwargs with fallback style
 58        kwargs = {**styles["default"], **styles[style], **kwargs}
 59        # Pop from kwargs. Remainder goes to countourf
 60        ax.set(**axprops(kwargs))
 61        cticks = kwargs.pop("cticks")
 62
 63        # Why extent=(0, Lx, 0, Ly), rather than merely changing ticks?
 64        # set_aspect("equal") and mouse hovering (reporting x,y).
 65        if "rel" in coord_type:
 66            Lx, Ly = 1, 1
 67        elif "abs" in coord_type:
 68            Lx, Ly = self.Lx, self.Ly
 69        elif "ind" in coord_type:
 70            Lx, Ly = self.Nx, self.Ny
 71        else:
 72            raise ValueError(f"Unsupported coord_type: {coord_type}")
 73
 74        # Apply transform
 75        Z = np.asarray(Z)
 76        Z = kwargs.pop("transf")(Z)
 77
 78        # Need to transpose coz orientation is model.shape==(Nx, Ny),
 79        # while contour() displays the same orientation as array printing.
 80        Z = Z.reshape(self.shape).T
 81        # Mask the inactive cells (both `contourf` and `pcolormesh` leave them blank)
 82        Z = np.ma.masked_where(~self.active.T, Z)
 83
 84        # Did we bother to specify set_over/set_under/set_bad ?
 85        has_out_of_range = getattr(kwargs["cmap"], "_rgba_over", None) is not None
 86        extend = "both" if has_out_of_range else "neither"
 87
 88        if cellwise:
 89            # Discretize the colours by the same levels as `contourf` would
 90            levels = kwargs.pop("levels")
 91            locator = kwargs.pop("locator")
 92            if np.ndim(levels) == 0:
 93                locator = locator or MaxNLocator(levels + 1)
 94                levels = locator.tick_values(Z.min(), Z.max())
 95            cmap = plt.get_cmap(kwargs.pop("cmap"))
 96            norm = BoundaryNorm(levels, cmap.N, extend=extend)
 97            collections = ax.pcolormesh(
 98                np.linspace(0, Lx, self.Nx + 1),
 99                np.linspace(0, Ly, self.Ny + 1),
100                Z,
101                cmap=cmap,
102                norm=norm,
103                **kwargs,
104            )
105        else:
106            # Unlike `ax.imshow(Z[::-1])`, `contourf` does not simply fill pixels/cells
107            # (but it does provide nice interpolation!) so there will be whitespace on
108            # the margins. No fix is needed, and anyway it would not be trivial/fast,
109            # ref https://github.com/matplotlib/basemap/issues/406 .
110            collections = ax.contourf(
111                Z,
112                **kwargs,
113                # origin=None,  # ⇒ NB: falsely stretches the field!!!
114                origin="lower",
115                extent=(0, Lx, 0, Ly),
116                extend=extend,
117            )
118
119        # Contourf does not plot (at all) the bad regions. "Fake it" by facecolor
120        if has_out_of_range:
121            ax.set_facecolor(getattr(kwargs["cmap"], "_rgba_bad", "w"))
122
123        # Grid (reflecting the model grid)
124        # NB: If not showing grid, then don't locate ticks on grid, because they're
125        #     generally uglier that mpl's default/automatic tick location. But, it
126        #     should be safe to go with 'g' format instead of 'f'.
127        ax.xaxis.set_major_formatter("{x:g}")
128        ax.yaxis.set_major_formatter("{x:g}")
129        ax.tick_params(which="minor", length=0, color="r")
130        ax.tick_params(which="major", width=1.5, direction="in")
131        if grid:
132            n1 = 10
133            xStep = 1 + self.Nx // n1
134            yStep = 1 + self.Ny // n1
135            ax.xaxis.set_major_locator(MultipleLocator(self.hx * xStep))
136            ax.yaxis.set_major_locator(MultipleLocator(self.hy * yStep))
137            ax.xaxis.set_minor_locator(MultipleLocator(self.hx))
138            ax.yaxis.set_minor_locator(MultipleLocator(self.hy))
139            ax.grid(True, which="both")
140
141        # Axis lims
142        ax.set_xlim((0, Lx))
143        ax.set_ylim((0, Ly))
144        ax.set_aspect("equal")
145        # Axis labels
146        if labels:
147            if "abs" in coord_type:
148                ax.set_xlabel("x")
149                ax.set_ylabel("y")
150            else:
151                ax.set_xlabel(f"x ({coord_type})")
152                ax.set_ylabel(f"y ({coord_type})")
153
154        # Add well markers, grouped (and numbered) by the sign of their rates,
155        # ref `minires.wells.Wells.signs`. The producers come first, so
156        # that their numbers and colors match those of `plt_production`.
157        if wells and self.wells.nComp:
158            sgn = self.wells.signs
159            # Label the completions by their well's name, if there are any
160            names = None
161            if self.wells.names is not None:
162                # Expand to one per completion (ref `minires.wells.Wells.group`)
163                group = self.wells.group
164                group = np.arange(self.wells.nComp) if group is None else group
165                names = np.asarray(self.wells.names)[group]
166            producer_colors = wells == "color"
167            if producer_colors:
168                # Colors matching `plt_production` of the producers
169                wells = cast(
170                    dict, {"color": [f"C{i}" for i in range(int(np.sum(sgn < 0)))]}
171                )
172            elif wells in [True, 1]:
173                wells = {}
174            else:
175                wells = dict(wells)  # NB: copy -- popped from below
176            # Colors given per well (`{name: color}`) become one per completion
177            if isinstance(by_name := wells.get("color"), dict):
178                cols = np.full(self.wells.nComp, None, object)
179                for nm, c in by_name.items():
180                    which = self.wells.which(nm)
181                    assert len(which), f"No well named {nm}."
182                    cs = len(which) * [c] if isinstance(c, str) else list(c)
183                    assert len(cs) == len(which), f"{nm} has {len(which)} completions."
184                    for i, ci in zip(which, cs):
185                        cols[i] = ci
186                wells["color"] = cols
187            # Hide the wells named by `exclude` (an aquifer's ring of contacts, say)
188            shown = np.ones(self.wells.nComp, bool)
189            if (exclude := wells.pop("exclude", None)) is not None:
190                shown[self.wells.which(*np.ravel(exclude))] = False
191            for s in [-1, +1, 0]:
192                sel = (sgn == s) & shown
193                if np.any(sel):  # NB: skip, lest empty artists upset the layout
194                    kws = dict(wells)  # NB: copy -- the labels are per sign
195                    if names is not None:
196                        kws.setdefault("text", names[sel])
197                    if isinstance(kws.get("color"), np.ndarray):
198                        kws["color"] = list(kws["color"][sel])
199                    self.well_scatter(ax, self.wells.xy[sel], s, **kws)
200                if producer_colors:
201                    wells.pop("color", None)  # producers only
202
203        # Add argmax marker
204        if argmax:
205            idx = Z.T.argmax()  # reverse above transpose
206            xy = self.ind2xy(idx)
207            for c, ms in zip(["b", "r", "y"], [10, 6, 3]):
208                ax.plot(*xy, "o", c=c, ms=ms, label="max", zorder=98)
209
210        # Add colorbar
211        if colorbar:
212            if isinstance(colorbar, type(ax)):
213                cax = dict(cax=colorbar)
214            else:
215                cax = dict(ax=ax, shrink=0.8)
216            ax.figure.colorbar(collections, **cax, ticks=cticks)
217
218        tight_show(ax.figure, finalize)
219        return collections
220
221    def plt_faces(
222        self: "ResSim", ax: Any, xy: Any, faces: str = "WESN", **kws: Any
223    ) -> Any:
224        """Stroke the boundary faces of the cells at `xy`, onto a `plt_field`.
225
226        I.e. the faces `minires.wells.boundary_faces` finds -- the contact
227        of an aquifer, say (ref `minires.wells.aquifer_WI`), whose ring of
228        well markers this replaces (hide those with `wells=dict(exclude=...)`).
229        `kws` go to the `LineCollection` (`color`, `lw`, ...).
230        """
231        from matplotlib.collections import LineCollection
232
233        from minires.wells import boundary_faces
234
235        xy = np.asarray(xy, float).reshape((-1, 2))
236        xy = self.sub2xy(*self.xy2sub(*xy.T)).T  # snap to the cell centres
237        hx, hy = self.hx / 2, self.hy / 2
238        # Each face as a segment from the centre: W, E, S, N
239        ends = np.array([[[-hx, -hy], [-hx, +hy]], [[+hx, -hy], [+hx, +hy]],
240                         [[-hx, -hy], [+hx, -hy]], [[-hx, +hy], [+hx, +hy]]])  # fmt: skip
241        segments = (xy[:, None, None, :] + ends)[boundary_faces(self, xy, faces)]
242        # fmt: off
243        if   "rel" in coord_type: s = 1/self.Lx, 1/self.Ly                     # noqa
244        elif "abs" in coord_type: s = 1, 1                                     # noqa
245        elif "ind" in coord_type: s = self.Nx/self.Lx, self.Ny/self.Ly         # noqa
246        else: raise ValueError("Unsupported coordinate type: %s" % coord_type) # noqa
247        # fmt: on
248        opts: dict = dict(color="C0", lw=4, capstyle="projecting", zorder=1.4) | kws
249        lc = LineCollection(segments * s, **opts)
250        ax.add_collection(lc)
251        return lc
252
253    def well_scatter(
254        self: "ResSim",
255        ax: Any,
256        ww: np.ndarray,
257        sgn: int = 1,
258        text: Any = None,
259        color: Any = None,  # e.g. "k", or a list of colors (one per completion)
260        size: float = 1,
261        dot: Any = True,
262    ) -> Any:
263        """Scatter-plot the wells of `ww` onto a `Plot2D.plt_field`.
264
265        The marker reflects `sgn`: injector (`+1`), producer (`-1`),
266        or neutral (`0`, i.e. of undecided sign, ref `minires.wells.Wells.signs`).
267
268        The label, `text`, is either one string for all of them, one *per* well
269        of `ww` (a list), or `False` for none.
270
271        `color` is one colour for all of them, or a list of one per row of `ww`
272        (wherein `None` means the default for `sgn`). `dot` is the little dot at
273        the marker's centre: `False` for none, or a colour (`True` means red).
274
275        .. note:: The default labels are indices *within* `ww`, not global ones.
276
277            I.e. with `text=None`, and since `plt_field` calls this once per
278            sign, the producers are numbered as `plt_production` numbers them --
279            separately from the injectors, and not as in the unified
280            `minires.wells.Wells.xy`. But `plt_field` supplies the names of
281            `minires.wells.Wells.names`, if the wells have been given any.
282        """
283        # Well coordinates
284        ww = self.sub2xy(*self.xy2sub(*ww.T)).T
285        # NB: make sure ww array data is not overwritten (avoid in-place)
286        # fmt: off
287        if   "rel" in coord_type: s = 1/self.Lx, 1/self.Ly                     # noqa
288        elif "abs" in coord_type: s = 1, 1                                     # noqa
289        elif "ind" in coord_type: s = self.Nx/self.Lx, self.Ny/self.Ly         # noqa
290        else: raise ValueError("Unsupported coordinate type: %s" % coord_type) # noqa
291        # fmt: on
292        ww = ww * s
293
294        # Style
295        if sgn > 0:
296            c = "#036dc9"
297            ec = "gray"
298            d = "w"
299            m = "v"
300        elif sgn < 0:
301            c = "#c83b03"
302            ec = "gray"
303            d = "w"
304            m = "^"
305        else:
306            c = "lightgray"
307            ec = "gray"
308            d = "k"
309            m = "o"
310
311        if isinstance(color, str):
312            c = color
313        elif color is not None:
314            # NB: a list may leave wells to the default, `c`, by `None`
315            c = [c if x is None else x for x in color]
316
317        # Markers
318        if dot is not False:
319            ax.plot(*ww.T, ".", c="k" if dot is True else dot, ms=3, clip_on=False)
320        sh = ax.scatter(
321            *ww.T,
322            s=(size * 26) ** 2,
323            c=c,
324            marker=m,
325            ec=ec,
326            clip_on=False,
327            zorder=1.5,  # required on Jupypter
328        )
329
330        # Text labels
331        if text is not False:
332            if text is None:
333                labels: Any = range(len(ww))
334            elif isinstance(text, str):
335                labels = len(ww) * [text]
336            else:
337                labels = text
338            for lbl, w in zip(labels, ww):
339                if sgn < 0:
340                    w[1] -= 0.01
341                ax.text(
342                    *w[:2],
343                    lbl,
344                    color=d,
345                    fontsize=size * 12,
346                    ha="center",
347                    va="center",
348                )
349
350        return sh
351
352    def plt_production(
353        self: "ResSim",
354        ax: Any,
355        production: np.ndarray,
356        obs: np.ndarray | None = None,
357        legend_outside: bool = True,
358        finalize: bool = True,
359        labels: Any = None,
360    ) -> list:
361        """Production time series. Multiple wells in 1 axes => not ensemble compat.
362
363        The curves are labelled by their index in `production`, unless `labels`
364        (e.g. a selection of `minires.wells.Wells.names`) says otherwise.
365        """
366        hh = []
367        tt = 1 + np.arange(len(production))
368        for i, p in enumerate(1 - production.T):
369            hh += ax.plot(tt, p, "-", label=i if labels is None else labels[i])
370
371        if obs is not None:
372            for i, y in enumerate(1 - obs.T):
373                ax.plot(tt, y, "*", c=hh[i].get_color())
374
375        # Add legend
376        if legend_outside:
377            kws = dict(
378                bbox_to_anchor=(1, 1),
379                loc="upper left",
380                ncol=1 + len(production.T) // 10,
381            )
382        else:
383            kws = dict(loc="lower left")
384        ax.legend(title="Well" if labels is not None else "Well #.", **kws)
385
386        ax.set_title("Oil saturation in producers")
387        ax.set_xlabel("Time index")
388        # ax.set_ylim(-0.01, 1.01)
389        ax.axhline(0, c="xkcd:light grey", ls="--", zorder=1.8)
390        ax.axhline(1, c="xkcd:light grey", ls="--", zorder=1.8)
391
392        tight_show(ax.figure, finalize)
393        return hh
394
395    def anim(
396        self: "ResSim",
397        S: np.ndarray,
398        prod: np.ndarray | None = None,
399        title: str = "",
400        figsize: tuple = (10, 3.5),
401        pause: int = 200,
402        animate: bool = True,
403        **kwargs,
404    ) -> Any:
405        """Animate the saturation field, alongside the oil saturation at the producers.
406
407        `S` is the saturation trajectory that `sim` returns, `(nSteps + 1, Nxy)`.
408        `prod`, the water saturation at the producers over time, `(nSteps, nPrd)`,
409        defaults to `S[1:]` at the cells of the wells with negative rates (ref
410        `minires.wells.Wells.signs`) -- `S[0]` being the initial condition, before
411        any production. Pass it to show another series there, e.g. the observed one.
412
413        The returned `Animation` displays itself as a cell's output in a notebook,
414        where the figure is closed, so that its final frame is not also displayed
415        as a static figure. `kwargs` go to `plt_field`.
416        """
417        if prod is None:
418            xy = self.wells.xy[self.wells.signs < 0]
419            prod = S[1:, self.xy2ind(*xy.T)]
420
421        # Create figure and axes
422        title = "Animation" + ("-- " + title if title else "")
423        fig, (ax1, ax2) = plt.subplots(
424            num=title, clear=True, ncols=2, figsize=figsize,
425            gridspec_kw=dict(width_ratios=(2, 3)),
426        )
427        fig.suptitle(title)  # coz animation never (any backend) displays title
428        # Saturations
429        kwargs.update(wells="color", colorbar=True, finalize=False)
430        ax2.cc = self.plt_field(ax2, S[-1], "oil", **kwargs)
431        # Production
432        hh = self.plt_production(ax1, prod, legend_outside=False, finalize=False)
433        fig.tight_layout()
434
435        if animate:
436            tt = np.arange(len(S))
437
438            def update_fig(iT):
439                # Update field.
440                # NB: since mpl 3.8 the `ContourSet` is itself an artist
441                # (it no longer has a `.collections` attribute).
442                try:
443                    ax2.cc.remove()
444                except ValueError:
445                    pass  # occurs when re-running script
446                kwargs.update(wells=False, colorbar=False)
447                ax2.cc = self.plt_field(ax2, S[iT], "oil", **kwargs)
448
449                # Update production lines
450                if iT >= 1:
451                    for h, p in zip(hh, prod.T):
452                        h.set_data(tt[1 : 1 + iT], 1 - p[:iT])
453
454            # `repeat=False`: else the busy/idle indicator of a notebook keeps
455            # flashing (despite `%%capture`, even after clearing the output).
456            ani = Animation(fig, update_fig, len(tt), blit=False, interval=pause,
457                            repeat=False)
458            if is_inline():
459                plt.close(fig)  # else displayed (statically) beside the animation
460            return ani
461
462
463class Animation(animation.FuncAnimation):
464    """A `FuncAnimation` that a notebook displays as its JavaScript player.
465
466    matplotlib's own `_repr_html_` consults `rcParams["animation.html"]`, whose
467    default (`"none"`) displays nothing but the `repr`; this one always renders
468    `to_jshtml`. Returned by `Plot2D.anim`.
469    """
470
471    def _repr_html_(self) -> str:
472        return self.to_jshtml()
473
474def show(block: Optional[bool] = None) -> None:
475    """Display the figures, whether run as script, in IPython, or in a notebook.
476
477    Reasons why a plain `plt.show()` does not suffice:
478
479    - In a script (`python script.py`) it does the right thing (blocks, thereby
480      running the GUI event loop) *only* if interactive mode (`plt.ion()`) is off.
481    - In IPython (e.g. `%run script.py`) the figures stay unpainted until the GUI
482      event loop gets to run, which requires the "input hook" installed by the
483      `%matplotlib` magic (which is why they only appear upon `ctrl-d`/exit).
484      Blocking here would be wrong: it would freeze the prompt.
485    - In Jupyter/Colab with the `inline` backend there is no event loop at all:
486      the (static) figures are rendered by `plt.show()` itself.
487    """
488    ip = _get_ipython()
489
490    if ip is not None and not is_inline():
491        # I.e. the `%matplotlib` magic: activate the event loop ("input hook")
492        # of the current backend. No-op if already active.
493        ip.run_line_magic("matplotlib", "")
494
495    if block is None:
496        # Run the event loop ourselves only if nobody else will.
497        block = not is_inline() and not _ipython_will_prompt(ip)
498
499    plt.show(block=block)
500
501
502# ╔══════════╗
503# ║ Settings ║
504# ╚══════════╝
505
506coord_type = "absolute"
507"""Define scaling of `Plot2D.plt_field` axes.
508- "relative": `(0, 1)  x (0, 1)`
509- "absolute": `(0, Lx) x (0, Ly)`
510- "index"   : `(0, Ny) x (0, Ny)`
511"""
512
513# Colormap for saturation
514# cm_ow = "managua"
515lin_cm = mpl.colors.LinearSegmentedColormap.from_list
516cm_ow = lin_cm("", [(0, "#1d9e97"), (0.3, "#b2e0dc"), (1, "#f48974")])
517# cOil, cWater = "red", "blue"        # Plain
518# cOil, cWater = "#d8345f", "#01a9b4" # Pastel/neon
519# cOil, cWater = "#e58a8a", "#086972" # Pastel
520# ccnvrt = lambda c: np.array(mpl.colors.colorConverter.to_rgb(c))
521# cMiddle = .3*ccnvrt(cWater) + .7*ccnvrt(cOil)
522# cm_ow = lin_cm("", [cWater, cMiddle, cOil])
523
524styles: dict = dict(
525    default=dict(
526        title="",
527        transf=lambda x: x,
528        cmap="viridis",
529        levels=10,
530        cticks=None,
531        # Note that providing vmin/vmax (and not a levels list) to mpl
532        # yields prettier colobar ticks, but destorys the consistency
533        # of the colorbars from one figure to another.
534        locator=None,
535    ),
536    oil=dict(
537        title="Oil saturation",
538        transf=lambda x: 1 - x,
539        cmap=cm_ow,
540        levels=np.linspace(0 - 1e-7, 1 + 1e-7, 20),
541        cticks=np.linspace(0, 1, 6),
542    ),
543)
544"""Default `Plot2D.plt_field` plot styling values."""
545
546
547# ╔═══════════════════╗
548# ║ Front-end helpers ║
549# ╚═══════════════════╝
550
551# NB: the two helpers below are all that is left of the `mpl-tools` dependency
552# (dropped so that the package need not concern itself with front-ends).
553# Its `freshfig` is not among them: `plt.subplots(num=..., clear=True)` is the
554# same thing, bar the screen placement, which belongs in a user's own tooling.
555
556# The axes properties that `plt_field` (and `styles`) let the caller set.
557_AXPROPS = ["title", "facecolor", "aspect"] + [
558    xyz + p for xyz in "xyz" for p in ["label", "ticks", "scale", "lim"]
559]
560
561
562def axprops(kwargs: dict) -> dict:
563    """Pop (and return) the `kwargs` that belong to `ax.set`, ref `_AXPROPS`."""
564    return {p: kwargs.pop(p) for p in _AXPROPS if p in kwargs}
565
566
567def is_inline() -> bool:
568    """Whether the active `mpl` backend is an IPython/Jupyter `inline` one.
569
570    NB: not the negation of "interactive" -- `Agg`, `PDF`, ... are neither.
571    """
572    return "inline" in mpl.get_backend()
573
574
575def tight_show(figure: Any, enabled: bool) -> None:
576    """`tight_layout`, then `plt.show()` -- except under marimo, whose backend renders
577    every `plt.show()` as console media, on top of the figure the cell outputs.
578    (The `inline` one displays *and closes* the figure, which callbacks of
579    `ipywidgets` -- with no cell end to flush the figures at -- rely on.)
580    """
581    if enabled:
582        figure.tight_layout()
583        if "marimo" not in mpl.get_backend():
584            plt.show()
585
586
587def _get_ipython() -> Any:
588    """Return the active IPython shell, or `None` (also if IPython isn't installed)."""
589    try:
590        from IPython import get_ipython
591    except ImportError:
592        return None
593    return get_ipython()
594
595
596def _ipython_will_prompt(ip: Any) -> bool:
597    """Whether IPython returns to its (event-loop running) prompt after this script.
598
599    False for `ipython script.py` (no `-i`), which exits immediately.
600    """
601    if ip is None:
602        return False
603    return bool(getattr(ip.parent, "interact", True))