def test_get_unpack_layer(names, dims_list, array: np.ndarray):
dataset = get_dataset(names, dims_list)
packer = ArrayPacker(SAMPLE_DIM, names)
packer.to_array(dataset) # must pack first to know dimension lengths
result = get_unpack_layer(packer, feature_dim=1)(array)
# to_dataset does not preserve coordinates
for name, array in zip(packer.pack_names, result):
if array.shape[1] == 1:
array = array[:, 0]
np.testing.assert_array_equal(array, dataset[name])
def test_unpack_matrix():
nz = 10
in_ = xr.Dataset({
"a": (["x", "z"], np.ones((1, nz))),
"b": (["x"], np.ones((1)))
})
out = xr.Dataset({
"c": (["x", "z"], np.ones((1, nz))),
"d": (["x", "z"], np.ones((1, nz)))
})
x_packer = ArrayPacker("x", pack_names=["a", "b"])
y_packer = ArrayPacker("x", pack_names=["c", "d"])
in_packed = x_packer.to_array(in_)
out_packed = y_packer.to_array(out)
matrix = np.outer(out_packed.squeeze(), in_packed.squeeze())
jacobian = unpack_matrix(x_packer, y_packer, matrix)
assert isinstance(jacobian[("a", "c")], xr.DataArray)
assert jacobian[("a", "c")].dims == ("c", "a")
assert isinstance(jacobian[("a", "d")], xr.DataArray)
assert jacobian[("a", "d")].dims == ("d", "a")
assert isinstance(jacobian[("b", "c")], xr.DataArray)
assert jacobian[("b", "c")].dims == ("c", "b")
assert isinstance(jacobian[("b", "d")], xr.DataArray)
assert jacobian[("b", "d")].dims == ("d", "b")
def test_array_packer_dump_and_load(tmpdir):
dataset = get_dataset(["var1"], [[SAMPLE_DIM, FEATURE_DIM]])
packer_config = PackerConfig({"var1": {"z": SliceConfig(None, 2)}})
packer = ArrayPacker(SAMPLE_DIM, list(dataset.data_vars), packer_config)
packer.to_array(dataset)
with open(str(tmpdir.join("packer.yaml")), "w") as f:
packer.dump(f)
with open(str(tmpdir.join("packer.yaml"))) as f:
loaded_packer = ArrayPacker.load(f)
assert packer._pack_names == loaded_packer._pack_names
assert packer._n_features == loaded_packer._n_features
assert packer._sample_dim_name == loaded_packer._sample_dim_name
assert packer._config == packer_config
for orig, loaded in zip(packer._feature_index,
loaded_packer._feature_index):
assert orig == loaded
def test_to_dataset(names, dims_list, array: np.ndarray):
dataset = get_dataset(names, dims_list)
packer = ArrayPacker(SAMPLE_DIM, names)
packer.to_array(dataset) # must pack first to know dimension lengths
result = packer.to_dataset(array)
xr.testing.assert_equal(result, dataset)
def test_to_array(names, dims_list, array: np.ndarray):
dataset = get_dataset(names, dims_list)
packer = ArrayPacker(SAMPLE_DIM, names)
result = packer.to_array(dataset)
np.testing.assert_array_equal(result, array)
def packer(names: Iterable[str]) -> ArrayPacker:
return ArrayPacker(SAMPLE_DIM, names)
def test_unpack_before_pack_raises(names, array: np.ndarray):
packer = ArrayPacker(SAMPLE_DIM, names)
with pytest.raises(RuntimeError):
packer.to_dataset(array)
def packer(names: Iterable[str], dataset: xr.Dataset) -> ArrayPacker:
packer = ArrayPacker(SAMPLE_DIM, names)
packer.to_array(dataset) # must let packer know about array shapes
return packer