Change dimension order in NetCDF output

Put ensemble dimension after height dimension (if height exists)

Author: tnipen

bris/outputs/netcdf.py

@@ -479,8 +479,12 @@ def _setup_prediction_vars(
                         shape = [len(times), len(y)]
 
                 if self.pm.num_members > 1:
-                    dims.insert(1, self.conv_name("ensemble_member"))
-                    shape.insert(1, self.pm.num_members)
+                    # We want the ensemble member dimension to be after height (if height exists)
+                    # I.e. we want (time, height, member, y, x) or (time, member, y, x)
+                    ens_dim_loc = 1 + (dim_name is not None)
+
+                    dims.insert(ens_dim_loc, self.conv_name("ensemble_member"))
+                    shape.insert(ens_dim_loc, self.pm.num_members)
 
                 ar = np.nan * np.zeros(shape, np.float32)
                 self.ds[ncname] = (dims, ar)
@@ -515,6 +519,7 @@ def _setup_prediction_vars(
                 # Accumulate over lead times
                 ar = np.cumsum(np.nan_to_num(ar, nan=0), axis=0)
 
+            # Rearrange to (time, member, y, [x]) or (time, y, [x])
             ar = np.moveaxis(ar, [-1], [1]) if self.pm.num_members > 1 else ar[..., 0]
 
             cfname = cf.get_metadata(variable)["cfname"]
@@ -527,10 +532,7 @@ def _setup_prediction_vars(
                 bris.units.convert(ar, from_units, to_units, inplace=True)
 
             if level_index is not None:
-                if self.pm.num_members > 1:
-                    self.ds[ncname][:, :, level_index, ...] = ar
-                else:
-                    self.ds[ncname][:, level_index, ...] = ar
+                self.ds[ncname][:, level_index, ...] = ar
             else:
                 self.ds[ncname][:] = ar
 

tests/test_outputs_netcdf.py

@@ -8,7 +8,7 @@
 from bris.predict_metadata import PredictMetadata
 
 
-def test_1():
+def test_deterministic():
     variables = ["u_800", "u_600", "2t", "v_500", "10u", "tp", "skt"]
     lats = np.array([1, 2])
     lons = np.array([2, 4])
@@ -25,6 +25,7 @@ def test_1():
     times = frt + leadtimes
 
     with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir = "./"
         pattern = os.path.join(temp_dir, "test_%Y%m%dT%HZ.nc")
         workdir = os.path.join(temp_dir, "test_gridded")
         attrs = {"creator": "met.no"}
@@ -61,6 +62,9 @@ def test_1():
                 assert "units" in var.attrs
                 assert "grid_mapping" in var.attrs
 
+            # (time, height, y, x)
+            assert file.variables["air_temperature_2m"].shape == (4, 1, 1, 2)
+
             height_dim = file.variables["air_temperature_2m"].dims[1]
             levels = file.variables[height_dim].values
             assert levels == [2]
@@ -95,6 +99,102 @@ def test_1():
             assert np.min(file.variables["altitude"]) == 100
 
 
+def test_ensemble():
+    variables = [
+        "u_500",
+        "u_800",
+        "u_600",
+        "2t",
+        "v_500",
+        "10u",
+        "tp",
+        "skt",
+        "unknown",
+    ]
+    lats = np.array([1, 2])
+    lons = np.array([2, 4])
+    altitudes = np.array([100, 200])
+    leadtimes = np.arange(0, 3600 * 4, 3600)
+    num_members = 2
+    field_shape = [1, 2]
+    pm = PredictMetadata(
+        variables, lats, lons, altitudes, leadtimes, num_members, field_shape
+    )
+
+    pred = np.random.rand(*pm.shape)
+    frt = 1672552800
+    times = frt + leadtimes
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir = "./"
+        pattern = os.path.join(temp_dir, "test_%Y%m%dT%HZ.nc")
+        workdir = os.path.join(temp_dir, "test_gridded")
+        attrs = {"creator": "met.no"}
+        output = Netcdf(pm, workdir, pattern, global_attributes=attrs)
+
+        for member in range(num_members):
+            output.add_forecast(times, member, pred)
+        output.finalize()
+
+        output_filename = os.path.join(temp_dir, "test_20230101T06Z.nc")
+
+        assert os.path.exists(output_filename)
+
+        with xr.open_dataset(output_filename) as file:
+            # Check that global attributes are written
+            for k, v in attrs.items():
+                assert file.attrs[k] == v
+
+            for variable in [
+                "altitude",
+                "air_temperature_2m",
+                "air_temperature_0m",
+                "x_wind_pl",
+            ]:
+                assert variable in file.variables, variable
+                var = file.variables[variable]
+                assert "units" in var.attrs
+                assert "grid_mapping" in var.attrs
+
+            # (time, height, member, y, x)
+            assert file.variables["air_temperature_2m"].shape == (4, 1, 2, 1, 2)
+
+            height_dim = file.variables["air_temperature_2m"].dims[1]
+            levels = file.variables[height_dim].values
+            assert levels == [2]
+
+            height_dim = file.variables["air_temperature_0m"].dims[1]
+            levels = file.variables[height_dim].values
+            assert levels == [0]
+
+            # (time, pressure, member, y, x)
+            assert file.variables["x_wind_pl"].shape == (4, 3, 2, 1, 2)
+
+            height_dim = file.variables["x_wind_pl"].dims[1]
+            levels = file.variables[height_dim].values
+            assert len(levels) == 3
+            assert levels[0] == 500
+            assert levels[1] == 600
+            assert levels[2] == 800
+            assert file.variables["x_wind_pl"].values.shape[2] == num_members
+
+    # Test interpolation
+    with tempfile.TemporaryDirectory() as temp_dir:
+        pattern = os.path.join(temp_dir, "test_%Y%m%dT%HZ.nc")
+        workdir = os.path.join(temp_dir, "test_gridded")
+        attrs = {"creator": "met.no"}
+        output = Netcdf(pm, workdir, pattern, interp_res=0.2)
+
+        for member in range(num_members):
+            output.add_forecast(times, member, pred)
+        output.finalize()
+
+        output_filename = os.path.join(temp_dir, "test_20230101T06Z.nc")
+        with xr.open_dataset(output_filename) as file:
+            # Check that altitude variable has attributes
+            assert "altitude" not in file.variables
+
+
 def test_domain_name():
     variables = ["u_800", "u_600", "2t", "v_500", "10u"]
     lats = np.array([1, 2])
@@ -122,4 +222,5 @@ def test_domain_name():
 
 if __name__ == "__main__":
     test_domain_name()
-    test_1()
+    test_deterministic()
+    test_ensemble()