Fix bug that was happening when slicing data

Author: pdevis

obsarray/err_corr.py

@@ -69,23 +69,68 @@ def form(self) -> str:
         """Form name"""
         pass
 
-    def expand_dim_matrix(self, submatrix, sli):
+    def expand_dim_matrix(self, submatrix, submatrix_dim, sli):
         return expand_errcorr_dims(
             in_corr=submatrix,
-            in_dim=self.dims,
+            in_dim=submatrix_dim,
             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_sizes=self.get_sliced_dim_sizes_uncvar(sli)
         )
 
+    def get_sliced_dim_sizes_uncvar(self, sli) -> dict:
+        """
+        return dictionary with sizes of sliced dimensions of unc variable, including all dimensions.
+
+        :param sli: slice (tuple with slice for each dimension)
+        :return: shape of included sliced dimensions
+        """
+        uncvar_dims = self._obj[self._unc_var_name][sli].dims
+        uncvar_shape = self._obj[self._unc_var_name][sli].shape
+        return {uncvar_dims[idim]: uncvar_shape[idim] for idim in range(len(uncvar_dims))}
+
+    def get_sliced_dim_sizes_errcorr(self, sli) -> dict:
+        """
+        return dictionary with sizes of sliced dimensions of unc variable, including only dimensions which are included in the current error correlation form.
+
+        :param sli: slice (tuple with slice for each dimension)
+        :return: shape of included sliced dimensions
+        """
+        uncvar_sizes=self.get_sliced_dim_sizes_uncvar(sli)
+        sliced_dims = self.get_sliced_dims_errcorr(sli)
+
+        return {dim: uncvar_sizes[dim] for dim in sliced_dims}
+
+    def get_sliced_dims_errcorr(self, sli) -> list:
+        """
+        return dimensions which are within the slice and included in the current error correlation form.
+
+        :param sli: slice (tuple with slice for each dimension)
+        :return: list with sliced dimensions
+        """
+        all_dims = self._obj[self._unc_var_name].dims
+        return [all_dims[idim] for idim in range(len(all_dims)) if (isinstance(sli[idim],slice) and all_dims[idim] in self.dims)]
+
+    def get_sliced_shape_errcorr(self, sli)->tuple:
+        """
+        return shape of sliced uncertainty variable, including only dimensions which are included in the current error correlation form.
+
+        :param sli: slice (tuple with slice for each dimension)
+        :return: shape of included sliced dimensions
+        """
+        uncvar_sizes=self.get_sliced_dim_sizes_uncvar(sli)
+        sliced_dims = self.get_sliced_dims_errcorr(sli)
+
+        return tuple([uncvar_sizes[dim] for dim in sliced_dims])
+
     def slice_full_cov(self, full_matrix, sli):
-        mask_array = np.ones(self._obj[self._unc_var_name].shape, dtype=bool)
+        return self.slice_flattened_matrix(full_matrix,self._obj[self._unc_var_name].shape,sli)
+
+    def slice_flattened_matrix(self, flattened_matrix, variable_shape, sli):
+        mask_array = np.ones(variable_shape, dtype=bool)
         mask_array[sli] = False
 
         return np.delete(
-            np.delete(full_matrix, mask_array.ravel(), 0), mask_array.ravel(), 1
+            np.delete(flattened_matrix, mask_array.ravel(), 0), mask_array.ravel(), 1
         )
 
     @abc.abstractmethod
@@ -100,6 +145,18 @@ def build_matrix(self, sli: Union[np.ndarray, tuple]) -> np.ndarray:
         """
         pass
 
+    def build_dot_matrix(self, sli: Union[np.ndarray, tuple]) -> np.ndarray:
+        """
+        Returns uncertainty effect error-correlation matrix, populated with error-correlation values defined
+        in this parameterisation
+
+        :param sli: slice of observation variable to return error-correlation matrix for
+
+        :return: populated error-correlation matrix
+        """
+
+        return self.expand_dim_matrix(self.build_matrix(sli), self.get_sliced_dims_errcorr(sli), sli)
+
 
 def register_err_corr_form(form_name: str) -> Callable:
     """
@@ -133,16 +190,12 @@ def build_matrix(self, sli):
         """
 
         # evaluate correlation over matrices in form defintion
-        dim_lens = [len(self._obj[dim]) for dim in self.dims]
+        dim_lens = self.get_sliced_shape_errcorr(sli)
         n_elems = int(np.prod(dim_lens))
 
-        dims_matrix = np.eye(n_elems)
-
-        # expand to correlation matrix over all variable dims
-        return self.expand_dim_matrix(dims_matrix, sli)
+        submatrix = np.eye(n_elems)
 
-        # # subset to slice
-        # return self.slice_full_cov(full_matrix, sli)
+        return submatrix
 
 
 @register_err_corr_form("systematic")
@@ -162,16 +215,12 @@ def build_matrix(self, sli):
         """
 
         # evaluate correlation over matrices in form defintion
-        dim_lens = [len(self._obj[dim]) for dim in self.dims]
+        dim_lens = self.get_sliced_shape_errcorr(sli)
         n_elems = int(np.prod(dim_lens))
 
-        dims_matrix = np.ones((n_elems, n_elems))
+        submatrix = np.ones((n_elems, n_elems))
 
-        # expand to correlation matrix over all variable dims
-        return self.expand_dim_matrix(dims_matrix, sli)
-
-        # subset to slice
-        # return self.slice_full_cov(full_matrix, sli)
+        return submatrix
 
 
 @register_err_corr_form("err_corr_matrix")
@@ -189,11 +238,16 @@ def build_matrix(self, sli):
         :return: populated error-correlation matrix
         """
 
-        # expand to correlation matrix over all variable dims
-        return self.expand_dim_matrix(self._obj[self.params[0]], sli)
+        all_dims = self._obj[self._unc_var_name].dims
+        all_dims_sizes = self._obj.sizes
+
+        sli_submatrix = tuple([sli[i] for i in range(len(all_dims)) if all_dims[i] in self.dims])
+
+        sliced_shape = tuple([all_dims_sizes[all_dims[i]] for i in range(len(all_dims)) if all_dims[i] in self.dims])
+
+        submatrix = self.slice_flattened_matrix(self._obj[self.params[0]],sliced_shape,sli_submatrix)
 
-        # # subset to slice
-        # return self.slice_full_cov(full_matrix, sli)
+        return submatrix
 
 
 @register_err_corr_form("ensemble")

obsarray/test/test_err_corr_forms.py

@@ -81,6 +81,55 @@ def test_slice_full_cov_full(self):
 
         np.testing.assert_equal(full_matrix, slice_matrix)
 
+    def test_get_sliced_dim_sizes_uncvar(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x"], [], []
+        )
+        dim_sizes=basicerrcorr.get_sliced_dim_sizes_uncvar((slice(None),0,slice(0,2,1)))
+        assert dim_sizes == {"x": 2, "time": 2}
+
+    def test_get_sliced_dim_sizes_errcorr(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x"], [], []
+        )
+        dim_sizes = basicerrcorr.get_sliced_dim_sizes_errcorr((slice(None), 0, slice(0, 2, 1)))
+        assert dim_sizes == {"x": 2}
+
+    def test_get_sliced_dims_errcorr(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x"], [], []
+        )
+        dims = basicerrcorr.get_sliced_dims_errcorr((slice(None), 0, slice(0, 2, 1)))
+        assert dims == ["x"]
+
+    def test_get_sliced_shape_errcorr(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x"], [], []
+        )
+        shape = basicerrcorr.get_sliced_shape_errcorr((slice(None), 0, slice(0, 2, 1)))
+        assert shape == (2,)
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x", "time"], [], []
+        )
+        shape = basicerrcorr.get_sliced_shape_errcorr((slice(None), 0, slice(0, 2, 1)))
+        assert shape == (2,2)
+
+    def test_slice_flattened_matrix(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x"], [], []
+        )
+
+        full_matrix = np.arange(144).reshape((12, 12))
+        slice_matrix = basicerrcorr.slice_flattened_matrix(
+            full_matrix, (2,2,3), (slice(None), slice(None), 0)
+        )
+
+        exp_slice_matrix = np.array(
+            [[0, 3, 6, 9], [36, 39, 42, 45], [72, 75, 78, 81], [108, 111, 114, 117]]
+        )
+
+        np.testing.assert_equal(slice_matrix, exp_slice_matrix)
+
     def test_slice_full_cov_slice(self):
         basicerrcorr = self.BasicErrCorrForm(
             self.ds, "u_ran_temperature", ["x"], [], []
@@ -97,22 +146,21 @@ def test_slice_full_cov_slice(self):
 
         np.testing.assert_equal(slice_matrix, exp_slice_matrix)
 
-
 class TestRandomUnc(unittest.TestCase):
     def setUp(self) -> None:
         self.ds = create_ds()
 
     def test_build_matrix_1stdim(self):
         rc = RandomCorrelation(self.ds, "u_ran_temperature", ["x"], [], [])
 
-        ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        ecrm = rc.build_dot_matrix((slice(None), slice(None), slice(None)))
 
         np.testing.assert_equal(ecrm, np.eye(12))
 
     def test_build_matrix_2nddim(self):
         rc = RandomCorrelation(self.ds, "u_ran_temperature", ["y"], [], [])
 
-        ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        ecrm = rc.build_dot_matrix((slice(None), slice(None), slice(None)))
 
         np.testing.assert_equal(ecrm, np.eye(12))
 
@@ -124,33 +172,33 @@ def setUp(self) -> None:
     def build_matrix_1stdim(self):
         rc = SystematicCorrelation(self.ds, "u_sys_temperature", ["x"], [], [])
 
-        ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        ecrm = rc.build_dot_matrix((slice(None), slice(None), slice(None)))
         # np.testing.assert_equal(ecrm, np.ones((12, 12)))
 
         return ecrm
 
     def build_matrix_2nddim(self):
         rc = SystematicCorrelation(self.ds, "u_sys_temperature", ["y"], [], [])
 
-        ecrm = rc.build_matrix((slice(None), slice(None), slice(None)))
+        ecrm = rc.build_dot_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)))
+        ecrm = rc.build_dot_matrix((slice(None), slice(None), slice(None)))
         # np.testing.assert_equal(ecrm, np.ones((12, 12)))
 
         return ecrm
 
-    def test_build_matrix(self):
+    def test_build_dot_matrix(self):
         x = self.build_matrix_1stdim()
         y = self.build_matrix_2nddim()
         time = self.build_matrix_3ddim()
+        print(x.dot(y),x,y)
         np.testing.assert_equal((x.dot(y)).dot(time), np.ones((12, 12)))
 
-
 if __name__ == "main":
     unittest.main()

obsarray/test/test_unc_accessor.py

@@ -117,6 +117,11 @@ def create_ds():
         },
     )
 
+    ds["err_corr_str_temperature"] = (
+        ["x", "x"],
+        np.eye(temperature.shape[0]),
+    )
+
     return ds
 
 
@@ -474,13 +479,15 @@ def test_systematic_unc(self, mock):
         self.ds.unc["temperature"][:, :, 0].systematic_unc()
         mock.assert_called_once_with(["u_sys_temperature"])
 
-    @patch(
-        "obsarray.unc_accessor.Uncertainty.err_corr_matrix",
-        return_value=xr.DataArray(np.ones((12, 12)), dims=["x.y.time", "x.y.time"]),
-    )
-    def test_total_err_corr_matrix(self, mock_err_corr_matrix):
-        pass
-        # tercm = self.ds.unc["temperature"].total_err_corr_matrix()
+    def test_total_err_corr_matrix(self, ):
+        tercm = self.ds.unc["temperature"].total_err_corr_matrix()
+        assert tercm.shape==(12,12)
+        tercm = self.ds.unc["temperature"][:,:,0].total_err_corr_matrix()
+        assert tercm.shape==(4,4)
+        tercm = self.ds.unc["temperature"][:,0,0:2].total_err_corr_matrix()
+        assert tercm.shape==(4,4)
+        tercm = self.ds.unc["temperature"][0,0,:].total_err_corr_matrix()
+        assert tercm.shape==(3,3)
 
     def test_structured_err_corr_matrix(self):
         pass
@@ -496,7 +503,7 @@ class TestUncertainty(unittest.TestCase):
     def setUp(self):
         self.ds = create_ds()
 
-    @patch("obsarray.unc_accessor.Uncertainty.expand_sli", return_value="slice")
+    @patch("obsarray.unc_accessor.Uncertainty._expand_sli", return_value="slice")
     def test___getitem__(self, m):
         self.assertEqual(
             self.ds.unc["temperature"]["u_ran_temperature"]["in_slice"]._sli, "slice"
@@ -513,19 +520,19 @@ def test_expand_slice_1d_full(self):
     def test_expand_slice_1d_None(self):
         self.ds["new"] = (["time"], np.ones(3), {})
         self.ds.unc["new"]["u_new"] = (["time"], np.ones(3), {})
-        sli = self.ds.unc["temperature"]["u_ran_temperature"]._expand_sli()
+        sli = self.ds.unc["new"]["u_new"]._expand_sli()
         self.assertEqual((slice(None),), sli)
 
     def test_expand_slice_full(self):
-        sli = self.ds.unc["temperature"]["u_ran_temperature"].expand_sli((1, 1, 1))
+        sli = self.ds.unc["temperature"]["u_ran_temperature"]._expand_sli((1, 1, 1))
         self.assertEqual((1, 1, 1), sli)
 
     def test_expand_slice_None(self):
-        sli = self.ds.unc["temperature"]["u_ran_temperature"].expand_sli()
+        sli = self.ds.unc["temperature"]["u_ran_temperature"]._expand_sli()
         self.assertEqual((slice(None), slice(None), slice(None)), sli)
 
     def test_expand_slice_first(self):
-        sli = self.ds.unc["temperature"]["u_ran_temperature"].expand_sli((0,))
+        sli = self.ds.unc["temperature"]["u_ran_temperature"]._expand_sli((0,))
         self.assertEqual((0, slice(None), slice(None)), sli)
 
     def test_err_corr(self):

obsarray/unc_accessor.py

@@ -320,14 +320,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:
-            if np.all(
-                [
-                    dim in self._obj[self._unc_var_name][self._sli].dims
-                    for dim in dim_err_corr[1].dims
-                ]
-            ):
+            sliced_dims=dim_err_corr[1].get_sliced_dims_errcorr(self._sli)
+            if len(dim_err_corr[1].get_sliced_dims_errcorr(self._sli))>0:
                 err_corr_matrix.values = err_corr_matrix.values.dot(
-                    dim_err_corr[1].build_matrix(self._sli)
+                    dim_err_corr[1].build_dot_matrix(self._sli)
                 )
 
         return err_corr_matrix