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()