Deltares / xugrid

Xarray and unstructured grids
https://deltares.github.io/xugrid/
MIT License
58 stars 8 forks source link

Investigate more effective preconditioner for laplace_interpolate #193

Closed Huite closed 7 months ago

Huite commented 8 months ago

laplace_interpolate currently use spilu as a preconditioner when not using a direct solution method:

https://github.com/Deltares/xugrid/blob/75c87e62d4dbfc27bebe88c206be925102be5e4e/xugrid/ugrid/interpolate.py#L120

In my experience, this scheme which is provided by SuperLU doesn't work very well (no clear documentation, better off reading the source to find what options do). The modified ILU scheme works well in MODFLOW6 (which also solves the laplace equation with the conjugate gradient method). The implementation can be found here:

https://github.com/MODFLOW-USGS/modflow6/blob/9ca0a13f04cf17cf7195682d1bf7a48fd89defce/src/Solution/LinearMethods/ims8base.f90#L947

It's not that complicated, I think it's worth porting it to numba, and then wrapping it as a scipy LinearOperator.

Huite commented 8 months ago

Using the preconditioner matrix (technically multiple I guess, U and L) generated by MODFLOW 6, this actually converges unlike a regular scipy ILU.

@nb.njit
def _solve(ilu: "ILU0", r: np.ndarray):
    # forward
    for n in range(ilu.n):
        tv = r[n]
        ic0 = ilu.ia[n]
        iu = ilu.ja[n]
        for j in range(ic0, iu):
            jcol = ilu.ja[j]
            tv -= ilu.apc[j] * ilu.work[jcol]
        ilu.work[n] = tv

    # backward
    for n in range(ilu.n - 1, -1, -1):
        ic1 = ilu.ia[n + 1]
        iu = ilu.ja[n]
        tv = ilu.work[n]
        for j in range(iu, ic1):
            jcol = ilu.ja[j]
            tv -= ilu.apc[j] * ilu.work[jcol]
        ilu.work[n] = tv * ilu.apc[n]

    return

class ILU0(NamedTuple):
    n: int
    m: int
    ia: np.ndarray
    ja: np.ndarray
    apc: np.ndarray
    work: np.ndarray

    def solve(self, r):
        _solve(self, r)
        return self.work

    @property
    def shape(self):
        return (self.n, self.m)

    @property
    def dtype(self):
        return self.apc.dtype

M = sparse.linalg.LinearOperator((n, n), ilu0.solve)

The only part left is to generate it directly.