Merge pull request #4 from comet-toolkit/small_fixes

Small fixes

Author: shunt16

obsarray/err_corr.py

@@ -75,7 +75,8 @@ def expand_dim_matrix(self, submatrix, sli):
             in_dim=self.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][sli].dims
+                dim: self._obj.dims[dim]
+                for dim in self._obj[self._unc_var_name][sli].dims
             },
         )
 
@@ -167,7 +168,7 @@ def build_matrix(self, sli):
         dims_matrix = np.ones((n_elems, n_elems))
 
         # expand to correlation matrix over all variable dims
-        return self.expand_dim_matrix(dims_matrix,sli)
+        return self.expand_dim_matrix(dims_matrix, sli)
 
         # subset to slice
         # return self.slice_full_cov(full_matrix, sli)
@@ -189,7 +190,7 @@ def build_matrix(self, sli):
         """
 
         # expand to correlation matrix over all variable dims
-        return self.expand_dim_matrix(self._obj[self.params[0]],sli)
+        return self.expand_dim_matrix(self._obj[self.params[0]], sli)
 
         # # subset to slice
         # return self.slice_full_cov(full_matrix, sli)

obsarray/templater/dswriter.py

@@ -3,7 +3,7 @@
 """
 
 import os
-
+from obsarray.templater.template_util import DatasetUtil
 
 __author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
 
@@ -73,6 +73,15 @@ def _write_netcdf(ds, path, compression_level=None):
         for var_name in ds.data_vars:
             var_encoding = dict(comp)
             var_encoding.update(ds[var_name].encoding)
+            if "dtype" in var_encoding.keys():
+                var_encoding.update(
+                    {
+                        "_FillValue": DatasetUtil.get_default_fill_value(
+                            var_encoding["dtype"]
+                        )
+                    }
+                )
+                ds[var_name].attrs.pop("_FillValue")
             encoding.update({var_name: var_encoding})
 
         ds.to_netcdf(path, format="netCDF4", engine="netcdf4", encoding=encoding)

obsarray/templater/template_util.py

@@ -129,7 +129,7 @@ def _create_var(
             )
 
         # Create variable and add to dataset
-        if isinstance(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

@@ -146,10 +146,11 @@ def build_matrix_3ddim(self):
         return ecrm
 
     def test_build_matrix(self):
-        x=self.build_matrix_1stdim()
-        y=self.build_matrix_2nddim()
-        time=self.build_matrix_3ddim()
+        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/unc_accessor.py

@@ -276,6 +276,55 @@ def is_systematic(self) -> bool:
 
         return all(e[1].is_systematic is True for e in self.err_corr)
 
+    def err_corr_dict(self) -> dict:
+        """
+        Error-correlation dictionary for uncertainty effect.
+
+        :return: dictionary with error-correlation for each dimension
+        """
+
+        # initialise error-correlation dictionary
+        err_corr_dict = {}
+
+        # populate with error-correlation matrices built be each error-correlation
+        # parameterisation object
+        for dim_err_corr in self.err_corr:
+            if np.all(
+                [
+                    dim in self._obj[self._unc_var_name][self._sli].dims
+                    for dim in dim_err_corr[1].dims
+                ]
+            ):
+                if dim_err_corr[1].form in ["random", "systematic"]:
+                    err_corr_dict[dim_err_corr[0]] = dim_err_corr[1].form
+
+                elif dim_err_corr[1].form == "err_corr_matrix":
+                    err_corr_dict[dim_err_corr[0]] = self._obj[
+                        dim_err_corr[1].params[0]
+                    ].values
+
+                else:
+                    raise NotImplementedError(
+                        "this correlation form is not implemented for err_corr_dict()"
+                    )
+        return err_corr_dict
+
+    def err_corr_dict_numdim(self) -> dict:
+        """
+        Error-correlation dictionary for uncertainty effect, where the keys are the dimension index rather than dimension name.
+
+        :return: dictionary with error-correlation for each dimension
+        """
+        # initialise error-correlation dictionary
+        err_corr_dict = self.err_corr_dict()
+        err_corr_dict_numdim = {}
+
+        for idim, dim in enumerate(self._obj.dims):
+            if dim in err_corr_dict.keys():
+                err_corr_dict_numdim[str(idim)] = err_corr_dict[dim]
+
+        return err_corr_dict_numdim
+
     def err_corr_matrix(self) -> xr.DataArray:
         """
         Error-correlation matrix for uncertainty effect.
@@ -291,7 +340,12 @@ 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:
-            if np.all([dim in self._obj[self._unc_var_name][self._sli].dims for dim in dim_err_corr[1].dims]):
+            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)
                 )
@@ -583,7 +637,8 @@ def structured_err_cov_matrix(self):
 
         covs_sum = np.zeros(structured_err_cov_matrix.shape)
         for unc in self:
-            covs_sum += unc[self._sli].err_cov_matrix().values
+            if unc.is_structured:
+                covs_sum += unc[self._sli].err_cov_matrix().values
 
         structured_err_cov_matrix.values = covs_sum