minires.geostat

  1"""Gaussian random fields, for synthetic permeabilities and priors.
  2
  3A stationary, isotropic Gaussian field over the given points -- typically the cell
  4centres of a `minires.Grid2D.mesh` -- defined by its *variogram*, $ γ(h) $, i.e.
  5the covariance $ c(h) = 1 - γ(h) $ as a function of the distance $ h $. Here the
  6Gaussian one, `variogram_gauss`, whose one parameter of note is the range, `r`
  7(the correlation drops to $ e^{-3} $ at $ h = r $, for the default `a`).
  8`cov_sqrt` forms the covariance matrix of the field over the points and returns
  9its (upper) Cholesky factor, $ \\mathbf{C} = \\mathbf{C}_{1/2}^T \\mathbf{C}_{1/2} $,
 10so that `gaussian_fields` samples the field as $ \\mathbf{z} \\, \\mathbf{C}_{1/2} $,
 11with $ \\mathbf{z} $ standard normal -- `N` of them at a time, each of unit
 12variance and zero mean.
 13
 14>>> from minires import Grid2D
 15>>> grid = Grid2D(Lx=1, Ly=1, Nx=8, Ny=8)
 16>>> C12 = cov_sqrt(grid.mesh, r=.3)
 17>>> C12.shape
 18(64, 64)
 19>>> Cov = 1 - variogram_gauss(dist_euclid(vectorize(*grid.mesh)), r=.3)
 20>>> bool(np.allclose(C12.T @ C12, Cov))
 21True
 22>>> fields = gaussian_fields(grid.mesh, N=3, r=.3, rng=np.random.default_rng(0))
 23>>> fields.shape
 24(3, 64)
 25
 26The covariance matrix is *dense*, `(nPts, nPts)`, and factorized outright (with a
 27tiny jitter on the diagonal, the Gaussian covariance being nearly singular), so
 28this is for the toy grids of this package -- a few thousand cells at most --
 29and not a geostatistics library. The fields being Gaussian, they serve as
 30*log*-permeabilities (`K = exp(logK)`), or as the Gaussian prior of an ensemble
 31method (ref `examples.random_fields`, and the tutorials of
 32[HistoryMatching](https://github.com/patnr/HistoryMatching), which also use
 33`cov_sqrt` as the change of variables of a variational method).
 34`funm_psd` is a general matrix function for such (positive semi-definite)
 35covariances, e.g. the symmetric square root.
 36"""
 37
 38import numpy as np
 39import scipy.linalg as sla
 40
 41
 42def variogram_gauss(xx: np.ndarray, r: float, n: float = 0, a: float = 1 / 3) -> np.ndarray:
 43    """Gaussian variogram at the distances `xx`.
 44
 45    Params: range `r`, nugget `n`, and `a`, which scales the range
 46    (for the default, the covariance $ 1 - γ $ is $ e^{-3} $ at `r`).
 47    Ref: <https://en.wikipedia.org/wiki/Variogram#Variogram_models>
 48
 49    >>> xx = np.array([0, 1, 2])
 50    >>> variogram_gauss(xx, 1, n=0.1, a=1)
 51    array([0.        , 0.6689085 , 0.98351593])
 52    """
 53    # Gauss
 54    gamma = 1 - np.exp(-(xx**2) / r**2 / a)
 55    # Sill (=1)
 56    gamma *= 1 - n
 57    # Nugget
 58    gamma[xx != 0] += n
 59    return gamma
 60
 61
 62def vectorize(*XYZ: np.ndarray) -> np.ndarray:
 63    """Reshape coordinate arrays (`nDim` of them, of equal `shape`) to `(nPts, nDim)`."""
 64    return np.stack(XYZ).reshape((len(XYZ), -1)).T
 65
 66
 67def dist_euclid(X: np.ndarray) -> np.ndarray:
 68    """Pairwise distances between the rows of `X`, like `squareform(pdist(X))`."""
 69    diff = X[:, None, :] - X
 70    d2 = np.sum(diff**2, axis=-1)
 71    return np.sqrt(d2)
 72
 73
 74def funm_psd(C, fun, rk=None, rtol=1e-8, sym_square=True, **kwargs):
 75    """Matrix function evaluation for a positive semi-definite matrix.
 76
 77    Adapted from the `scipy.linalg.funm` doc.
 78
 79    Note: small `rk` and `driver="evx"` should be faster,
 80    but in my (simple but hopefully relevant) trials
 81    sticking with the default "evr" and `rk=None` is usually faster.
 82    Alternatively, could try iterative algorithms from `sparse.linalg`.
 83
 84    Example
 85    -------
 86    >>> def sqrtm(C):
 87    ...     return funm_psd(C, np.sqrt)
 88    """
 89    # EVD -- possibly truncated (for speed)
 90    idx = [max(0, len(C) - rk), len(C) - 1] if rk else None
 91    ews, V = sla.eigh(C, subset_by_index=idx, **kwargs)
 92
 93    # Truncate (for stability) -- NB: ordering low-->high!
 94    nNull = sum(ews <= rtol * ews.max())
 95    ews = ews[nNull:]
 96    V = V[:, nNull:]
 97
 98    # Apply
 99    ews = fun(ews)
100
101    # Reconstruct
102    funC = V * ews
103    if sym_square:
104        # Optional, since not necessary e.g. for cholesky factors (for sampling)
105        # But, without it the sqrtm is beholden to the "vagaries" of lapack version?
106        funC = funC @ V.T
107    return funC
108
109
110def cov_sqrt(pts: tuple, r: float = 0.2) -> np.ndarray:
111    """Cholesky factor, `C12`, of the covariance (Gaussian variogram) over `pts`.
112
113    `pts` are coordinate arrays, e.g. `Grid2D.mesh` (or `(xx,)` in 1D).
114    The factor is upper triangular, such that `Cov == C12.T @ C12`,
115    and so `z @ C12` has covariance `Cov` if `z` is standard normal.
116    """
117    dists = dist_euclid(vectorize(*pts))
118    Cov = 1 - variogram_gauss(dists, r)
119    # C12    = sla.sqrtm(Cov).real.T                      # unstable for n >≈ 20
120    # C12    = funm_psd(Cov, np.sqrt, sym_square=True).T  # too slow for n >= 50^2
121    C12 = sla.cholesky(Cov + 1e-10 * np.eye(len(Cov)))
122    return C12
123
124
125def gaussian_fields(pts: tuple, N: int = 1, r: float = 0.2, rng=None) -> np.ndarray:
126    """Sample `N` Gaussian random fields over `pts`, as `(N, nPts)`.
127
128    Each of zero mean and unit variance, with the covariance of `cov_sqrt`
129    (Gaussian variogram of range `r`). The draws come from `rng`
130    (a `numpy.random.Generator`), or the global `numpy.random` state if `None`.
131    """
132    rng = np.random if rng is None else rng
133    C12 = cov_sqrt(pts, r)
134    fields = rng.standard_normal((N, len(C12))) @ C12
135    return fields