def xs_lpyprocess(name, lpyfile, graphicalparameters = None, modulestoconsider = None, globaldependencies = {}, propertymapping = {}):
""" Generate the xs process class under the given name with adequate properties from the lpy file. """
externs, modules = parse_extern_modules(lpyfile)
properties = gen_properties(externs, modules, modulestoconsider, globaldependencies, propertymapping)
properties['lpyfile'] = lpyfile
properties['graphicalparameters'] = graphicalparameters
properties['lscene'] = xs.any_object()
properties['lstring'] = xs.any_object()
process = xs.process(type(name, (AbstractLpyProcess,), properties))
get_caller_namespace()[name] = process
return process
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 FastscapelibContext:
"""This process takes care of proper initialization,
update and clean-up of fastscapelib-fortran internal
state.
"""
shape = xs.foreign(UniformRectilinearGrid2D, 'shape')
length = xs.foreign(UniformRectilinearGrid2D, 'length')
ibc = xs.foreign(BorderBoundary, 'ibc')
context = xs.any_object(
description='accessor to fastscapelib-fortran internal variables')
def initialize(self):
fs.fastscape_init()
fs.fastscape_set_nx_ny(*np.flip(self.shape))
fs.fastscape_setup()
fs.fastscape_set_xl_yl(*np.flip(self.length))
fs.fastscape_set_bc(self.ibc)
self.context = SerializableFastscapeContext()
@xs.runtime(args='step_delta')
def run_step(self, dt):
# fastscapelib-fortran runtime routines use dt from context
self.context["dt"] = dt
def finalize(self):
fs.fastscape_destroy()
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
class Backend:
"""this object contains the backend model and is modified or read by all other components"""
solver_type = xs.variable(intent='in')
m = xs.any_object(description='model backend instance is stored here')
def initialize(self):
print('initializing model backend')
self.m = PhydraCore(self.solver_type)
def finalize(self):
print('finalizing: cleanup')
self.m.cleanup() # for now only affects gekko solve
class P:
obj = xs.any_object()
class ProfileFix(Profile):
# limitation of using distributed for single-model parallelism
# internal instance attributes created and used in multiple stage
# methods are not supported.
u_change = xs.any_object()
class TestSimlabAccessor:
_clock_key = xr_accessor.SimlabAccessor._clock_key
_main_clock_key = xr_accessor.SimlabAccessor._main_clock_key
_output_vars_key = xr_accessor.SimlabAccessor._output_vars_key
def test_clock_coords(self):
ds = xr.Dataset(
coords={
"mclock": (
"mclock",
[0, 1, 2],
{
self._clock_key: 1,
self._main_clock_key: 1
},
),
"sclock": ("sclock", [0, 2], {
self._clock_key: 1
}),
"no_clock": ("no_clock", [3, 4]),
})
assert set(ds.xsimlab.clock_coords) == {"mclock", "sclock"}
def test_main_clock_coords(self):
ds = xr.Dataset(
coords={
"mclock": (
"mclock",
[0, 1, 2],
{
self._clock_key: 1,
self._main_clock_key: 1
},
),
"sclock": ("sclock", [0, 2], {
self._clock_key: 1
}),
"no_clock": ("no_clock", [3, 4]),
})
xr.testing.assert_equal(ds.xsimlab.main_clock_coord, ds.mclock)
def test_master_clock_coords_warning(self):
ds = xr.Dataset(
coords={
"mclock": (
"mclock",
[0, 1, 2],
{
self._clock_key: 1,
self._main_clock_key: 1
},
),
"sclock": ("sclock", [0, 2], {
self._clock_key: 1
}),
"no_clock": ("no_clock", [3, 4]),
})
with pytest.warns(
FutureWarning,
match=
"master_clock_coord is to be deprecated in favour of main_clock",
):
ds.xsimlab.master_clock_coord
def test_clock_sizes(self):
ds = xr.Dataset(
coords={
"clock1": ("clock1", [0, 1, 2], {
self._clock_key: 1
}),
"clock2": ("clock2", [0, 2], {
self._clock_key: 1
}),
"no_clock": ("no_clock", [3, 4]),
})
assert ds.xsimlab.clock_sizes == {"clock1": 3, "clock2": 2}
def test_main_clock_dim(self):
attrs = {self._clock_key: 1, self._main_clock_key: 1}
ds = xr.Dataset(coords={"clock": ("clock", [1, 2], attrs)})
assert ds.xsimlab.main_clock_dim == "clock"
assert ds.xsimlab._main_clock_dim == "clock" # cache
assert ds.xsimlab.main_clock_dim == "clock" # get cached value
ds = xr.Dataset()
assert ds.xsimlab.main_clock_dim is None
def test_master_clock_dim_warning(self):
attrs = {self._clock_key: 1, self._main_clock_key: 1}
ds = xr.Dataset(coords={"clock": ("clock", [1, 2], attrs)})
with pytest.warns(
FutureWarning,
match=
"master_clock is to be deprecated in favour of main_clock",
):
assert ds.xsimlab.master_clock_dim == "clock"
# internally, _main_clock_dim is used
assert ds.xsimlab._main_clock_dim == "clock" # cache
with pytest.warns(
FutureWarning,
match=
"master_clock is to be deprecated in favour of main_clock",
):
assert ds.xsimlab.master_clock_dim == "clock" # get cached value
ds = xr.Dataset()
with pytest.warns(
FutureWarning,
match=
"master_clock is to be deprecated in favour of main_clock",
):
assert ds.xsimlab.master_clock_dim is None
def test_nsteps(self):
attrs = {self._clock_key: 1, self._main_clock_key: 1}
ds = xr.Dataset(coords={"clock": ("clock", [1, 2, 3], attrs)})
assert ds.xsimlab.nsteps == 2
ds = xr.Dataset()
assert ds.xsimlab.nsteps == 0
def test_get_output_save_steps(self):
attrs = {self._clock_key: 1, self._main_clock_key: 1}
ds = xr.Dataset(
coords={
"clock": ("clock", [0, 1, 2, 3, 4], attrs),
"clock1": ("clock1", [0, 2, 4], {
self._clock_key: 1
}),
"clock2": ("clock2", [0, 4], {
self._clock_key: 1
}),
})
expected = xr.Dataset(
coords={"clock": ("clock", [0, 1, 2, 3, 4], attrs)},
data_vars={
"clock1": ("clock", [True, False, True, False, True]),
"clock2": ("clock", [True, False, False, False, True]),
},
)
xr.testing.assert_identical(ds.xsimlab.get_output_save_steps(),
expected)
def test_set_input_vars(self, model, in_dataset):
in_vars = {
("init_profile", "n_points"): 5,
("roll", "shift"): 1,
("add", "offset"): ("clock", [1, 2, 3, 4, 5]),
}
ds = xr.Dataset(coords={"clock": [0, 2, 4, 6, 8]})
ds.xsimlab._set_input_vars(model, in_vars)
for vname in ("init_profile__n_points", "roll__shift", "add__offset"):
# xr.testing.assert_identical also checks attrs of coordinates
# (not needed here)
xr.testing.assert_equal(ds[vname], in_dataset[vname])
assert ds[vname].attrs == in_dataset[vname].attrs
# test errors
in_vars[("not_an", "input_var")] = None
with pytest.raises(KeyError) as excinfo:
ds.xsimlab._set_input_vars(model, in_vars)
assert "not valid key(s)" in str(excinfo.value)
# test implicit dimension label
in_vars = {("add", "offset"): [1, 2, 3, 4, 5]}
ds.xsimlab._set_input_vars(model, in_vars)
assert ds["add__offset"].dims == ("x", )
# test implicit dimension label error
in_vars = {("roll", "shift"): [1, 2]}
with pytest.raises(TypeError,
match=r"Could not get dimension labels.*"):
ds.xsimlab._set_input_vars(model, in_vars)
def test_update_clocks(self, model):
ds = xr.Dataset()
with pytest.raises(ValueError, match="Cannot determine which clock.*"):
ds.xsimlab.update_clocks(model=model, clocks={})
ds = xr.Dataset()
with pytest.raises(ValueError, match="Cannot determine which clock.*"):
ds.xsimlab.update_clocks(model=model,
clocks={
"clock": [0, 1, 2],
"out": [0, 2]
})
ds = xr.Dataset()
with pytest.raises(KeyError, match="Main clock dimension name.*"):
ds.xsimlab.update_clocks(
model=model,
clocks={"clock": [0, 1, 2]},
main_clock="non_existing_clock_dim",
)
ds = xr.Dataset()
with pytest.raises(ValueError, match="Invalid dimension.*"):
ds.xsimlab.update_clocks(
model=model,
clocks={"clock": ("x", [0, 1, 2])},
)
ds = xr.Dataset()
with pytest.raises(ValueError, match=".*not synchronized.*"):
ds.xsimlab.update_clocks(
model=model,
clocks={
"clock": [0, 1, 2],
"out": [0, 0.5, 2]
},
main_clock="clock",
)
ds = xr.Dataset()
ds = ds.xsimlab.update_clocks(model=model, clocks={"clock": [0, 1, 2]})
assert ds.xsimlab.main_clock_dim == "clock"
ds.clock.attrs[self._output_vars_key] = "profile__u"
ds = ds.xsimlab.update_clocks(
model=model,
clocks={"clock": [0, 1, 2]},
main_clock={
"dim": "clock",
"units": "days since 1-1-1 0:0:0",
"calendar": "365_days",
},
)
np.testing.assert_array_equal(ds.clock.values, [0, 1, 2])
assert "units" in ds.clock.attrs
assert "calendar" in ds.clock.attrs
assert ds.clock.attrs[self._output_vars_key] == "profile__u"
new_ds = ds.xsimlab.update_clocks(
model=model,
clocks={"clock2": [0, 0.5, 1, 1.5, 2]},
main_clock="clock2",
)
assert new_ds.xsimlab.main_clock_dim == "clock2"
new_ds = ds.xsimlab.update_clocks(model=model, clocks={"out2": [0, 2]})
assert new_ds.xsimlab.main_clock_dim == "clock"
new_ds = ds.xsimlab.update_clocks(model=model,
clocks={"clock": [0, 2, 4]})
assert new_ds.xsimlab.main_clock_dim == "clock"
np.testing.assert_array_equal(new_ds.clock.values, [0, 2, 4])
def test_update_clocks_master_clock_warning(self, model):
ds = xr.Dataset()
ds = ds.xsimlab.update_clocks(model=model, clocks={"clock": [0, 1, 2]})
assert ds.xsimlab.main_clock_dim == "clock"
ds.clock.attrs[self._output_vars_key] = "profile__u"
# assert that a warning is raised with correct use of update master clock
with pytest.warns(
FutureWarning,
match=
"master_clock is to be deprecated in favour of main_clock",
):
ds = ds.xsimlab.update_clocks(
model=model,
clocks={"clock": [0, 1, 2]},
master_clock={
"dim": "clock",
"units": "days since 1-1-1 0:0:0",
"calendar": "365_days",
},
)
np.testing.assert_array_equal(ds.clock.values, [0, 1, 2])
assert "units" in ds.clock.attrs
assert "calendar" in ds.clock.attrs
assert ds.clock.attrs[self._output_vars_key] == "profile__u"
with pytest.warns(
FutureWarning,
match=
"master_clock is to be deprecated in favour of main_clock",
):
new_ds = ds.xsimlab.update_clocks(
model=model,
clocks={"clock2": [0, 0.5, 1, 1.5, 2]},
master_clock="clock2",
)
assert new_ds.xsimlab.main_clock_dim == "clock2"
def test_update_vars(self, model, in_dataset):
ds = in_dataset.xsimlab.update_vars(
model=model,
input_vars={("roll", "shift"): 2},
output_vars={("profile", "u"): "out"},
)
assert not ds["roll__shift"].equals(in_dataset["roll__shift"])
assert not ds["out"].identical(in_dataset["out"])
def test_update_vars_promote_to_coords(self, model, in_dataset):
# It should be possible to update an input variable with a dimension
# label that cooresponds to its name (turned into a coordinate). This
# should not raise any merge conflict error
ds = in_dataset.xsimlab.update_vars(
model=model,
input_vars={"roll__shift": ("roll__shift", [1, 2])},
)
assert "roll__shift" in ds.coords
def test_reset_vars(self, model, in_dataset):
# add new variable
ds = xr.Dataset().xsimlab.reset_vars(model)
assert ds["roll__shift"] == 2
# overwrite existing variable
reset_ds = in_dataset.xsimlab.reset_vars(model)
assert reset_ds["roll__shift"] == 2
def test_filter_vars(self, simple_model, in_dataset):
in_dataset["not_a_xsimlab_model_input"] = 1
filtered_ds = in_dataset.xsimlab.filter_vars(model=simple_model)
assert "add__offset" not in filtered_ds
assert "not_a_xsimlab_model_input" not in filtered_ds
assert sorted(filtered_ds.xsimlab.clock_coords) == ["clock", "out"]
assert filtered_ds.out.attrs[self._output_vars_key] == "roll__u_diff"
# test unchanged attributes in original dataset
assert in_dataset.out.attrs[
self._output_vars_key] == "roll__u_diff,add__u_diff"
assert in_dataset.attrs[self._output_vars_key] == "profile__u_opp"
def test_set_output_vars(self, model):
ds = xr.Dataset()
ds["clock"] = (
"clock",
[0, 2, 4, 6, 8],
{
self._clock_key: 1,
self._main_clock_key: 1
},
)
ds["out"] = ("out", [0, 4, 8], {self._clock_key: 1})
ds["not_a_clock"] = ("not_a_clock", [0, 1])
with pytest.raises(KeyError, match=r".*not valid key.*"):
ds.xsimlab._set_output_vars(model, {("invalid", "var"): None})
ds.xsimlab._set_output_vars(model, {("profile", "u_opp"): None})
assert ds.attrs[self._output_vars_key] == "profile__u_opp"
ds.xsimlab._set_output_vars(model, {
("roll", "u_diff"): "out",
("add", "u_diff"): "out"
})
expected = "roll__u_diff,add__u_diff"
assert ds["out"].attrs[self._output_vars_key] == expected
with pytest.raises(ValueError, match=r".not a valid clock.*"):
ds.xsimlab._set_output_vars(model,
{("profile", "u"): "not_a_clock"})
with pytest.warns(FutureWarning):
ds.xsimlab._set_output_vars(model, {None: ("profile", "u_opp")})
with pytest.warns(FutureWarning):
ds.xsimlab._set_output_vars(model, {"out": ("profile", "u_opp")})
@pytest.mark.parametrize(
"field",
[xs.any_object(), xs.group("g"),
xs.group_dict("g")])
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_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_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_run(self, model, in_dataset, out_dataset, parallel, scheduler):
@xs.process
class ProfileFix(Profile):
# limitation of using distributed for single-model parallelism
# internal instance attributes created and used in multiple stage
# methods are not supported.
u_change = xs.any_object()
m = model.update_processes({"profile": ProfileFix})
out_ds = in_dataset.xsimlab.run(model=m,
parallel=parallel,
scheduler=scheduler)
xr.testing.assert_equal(out_ds.load(), out_dataset)
def test_run_safe_mode(self, model, in_dataset):
# safe mode True: ensure model is cloned (empty state)
_ = in_dataset.xsimlab.run(model=model, safe_mode=True)
assert model.state == {}
# safe mode False: model not cloned (non empty state)
_ = in_dataset.xsimlab.run(model=model, safe_mode=False)
assert model.state != {}
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_run_validate(self, model, in_dataset):
in_dataset["roll__shift"] = 2.5
# no input validation -> raises within np.roll()
with pytest.raises(TypeError,
match=r"slice indices must be integers.*"):
in_dataset.xsimlab.run(model=model, validate=None)
# input validation at initialization -> raises within attr.validate()
with pytest.raises(TypeError, match=r".*'int'.*"):
in_dataset.xsimlab.run(model=model, validate="inputs")
in_dataset["roll__shift"] = ("clock", [1, 2.5, 1, 1, 1])
# input validation at runtime -> raises within attr.validate()
with pytest.raises(TypeError, match=r".*'int'.*"):
in_dataset.xsimlab.run(model=model, validate="inputs")
@xs.process
class SetRollShift:
shift = xs.foreign(Roll, "shift", intent="out")
def initialize(self):
self.shift = 2.5
m = model.update_processes({"set_shift": SetRollShift})
# no internal validation -> raises within np.roll()
with pytest.raises(TypeError,
match=r"slice indices must be integers.*"):
in_dataset.xsimlab.run(model=m, validate="inputs")
# internal validation -> raises within attr.validate()
with pytest.raises(TypeError, match=r".*'int'.*"):
in_dataset.xsimlab.run(model=m, validate="all")
@pytest.mark.parametrize(
"dims,data,clock",
[
("batch", [1, 2], None),
(("batch", "clock"), [[1, 1, 1], [2, 2, 2]], "clock"),
(("batch", "x"), [[1, 1], [2, 2]], None),
],
)
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)
@pytest.mark.parametrize(
"decoding,expected",
[
(None, [nan, nan]), # mask_and_scale=True by default
({
"mask_and_scale": False
}, [-1, -1]),
],
)
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)
class Foo:
var = xs.any_object()
