def test_constructor(self):
with pytest.raises(KeyError) as excinfo:
xs.Model({"init_profile": InitProfile})
assert "Process class 'Profile' missing" in str(excinfo.value)
# test empty model
assert len(xs.Model({})) == 0
def test_context_manager(self):
m1 = xs.Model({})
m2 = xs.Model({})
with pytest.raises(ValueError, match=r"There is already a model object.*"):
with m1, m2:
pass
def test_global_variable():
@xs.process
class Foo:
var = xs.variable(global_name="global_var")
idx = xs.index(dims="x", global_name="global_idx")
obj = xs.any_object(global_name="global_obj")
def initialize(self):
self.idx = np.array([1, 1])
self.obj = 2
@xs.process
class Bar:
var = xs.global_ref("global_var")
idx = xs.global_ref("global_idx")
obj = xs.global_ref("global_obj")
actual = xs.variable(intent="out")
def initialize(self):
self.actual = self.var + self.obj * np.sum(self.idx)
@xs.process
class Baz:
# foreign pointing to global reference Bar.var
# --> must pass through and actually points to Foo.var
var = xs.foreign(Bar, "var", intent="out")
def initialize(self):
self.var = 1
model = xs.Model({"foo": Foo, "bar": Bar, "baz": Baz})
model.execute("initialize", {})
assert model.state[("foo", "var")] == 1
assert model.state[("bar", "actual")] == 5
# -- test errors
@xs.process
class NotFound:
var = xs.global_ref("missing")
with pytest.raises(KeyError,
match="No variable with global name 'missing' found.*"):
xs.Model({"foo": Foo, "not_found": NotFound})
@xs.process
class Duplicate:
var = xs.variable(global_name="global_var")
with pytest.raises(ValueError,
match="Found multiple variables with global name.*"):
xs.Model({"foo": Foo, "bar": Bar, "dup": Duplicate})
def test_process_executor():
@xs.process
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
return xs.RuntimeSignal.BREAK
@od_var.compute
def _dummy(self):
return None
m = xs.Model({"p": P})
executor = m.p.__xsimlab_executor__
state = {("p", "in_var"): 1}
state_out, signal_out = executor.execute(m.p,
SimulationStage.RUN_STEP, {},
state=state)
assert state_out == {("p", "out_var"): 2}
assert signal_out == xs.RuntimeSignal.BREAK
state_out, signal_out = executor.execute(m.p,
SimulationStage.INITIALIZE, {},
state=state)
assert state_out == {}
assert signal_out == xs.RuntimeSignal.NONE
def test_group_dict_variable():
@xs.process
class Foo:
a = xs.variable(groups="g", intent="out")
def initialize(self):
self.a = 1
@xs.process
class Bar:
b = xs.variable(groups="g", intent="out")
def initialize(self):
self.b = 2
@xs.process
class Baz:
c = xs.group_dict("g")
actual = xs.variable(intent="out")
def initialize(self):
self.actual = self.c
model = xs.Model({"foo": Foo, "bar": Bar, "baz": Baz})
model.execute("initialize", {})
assert model.state[("baz", "actual")] == Frozen({
("foo", "a"): 1,
("bar", "b"): 2
})
def model():
return xs.Model({
'roll': Roll,
'add': AddOnDemand,
'profile': Profile,
'init_profile': InitProfile
})
def model():
return xs.Model({
"roll": Roll,
"add": AddOnDemand,
"profile": Profile,
"init_profile": InitProfile,
})
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_object_variable(self):
@xs.process
class P:
obj = xs.any_object()
m = xs.Model({"p": P})
assert m.p.__xsimlab_state_keys__["obj"] == ("p", "obj")
def test_set_output_object_vars(self):
@xs.process
class P:
obj = xs.any_object()
m = xs.Model({"p": P})
ds = xr.Dataset()
with pytest.raises(ValueError,
match=r"Object variables can't be set.*"):
ds.xsimlab._set_output_vars(m, {("p", "obj"): None})
def test_set_output_object_or_group_vars(self, field):
@xs.process
class P:
var = field
m = xs.Model({"p": P})
ds = xr.Dataset()
with pytest.raises(ValueError,
match=r"Object or group variables can't be set.*"):
ds.xsimlab._set_output_vars(m, {("p", "var"): None})
def test_sort_processes_cycle(self, model):
@xs.process
class Foo:
in_var = xs.variable()
out_var = xs.variable(intent="out")
@xs.process
class Bar:
in_foreign = xs.foreign(Foo, "out_var")
out_foreign = xs.foreign(Foo, "in_var", intent="out")
with pytest.raises(RuntimeError, match=r"Cycle detected.*"):
xs.Model({"foo": Foo, "bar": Bar})
def test_multiple_groups(self):
@xs.process
class A:
v = xs.variable(groups=["g1", "g2"])
@xs.process
class B:
g1 = xs.group("g1")
g2 = xs.group("g2")
m = xs.Model({"a": A, "b": B})
assert m.b.__xsimlab_state_keys__["g1"] == [("a", "v")]
assert m.b.__xsimlab_state_keys__["g2"] == [("a", "v")]
def test_sort_processes_cycle(self, model):
@xs.process
class Foo(object):
in_var = xs.variable()
out_var = xs.variable(intent='out')
@xs.process
class Bar(object):
in_foreign = xs.foreign(Foo, 'out_var')
out_foreign = xs.foreign(Foo, 'in_var', intent='out')
with pytest.raises(RuntimeError) as excinfo:
xs.Model({'foo': Foo, 'bar': Bar})
assert "Cycle detected" in str(excinfo.value)
def test_on_demand_cache():
@xs.process
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)
@xs.process
class P2:
var = xs.foreign(P1, "var")
cached_var = xs.foreign(P1, "cached_var")
view = xs.variable(dims="x", intent="out")
cached_view = xs.variable(dims="x", intent="out")
def run_step(self):
self.view = self.var
self.cached_view = self.cached_var
@xs.process
class P3:
p1_view = xs.foreign(P1, "var")
p1_cached_view = xs.foreign(P1, "cached_var")
p2_view = xs.foreign(P2, "view")
p2_cached_view = xs.foreign(P2, "cached_view")
def initialize(self):
self._p1_cached_view_init = self.p1_cached_view
def run_step(self):
# P1.var's compute method called twice
assert not np.all(self.p1_view == self.p2_view)
# P1.cached_var's compute method called once
assert self.p1_cached_view is self.p2_cached_view
# check cache cleared between simulation stages
assert not np.all(self.p1_cached_view == self._p1_cached_view_init)
model = xs.Model({"p1": P1, "p2": P2, "p3": P3})
model.execute("initialize", {})
model.execute("run_step", {})
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_signal_process_level(step, trigger, signal, break_bar, expected_v1,
expected_v2, parallel):
@xs.process
class Foo:
v1 = xs.variable(intent="out")
v2 = xs.variable()
def initialize(self):
self.v1 = 0
def run_step(self):
self.v1 = 1
@xs.process
class Bar:
v2 = xs.foreign(Foo, "v2", intent="out")
def initialize(self):
self.v2 = 0
def run_step(self):
self.v2 = 2
if break_bar:
return xs.RuntimeSignal.BREAK
def hook_func(model, context, state):
if context["step"] == 1:
return signal
hook_dict = {SimulationStage.RUN_STEP: {"process": {trigger: [hook_func]}}}
model = xs.Model({"foo": Foo, "bar": Bar})
model.execute("initialize", {})
model.execute("run_step", {"step": step},
hooks=hook_dict,
parallel=parallel)
assert model.state[("foo", "v1")] == expected_v1
assert model.state[("foo", "v2")] == expected_v2
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_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_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_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_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_process_executor():
@xs.process
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
m = xs.Model({"p": P})
executor = m.p.__xsimlab_executor__
state = {("p", "in_var"): 1}
expected = {("p", "out_var"): 2}
actual = executor.execute(m.p, SimulationStage.RUN_STEP, {}, state=state)
assert actual == expected
assert executor.execute(m.p, SimulationStage.INITIALIZE, {},
state=state) == {}
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)
def test_model_repr(simple_model, simple_model_repr):
assert repr_model(simple_model) == simple_model_repr
expected = "<xsimlab.Model (0 processes, 0 inputs)>\n"
assert repr(xs.Model({})) == expected
import xsimlab as xs
import phydra
from phydra.components.main import Backend
# ``Bootstrap model`` has the minimal set of components required to
# simulate a
Bootstrap_model = xs.Model({
'Grid': Backend
})
# ``NPZD Slab model`` has the minimal set of components required to
# simulate a
ZeroD_NPZD_Slab_model = xs.Model({
'Grid': Backend
})
# ``ZeroD NPxZx Chemostat model`` has the minimal set of components required to
# simulate a
ZeroD_NPxZx_Chemostat_model = xs.Model({
'Grid': Backend
})
def no_init_model():
return xs.Model({"roll": Roll, "add": Add, "profile": Profile})
def simple_model():
return xs.Model({"roll": Roll, "profile": Profile})
def initialize(self):
self.x = np.arange(0, self.length, self.spacing)
self.u = np.exp(-1 / self.scale ** 2 * (self.x - self.loc) ** 2)
@xs.runtime(args="step_delta")
def run_step(self, dt):
factor = (self.v * dt) / (2 * self.spacing)
u_left = np.roll(self.u, 1)
u_right = np.roll(self.u, -1)
self.u1 = 0.5 * (u_right + u_left) - factor * (u_right - u_left)
def finalize_step(self):
self.u = self.u1
advect_model_raw = xs.Model({"advect": AdvectionLax1D})
@xs.process
class UniformGrid1D:
"""Create a 1-dimensional, equally spaced grid."""
spacing = xs.variable(description="uniform spacing", static=True)
length = xs.variable(description="total length", static=True)
x = xs.index(dims="x")
def initialize(self):
self.x = np.arange(0, self.length, self.spacing)
@xs.process