import os
import gc
import sys
import petsc4py
from gospl.tools.petscgc import safe_garbage_cleanup
import numpy as np
import numpy_indexed as npi
from mpi4py import MPI
from time import process_time
from gospl.tools.constants import BEDROCK_EXPOSED, ICE_COVER_MIN
if "READTHEDOCS" not in os.environ:
from gospl._fortran import fctcoeff
from gospl._fortran import local_spl
from gospl._fortran import jacobiancoeff
MPIrank = petsc4py.PETSc.COMM_WORLD.Get_rank()
[docs]
class soilSPL(object):
"""
The class computes river incision expressed using a **stream power formulation** function of river discharge and slope also **accounting for soil production**.
A non-linear diffusion of soil based on soil thickness is also implemented in this class.
If the user has turned-on the sedimentation capability, this class will solve implicitly the **stream power formulation** accounting for a sediment transport/deposition term (`Yuan et al, 2019 <https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2018JF004867>`_).
"""
def __init__(self, *args, **kwargs):
"""
Initialisation of `soilSPL` class.
"""
# Regolith soil mode is only fully meaningful with stratigraphy on: it
# keeps deposited sediment OUT of Lsoil, expecting the stratigraphy
# (stratK) to carry that sediment's (soft) erodibility. Stratigraphy is
# enabled purely by the stratal time step `time: strat:` (which sets
# stratNb>0) — NOT by a `strata:` block (that is only for initial layers).
# Without it `_surfaceK` returns 1.0, so a fresh deposit erodes at raw
# bedrock K (harder than lumped mode would make it). Warn rather than
# hard-fail: erosional/low-deposition runs can still use regolith mode
# sensibly. See DESIGN_SOIL_REGOLITH.md §5.
if (
getattr(self, "regolithSoil", False)
and getattr(self, "cptSoil", False)
and getattr(self, "stratNb", 0) == 0
and MPIrank == 0
):
print(
"[soil] WARNING: mode 'regolith' without stratigraphy — deposited "
"sediment has no soft-erodibility home and will erode as bedrock. "
"Define a stratal time step ('time: strat:', which sets stratNb>0) "
"so deposits are recorded with soft stratK, or use mode 'lumped'.",
flush=True,
)
# Soil-SPL SNES controls. These are normally set from the YAML soil
# block by the input parser; fall back to robust defaults when soilSPL
# is built standalone (e.g. unit tests).
self.soil_rtol = getattr(self, "soil_rtol", 1.0e-6)
self.soil_atol = getattr(self, "soil_atol", 1.0e-6)
self.soil_maxit = getattr(self, "soil_maxit", 500)
self.soil_pc = getattr(self, "soil_pc", "hypre")
self.soil_solver = getattr(self, "soil_solver", "qn")
# Cached SNES + helper vectors for the soil-aware SPL solver; created
# lazily on first call and reused across timesteps.
self._snes_soil = None
self._snes_soil_f = None
self._snes_soil_x = None
# Robustness fallback SNES (L-BFGS quasi-Newton); created lazily and
# only used on timesteps where the primary solver fails to converge.
self._snes_soil_fb = None
self._snes_soil_fb_f = None
# Cached TS for the non-linear soil diffusion (rosw scheme); created
# lazily on first call and reused across timesteps.
self._ts_soil = None
self._ts_soil_x = None
self._ts_soil_f = None
# Maximum soil (regolith) thickness — the depth at which soil production
# has decayed to a fraction `Sperc` (= soil.bedrockConv) of its surface
# rate, i.e. `-ln(Sperc)*Hs`. Applied as a clip in the two Lsoil
# write-backs. When `Sperc == 0` (bedrockConv: 0, meaning "no bedrock-
# conversion depth") there is NO maximum: use +inf so the clips
# (`nHsoil > soil_transition`) are no-ops — a true no-cap. (Was 100.0 m,
# which silently imposed a large-but-finite cap on a `bedrockConv: 0` run.)
if self.Sperc > 0:
self.soil_transition = -np.log(self.Sperc) * self.Hs
else:
self.soil_transition = np.inf
self.Gsoil = self.hGlobal.duplicate()
self.Lsoil = self.hLocal.duplicate()
self.lHbed = self.hLocal.duplicate()
self.gHbed = self.hGlobal.duplicate()
# Allocate initial soil thicknesses
if self.soilFile is not None:
loadData = np.load(self.soilFile)
soilH = loadData[self.soilData]
self.Lsoil.setArray(soilH[self.locIDs])
self.dm.localToGlobal(self.Lsoil, self.Gsoil)
else:
# Create an uniform soil thickness distribution
self.Gsoil.set(self.cstSoilH)
self.Lsoil.set(self.cstSoilH)
# Initialise the bedrock elevation (base of the soil/regolith mantle,
# lHbed = z − Lsoil) from the start. It is refreshed each soil step
# (diffuseSoil), but the groundwater `aquifer_base: from_soil` coupling
# reads it on the first step — before the first soil update — so it must
# already hold a valid value here rather than an empty duplicate.
self.lHbed.waxpy(-1.0, self.Lsoil, self.hLocal)
self.gHbed.waxpy(-1.0, self.Gsoil, self.hGlobal)
# If temperatures dataset is provided then compute the corresponding soil production rate
if self.tempFile is not None:
Tref = self.tempRef + 273.15
loadData = np.load(self.tempFile)
# Subset the full-mesh temperature map to this rank's local nodes
# (mirrors the soilFile branch above). Without [self.locIDs] the
# derived prodSoil stays global (mpoints) and broadcasts against
# the local hSoil/rainVal arrays only when MPIsize==1; in parallel
# (lpoints < mpoints) it raises a ValueError shape mismatch.
T = loadData[self.tempData][self.locIDs] + 273.15 # Conversion to Kelvin
# Compute Arrhenius term, including Ea / R T0 term
R = 8.314 # Gas constant (J/mol/K)
Arr_terms = self.energyAct * (1./Tref - 1./T) / R
self.prodSoil = self.P0 * np.exp(Arr_terms)
else:
self.prodSoil = self.P0 * np.ones(len(self.locIDs))
return
[docs]
def _monitorsoil(self, snes, its, norm):
"""
Non-linear SPL with soil production solver convergence evaluation.
"""
if MPIrank == 0 and its % 10 == 0:
print(f" --- Non-linear soil SPL solver iteration {its}, Residual norm: {norm}", flush=True)
[docs]
def _build_soil_snes(self, primary=True):
"""
Construct (and return) a soil-SPL SNES solver and its residual vector.
Two configurations are available:
* **primary** (``primary=True``) -- a Nonlinear GMRES accelerator
(``ngmres``) right-preconditioned by Nonlinear Richardson
(``nrichardson``). The Richardson sweep is what actually applies the
Krylov solve (``cg``) and the multigrid preconditioner (HYPRE
BoomerAMG by default, or ``self.soil_pc``); a *bare* ``ngmres`` ignores
the KSP/PC entirely, which is why the previous setup stalled to
``SNES_DIVERGED_MAX_IT`` on the stiff soil-production residual.
* **fallback** (``primary=False``) -- a limited-memory quasi-Newton
solver (``qn``, L-BFGS) with a matrix-free critical-point line search.
It builds an approximate Jacobian from secant updates (no analytic
Jacobian required) and is markedly more robust for stiff problems; it
runs only when the primary solver fails to converge.
Each SNES gets its own PETSc options prefix so per-solver options (e.g.
the line-search type) do not leak into the model's other SNES/KSP
objects.
:arg primary: select the primary (True) or fallback (False) solver.
:return: the configured ``(SNES, residual Vec)`` pair.
"""
opts = petsc4py.PETSc.Options()
snes = petsc4py.PETSc.SNES().create(comm=petsc4py.PETSc.COMM_WORLD)
f = self.hGlobal.duplicate()
snes.setFunction(self._form_residual_soil, f)
if self.verbose:
snes.setMonitor(self._monitorsoil)
if primary:
prefix = "soilspl_"
solver = self.soil_solver
snes.setTolerances(rtol=self.soil_rtol, atol=self.soil_atol,
max_it=self.soil_maxit)
else:
prefix = "soilsplfb_"
# The fallback only runs when the primary has already stalled, so
# use the *other* method (the two are complementary: a fast
# quasi-Newton vs the globally-convergent multigrid-preconditioned
# accelerator), with a relaxed relative tolerance and a larger
# iteration budget as a last resort.
solver = "ngmres" if self.soil_solver == "qn" else "qn"
snes.setTolerances(rtol=max(self.soil_rtol * 100.0, 1.0e-4),
atol=self.soil_atol,
max_it=max(2 * self.soil_maxit, 200))
snes.setOptionsPrefix(prefix)
if solver == "ngmres":
snes.setType("ngmres")
# Nonlinear right-preconditioner: one Richardson sweep per outer
# iteration is what engages the Krylov solve + multigrid PC.
npc = snes.getNPC()
npc.setType("nrichardson")
npc.setTolerances(max_it=1)
ksp = npc.getKSP()
ksp.setType("cg")
pc = ksp.getPC()
pc.setType(self.soil_pc)
if self.soil_pc == "hypre":
opts["pc_hypre_type"] = "boomeramg"
pc.setFromOptions()
ksp.setTolerances(rtol=1.0e-6)
else:
# Limited-memory quasi-Newton (L-BFGS): builds curvature from secant
# updates (no analytic Jacobian) and typically converges in far
# fewer iterations than the ngmres accelerator on this stiff
# residual. Critical-point line search is matrix-free (the 'bt'
# backtracking search requires a Jacobian, which qn does not form).
snes.setType("qn")
opts[prefix + "snes_qn_type"] = "lbfgs"
opts[prefix + "snes_linesearch_type"] = "cp"
snes.setFromOptions()
return snes, f
[docs]
def _soilErodibility(self, PA, surfK):
"""
(Re)build the dt-scaled stream-power erodibility coefficients ``Kbr``
(bedrock) and ``K_soil`` (soil layer) for the CURRENT ``self.dt`` and
``self.hOldArray``. The erosion limiter follows the current elevation
(drop to the receiver), so this is recomputed per sub-step by the
adaptive sub-stepper as well as once for the full step.
"""
dh = (self.hOldArray[:, None] - self.hOldArray[self.rcvIDi]).max(axis=1)
elimiter = np.divide(dh, dh + 1.0e-2, out=np.zeros_like(dh),
where=dh != 0)
if self.sedfacVal is not None:
self.Kbr = self.K * surfK * self.sedfacVal * (self.rainVal ** self.coeffd)
else:
self.Kbr = self.K * surfK * (self.rainVal ** self.coeffd)
self.Kbr *= self.dt * (PA ** self.spl_m) * elimiter
self.Kbr[self.seaID] = 0.0
self.K_soil = self.Ksoil * self.dt * (PA ** self.spl_m) * elimiter
self.K_soil[self.seaID] = 0.0
[docs]
def _soilSNESsolve(self, x, guess):
"""
Solve the soil-SPL SNES into ``x`` from initial guess ``guess``: the
primary solver, then the complementary quasi-Newton / ngmres fallback on
failure (same primary→fallback net as the flow KSP). Returns the final
converged reason (``>= 0`` converged, ``< 0`` both failed). Prints only
the informational "fallback recovered" line; the caller decides what to
do on an outright failure (sub-step, then discard).
"""
guess.copy(result=x)
self._snes_soil.solve(None, x)
r = self._snes_soil.getConvergedReason()
if r < 0:
if self._snes_soil_fb is None:
self._snes_soil_fb, self._snes_soil_fb_f = self._build_soil_snes(
primary=False
)
r0, it0 = r, self._snes_soil.getIterationNumber()
guess.copy(result=x)
self._snes_soil_fb.solve(None, x)
r = self._snes_soil_fb.getConvergedReason()
if MPIrank == 0 and r >= 0:
pn = self.soil_solver
fn = "ngmres" if self.soil_solver == "qn" else "qn"
print(
"Soil SPL: primary (%s) stalled (reason %d after %d its); "
"%s fallback converged (reason %d, %d its)."
% (pn, r0, it0, fn, r, self._snes_soil_fb.getIterationNumber()),
flush=True,
)
return r
[docs]
def _adaptiveSubstepSoil(self, x, PA, surfK):
"""
Adaptive sub-stepping for the soil-SPL solve when the full-Δt SNES
diverges. The divergence is a stiffness problem: a de-armoured node that
has just captured a large discharge demands an enormous single-step
incision (``Kbr·Sⁿ`` with ``Kbr ∝ Δt·A^m``), so the slope term overshoots
and the residual grows. Because ``Kbr ∝ Δt``, splitting the step into
``N`` increments of ``Δt/N`` cuts the demanded per-solve incision by
``N`` until the residual is smooth enough to converge; summing the
increments recovers the FULL step's erosion (unlike simply discarding).
Escalates ``N`` over ``_soilSubstepN`` (default 4, 8, 16); accepts the
first ``N`` whose every sub-step converges. If none does, the step is
discarded (revert to the prior elevation) as a last resort. On success
the step-start bookkeeping (``hOld``/``hOldArray``/``soilH``/``dt``) is
restored and one final residual evaluation rebuilds ``nsoilH`` /
``_soilDepoGrowth`` consistently with the full step (used by the soil
update). Returns the final reason (``0`` sub-stepped OK, ``-1``
discarded).
"""
dt_full = self.dt
hstart = self.hOld.getArray().copy() # step-start (owned) elevation
hOldArray0 = self.hOldArray.copy()
soilH0 = self.soilH.copy()
converged, used = False, 0
for N in getattr(self, "_soilSubstepN", (4, 8, 16)):
x.setArray(hstart) # running elevation ← step start
self.soilH = soilH0.copy()
self.dt = dt_full / N
ok = True
for _ in range(N):
x.copy(result=self.hOld) # sub-step start = current elevation
self.dm.globalToLocal(self.hOld, self.hOldLocal)
self.hOldArray = self.hOldLocal.getArray().copy()
self._soilErodibility(PA, surfK) # Kbr/K_soil at Δt/N + current elev
if self._soilSNESsolve(x, self.hOld) < 0:
ok = False
break
self.soilH = self.nsoilH.copy() # soil carried across sub-steps
if ok:
converged, used = True, N
break
# Restore full-step bookkeeping (step-start state, full Δt) so the eroded
# thickness `x − hOld` and the soil update reflect the whole step.
self.dt = dt_full
self.hOld.setArray(hstart)
self.dm.globalToLocal(self.hOld, self.hOldLocal)
self.hOldArray = hOldArray0
self.soilH = soilH0
if not converged:
x.setArray(hstart) # last resort: no erosion this step
# Rebuild nsoilH / _soilDepoGrowth (Kbr-independent) at the final
# elevation with the full-Δt production convention.
self._form_residual_soil(self._snes_soil, x, self.tmp)
if MPIrank == 0:
if converged:
print(
" Soil SPL: full-Δt solve diverged — converged via %d "
"sub-steps (Δt/%d); fluvial erosion retained." % (used, used),
flush=True,
)
else:
print(
" Soil SPL: still diverged after sub-stepping — discarded "
"(no fluvial erosion this step).",
flush=True,
)
return 0 if converged else -1
[docs]
def _solveSoil(self):
"""
Solves the non-linear stream power law for the transport limited and soil case. This calls the following *private function*:
- _form_residual_soil
.. note::
PETSc SNES approach is used to solve the nonlinear equation without forming an analytic Jacobian. The primary solver is a Nonlinear GMRES accelerator (``ngmres``) right-preconditioned by Nonlinear Richardson (``nrichardson``), whose Krylov solve uses a Preconditioned Conjugate Gradient (``cg``) method with a multi-grid preconditioner (HYPRE BoomerAMG by default). If the primary solver stalls on the stiff soil-production residual, a limited-memory quasi-Newton fallback (``qn``, L-BFGS) with a critical-point line search is used. The iteration budget, tolerances and preconditioner are configurable through the YAML ``soil`` block (``maxIter``, ``rtol``, ``atol``, ``pcType``).
"""
self.hOldArray = self.hLocal.getArray().copy()
# Get soil thickness from previous time step
self.soilH = self.Lsoil.getArray().copy()
# Consider bedrock exposed when soil thickness is below 10 cm
self.soilH[self.soilH < BEDROCK_EXPOSED] = 0.0
# Upstream-averaged mean annual precipitation rate based on drainage area
PA = self.FAL.getArray()
# Per-node erodibility multiplier from the top of the local
# stratigraphic column (1.0 = use self.K as-is). Only scales the
# *bedrock* SPL coefficient; the soil-layer K is governed by
# `self.Ksoil` and is left unchanged.
# Fold in the dual-lithology erodibility blend (1.0 everywhere when
# single-fraction, so behaviour is unchanged): K_eff = K*surfK*litK.
surfK = self._surfaceK() * self._surfaceLithoK()
# Dimensionless depositional coefficient fDep = fDepa*area / PA (step-
# constant — depends on drainage area, not on Δt or the evolving
# elevation). Floor the denominator at the cap value (num/0.99) so a
# denormal-tiny PA — more likely with evaporation reducing the discharge —
# does not overflow to inf before the 0.99 cap; bit-identical to dividing
# then clamping.
num = self.fDepa * self.larea
self.fDep = np.divide(
num, np.maximum(PA, num / 0.99),
out=np.zeros_like(PA), where=PA != 0,
)
self.fDep[self.seaID] = 0.
self.fDep[self.fDep > 0.99] = 0.99
if self.flatModel:
self.fDep[self.outletIDs] = 0.
# dt-scaled erodibility (erosion limiter + Kbr / K_soil) for the full step.
self._soilErodibility(PA, surfK)
if self._snes_soil is None:
self._snes_soil, self._snes_soil_f = self._build_soil_snes(primary=True)
self._snes_soil_x = self.hGlobal.duplicate()
x = self._snes_soil_x
r = self._soilSNESsolve(x, self.hGlobal)
# When the full-Δt solve diverges (both primary and fallback), retry with
# adaptive sub-stepping — Δt/N increments make the demanded per-solve
# incision small enough to converge, and summing them recovers the full
# step's erosion. Only a still-diverging step (after sub-stepping) is
# discarded. See `_adaptiveSubstepSoil`.
if r < 0:
self._adaptiveSubstepSoil(x, PA, surfK)
# Get eroded sediment thicknesses
self.tmp.waxpy(-1.0, self.hOld, x)
# Update soil thicknesses
nHsoil = self.nsoilH.copy()
# Regolith mode: remove the fluvial transport-limited DEPOSITION growth
# (post-solve, so the SNES residual — and its smoothness — is untouched).
# Lsoil then reflects weathering production + erosion only; the deposit is
# tracked by the stratigraphy. Erosion (negative elevation change) still
# strips soil via nsoilH. See DESIGN_SOIL_REGOLITH.md §5.
if self.regolithSoil:
nHsoil = nHsoil - self._soilDepoGrowth
# No subaerial soil under standing water (marine seaID OR a ponded
# continental lake). Extends the former marine-only mask to lakes for a
# coherent subaerial gate — the rainfall-scaled production term would
# otherwise leave a spurious cover on submerged nodes.
nHsoil[self._subaqueousMask(self.hOldArray)] = 0.0
nHsoil[nHsoil < BEDROCK_EXPOSED] = 0.
# Limit soil thickness
nHsoil[nHsoil > self.soil_transition] = self.soil_transition
# Ice-covered land: freeze the regolith inert (preserve the prior column,
# no production). Applied LAST so the preserved value is not re-clipped.
ice = self._iceFrozenMask(self.hOldArray)
nHsoil[ice] = self.Lsoil.getArray()[ice]
self.Lsoil.setArray(nHsoil)
self.dm.localToGlobal(self.Lsoil, self.Gsoil)
safe_garbage_cleanup()
return
[docs]
def _getEroDepRateSoil(self):
"""
This function computes erosion deposition rates in metres per year and associated soil evolution. This is done on the filled elevation.
The approach is based on **BasicHySa** governing equations from Terrainbento (as described in Appendix B20 from `Barnhart et al. (2019) <https://gmd.copernicus.org/articles/12/1267/2019/gmd-12-1267-2019.pdf>`_).
.. note::
The approach uses a continuous layer of soil-alluvium, which influences both hillslope and river-induced erosion. It relies on the SPACE algorithm of `Shobe et al. (2017) <https://gmd.copernicus.org/articles/10/4577/2017/>`_.
"""
t0 = process_time()
# Build the SPL erosion arrays. Snapshot the flexure load reference only
# at the start of a flexure interval so the load accumulates across
# skipped steps (flex_interval).
if self.flexOn and self.flexCount % self.flex_interval == 0:
self.hLocal.copy(result=self.hOldFlex)
self._solveSoil()
# Update erosion/deposition rate (thickness convention: positive
# for deposition, negative for incision; same sign as cumED and
# the on-disk EDrate field). See SPL.py for the convention note.
E = self.tmp.getArray().copy()
E = np.divide(E, self.dt)
self.Eb.setArray(E)
self.dm.globalToLocal(self.Eb, self.EbLocal)
E = self.EbLocal.getArray().copy()
if self.flatModel:
E[self.outletIDs] = 0.0
E[self.lsink] = 0.0
self.EbLocal.setArray(E)
self.dm.localToGlobal(self.EbLocal, self.Eb)
if MPIrank == 0 and self.verbose:
print(
"Finalise erosion deposition rates (%0.02f seconds)" % (process_time() - t0),
flush=True,
)
return
def _subaqueousMask(self, hl):
"""
Boolean mask (``lpoints``) of **subaqueous** nodes — under standing water —
where subaerial soil is suppressed. ``self.Lsoil`` is a *subaerial* regolith
cover, so it is held at 0 on this mask; sediment deposited under water is
tracked by the stratigraphy, not as soil.
Two contributions:
- **marine** — ``self.seaID`` (filled level at/below sea level);
- **ponded continental lake** — a depression node (``pitIDs > -1``) whose
fill/spill level sits above the bed (``lFill > hl``), i.e. its
accommodation is still water-filled (a lake fills to its spillover
before emerging; see ``DESIGN_SOIL_REGOLITH.md`` §3).
Purely local (per-node) — no collective. Falls back to a marine-only mask
if the pit fields are not yet populated (bare ``STRAMesh.__new__`` stubs).
"""
sub = np.zeros(self.lpoints, dtype=bool)
sub[self.seaID] = True
pitIDs = getattr(self, "pitIDs", None)
lFill = getattr(self, "lFill", None)
if pitIDs is not None and lFill is not None:
sub |= (pitIDs > -1) & (lFill > hl)
return sub
def _iceFrozenMask(self, hl):
"""
Boolean mask (``lpoints``) of ice-covered LAND where soil is held **frozen
inert** — the pre-existing regolith is *preserved* (not zeroed) and no new
pedogenic soil is produced, since subaerial weathering does not operate
beneath ice (frozen, insulated from the atmosphere, no biota / rain
infiltration). This differs from the subaqueous case (marine / ponded lake),
where soil is held at 0: ice can preserve a buried regolith for a long time,
so glaciation freezes the soil column rather than removing it. Glacial
erosion / till are handled separately by the ice model.
Ice-covered ⇔ ``iceOn`` and ``iceHL > ICE_COVER_MIN``, restricted to LAND
(subaqueous cells are excluded — an ice shelf over sea/lake stays a
subaqueous, soil-free cell). Purely local; empty when ice is off.
"""
ice = np.zeros(self.lpoints, dtype=bool)
if not getattr(self, "iceOn", False):
return ice
iceHL = getattr(self, "iceHL", None)
if iceHL is None:
return ice
ice = iceHL.getArray() > ICE_COVER_MIN
ice &= ~self._subaqueousMask(hl) # subaqueous (zero) wins over ice
return ice
[docs]
def updateSoilThickness(self, deposition=True):
"""
Updates soil thickness through time from the increment in ``self.tmp``.
``Lsoil`` is a **subaerial** regolith cover, so the increment is not
retained at subaqueous nodes (marine or ponded lake) — see
``_subaqueousMask`` — and is frozen under ice (``_iceFrozenMask``). These
gates run in every mode, so callers invoke this unconditionally.
``deposition`` distinguishes the two kinds of increment:
- **deposition=True** (lake/pit, marine): in **regolith mode** the deposit
is stratigraphy, not soil, so the increment is **skipped** (the gates
still run — a cell newly ponded by this step's deposition is re-zeroed
consistently with lumped mode). In lumped mode it is added.
- **deposition=False** (soil creep / hillslope transport): always added —
creep moves the weathering regolith itself, in both modes.
"""
self.dm.globalToLocal(self.tmp, self.tmpL)
prevL = self.Lsoil.getArray().copy()
if deposition and getattr(self, "regolithSoil", False):
nHsoil = prevL.copy() # regolith: deposit is strata, not soil
else:
nHsoil = prevL + self.tmpL.getArray()
hl = self.hLocal.getArray()
# No subaerial soil under standing water (marine or ponded lake).
nHsoil[self._subaqueousMask(hl)] = 0.0
# Limit soil thickness
nHsoil[nHsoil < 0.] = 0.
nHsoil[nHsoil > self.soil_transition] = self.soil_transition
# Ice-covered land: freeze the regolith inert (no deposition-into-soil
# increment under ice — glacial till is tracked by the stratigraphy).
ice = self._iceFrozenMask(hl)
nHsoil[ice] = prevL[ice]
self.Lsoil.setArray(nHsoil)
self.dm.localToGlobal(self.Lsoil, self.Gsoil)
return
[docs]
def erodepSPLsoil(self):
"""
Modified **stream power law** model used to represent erosion by rivers also taking into account the role played by sediments in modulating erosion and deposition rate, considering **non-linear slope dependency** and accounting for soil production.
It calls the private function `_getEroDepRateSoil` described above. Once erosion/deposition rates have been calculated, the function computes local thicknesses and soil evolution for the considered time step and update local elevation and cumulative erosion, deposition values.
"""
t0 = process_time()
# Computes the erosion deposition rates based on flow accumulation
self.Eb.set(0.0)
self.hGlobal.copy(result=self.hOld)
self.dm.globalToLocal(self.hOld, self.hOldLocal)
self._getEroDepRateSoil()
self._glacialAbrasion()
# Get erosion / deposition thicknesses (Eb is in thickness rate
# convention: positive deposition, negative incision). See SPL.py.
Eb = self.Eb.getArray().copy()
self.tmp.setArray(Eb * self.dt)
self.cumED.axpy(1.0, self.tmp)
self.dm.globalToLocal(self.cumED, self.cumEDLocal)
self.hGlobal.axpy(1.0, self.tmp)
self.dm.globalToLocal(self.hGlobal, self.hLocal)
self.tmp1.pointwiseMult(self.tmp, self.areaGlobal)
# Update stratigraphic layers
if self.stratNb > 0:
self.erodeStrat()
self.deposeStrat()
# Update erosion/deposition rates
self.dm.globalToLocal(self.tmp, self.tmpL)
add_rate = self.tmpL.getArray() / self.dt
self.EbLocal.setArray(add_rate)
# Destroy flow matrices
self.fMati.destroy()
self.fMat.destroy()
if MPIrank == 0 and self.verbose:
print(
"Get Erosion Deposition values (%0.02f seconds)" % (process_time() - t0),
flush=True,
)
if self.memclear:
del Eb
gc.collect()
return
[docs]
def _evalFunctionSoil(self, ts, t, x, xdot, f):
"""
The non-linear system for soil diffusion is solved iteratively using PETSc time stepping and SNES solution and is based on Rosenbrock W-scheme (``rosw``).
Here again, we evaluate the residual function on a DMPlex for an implicit time-stepping method.
Parameters:
-----------
ts : PETSc.TS: The time-stepper object.
t : float: The current time.
x : PETSc.Vec: The current solution vector (h^{n+1}) at the new time step.
xdot : PETSc.Vec: The time derivative approximation (h^{n+1} - h^n) / dt.
f : PETSc.Vec: The residual vector to be filled.
"""
self.dm.globalToLocal(x, self.hl)
with self.hl as hl, self.lHbed as zb, xdot as hdot:
dh = hl - zb
dh[dh < 0.1] = 0.0
Cd = self.minDiff + np.multiply(self.Cd, (1.0 - np.exp(-dh / self.H0)))
nlvec = fctcoeff(hl, Cd)
f.setArray(hdot + nlvec[self.glIDs])
return
[docs]
def _evalJacobianSoil(self, ts, t, x, xdot, a, A, B):
"""
The non-linear system for soil diffusion is solved iteratively using PETSc time stepping and SNES solution and is based on Rosenbrock W-scheme (``rosw``).
Here, we define the Jacobian matrix A and the preconditioner matrix B on a DMPlex.
Parameters:
-----------
ts : PETSc.TS: The time-stepper object.
t : float: The current time.
x : PETSc.Vec: The current solution vector (h^{n+1}) at the new time step.
xdot : PETSc.Vec: The time derivative approximation (h^{n+1} - h^n) / dt.
a : float: The shift factor for implicit methods.
A : PETSc.Mat: The Jacobian matrix to be filled.
B : PETSc.Mat: The preconditioner matrix to be filled.
"""
self.dm.globalToLocal(x, self.hl)
with self.hl as hl, self.lHbed as zb:
dh = hl - zb
dh[dh < 0.1] = 0.0
Cd = self.minDiff + np.multiply(self.Cd, (1.0 - np.exp(-dh / self.H0)))
# Coefficient derivatives
Cp = np.multiply(self.Cd, np.exp(-dh / self.H0) / self.H0)
nlC = jacobiancoeff(hl, Cd, Cp)
# Assemble ONLY the rows this rank OWNS (self.glIDs). Looping over
# all self.lpoints rows would also set the GHOST rows, whose stencil
# is computed here from an INCOMPLETE neighbour set (the ghost's full
# neighbourhood is not all present on this rank). Those off-process
# ghost-row values collide under INSERT_VALUES with the owning rank's
# correct values, leaving the boundary rows of the Jacobian
# partition-dependent -- which drove an isolated elevation spike at
# sub-domain boundaries (np=1 clean, np>1 spiked). Owned-row-only
# assembly (mirroring _evalJacobianMardDiff) makes the matrix
# partition-invariant. The diagonal column for owned row r is r
# itself; its off-diagonal columns are r's FV neighbours.
ngb_cols = self.FVmesh_ngbID[self.glIDs, :]
cols_2d = np.column_stack([self.glIDs[:, None], ngb_cols]).astype(
petsc4py.PETSc.IntType
)
vals_2d = np.column_stack(
[(a + nlC[self.glIDs, 0])[:, None], nlC[self.glIDs, 1:]]
)
for i, row in enumerate(self.glIDs):
B.setValuesLocal(row, cols_2d[i], vals_2d[i])
B.assemble()
if A != B:
A.assemble()
return True
[docs]
def _evalSolutionSoil(self, t, x):
"""
Evaluate the initial solution of the SNES system.
"""
assert t == 0.0, "only for t=0.0"
x.setArray(self.h.getArray())
return
[docs]
def diffuseSoil(self):
r"""
For river-transported sediments reaching the marine realm, this function computes the related marine deposition diffusion. It is based on a non-linear diffusion approach.
.. math::
\frac{\partial h}{\partial t}= \nabla \cdot \left( C_d \times (1.0 - e^{-h_s/H_0} \nabla h \right)
It calls the following *private functions*:
- _evalFunctionSoil
- _evalJacobianSoil
- _evalSolutionSoil
.. note::
PETSc SNES and time stepping TS approaches are used to solve the non-linear equation above over the considered time step.
"""
t0 = process_time()
# Get diffusion soil coefficient
self.Cd = np.full(self.lpoints, self.Cda, dtype=np.float64)
self.Cd[self.seaID] = self.Cdm
# Dual-lithology (Phase 7): scale soil diffusivity by the surface
# composition so fine-rich soil diffuses faster (neutral when single-
# fraction / no contrast).
if self.stratLith:
self.Cd = self.Cd * self._surfaceLithoD()
# Remove the soil thickness from the elevation
self.hLocal.copy(result=self.hl)
self.dm.localToGlobal(self.hl, self.h)
self.gHbed.waxpy(-1.0, self.Gsoil, self.hGlobal)
self.lHbed.waxpy(-1.0, self.Lsoil, self.hLocal)
# Time stepping definition (cached across timesteps)
if self._ts_soil is None:
ts = petsc4py.PETSc.TS().create(comm=petsc4py.PETSc.COMM_WORLD)
# arkimex: IMEX Runge-Kutta schemes | rosw: Rosenbrock W-schemes
ts.setType("rosw")
ts.setIFunction(self._evalFunctionSoil, self.tmp1)
ts.setIJacobian(self._evalJacobianSoil, self.mat)
ts.setExactFinalTime(petsc4py.PETSc.TS.ExactFinalTime.MATCHSTEP)
# Allow an unlimited number of failures (step rejected and retried)
ts.setMaxSNESFailures(-1)
# SNES nonlinear solver
snes = ts.getSNES()
snes.setTolerances(max_it=10)
# KSP linear solver
ksp = snes.getKSP()
ksp.setType("preonly")
pc = ksp.getPC()
pc.setType("gasm")
ts.setFromOptions()
self._ts_soil = ts
self._ts_soil_x = self.tmp1.duplicate()
self._ts_soil_f = self.tmp1.duplicate()
ts = self._ts_soil
x = self._ts_soil_x
# Soil thicknesses are meters; mm-level absolute tolerance is plenty.
ts.setTolerances(atol=1e-3, rtol=1e-3)
ts.setTime(0.0)
# Reset the step COUNTER (setTime only resets the clock). The cached TS
# is reused and getStepNumber() is not reset by setTime, so without this
# setMaxSteps below becomes a *cumulative* cap — after ~tsStep total
# substeps it is exceeded on entry and TSSolve returns immediately,
# leaving the soil column un-diffused (same bug as hillslope marine TS).
ts.setStepNumber(0)
# Larger initial step (was self.dt / 1000.0).
ts.setTimeStep(self.dt / 100.0)
ts.setMaxTime(self.dt)
ts.setMaxSteps(self.tsStep)
tstart = ts.getTime()
self._evalSolutionSoil(tstart, x)
# Solve nonlinear equation
ts.solve(x)
if MPIrank == 0 and self.verbose:
print(
"Nonlinear soil diffusion solution (%0.02f seconds)" % (process_time() - t0),
flush=True,
)
print(
"steps %d (%d rejected, %d SNES fails), nonlinear its %d, linear its %d"
% (
ts.getStepNumber(),
ts.getStepRejections(),
ts.getSNESFailures(),
ts.getSNESIterations(),
ts.getKSPIterations(),
),
flush=True,
)
# Get diffused sediment thicknesses
self.dh.waxpy(-1.0, self.hGlobal, x)
self.dm.globalToLocal(self.dh, self.tmpL)
chgSoil = self.tmpL.getArray().copy()
self.tmpL.setArray(chgSoil)
self.dm.localToGlobal(self.tmpL, self.tmp)
safe_garbage_cleanup()
# Update cumulative erosion and deposition as well as elevation
self.cumED.axpy(1.0, self.tmp)
self.dm.globalToLocal(self.cumED, self.cumEDLocal)
self.hGlobal.axpy(1.0, self.tmp)
self.dm.globalToLocal(self.hGlobal, self.hLocal)
# Update soil thickness. Soil creep TRANSPORTS the weathering regolith
# itself (not a deposit), so it is applied in both modes — deposition=False.
self.updateSoilThickness(deposition=False)
# Update erosion/deposition rates
self.dm.globalToLocal(self.tmp, self.tmpL)
add_rate = self.tmpL.getArray() / self.dt
self.tmpL.setArray(add_rate)
self.EbLocal.axpy(1.0, self.tmpL)
# Update stratigraphic layer parameters
if self.stratNb > 0:
self.deposeStrat()
if MPIrank == 0 and self.verbose:
print(
"Diffuse Soil Sediments (%0.02f seconds)"
% (process_time() - t0),
flush=True,
)
return
