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