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