Merge pull request #2 from comet-toolkit/small_fixes

Small fixes for use in punpy

Author: shunt16

notebook/test_register_err_corr_form.py

@@ -14,32 +14,32 @@ def build_matrix(self, idx: np.ndarray) -> np.ndarray:
         return "abc"
 
 
-ds = xr.open_dataset("obs_example.nc")
-
-ds.unc["temperature"]["u_sys_temperature"] = (
-    ["x", "y", "time"],
-    ds.temperature * 0.03,
-    {
-        "err_corr": [
-            {
-                "dim": "x",
-                "form": "new",
-                "params": [],
-            },
-            {
-                "dim": "y",
-                "form": "systematic",
-                "params": [],
-            },
-            {
-                "dim": "time",
-                "form": "systematic",
-                "params": [],
-            },
-        ]
-    },
-)
-
-p = ds.unc["temperature"]["u_sys_temperature"]
-
-pass
+# ds = xr.open_dataset("obs_example.nc")
+#
+# ds.unc["temperature"]["u_sys_temperature"] = (
+#     ["x", "y", "time"],
+#     ds.temperature * 0.03,
+#     {
+#         "err_corr": [
+#             {
+#                 "dim": "x",
+#                 "form": "new",
+#                 "params": [],
+#             },
+#             {
+#                 "dim": "y",
+#                 "form": "systematic",
+#                 "params": [],
+#             },
+#             {
+#                 "dim": "time",
+#                 "form": "systematic",
+#                 "params": [],
+#             },
+#         ]
+#     },
+# )
+#
+# p = ds.unc["temperature"]["u_sys_temperature"]
+#
+# pass

obsarray/err_corr.py

@@ -49,7 +49,7 @@ class BaseErrCorrForm(abc.ABC):
 
     def __init__(self, xarray_obj, unc_var_name, dims, params, units):
         self._obj = xarray_obj
-        self.unc_var_name = unc_var_name
+        self._unc_var_name = unc_var_name
         self.dims = dims if isinstance(dims, list) else [dims]
         self.params = params if isinstance(params, list) else [params]
         self.units = units
@@ -69,18 +69,18 @@ def form(self) -> str:
         """Form name"""
         pass
 
-    def expand_dim_matrix(self, submatrix):
+    def expand_dim_matrix(self, submatrix, sli):
         return expand_errcorr_dims(
             in_corr=submatrix,
             in_dim=self.dims,
-            out_dim=list(self._obj[self.unc_var_name].dims),
+            out_dim=list(self._obj[self._unc_var_name][sli].dims),
             dim_sizes={
-                dim: self._obj.dims[dim] for dim in self._obj[self.unc_var_name].dims
+                dim: self._obj.dims[dim] for dim in self._obj[self._unc_var_name][sli].dims
             },
         )
 
     def slice_full_cov(self, full_matrix, sli):
-        mask_array = np.ones(self._obj[self.unc_var_name].shape, dtype=bool)
+        mask_array = np.ones(self._obj[self._unc_var_name].shape, dtype=bool)
         mask_array[sli] = False
 
         return np.delete(
@@ -138,10 +138,10 @@ def build_matrix(self, sli):
         dims_matrix = np.eye(n_elems)
 
         # expand to correlation matrix over all variable dims
-        full_matrix = self.expand_dim_matrix(dims_matrix)
+        return self.expand_dim_matrix(dims_matrix, sli)
 
-        # subset to slice
-        return self.slice_full_cov(full_matrix, sli)
+        # # subset to slice
+        # return self.slice_full_cov(full_matrix, sli)
 
 
 @register_err_corr_form("systematic")
@@ -167,10 +167,10 @@ def build_matrix(self, sli):
         dims_matrix = np.ones((n_elems, n_elems))
 
         # expand to correlation matrix over all variable dims
-        full_matrix = self.expand_dim_matrix(dims_matrix)
+        return self.expand_dim_matrix(dims_matrix,sli)
 
         # subset to slice
-        return self.slice_full_cov(full_matrix, sli)
+        # return self.slice_full_cov(full_matrix, sli)
 
 
 @register_err_corr_form("err_corr_matrix")
@@ -189,10 +189,10 @@ def build_matrix(self, sli):
         """
 
         # expand to correlation matrix over all variable dims
-        full_matrix = self.expand_dim_matrix(self._obj[self.params[0]])
+        return self.expand_dim_matrix(self._obj[self.params[0]],sli)
 
-        # subset to slice
-        return self.slice_full_cov(full_matrix, sli)
+        # # subset to slice
+        # return self.slice_full_cov(full_matrix, sli)
 
 
 @register_err_corr_form("ensemble")

obsarray/templater/template_util.py

@@ -129,7 +129,7 @@ def _create_var(
             )
 
         # Create variable and add to dataset
-        if dtype == str:
+        if isinstance(dtype,str):
             if dtype == "flag":
                 flag_meanings = attributes.pop("flag_meanings")
                 variable = du.create_flags_variable(

obsarray/test/test_err_corr_forms.py

@@ -121,20 +121,35 @@ class TestSystematicUnc(unittest.TestCase):
     def setUp(self) -> None:
         self.ds = create_ds()
 
-    def test_build_matrix_1stdim(self):
+    def build_matrix_1stdim(self):
         rc = SystematicCorrelation(self.ds, "u_sys_temperature", ["x"], [], [])
 
         ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        # np.testing.assert_equal(ecrm, np.ones((12, 12)))
 
-        np.testing.assert_equal(ecrm, np.ones((12, 12)))
+        return ecrm
 
-    def test_build_matrix_2nddim(self):
+    def build_matrix_2nddim(self):
         rc = SystematicCorrelation(self.ds, "u_sys_temperature", ["y"], [], [])
 
         ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        # np.testing.assert_equal(ecrm, np.ones((12, 12)))
+
+        return ecrm
+
+    def build_matrix_3ddim(self):
+        rc = SystematicCorrelation(self.ds, "u_sys_temperature", ["time"], [], [])
+
+        ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        # np.testing.assert_equal(ecrm, np.ones((12, 12)))
 
-        np.testing.assert_equal(ecrm, np.ones((12, 12)))
+        return ecrm
 
+    def test_build_matrix(self):
+        x=self.build_matrix_1stdim()
+        y=self.build_matrix_2nddim()
+        time=self.build_matrix_3ddim()
+        np.testing.assert_equal((x.dot(y)).dot(time), np.ones((12, 12)))
 
 if __name__ == "main":
     unittest.main()

obsarray/test/test_unc_accessor.py

@@ -31,7 +31,7 @@ def compare_err_corr_form(self, form, exp_form):
     self.assertEqual(form.form, exp_form.form)
     self.assertCountEqual(form.params, exp_form.params)
     self.assertCountEqual(form.units, exp_form.units)
-    self.assertCountEqual(form.unc_var_name, exp_form.unc_var_name)
+    self.assertCountEqual(form._unc_var_name, exp_form._unc_var_name)
 
 
 def create_ds():
@@ -99,7 +99,7 @@ def create_ds():
             "err_corr": [
                 {
                     "dim": "x",
-                    "form": "custom",
+                    "form": "err_corr_matrix",
                     "params": ["err_corr_str_temperature"],
                 },
                 {

obsarray/unc_accessor.py

@@ -291,9 +291,10 @@ def err_corr_matrix(self) -> xr.DataArray:
         # populate with error-correlation matrices built be each error-correlation
         # parameterisation object
         for dim_err_corr in self.err_corr:
-            err_corr_matrix.values = err_corr_matrix.values.dot(
-                dim_err_corr[1].build_matrix(self._sli)
-            )
+            if np.all([dim in self._obj[self._unc_var_name][self._sli].dims for dim in dim_err_corr[1].dims]):
+                err_corr_matrix.values = err_corr_matrix.values.dot(
+                    dim_err_corr[1].build_matrix(self._sli)
+                )
 
         return err_corr_matrix
 
@@ -563,7 +564,7 @@ def total_err_cov_matrix(self) -> xr.DataArray:
         )
         covs_sum = np.zeros(total_err_cov_matrix.shape)
         for unc in self:
-            covs_sum += unc.err_cov_matrix().values
+            covs_sum += unc[self._sli].err_cov_matrix().values
 
         total_err_cov_matrix.values = covs_sum
 
@@ -582,7 +583,7 @@ def structured_err_cov_matrix(self):
 
         covs_sum = np.zeros(structured_err_cov_matrix.shape)
         for unc in self:
-            covs_sum += unc.err_cov_matrix().values
+            covs_sum += unc[self._sli].err_cov_matrix().values
 
         structured_err_cov_matrix.values = covs_sum