TPFA_ResSim.grid

  1"""Tools for working with model grid coordinates.
  2
  3Most functions here are barely in use, mostly serving as reference.
  4After all, it is surprisingly hard to remember
  5which direction and index is for x and which is for y.
  6
  7The index ordering is "C-style" (numpy default).
  8This choice means that `x` is the 1st coord., `y` is 2nd,
  9and is hardcoded in the reservoir simulator model code
 10(in what takes place **between** `np.ravel` and `np.reshape`,
 11both of which are configured to use row-major index ordering.
 12"F-style" (column-major) indexing implementation is perfectly possible,
 13but would imply an undue amount hassle).
 14Conveniently, it also means that `x` and `y` tend to occur in alphabetic order.
 15Thus, in printing a matrix of a field, the `x` coordinate corresponds to the row index.
 16By contrast, the plotting module depicts `x` from left to right, `y` from bottom to top.
 17"""
 18
 19from dataclasses import dataclass
 20from typing import NamedTuple, overload
 21
 22import numpy as np
 23import numpy.typing as npt
 24
 25
 26class Fluxes(NamedTuple):
 27    """Data container with dot (attr) access for cell face fluxes. Mimicks matlab code.
 28
 29    Positive is in the direction of increasing index. The fluxes through the
 30    *boundary* faces are `0`: the reservoir is closed (no-flow) all around.
 31    """
 32
 33    x: np.ndarray
 34    """Fluxes through the x-normal faces. Shape `(Nx+1, Ny)`."""
 35    y: np.ndarray
 36    """Fluxes through the y-normal faces. Shape `(Nx, Ny+1)`."""
 37
 38
 39@dataclass
 40class Grid2D:
 41    """Defines a 2D rectangular grid.
 42
 43    Example (2 x-nodes, 5 y-nodes):
 44    >>> grid = Grid2D(Lx=6, Ly=10, Nx=3, Ny=5)
 45
 46    The nodes are centered in the cells:
 47    >>> X, Y = grid.mesh
 48    >>> X
 49    array([[1., 1., 1., 1., 1.],
 50           [3., 3., 3., 3., 3.],
 51           [5., 5., 5., 5., 5.]])
 52
 53    You can compute cell boundaries (i.e. non-central nodes) by adding or subtracting
 54    `hx`/2 and `hy`/2 (i.e. you will miss either boundary at 0 or `Lx` or `Ly`).
 55
 56    Test of round-trip capability of grid mapping computations:
 57    >>> ij = (0, 4)
 58    >>> grid.xy2sub(X[ij], Y[ij]) == ij
 59    array([ True,  True])
 60
 61    >>> grid.sub2xy(*ij) == (X[ij], Y[ij])
 62    array([ True,  True])
 63    """
 64
 65    Lx: float = 1.0
 66    """Physical x-length of domain."""
 67    Ly: float = 1.0
 68    """Physical y-length of domain."""
 69    Nx: int = 32
 70    """Number of grid cells (and their centres) in x dir."""
 71    Ny: int = 32
 72    """Number of grid cells (and their centres) in y dir."""
 73
 74    @property
 75    def shape(self) -> tuple:
 76        """`(Nx, Ny)`"""
 77        return self.Nx, self.Ny
 78
 79    @property
 80    def size(self) -> int:
 81        """Total number of elements."""
 82        return int(np.prod(self.shape))
 83
 84    @property
 85    def domain(self) -> tuple:
 86        """`((0, 0), (Lx, Ly))`"""
 87        return ((0, 0), (self.Lx, self.Ly))
 88
 89    @property
 90    def Nxy(self) -> int:
 91        """`Nx` * `Ny`"""
 92        return int(np.prod(self.shape))
 93
 94    @property
 95    def hx(self) -> float:
 96        """x-length of cells"""
 97        return self.Lx / self.Nx
 98
 99    @property
100    def hy(self) -> float:
101        """y-length of cells"""
102        return self.Ly / self.Ny
103
104    @property
105    def h2(self) -> float:
106        """`hx` * `hy`"""
107        return self.hx * self.hy
108
109    @property
110    def mesh(self) -> tuple:
111        """Generate 2D coordinate grid of cell centres."""
112        xx = np.linspace(0, self.Lx, self.Nx, endpoint=False) + self.hx / 2
113        yy = np.linspace(0, self.Ly, self.Ny, endpoint=False) + self.hy / 2
114        return np.meshgrid(xx, yy, indexing="ij")
115
116    def sub2ind(
117        self, ix: int | np.ndarray, iy: int | np.ndarray
118    ) -> np.intp | np.ndarray:
119        """Convert index `(ix, iy)` to index in flattened array."""
120        idx = np.ravel_multi_index((ix, iy), self.shape)
121        return idx
122
123    def ind2sub(self, ind: int | np.intp | np.ndarray) -> np.ndarray:
124        """Inv. of `self.sub2ind`."""
125        ix, iy = np.unravel_index(ind, self.shape)
126        return np.asarray([ix, iy])
127
128    def xy2sub(self, x: npt.ArrayLike, y: npt.ArrayLike) -> np.ndarray:
129        """Convert physical coordinate tuple to `(ix, iy)`, ix ∈ {0, ..., Nx-1}.
130
131        .. warning:: `xy2sub` and `xy2ind` *round* to nearest cell center.
132
133            I.e. they are not injective.
134            The alternative would be to return some kind
135            of interpolation weights distributing `(x, y)` over multiple nodes.
136            This was tried and rejected (not worth it), ref "Missing features"
137            section of docs.
138        """
139        x = np.asarray(x)
140        y = np.asarray(y)
141        # Don't silence errors! Validation is useful in optimisation (e.g.)
142        assert np.all(x <= self.Lx)
143        assert np.all(y <= self.Ly)
144        # Set upper border values to slightly interior.
145        # NB: the nudge is *relative*, `Lx` being of whatever magnitude the
146        # units imply (ref `TPFA_ResSim.ResSim.cdarcy`).
147        x = x.clip(max=self.Lx * (1 - 1e-12))
148        y = y.clip(max=self.Ly * (1 - 1e-12))
149        ix = np.floor(x / self.Lx * self.Nx).astype(int)
150        iy = np.floor(y / self.Ly * self.Ny).astype(int)
151        return np.asarray([ix, iy])
152
153    # Overloaded so that the (array-valued) well lookups of `ResSim` type-check
154    @overload
155    def xy2ind(self, x: float, y: float) -> np.intp: ...
156    @overload
157    def xy2ind(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: ...
158
159    def xy2ind(self, x: npt.ArrayLike, y: npt.ArrayLike) -> np.intp | np.ndarray:
160        """Convert physical coord to flat indx."""
161        i, j = self.xy2sub(x, y)
162        return self.sub2ind(i, j)
163
164    def sub2xy(self, ix: npt.ArrayLike, iy: npt.ArrayLike) -> np.ndarray:
165        """Inverse of `self.xy2sub`."""
166        x = (np.asarray(ix) + 0.5) * self.hx
167        y = (np.asarray(iy) + 0.5) * self.hy
168        return np.asarray([x, y])
169
170    def ind2xy(self, ind: int | np.intp | np.ndarray) -> np.ndarray:
171        """Inverse of `self.xy2ind`."""
172        i, j = self.ind2sub(ind)
173        return self.sub2xy(i, j)
174
175    def _crossings(self, p0: np.ndarray, d: np.ndarray) -> np.ndarray:
176        """Parameters $t ∈ [0, 1]$ at which `p0 + t*d` crosses a cell boundary."""
177        ts = [0.0, 1.0]
178        for ax, h in enumerate([self.hx, self.hy]):
179            if d[ax] == 0:
180                continue
181            lo, hi = sorted([p0[ax], p0[ax] + d[ax]])
182            ts += [
183                (i * h - p0[ax]) / d[ax]
184                for i in range(int(np.floor(lo / h)) + 1, int(np.ceil(hi / h)))
185            ]
186        return np.unique(np.clip(ts, 0, 1))  # NB: `unique` also sorts