class TotalVerticalMotion: """Sum up all vertical motions of bedrock and topographic surface, respectively. Vertical motions may result from external forcing, erosion and/or feedback of erosion on tectonics (isostasy). """ #TODO: remove any_upward_vars # see https://github.com/benbovy/xarray-simlab/issues/64 any_upward_vars = xs.group('any_upward') bedrock_upward_vars = xs.group('bedrock_upward') surface_upward_vars = xs.group('surface_upward') surface_downward_vars = xs.group('surface_downward') bedrock_upward = xs.variable( dims=('y', 'x'), intent='out', description='bedrock motion in upward direction') surface_upward = xs.variable( dims=('y', 'x'), intent='out', description='topographic surface motion in upward direction') def run_step(self): sum_any = sum(self.any_upward_vars) self.bedrock_upward = sum_any + sum(self.bedrock_upward_vars) self.surface_upward = (sum_any + sum(self.surface_upward_vars) - sum(self.surface_downward_vars))
class TotalVerticalMotion: """Sum up all vertical motions of bedrock and topographic surface, respectively. Vertical motions may result from external forcing, erosion and/or feedback of erosion on tectonics (isostasy). """ bedrock_upward_vars = xs.group('bedrock_upward') surface_upward_vars = xs.group('surface_upward') surface_downward_vars = xs.group('surface_downward') bedrock_upward = xs.variable( dims=('y', 'x'), intent='out', description='bedrock motion in upward direction' ) surface_upward = xs.variable( dims=('y', 'x'), intent='out', description='topographic surface motion in upward direction' ) def run_step(self): self.bedrock_upward = sum(self.bedrock_upward_vars) self.surface_upward = (sum(self.surface_upward_vars) - sum(self.surface_downward_vars))
class ThirdInit(Context): """ Inherited by all other model components to access backend""" firstinit = xs.group('FirstInit') secondinit = xs.group('SecondInit') group = xs.variable(intent='out', groups='ThirdInit') def initialize(self): super(ThirdInit, self).initialize() self.group = 3
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
def _make_phydra_flux(label, variable): """ """ xs_var_dict = defaultdict() xs_var_dict[label + '_value'] = _convert_2_xsimlabvar( var=variable, intent='out', value_store=True, description_label='output of flux value / ') group = variable.metadata.get('group') group_to_arg = variable.metadata.get('group_to_arg') #if group and group_to_arg: # raise Exception("A flux can be either added to group or take a group as argument, not both.") if group: xs_var_dict[label + '_label'] = _convert_2_xsimlabvar( var=variable, intent='out', groups=group, var_dims=(), description_label='label reference with group / ', attrs=False) if group_to_arg: xs_var_dict[group_to_arg] = xs.group(group_to_arg) return xs_var_dict
class ProfileZ: """Compute the evolution of the elevation (z) profile""" h_vars = xs.group( "h_vars" ) #allows for multiple processes influencing; say diffusion and subsidence z = xs.variable(dims="x", intent="inout", description="surface elevation z", attrs={"units": "m"}) br = xs.variable(dims=[(), "x"], intent="in", description="bedrock_elevation", attrs={"units": "m"}) h = xs.variable(dims="x", intent="inout", description="sed_thickness", attrs={"units": "m"}) def run_step(self): #self._delta_br = sum((br for br in self.br_vars)) self._delta_h = sum((h for h in self.h_vars)) def finalize_step(self): #self.br += self._delta_br #update bedrock surface self.h += self._delta_h #update sediment thickness self.z = self.br + self.h #add sediment to bedrock to get topo elev.
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 Solver(Context): """ Solver process executed last """ firstinit = xs.group('FirstInit') secondinit = xs.group('SecondInit') thirdinit = xs.group('ThirdInit') fourthinit = xs.group('FourthInit') fifthinit = xs.group('FifthInit') def initialize(self): """""" print("assembling model") print("SOLVER :", self.m.Solver) self.m.assemble() @xs.runtime(args="step_delta") def run_step(self, dt): self.m.solve(dt)
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 ProfileU: """Compute the evolution of the profile of quantity `u`.""" u_vars = xs.group("u_vars") u = xs.variable( dims="x", intent="inout", description="quantity u", attrs={"units": "m"} ) def run_step(self): self._delta_u = sum((v for v in self.u_vars)) def finalize_step(self): self.u += self._delta_u
class ProfileU(object): """Compute the evolution of the profile of quantity `u`.""" u_vars = xs.group('u_vars') u = xs.variable(dims='x', intent='inout', description='quantity u', attrs={'units': 'm'}) def run_step(self, *args): self._delta_u = sum((v for v in self.u_vars)) def finalize_step(self): self.u += self._delta_u
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 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 B: g1 = xs.group("g1") g2 = xs.group("g2")
class Foo: bar = xs.group("g")
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)