Merge pull request #34 from comet-toolkit/33-suffix

Renaming of dimensions and variables

Author: shunt16

docs/content/user/api.rst

@@ -30,6 +30,8 @@ Uncertainty functions
    unc_accessor.UncAccessor.unc_vars
    unc_accessor.UncAccessor.__getitem__
    unc_accessor.UncAccessor.keys
+   unc_accessor.UncAccessor.rename
+   unc_accessor.UncAccessor.rename_dims
    unc_accessor.VariableUncertainty
    unc_accessor.VariableUncertainty.__getitem__
    unc_accessor.VariableUncertainty.__setitem__
@@ -83,4 +85,12 @@ Flag functions
    flag_accessor.Flag
    flag_accessor.Flag.__getitem__
    flag_accessor.Flag.__setitem__
-   flag_accessor.Flag.value
+   flag_accessor.Flag.value
+
+Utility functions
+=================
+
+.. autosummary::
+   :toctree: generated/
+
+   utils.append_names

docs/content/user/unc_accessor.rst

@@ -186,15 +186,16 @@ A component of uncertainty can be simply be deleted as,
    # Check uncertainties
    ds.unc["temperature"].keys()
 
-Renaming Variables
-------------------
+Renaming Variables and Dimensions
+---------------------------------
 
-The storage of uncertainty information is underpinned by variable attributes, which include referencing other variables (for example, which variables are the uncertainties associated with a particular observation variable). Because of this it is important, if renaming uncertainty variables, to use **obsarray**'s renaming functionality. This renames the uncertainty variable and safely updates attribute variable references. This is done as follows:
+The storage of uncertainty information is underpinned by variable attributes, which include referencing other variables/dimensions (for example, which variables are the uncertainties associated with a particular observation variable). Because of this it is important, if renaming uncertainty variables or dimensions, to use **obsarray**'s renaming functionality. This renames the uncertainty variable or dimension and safely updates attribute variable references. This is done as follows (mirroring the interface to `xarray renaming <https://docs.xarray.dev/en/latest/generated/xarray.Dataset.rename.html>`_):
 
 
 .. ipython:: python
    :okwarning:
 
    print(ds.unc["temperature"])
-   ds = ds.unc["temperature"]["u_ran_temperature"].rename("u_noise")
+   ds = ds.unc.rename({"u_ran_temperature": "u_noise"})
+   ds = ds.unc.rename_dims({"time": "t"})
    print(ds.unc["temperature"])

obsarray/__init__.py

@@ -9,6 +9,7 @@
 from obsarray.templater.template_util import create_ds
 from obsarray.templater.dstemplater import DSTemplater
 from obsarray.templater.dswriter import DSWriter
+from obsarray.utils import append_names
 
 __version__ = get_versions()["version"]
 del get_versions

obsarray/test/test_unc_accessor.py

@@ -34,7 +34,7 @@ def compare_err_corr_form(self, form, exp_form):
     self.assertCountEqual(form._unc_var_name, exp_form._unc_var_name)
 
 
-def create_ds():
+def create_ds(var_suffix="", dim_suffix="", coord_dim_suffix_extra=""):
     np.random.seed(0)
     temperature = 15 + 8 * np.random.randn(2, 2, 3)
     u_r_temperature = temperature * 0.02
@@ -47,42 +47,59 @@ def create_ds():
     reference_time = pd.Timestamp("2014-09-05")
 
     ds = xr.Dataset(
-        data_vars=dict(
-            temperature=(["x", "y", "time"], temperature, {"units": "K"}),
-        ),
-        coords=dict(
-            lon=(["x", "y"], lon),
-            lat=(["x", "y"], lat),
-            time=time,
-            reference_time=reference_time,
-        ),
+        data_vars={
+            "temperature"
+            + var_suffix: (
+                [
+                    "x" + dim_suffix,
+                    "y" + dim_suffix,
+                    "time" + dim_suffix + coord_dim_suffix_extra,
+                ],
+                temperature,
+                {"units": "K"},
+            ),
+        },
+        coords={
+            "lon" + var_suffix: (["x" + dim_suffix, "y" + dim_suffix], lon),
+            "lat" + var_suffix: (["x" + dim_suffix, "y" + dim_suffix], lat),
+            "time" + var_suffix: ("time" + dim_suffix + coord_dim_suffix_extra, time),
+            "reference_time": reference_time,
+        },
         attrs=dict(description="Weather related data."),
     )
 
-    ds.unc["temperature"]["u_ran_temperature"] = (
-        ["x", "y", "time"],
+    ds.unc["temperature" + var_suffix]["u_ran_temperature" + var_suffix] = (
+        [
+            "x" + dim_suffix,
+            "y" + dim_suffix,
+            "time" + dim_suffix + coord_dim_suffix_extra,
+        ],
         temperature * 0.05,
         {"units": "K", "pdf_shape": "gaussian"},
     )
 
-    ds.unc["temperature"]["u_sys_temperature"] = (
-        ["x", "y", "time"],
+    ds.unc["temperature" + var_suffix]["u_sys_temperature" + var_suffix] = (
+        [
+            "x" + dim_suffix,
+            "y" + dim_suffix,
+            "time" + dim_suffix + coord_dim_suffix_extra,
+        ],
         temperature * 0.03,
         {
             "units": "K",
             "err_corr": [
                 {
-                    "dim": "x",
+                    "dim": "x" + dim_suffix,
                     "form": "systematic",
                     "params": [],
                 },
                 {
-                    "dim": "y",
+                    "dim": "y" + dim_suffix,
                     "form": "systematic",
                     "params": [],
                 },
                 {
-                    "dim": "time",
+                    "dim": "time" + dim_suffix + coord_dim_suffix_extra,
                     "form": "systematic",
                     "params": [],
                 },
@@ -91,19 +108,26 @@ def create_ds():
         },
     )
 
-    ds.unc["temperature"]["u_str_temperature"] = (
-        ["x", "y", "time"],
+    ds.unc["temperature" + var_suffix]["u_str_temperature" + var_suffix] = (
+        [
+            "x" + dim_suffix,
+            "y" + dim_suffix,
+            "time" + dim_suffix + coord_dim_suffix_extra,
+        ],
         temperature * 0.1,
         {
             "units": "K",
             "err_corr": [
                 {
-                    "dim": ["x", "time"],
+                    "dim": [
+                        "x" + dim_suffix,
+                        "time" + dim_suffix + coord_dim_suffix_extra,
+                    ],
                     "form": "err_corr_matrix",
-                    "params": ["err_corr_str_temperature"],
+                    "params": ["err_corr_str_temperature" + var_suffix],
                 },
                 {
-                    "dim": "y",
+                    "dim": "y" + dim_suffix,
                     "form": "systematic",
                     "params": [],
                 },
@@ -112,8 +136,8 @@ def create_ds():
         },
     )
 
-    ds["err_corr_str_temperature"] = (
-        ["x.time", "x.time"],
+    ds["err_corr_str_temperature" + var_suffix] = (
+        ["x.time" + dim_suffix, "x.time" + dim_suffix],
         np.ones(
             (
                 temperature.shape[0] * temperature.shape[2],
@@ -740,6 +764,44 @@ def test_err_cov_matrix(self, mock_cr2cv, mock_ecrm, mock_value):
         exp_ecm = xr.DataArray(np.ones((12, 12)), dims=["x.y.time", "x.y.time"])
         xr.testing.assert_equal(ecm, exp_ecm)
 
+    def test_rename_vars(self):
+        var_suffix = "_test"
+        input_ds = create_ds()
+
+        ds = input_ds.unc.rename(
+            {
+                "temperature": "temperature" + var_suffix,
+                "lon": "lon" + var_suffix,
+                "lat": "lat" + var_suffix,
+                "time": "time" + var_suffix,
+                "u_ran_temperature": "u_ran_temperature" + var_suffix,
+                "u_str_temperature": "u_str_temperature" + var_suffix,
+                "u_sys_temperature": "u_sys_temperature" + var_suffix,
+                "err_corr_str_temperature": "err_corr_str_temperature" + var_suffix,
+            }
+        )
+
+        exp_ds = create_ds(var_suffix=var_suffix, coord_dim_suffix_extra=var_suffix)
+
+        xr.testing.assert_identical(ds, exp_ds)
+
+    def test_rename_dims(self):
+        dim_suffix = "_test"
+        input_ds = create_ds()
+
+        ds = input_ds.unc.rename_dims(
+            {
+                "x": "x" + dim_suffix,
+                "y": "y" + dim_suffix,
+                "time": "time" + dim_suffix,
+                "x.time": "x.time_test",
+            }
+        )
+
+        exp_ds = create_ds(dim_suffix=dim_suffix)
+
+        xr.testing.assert_identical(ds, exp_ds)
+
 
 if __name__ == "__main__":
     unittest.main()

obsarray/test/test_utils.py

@@ -0,0 +1,169 @@
+"""test_utils - tests for obsarray.utils"""
+
+import unittest
+import numpy as np
+from obsarray import append_names, create_ds
+import xarray as xr
+
+__author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
+__all__ = []
+
+
+def create_test_ds(suffix):
+    # define ds variables
+    template = {
+        "temperature"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {
+                "units": "K",
+                "unc_comps": [
+                    "u_ran_temperature" + suffix,
+                    "u_sys_temperature" + suffix,
+                ],
+            },
+        },
+        "u_ran_temperature"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {
+                "units": "K",
+                "err_corr": [
+                    {"dim": "x" + suffix, "form": "random", "params": [], "units": []},
+                    {"dim": "y" + suffix, "form": "random", "params": [], "units": []},
+                    {
+                        "dim": "time" + suffix,
+                        "form": "random",
+                        "params": [],
+                        "units": [],
+                    },
+                ],
+            },
+        },
+        "u_sys_temperature"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {
+                "units": "K",
+                "err_corr": [
+                    {
+                        "dim": "x" + suffix,
+                        "form": "systematic",
+                        "params": [],
+                        "units": [],
+                    },
+                    {
+                        "dim": "y" + suffix,
+                        "form": "systematic",
+                        "params": [],
+                        "units": [],
+                    },
+                    {
+                        "dim": "time" + suffix,
+                        "form": "systematic",
+                        "params": [],
+                        "units": [],
+                    },
+                ],
+            },
+        },
+        "pressure"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {"units": "Pa", "unc_comps": ["u_str_pressure" + suffix]},
+        },
+        "u_str_pressure"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {
+                "units": "Pa",
+                "err_corr": [
+                    {"dim": "x" + suffix, "form": "random", "params": [], "units": []},
+                    {
+                        "dim": "y" + suffix,
+                        "form": "err_corr_matrix",
+                        "params": "err_corr_str_pressure_y",
+                        "units": [],
+                    },
+                    {
+                        "dim": "time" + suffix,
+                        "form": "systematic",
+                        "params": [],
+                        "units": [],
+                    },
+                ],
+            },
+        },
+        "err_corr_str_pressure_y"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["y" + suffix, "y" + suffix],
+            "attributes": {"units": ""},
+        },
+        "n_moles"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {"units": "", "unc_comps": ["u_ran_n_moles" + suffix]},
+        },
+        "u_ran_n_moles"
+        + suffix: {
+            "dtype": np.float32,
+            "dim": ["x" + suffix, "y" + suffix, "time" + suffix],
+            "attributes": {
+                "units": "",
+                "err_corr": [
+                    {"dim": "x" + suffix, "form": "random", "params": [], "units": []},
+                    {"dim": "y" + suffix, "form": "random", "params": [], "units": []},
+                    {
+                        "dim": "time" + suffix,
+                        "form": "random",
+                        "params": [],
+                        "units": [],
+                    },
+                ],
+            },
+        },
+    }
+
+    # define dim_size_dict to specify size of arrays
+    dim_sizes = {"x" + suffix: 20, "y" + suffix: 30, "time" + suffix: 6}
+
+    # create dataset template
+    ds = create_ds(template, dim_sizes)
+
+    # populate with example data
+    ds["temperature" + suffix].values = 293 * np.ones((20, 30, 6))
+    ds["u_ran_temperature" + suffix].values = 1 * np.ones((20, 30, 6))
+    ds["u_sys_temperature" + suffix].values = 0.4 * np.ones((20, 30, 6))
+    ds["pressure" + suffix].values = 10**5 * np.ones((20, 30, 6))
+    ds["u_str_pressure" + suffix].values = 10 * np.ones((20, 30, 6))
+    ds["err_corr_str_pressure_y" + suffix].values = 0.5 * np.ones(
+        (30, 30)
+    ) + 0.5 * np.eye(30)
+    ds["n_moles" + suffix].values = 40 * np.ones((20, 30, 6))
+    ds["u_ran_n_moles" + suffix].values = 1 * np.ones((20, 30, 6))
+
+    ds.attrs["attr" + suffix] = "val"
+
+    return ds
+
+
+class TestAppendNames(unittest.TestCase):
+    def test_append_names(self):
+
+        input_ds = create_test_ds(suffix="")
+        ds = append_names(input_ds, "_test")
+
+        exp_ds = create_test_ds(suffix="_test")
+
+        xr.testing.assert_identical(ds, exp_ds)
+
+
+if __name__ == "__main__":
+    unittest.main()