Skip to content

Commit

Permalink
add test fixture to create a datatree using the same variable, dim an…
Browse files Browse the repository at this point in the history
…d coord names, but different values and sizes, and nothing in the root group to inherit
  • Loading branch information
acrosby committed Jul 13, 2024
1 parent 0eac740 commit 97e971e
Showing 1 changed file with 115 additions and 0 deletions.
115 changes: 115 additions & 0 deletions xarray/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,3 +198,118 @@ def simple_datatree(create_test_datatree):
Returns a DataTree.
"""
return create_test_datatree()


@pytest.fixture(scope="module")
def create_test_multidataset_withoutroot_datatree():
"""
Create a test datatree with this structure:
<xarray.DataTree>
Group: /
├── Group: /Main
│ Dimensions: (yi: 50, xi: 60, t: 30)
│ Coordinates:
│ * t (t) int64 240B 0 1 2 3 4 5 6 7 8 9 ... 21 22 23 24 25 26 27 28 29
│ Dimensions without coordinates: yi, xi
│ Data variables:
│ x (yi, xi) int64 24kB 0 0 0 0 0 0 0 0 0 ... 49 49 49 49 49 49 49 49
│ y (yi, xi) int64 24kB 0 1 2 3 4 5 6 7 8 ... 52 53 54 55 56 57 58 59
│ u (t, yi, xi) float64 720kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
│ v (t, yi, xi) float64 720kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
│ p (t, yi, xi) float64 720kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
├── Group: /ovr1
│ Dimensions: (yi: 51, xi: 51, t: 10)
│ Coordinates:
│ * t (t) int64 80B 10 11 12 13 14 15 16 17 18 19
│ Dimensions without coordinates: yi, xi
│ Data variables:
│ x (yi, xi) float64 21kB 150.0 150.0 150.0 150.0 ... 200.0 200.0 200.0
│ y (yi, xi) int64 21kB 0 1 2 3 4 5 6 7 8 ... 43 44 45 46 47 48 49 50
│ u (t, yi, xi) float64 208kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
│ v (t, yi, xi) float64 208kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
│ p (t, yi, xi) float64 208kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
│ cx (t) int64 80B 10 11 12 13 14 15 16 17 18 19
│ cy (t) int64 80B 10 11 12 13 14 15 16 17 18 19
└── Group: /ovr2
Dimensions: (yi: 51, xi: 51, t: 20)
Coordinates:
* t (t) int64 160B 10 11 12 13 14 15 16 17 ... 22 23 24 25 26 27 28 29
Dimensions without coordinates: yi, xi
Data variables:
x (yi, xi) int64 21kB 0 0 0 0 0 0 0 0 0 ... 50 50 50 50 50 50 50 50
y (yi, xi) int64 21kB 0 1 2 3 4 5 6 7 8 ... 43 44 45 46 47 48 49 50
u (t, yi, xi) float64 416kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
v (t, yi, xi) float64 416kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
p (t, yi, xi) float64 416kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
cx (t) int64 160B 10 11 12 13 14 15 16 17 ... 22 23 24 25 26 27 28 29
cy (t) int64 160B 10 11 12 13 14 15 16 17 ... 22 23 24 25 26 27 28 29
The structure reflects a set of groups that share identical dimension,
coordinate and variable names which do not exist in the root group, so
that there should be no problem with the cascading inheritance.
The 'u','v' and 'p' variables in each group depend on their group-specific
coordinates 't', 'x' and 'y', and dimensions 't', 'xi' and 'yi'. The 'cx'
and 'cy' variables depend on only the 't' coordinate/dimension for the
respective group. The 2D 'x' and 'y' variables could be treated like coordinates
for the arrays that share their 'xi' and 'yi' dimensions.
"""

def _create_test_multidataset_withoutroot_datatree(
modify=lambda ds: ds,
):
main_data = modify(
xr.Dataset(
{
"x": (["yi", "xi"], np.meshgrid(np.arange(60), np.arange(50))[1]),
"y": (["yi", "xi"], np.meshgrid(np.arange(60), np.arange(50))[0]),
"t": ("t", np.arange(30)),
"u": (["t", "yi", "xi"], np.zeros((30, 50, 60))),
"v": (["t", "yi", "xi"], np.zeros((30, 50, 60))),
"p": (["t", "yi", "xi"], np.zeros((30, 50, 60))),
}
)
)
ovr1_data = modify(
xr.Dataset(
{
"x": (
["yi", "xi"],
np.meshgrid(np.arange(51), np.arange(51) + 150.0)[1],
),
"y": (["yi", "xi"], np.meshgrid(np.arange(51), np.arange(51))[0]),
"t": ("t", np.arange(10, 20)),
"u": (["t", "yi", "xi"], np.zeros((10, 51, 51))),
"v": (["t", "yi", "xi"], np.zeros((10, 51, 51))),
"p": (["t", "yi", "xi"], np.zeros((10, 51, 51))),
"cx": ("t", np.arange(10, 20)),
"cy": ("t", np.arange(10, 20)),
}
)
)
ovr2_data = modify(
xr.Dataset(
{
"x": (["yi", "xi"], np.meshgrid(np.arange(51), np.arange(51))[1]),
"y": (["yi", "xi"], np.meshgrid(np.arange(51), np.arange(51))[0]),
"t": ("t", np.arange(10, 30)),
"u": (["t", "yi", "xi"], np.zeros((20, 51, 51))),
"v": (["t", "yi", "xi"], np.zeros((20, 51, 51))),
"p": (["t", "yi", "xi"], np.zeros((20, 51, 51))),
"cx": ("t", np.arange(10, 30)),
"cy": ("t", np.arange(10, 30)),
}
)
)
root_data = modify(xr.Dataset())

# Avoid using __init__ so we can independently test it
root: DataTree = DataTree(data=root_data)
main: DataTree = DataTree(name="Main", parent=root, data=main_data)
ovr1: DataTree = DataTree(name="ovr1", parent=root, data=ovr1_data)
ovr2: DataTree = DataTree(name="ovr2", parent=root, data=ovr2_data)

return root

return _create_test_multidataset_withoutroot_datatree

0 comments on commit 97e971e

Please sign in to comment.