class TerrainDerivatives:
"""Compute, on demand, terrain derivatives such as slope or
curvature.
"""
shape = xs.foreign(UniformRectilinearGrid2D, 'shape')
spacing = xs.foreign(UniformRectilinearGrid2D, 'spacing')
elevation = xs.foreign(SurfaceTopography, 'elevation')
slope = xs.on_demand(dims=('y', 'x'), description='terrain local slope')
curvature = xs.on_demand(dims=('y', 'x'),
description='terrain local curvature')
@slope.compute
def _slope(self):
slope = np.empty_like(self.elevation)
ny, nx = self.shape
dy, dx = self.spacing
fs.slope(self.elevation.ravel(), slope.ravel(), nx, ny, dx, dy)
return slope
@curvature.compute
def _curvature(self):
curv = np.empty_like(self.elevation)
ny, nx = self.shape
dy, dx = self.spacing
fs.curvature(self.elevation.ravel(), curv.ravel(), nx, ny, dx, dy)
return curv
class P1:
var = xs.on_demand(dims="x")
cached_var = xs.on_demand(dims="x")
@var.compute
def _compute_var(self):
return np.random.rand(10)
@cached_var.compute(cache=True)
def _compute_cached_var(self):
return np.random.rand(10)
class Profile:
u = xs.variable(
dims="x",
description="quantity u",
intent="inout",
encoding={"fill_value": np.nan},
)
u_diffs = xs.group("diff")
u_opp = xs.on_demand(dims="x")
def initialize(self):
self.u_change = np.zeros_like(self.u)
def run_step(self):
self.u_change[:] = sum((d for d in self.u_diffs))
def finalize_step(self):
self.u += self.u_change
def finalize(self):
self.u[:] = 0.0
@u_opp.compute
def _get_u_opposite(self):
return -self.u
class ExampleProcess:
"""A process with complete interface for testing."""
in_var = xs.variable(dims=["x", ("x", "y")], description="input variable")
out_var = xs.variable(groups="example_group", intent="out")
inout_var = xs.variable(intent="inout", converter=int)
od_var = xs.on_demand()
obj_var = xs.any_object(description="arbitrary object")
in_foreign_var = xs.foreign(SomeProcess, "some_var")
in_foreign_var2 = xs.foreign(AnotherProcess, "some_var")
out_foreign_var = xs.foreign(AnotherProcess, "another_var", intent="out")
in_foreign_od_var = xs.foreign(SomeProcess, "some_od_var")
in_global_var = xs.global_ref("some_global_var")
out_global_var = xs.global_ref("another_global_var", intent="out")
group_var = xs.group("some_group")
group_dict_var = xs.group_dict("some_group")
other_attrib = attr.attrib(init=False, repr=False)
other_attr = "this is not a xsimlab variable attribute"
@od_var.compute
def compute_od_var(self):
return 0
class Bedrock:
"""Update the elevation of bedrock (i.e., land and/or submarine
basement).
"""
elevation = xs.variable(dims=('y', 'x'),
intent='inout',
description='bedrock elevation')
depth = xs.on_demand(dims=('y', 'x'),
description='bedrock depth below topographic surface')
bedrock_motion_up = xs.foreign(TotalVerticalMotion, 'bedrock_upward')
surface_motion_up = xs.foreign(TotalVerticalMotion, 'surface_upward')
surface_elevation = xs.foreign(SurfaceTopography, 'elevation')
@depth.compute
def _depth(self):
return self.surf_elevation - self.elevation
def initialize(self):
if np.any(self.elevation > self.surface_elevation):
raise ValueError("Encountered bedrock elevation higher than "
"topographic surface elevation.")
def run_step(self):
self._elevation_next = np.minimum(
self.elevation + self.bedrock_motion_up,
self.surface_elevation + self.surface_motion_up)
def finalize_step(self):
self.elevation = self._elevation_next
class SingleFlowRouter(FlowRouter):
"""Single direction (convergent) flow router."""
slope = xs.on_demand(dims='node', description='out flow path slope')
def initialize(self):
# for compatibility
self.nb_receivers = np.ones_like(self.fs_context.rec)
self.weights = np.ones_like(self.fs_context.length)
def route_flow(self):
fs.flowroutingsingleflowdirection()
# Fortran 1 vs Python 0 index
self.stack = self.fs_context.stack.astype('int') - 1
self.receivers = self.fs_context.rec - 1
self.lengths = self.fs_context.length
@slope.compute
def _slope(self):
elev_flat = self.elevation.ravel()
elev_flat_diff = elev_flat - elev_flat[self.receivers]
# skip base levels
slope = np.zeros_like(self.lengths)
idx = np.argwhere(self.lengths > 0)
slope[idx] = elev_flat_diff[idx] / self.lengths[idx],
return slope
class AddOnDemand(object):
offset = xs.variable(description='offset added to profile u')
u_diff = xs.on_demand(group='diff')
@u_diff.compute
def _compute_u_diff(self):
return self.offset
class OrographicDrainageDischarge(FlowAccumulator):
"""Accumulate orographic precipitation from upstream to downstream.
For use in the context of landscape evolution modeling, ``flowacc`` values
are converted from mm^3 h^-1 to m^3 yr^-1. For convenience, the ``discharge``
on demand variable still returns the values in mm^3 h^-1.
"""
runoff = xs.foreign(OrographicPrecipitation, 'precip_rate')
rainfall_frequency = xs.foreign(OrographicPrecipitation,
'rainfall_frequency')
discharge = xs.on_demand(
dims=('y', 'x'),
description='discharge from orographic precipitation',
attrs={"units": "mm^3/h"})
def run_step(self):
super().run_step()
# scale mm m^2/h to m^3/yr assuming that
# the rainfall frequency is the number of storm of 1 hour duration per day
self.flowacc *= 8.76 * self.rainfall_frequency
@discharge.compute
def _discharge(self):
return self.flowacc / (8.76 * self.rainfall_frequency)
class SomeProcess(object):
"""Just used for foreign variables in ExampleProcess."""
some_var = xs.variable(group='some_group', intent='out')
some_od_var = xs.on_demand(group='some_group')
@some_od_var.compute
def compute_some_od_var(self):
return 1
class AddOnDemand:
offset = xs.variable(dims=[(), "x"],
description="offset added to profile u")
u_diff = xs.on_demand(dims=[(), "x"], groups="diff")
@u_diff.compute
def _compute_u_diff(self):
return self.offset
class SomeProcess:
"""Just used for foreign variables in ExampleProcess."""
some_var = xs.variable(groups="some_group", intent="out")
some_od_var = xs.on_demand(groups="some_group")
@some_od_var.compute
def compute_some_od_var(self):
return 1
class P:
v1 = xs.variable(dims="x",
intent="out",
encoding={"dtype": np.int32})
v2 = xs.on_demand(dims="x", encoding={"fill_value": 0})
v3 = xs.index(dims="x")
@v2.compute
def _get_v2(self):
return [0]
class AddOnDemand:
offset = xs.variable(dims=[(), "x"],
description="offset added to profile u")
u_diff = xs.on_demand(dims=[(), "x"],
groups="diff",
encoding={"fill_value": np.nan})
@u_diff.compute
def _compute_u_diff(self):
return self.offset * 1.0
class StreamPowerChannel(ChannelErosion):
"""Stream-Power channel erosion."""
k_coef = xs.variable(dims=[(), ('y', 'x')],
description='bedrock channel incision coefficient')
area_exp = xs.variable(description='drainage area exponent')
slope_exp = xs.variable(description='slope exponent')
shape = xs.foreign(UniformRectilinearGrid2D, 'shape')
elevation = xs.foreign(FlowRouter, 'elevation')
receivers = xs.foreign(FlowRouter, 'receivers')
flowacc = xs.foreign(FlowAccumulator, 'flowacc')
fs_context = xs.foreign(FastscapelibContext, 'context')
chi = xs.on_demand(dims=('y', 'x'),
description='integrated drainage area (chi)')
def _set_g_in_context(self):
# transport/deposition feature is exposed in subclasses
self.fs_context.g1 = 0.
self.fs_context.g2 = 0.
def run_step(self):
kf = np.broadcast_to(self.k_coef, self.shape).flatten()
self.fs_context.kf = kf
# we don't use kfsed fastscapelib-fortran feature directly
self.fs_context.kfsed = -1.
self._set_g_in_context()
self.fs_context.m = self.area_exp
self.fs_context.n = self.slope_exp
# bypass fastscapelib_fortran global state
self.fs_context.h = self.elevation.flatten()
# TODO: https://github.com/fastscape-lem/fastscapelib-fortran/pull/25
# this has no effect yet.
self.fs_context.a = self.flowacc.flatten()
if self.receivers.ndim == 1:
fs.streampowerlawsingleflowdirection()
else:
fs.streampowerlaw()
erosion_flat = self.elevation.ravel() - self.fs_context.h
self.erosion = erosion_flat.reshape(self.shape)
@chi.compute
def _chi(self):
chi_arr = np.empty_like(self.elevation, dtype='d')
self.fs_context.copychi(chi_arr.ravel())
return chi_arr
class P:
in_var = xs.variable()
out_var = xs.variable(intent="out")
od_var = xs.on_demand()
def run_step(self):
self.out_var = self.in_var * 2
@od_var.compute
def _dummy(self):
return None
class TotalErosion:
"""Sum up all erosion processes."""
erosion_vars = xs.group('erosion')
cumulative_height = xs.variable(
dims=[(), ('y', 'x')],
intent='inout',
description='erosion height accumulated over time')
height = xs.variable(dims=[(), ('y', 'x')],
intent='out',
description='total erosion height at current step',
groups='surface_downward')
rate = xs.on_demand(dims=[(), ('y', 'x')],
description='total erosion rate at current step')
grid_area = xs.foreign(UniformRectilinearGrid2D, 'area')
domain_rate = xs.on_demand(
description='domain-integrated volumetric erosion rate')
@xs.runtime(args='step_delta')
def run_step(self, dt):
self._dt = dt
self.height = sum(self.erosion_vars)
self.cumulative_height += self.height
@rate.compute
def _rate(self):
return self.height / self._dt
@domain_rate.compute
def _domain_rate(self):
return np.sum(self.height) * self.grid_area / self._dt
class P:
v1 = xs.variable(dims="x",
intent="out",
encoding={"dtype": np.int32})
v2 = xs.on_demand(dims="x", encoding={"fill_value": 0})
v3 = xs.index(dims="x")
v4 = xs.variable(
dims="x",
intent="out",
encoding={
"dtype": object,
"object_codec": zarr.codecs.Pickle()
},
)
@v2.compute
def _get_v2(self):
return [0]
class Profile(object):
u = xs.variable(dims='x', description='quantity u', intent='inout')
u_diffs = xs.group('diff')
u_opp = xs.on_demand(dims='x')
def initialize(self):
self.u_change = np.zeros_like(self.u)
def run_step(self, *args):
self.u_change[:] = np.sum((d for d in self.u_diffs))
def finalize_step(self):
self.u += self.u_change
def finalize(self):
self.u[:] = 0.
@u_opp.compute
def _get_u_opposite(self):
return -self.u
class ExampleProcess(object):
"""A process with complete interface for testing."""
in_var = xs.variable(dims=['x', ('x', 'y')], description='input variable')
out_var = xs.variable(group='example_group', intent='out')
inout_var = xs.variable(intent='inout')
od_var = xs.on_demand()
in_foreign_var = xs.foreign(SomeProcess, 'some_var')
in_foreign_var2 = xs.foreign(AnotherProcess, 'some_var')
out_foreign_var = xs.foreign(AnotherProcess, 'another_var', intent='out')
in_foreign_od_var = xs.foreign(SomeProcess, 'some_od_var')
group_var = xs.group('some_group')
other_attrib = attr.attrib(init=False, cmp=False, repr=False)
other_attr = "this is not a xsimlab variable attribute"
@od_var.compute
def compute_od_var(self):
return 0
class AdaptiveFlowRouter(MultipleFlowRouter):
"""Multiple direction (convergent/divergent) flow router where the
slope exponent is itself a function of slope.
slope_exp = 0.5 + 0.6 * slope
"""
slope_exp = xs.on_demand(description='MFD partioner slope exponent')
def initialize(self):
# this is defined like that in fastscapelib-fortran
self.fs_context.p = -1.
@slope_exp.compute
def _slope_exp(self):
# see https://github.com/fastscape-lem/fastscapelib-fortran/issues/24
warnings.warn(
"'AdaptiveFlowRouter.slope_exp' "
"has no meaningful value.", UserWarning)
return -1
class TectonicForcing:
"""Sum up all tectonic forcing processes and their effect on the
vertical motion of the bedrock surface and the topographic
surface, respectively.
"""
#TODO: remove any_forcing_vars
# see https://github.com/benbovy/xarray-simlab/issues/64
any_forcing_vars = xs.group('any_forcing_upward')
bedrock_forcing_vars = xs.group('bedrock_forcing_upward')
surface_forcing_vars = xs.group('surface_forcing_upward')
bedrock_upward = xs.variable(
dims=[(), ('y', 'x')],
intent='out',
group='bedrock_upward',
description='imposed vertical motion of bedrock surface')
surface_upward = xs.variable(
dims=[(), ('y', 'x')],
intent='out',
group='surface_upward',
description='imposed vertical motion of topographic surface')
grid_area = xs.foreign(UniformRectilinearGrid2D, 'area')
domain_rate = xs.on_demand(
description='domain-integrated volumetric tectonic rate')
@xs.runtime(args='step_delta')
def run_step(self, dt):
self._dt = dt
sum_any = sum(self.any_forcing_vars)
self.bedrock_upward = sum_any + sum(self.bedrock_forcing_vars)
self.surface_upward = sum_any + sum(self.surface_forcing_vars)
@domain_rate.compute
def _domain_rate(self):
return np.sum(self.surface_upward) * self.grid_area / self._dt
class Process3:
var = xs.on_demand()
class FlowRouter:
"""Base process class to route flow on the topographic surface.
Do not use this base class directly in a model! Use one of its
subclasses instead.
However, if you need one or several of the variables declared here
in another process, it is preferable to pass this base class in
:func:`xsimlab.foreign`.
"""
shape = xs.foreign(UniformRectilinearGrid2D, 'shape')
elevation = xs.foreign(SurfaceToErode, 'elevation')
fs_context = xs.foreign(FastscapelibContext, 'context')
stack = xs.variable(dims='node',
intent='out',
description='DFS ordered grid node indices')
nb_receivers = xs.variable(dims='node',
intent='out',
description='number of flow receivers')
receivers = xs.variable(dims=['node', ('node', 'nb_rec_max')],
intent='out',
description='flow receiver node indices')
lengths = xs.variable(dims=['node', ('node', 'nb_rec_max')],
intent='out',
description='out flow path length')
weights = xs.variable(dims=['node', ('node', 'nb_rec_max')],
intent='out',
description='flow partition weights')
nb_donors = xs.variable(dims='node',
intent='out',
description='number of flow donors')
donors = xs.variable(dims=('node', 'nb_don_max'),
intent='out',
description='flow donors node indices')
basin = xs.on_demand(dims=('y', 'x'), description='river catchments')
lake_depth = xs.on_demand(dims=('y', 'x'), description='lake depth')
def route_flow(self):
# must be implemented in sub-classes
pass
def run_step(self):
# bypass fastscapelib_fortran global state
self.fs_context.h = self.elevation.ravel()
self.route_flow()
self.nb_donors = self.fs_context.ndon.astype('int')
# Fortran 1 vs Python 0 index
self.donors = self.fs_context.don.astype('int') - 1
@basin.compute
def _basin(self):
catch = self.fs_context.catch.reshape(self.shape)
# storing basin ids as integers is safer
return (catch * catch.size).astype(np.int)
@lake_depth.compute
def _lake_depth(self):
return self.fs_context.lake_depth.reshape(self.shape).copy()
class Foo:
var = xs.on_demand()
