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)
