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