def test_create_setup(model, in_dataset):
expected = xr.Dataset()
actual = create_setup(model=model)
xr.testing.assert_identical(actual, expected)
ds = create_setup(model=model,
input_vars={
'init_profile': {
'n_points': 5
},
('roll', 'shift'): 1,
'add__offset': ('clock', [1, 2, 3, 4, 5])
},
clocks={
'clock': [0, 2, 4, 6, 8],
'out': [0, 4, 8],
},
master_clock='clock',
output_vars={
'clock': 'profile__u',
'out': [('roll', 'u_diff'), ('add', 'u_diff')],
None: {
'profile': 'u_opp'
}
})
xr.testing.assert_identical(ds, in_dataset)
def test_create_setup(model, in_dataset):
expected = xr.Dataset()
actual = create_setup(model=model, fill_default=False)
xr.testing.assert_identical(actual, expected)
expected = xr.Dataset({"roll__shift": 2})
actual = create_setup(model=model, fill_default=True)
xr.testing.assert_equal(actual, expected)
ds = create_setup(
model=model,
input_vars={
"init_profile": {
"n_points": 5
},
("roll", "shift"): 1,
"add__offset": ("clock", [1, 2, 3, 4, 5]),
},
clocks={
"clock": [0, 2, 4, 6, 8],
"out": [0, 4, 8]
},
main_clock="clock",
output_vars={
"profile__u": "clock",
("roll", "u_diff"): "out",
("add", "u_diff"): "out",
"profile": {
"u_opp": None
},
},
)
xr.testing.assert_identical(ds, in_dataset)
def test_runtime_hook_instance(model, given_as):
flag = [False]
@runtime_hook("run_step", "model", "pre")
def test_hook(_model, context, state):
flag[0] = True
rd = RuntimeHook(test_hook)
in_ds = xs.create_setup(model=model, clocks={"c": [0, 1]})
if given_as == "argument":
in_ds.xsimlab.run(model=model, hooks=[rd])
elif given_as == "context":
with rd:
in_ds.xsimlab.run(model=model)
elif given_as == "register":
rd.register()
in_ds.xsimlab.run(model=model)
assert flag[0] is True
if given_as == "register":
flag[0] = False
rd.unregister()
in_ds.xsimlab.run(model=model)
assert flag[0] is False
def test_write_global_vars(self):
# ensure that variable metadata (dims, etc.) is properly accessed for global references
@xs.process
class Foo:
var = xs.variable(dims="x", global_name="global_var", intent="out")
@xs.process
class Bar:
var = xs.global_ref("global_var")
model = xs.Model({"foo": Foo, "bar": Bar})
in_ds = xs.create_setup(
model=model,
clocks={"clock": [0, 1]},
output_vars={"bar__var": None},
)
store = ZarrSimulationStore(in_ds, model)
model.state[("foo", "var")] = np.array([1, 2, 3])
store.write_output_vars(-1, -1)
ztest = zarr.open_group(store.zgroup.store, mode="r")
np.testing.assert_array_equal(ztest.bar__var, np.array([1, 2, 3]))
def test_encoding(self):
@xs.process
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]
model = xs.Model({"p": P})
in_ds = xs.create_setup(
model=model,
clocks={"clock": [0]},
output_vars={
"p__v1": None,
"p__v2": None,
"p__v3": None,
"p__v4": None
},
)
store = ZarrSimulationStore(
in_ds,
model,
encoding={
"p__v2": {
"fill_value": -1
},
"p__v3": {
"chunks": (10, )
}
},
)
model.state[("p", "v1")] = [0]
model.state[("p", "v3")] = [0]
model.state[("p", "v4")] = [{"foo": "bar"}]
store.write_output_vars(-1, -1)
ztest = zarr.open_group(store.zgroup.store, mode="r")
assert ztest.p__v1.dtype == np.int32
# test encoding precedence ZarrSimulationStore > model variable
assert ztest.p__v2.fill_value == -1
assert ztest.p__v3.chunks == (10, )
assert ztest.p__v4[0] == {"foo": "bar"}
def test_output_vars_by_clock(self, model):
o_vars = {("roll", "u_diff"): "clock", ("add", "u_diff"): None}
ds = xs.create_setup(
model=model,
clocks={"clock": [0, 2, 4, 6, 8]},
output_vars=o_vars,
)
expected = {"clock": [("roll", "u_diff")], None: [("add", "u_diff")]}
assert ds.xsimlab.output_vars_by_clock == expected
def test_runtime_hook_subclass(model):
flag = [False]
class TestHook(RuntimeHook):
@runtime_hook("run_step", "model", "pre")
def test_hook(self, _model, context, state):
flag[0] = True
in_ds = xs.create_setup(model=model, clocks={"c": [0, 1]})
with TestHook():
in_ds.xsimlab.run(model=model)
assert flag[0] is True
def test_create_setup_masterclock_warning(model, in_dataset):
expected = xr.Dataset()
actual = create_setup(model=model, fill_default=False)
xr.testing.assert_identical(actual, expected)
expected = xr.Dataset({"roll__shift": 2})
actual = create_setup(model=model, fill_default=True)
xr.testing.assert_equal(actual, expected)
with pytest.warns(
FutureWarning,
match="master_clock is to be deprecated in favour of main_clock"):
ds = create_setup(
model=model,
input_vars={
"init_profile": {
"n_points": 5
},
("roll", "shift"): 1,
"add__offset": ("clock", [1, 2, 3, 4, 5]),
},
clocks={
"clock": [0, 2, 4, 6, 8],
"out": [0, 4, 8]
},
master_clock="clock",
output_vars={
"profile__u": "clock",
("roll", "u_diff"): "out",
("add", "u_diff"): "out",
"profile": {
"u_opp": None
},
},
)
xr.testing.assert_identical(ds, in_dataset)
def test_runtime_hook_call_frozen(model):
in_ds = xs.create_setup(model=model, clocks={"c": [0, 1]})
@runtime_hook("run_step", "model", "pre")
def change_context(model, context, state):
context["step"] = 0
with pytest.raises(TypeError, match=".*not support item assignment"):
in_ds.xsimlab.run(model=model, hooks=[change_context])
@runtime_hook("run_step", "model", "pre")
def change_state(model, context, state):
state[("p", "u")] = 0
with pytest.raises(TypeError, match=".*not support item assignment"):
in_ds.xsimlab.run(model=model, hooks=[change_state])
def test_runtime_hook_calls(model, event, expected_ncalls, expected_u):
inc = [0]
@runtime_hook(*event)
def test_hook(_model, context, state):
inc[0] += 1
assert "step" in context
assert _model is model
# TODO: get name of current process executed in context?
if expected_u is not None:
assert state[("p", "u")] == expected_u
in_ds = xs.create_setup(model=model, clocks={"c": [0, 1]})
# safe mode disabled so that we can assert _model is model above
in_ds.xsimlab.run(model=model, hooks=[test_hook], safe_mode=False)
assert inc[0] == expected_ncalls
def test_run_decoding(self, decoding, expected):
@xs.process
class P:
var = xs.variable(dims="x",
intent="out",
encoding={"fill_value": -1})
def initialize(self):
self.var = [-1, -1]
m = xs.Model({"p": P})
in_ds = xs.create_setup(
model=m,
clocks={"clock": [0, 1]},
output_vars={"p__var": None},
)
out_ds = in_ds.xsimlab.run(model=m, decoding=decoding)
np.testing.assert_array_equal(out_ds.p__var, expected)
def test_output_vars(self, model):
o_vars = {
("profile", "u_opp"): None,
("profile", "u"): "clock",
("roll", "u_diff"): "out",
("add", "u_diff"): "out",
}
ds = xs.create_setup(
model=model,
clocks={
"clock": [0, 2, 4, 6, 8],
"out": [0, 4, 8],
# snapshot clock with no output variable
"out2": [0, 8],
},
main_clock="clock",
output_vars=o_vars,
)
assert ds.xsimlab.output_vars == o_vars
def test_run_check_dims(self):
@xs.process
class P:
var = xs.variable(dims=["x", ("x", "y")])
m = xs.Model({"p": P})
arr = np.array([[1, 2], [3, 4]])
in_ds = xs.create_setup(
model=m,
clocks={"clock": [1, 2]},
input_vars={"p__var": (("y", "x"), arr)},
output_vars={"p__var": None},
)
out_ds = in_ds.xsimlab.run(model=m, check_dims=None)
actual = out_ds.p__var.values
np.testing.assert_array_equal(actual, arr)
with pytest.raises(ValueError, match=r"Invalid dimension.*"):
in_ds.xsimlab.run(model=m, check_dims="strict")
out_ds = in_ds.xsimlab.run(model=m,
check_dims="transpose",
safe_mode=False)
actual = out_ds.p__var.values
np.testing.assert_array_equal(actual, arr)
np.testing.assert_array_equal(m.p.var, arr.transpose())
in_ds2 = in_ds.xsimlab.update_vars(model=m,
output_vars={"p__var": "clock"})
# TODO: fix update output vars time-independet -> dependent
# currently need the workaround below
in_ds2.attrs = {}
out_ds = in_ds2.xsimlab.run(model=m, check_dims="transpose")
actual = out_ds.p__var.isel(clock=-1).values
np.testing.assert_array_equal(actual, arr)
def test_fill_values(self):
@xs.process
class Foo:
v_int64 = xs.variable(dims="x", intent="out")
v_float64 = xs.variable(dims="x", intent="out")
v_uint8 = xs.variable(dims="x",
intent="out",
encoding={"dtype": np.uint8})
v_string = xs.variable(dims="x", intent="out")
v_bool = xs.variable(dims="x", intent="out")
def initialize(self):
self.v_int64 = [0, np.iinfo("int64").max]
self.v_float64 = [0.0, np.nan]
self.v_uint8 = [0, 255]
self.v_string = ["hello", ""]
self.v_bool = [True, False]
model = xs.Model({"foo": Foo})
in_ds = xs.create_setup(
model=model,
clocks={"clock": [0, 1]},
output_vars={
"foo__v_int64": None,
"foo__v_float64": None,
"foo__v_uint8": None,
"foo__v_string": None,
"foo__v_bool": None,
},
)
out_ds = in_ds.xsimlab.run(model=model)
np.testing.assert_equal(out_ds["foo__v_int64"].data, [0, np.nan])
np.testing.assert_equal(out_ds["foo__v_float64"].data, [0.0, np.nan])
np.testing.assert_equal(out_ds["foo__v_uint8"].data, [0, np.nan])
# np.testing.assert_equal does not work for "object" dtypes, so test each value explicitly:
assert out_ds["foo__v_string"].data[0] == "hello"
assert np.isnan(out_ds["foo__v_string"].data[1])
assert out_ds["foo__v_bool"].data[0] == True
assert np.isnan(out_ds["foo__v_bool"].data[1])
def test_finalize_always_called():
@xs.process
class P:
var = xs.variable(intent="out")
def initialize(self):
self.var = "initialized"
raise RuntimeError()
def finalize(self):
self.var = "finalized"
model = xs.Model({"p": P})
in_dataset = xs.create_setup(model=model, clocks={"clock": [0, 1]})
driver = XarraySimulationDriver(in_dataset, model)
try:
driver.run_model()
except RuntimeError:
pass
assert model.state[("p", "var")] == "finalized"
def test_run_batch_dim(self, dims, data, clock, parallel, scheduler):
@xs.process
class P:
in_var = xs.variable(dims=[(), "x"])
out_var = xs.variable(dims=[(), "x"], intent="out")
idx_var = xs.index(dims="x")
def initialize(self):
self.idx_var = [0, 1]
def run_step(self):
self.out_var = self.in_var * 2
m = xs.Model({"p": P})
in_ds = xs.create_setup(
model=m,
clocks={"clock": [0, 1, 2]},
input_vars={"p__in_var": (dims, data)},
output_vars={"p__out_var": clock},
)
out_ds = in_ds.xsimlab.run(
model=m,
batch_dim="batch",
parallel=parallel,
scheduler=scheduler,
store=zarr.TempStore(),
)
if clock is None:
coords = {}
else:
coords = {"clock": in_ds["clock"]}
expected = xr.DataArray(data, dims=dims, coords=coords) * 2
xr.testing.assert_equal(out_ds["p__out_var"], expected)
def test_signal_model_level(trigger, signal, expected):
@xs.process
class Foo:
v = xs.variable(intent="out")
vv = xs.variable(intent="out")
def initialize(self):
self.v = 0.0
self.vv = 10.0
def run_step(self):
self.v += 1.0
def finalize_step(self):
self.v += 1.0
@xs.runtime_hook("run_step", level="model", trigger=trigger)
def hook_func(model, context, state):
if context["step"] == 1:
return signal
model = xs.Model({"foo": Foo})
ds_in = xs.create_setup(
model=model,
clocks={"clock": range(4)},
output_vars={
"foo__v": "clock",
"foo__vv": None
},
)
ds_out = ds_in.xsimlab.run(model=model, hooks=[hook_func])
np.testing.assert_equal(ds_out.foo__v.values, expected)
# ensure that clock-independent output variables are properly
# saved even when the simulation stops early
assert ds_out.foo__vv == 10.0
def test_resize_zarr_dataset(self):
@xs.process
class P:
arr = xs.variable(dims="x", intent="out")
model = xs.Model({"p": P})
in_ds = xs.create_setup(
model=model,
clocks={"clock": [0, 1, 2]},
output_vars={"p__arr": "clock"},
)
store = ZarrSimulationStore(in_ds, model)
for step, size in zip([0, 1, 2], [1, 3, 2]):
model.state[("p", "arr")] = np.ones(size)
store.write_output_vars(-1, step)
ztest = zarr.open_group(store.zgroup.store, mode="r")
expected = np.array([[1.0, np.nan, np.nan], [1.0, 1.0, 1.0],
[1.0, 1.0, np.nan]])
np.testing.assert_array_equal(ztest.p__arr, expected)
in_ds = xsimlab.create_setup(
model=custom_model,
clocks={
'time': time,
'out': time,
},
master_clock='time',
input_vars={
'grid__shape': ('shape_yx', [ny, nx]),
'grid__length': ('shape_yx', [lenghty, lenghtx]),
'boundary__status':
('border',
['fixed_value', 'fixed_value', 'fixed_value', 'fixed_value']),
'uplift_func': {
'uplift_rate': ('time', U),
'coeff': 1,
'axis': 2
},
'square_basement': {
'x_origin_position': 80,
'y_origin_position': 45,
'x_lenght': 100,
'y_lenght': 20,
'basement_k_coef': ('time', K_basement),
'rock_k_coef': 2e-6,
'basement_diff': 0.1,
'rock_diff': 0.1
},
# 'circle_basement': {
# 'x_origin_position': 100,
# 'y_origin_position': 60,
# 'radius': 25,
# 'basement_k_coef': ('time', K_basement),
# 'rock_k_coef': 2e-6,
# 'basement_diff': 0.1,
# 'rock_diff': 0.1
# },
'spl': {
'area_exp': 0.6,
'slope_exp': 1.5
},
},
output_vars={
'out': [
'topography__elevation', 'drainage__area', 'flow__basin',
'terrain__slope', 'spl__erosion'
],
None: ['boundary__border', 'grid__x', 'grid__y', 'grid__spacing'],
})
def test_main_clock_access():
@xs.process
class Foo:
a = xs.variable(intent="out", dims=xs.MAIN_CLOCK)
b = xs.variable(intent="out", dims=xs.MAIN_CLOCK)
@xs.runtime(args=["main_clock_values", "main_clock_dataarray"])
def initialize(self, clock_values, clock_array):
self.a = clock_values * 2
np.testing.assert_equal(self.a, [0, 2, 4, 6])
self.b = clock_array * 2
assert clock_array.dims[0] == "clock"
assert all(clock_array[clock_array.dims[0]].data == [0, 1, 2, 3])
@xs.runtime(args=["step_delta", "step"])
def run_step(self, dt, n):
assert self.a[n] == 2 * n
self.a[n] += 1
model = xs.Model({"foo": Foo})
ds_in = xs.create_setup(
model=model,
clocks={"clock": range(4)},
input_vars={},
output_vars={"foo__a": None},
)
ds_out = ds_in.xsimlab.run(model=model)
assert all(ds_out.foo__a.data == [1, 3, 5, 6])
# test for error when another dim has the same name as xs.MAIN_CLOCK
@xs.process
class DoubleMainClockDim:
a = xs.variable(intent="out", dims=("clock", xs.MAIN_CLOCK))
def initialize(self):
self.a = [[1, 2, 3], [3, 4, 5]]
def run_step(self):
self.a += self.a
model = xs.Model({"foo": DoubleMainClockDim})
with pytest.raises(ValueError, match=r"Main clock:*"):
xs.create_setup(
model=model,
clocks={
"clock": [0, 1, 2, 3]
},
input_vars={},
output_vars={
"foo__a": None
},
).xsimlab.run(model)
# test for error when trying to put xs.MAIN_CLOCK as a dim in an input var
with pytest.raises(
ValueError,
match="Do not pass xs.MAIN_CLOCK into input vars dimensions"):
a = xs.variable(intent="in", dims=xs.MAIN_CLOCK)
with pytest.raises(
ValueError,
match="Do not pass xs.MAIN_CLOCK into input vars dimensions"):
b = xs.variable(intent="in", dims=(xs.MAIN_CLOCK, ))
with pytest.raises(
ValueError,
match="Do not pass xs.MAIN_CLOCK into input vars dimensions"):
c = xs.variable(intent="in", dims=["a", ("a", xs.MAIN_CLOCK)])
def test_hook_arg_type(model):
in_ds = xs.create_setup(model=model, clocks={"c": [0, 1]})
with pytest.raises(TypeError, match=".*not a RuntimeHook.*"):
in_ds.xsimlab.run(model=model, hooks=[1])
def run_fastscape(k_sp,
k_diff,
u_rate,
x_size=601,
y_size=401,
spacing=200.,
time_step=1e5,
time_total=1e7):
"""
Run "Fastscape" landscape evolution model.
This version of Fastscape includes the following processes:
- uniform, block uplift at a constant rate through time ;
- bedrock channel erosion using the stream power law ;
- hillslope erosion using linear diffusion.
The total simulation time is 10 Myr.
The simulation start with a random, nearly flat topography.
Topographic elevation remains fixed at the domain boundary.
Parameters
----------
k_sp : float
Stream power law coefficient. The drainage area and slope
exponents of the stream power law are fixed to 0.4 and 1,
respectively.
k_diff : float
Hillslope diffusivity (units: m**2/yr).
u_rate : float
Uplift rate (units: m/yr).
Other parameters
----------------
x_size : int, optional
Grid size in x (i.e., number of columns).
y_size : int, optional
Grid size in y (i.e., number of rows).
spacing : float, optional
Uniform grid spacing in x and y (units: m).
time_step : float, optional
Simulation time step (units: yr).
time_total : float, optional
Total simulation duration (units: yr).
Returns
-------
elevation : numpy.ndarray
Topographic elevation at the end of the simulation (units: m).
Array shape is (`y_size`, `x_size`).
"""
in_ds = xsimlab.create_setup(
model=fastscape_base_model,
clocks={'time': {
'end': time_total,
'step': time_step
}},
input_vars={
'grid': {
'x_size': x_size,
'y_size': y_size,
'x_spacing': spacing,
'y_spacing': spacing
},
'flow_routing': {
'pit_method': 'mst_linear'
},
'block_uplift': {
'u_coef': u_rate
},
'spower': {
'k_coef': k_sp,
'm_exp': 0.4,
'n_exp': 1
},
'diffusion': {
'k_coef': k_diff
},
'topography': {
'elevation': (('y', 'x'), np.random.rand(y_size, x_size))
}
},
snapshot_vars={
None: {
'grid': ('x', 'y'),
'topography': 'elevation'
},
})
out_ds = (in_ds.xsimlab.run(model=fastscape_base_model).set_index(
y='grid__y', x='grid__x'))
return out_ds.topography__elevation.values
