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