Merge pull request #28 from comet-toolkit/small_fixes

Small fixes

Author: shunt16

obsarray/err_corr.py

@@ -69,22 +69,23 @@ def form(self) -> str:
         """Form name"""
         pass
 
-    def expand_dim_matrix(
-        self, submatrix: np.ndarray, submatrix_dim: Union[str, List[str]], sli: tuple
-    ):
-        return expand_errcorr_dims(
-            in_corr=submatrix,
-            in_dim=submatrix_dim,
-            out_dim=list(self._obj[self._unc_var_name][sli].dims),
-            dim_sizes=self.get_sliced_dim_sizes_uncvar(sli),
-        )
+    def get_varshape_errcorr(self):
+        """
+        return shape of uncertainty variable, including only dimensions which are included in the current error correlation form.
+
+        :return: shape of included dimensions
+        """
+        all_dims = self._obj[self._unc_var_name].dims
+        all_dims_sizes = self._obj.sizes
+
+        return tuple([all_dims_sizes[dim] for dim in all_dims if dim in self.dims])
 
     def get_sliced_dim_sizes_uncvar(self, sli: tuple) -> 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
+        :return: dictionary with 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
@@ -98,7 +99,7 @@ def get_sliced_dim_sizes_errcorr(self, sli: tuple) -> dict:
         included in the current error correlation form.
 
         :param sli: slice (tuple with slice for each dimension)
-        :return: shape of included sliced dimensions
+        :return: dictionary with shape of included sliced dimensions
         """
         uncvar_sizes = self.get_sliced_dim_sizes_uncvar(sli)
         sliced_dims = self.get_sliced_dims_errcorr(sli)
@@ -131,17 +132,21 @@ def get_sliced_shape_errcorr(self, sli: tuple) -> tuple:
 
         return tuple([uncvar_sizes[dim] for dim in sliced_dims])
 
-    def slice_full_cov(self, full_matrix: np.ndarray, sli: tuple) -> np.ndarray:
-        return self.slice_flattened_matrix(
-            full_matrix, self._obj[self._unc_var_name].shape, sli
-        )
+    def slice_errcorr_matrix(self, err_corr_matrix, variable_shape, sli) -> np.ndarray:
+        """
+        Slice the provided error correlation matrix (typically the error correlation matrix of the
+        BaseErrCorrForm) using the
 
-    def slice_flattened_matrix(self, flattened_matrix, variable_shape, sli):
+        :param err_corr_matrix: error correlation matrix to be sliced
+        :param variable_shape: tuple with the length of the dimensions in the error correlation matrix (in correct order for flattening)
+        :param sli: slice of observation variable to return error-correlation matrix for
+        :return: sliced error correlation matrix
+        """
         mask_array = np.ones(variable_shape, dtype=bool)
         mask_array[sli] = False
 
         return np.delete(
-            np.delete(flattened_matrix, mask_array.ravel(), 0), mask_array.ravel(), 1
+            np.delete(err_corr_matrix, mask_array.ravel(), 0), mask_array.ravel(), 1
         )
 
     @abc.abstractmethod
@@ -158,16 +163,24 @@ def build_matrix(self, sli: Union[np.ndarray, tuple]) -> np.ndarray:
 
     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
+        Returns expanded error correlation matrix for use in dot product with error correlation
+        in other dimensions.
 
-        :param sli: slice of observation variable to return error-correlation matrix for
+        The (sliced) error correlation matrix for this BaseErrCorrForm is expanded from its current
+        (sliced) dimensions (which often don't include all dimensions of the associated uncertainty
+        variable) to the dimensions of the full (sliced) error correlation (i.e. all dimensions of
+        the uncertainty).
+        The returned matrix is not meaningfull unless combined in a dot product with the expanded
+        matrices of other error correlation matrices (together spanning all uncertainty dimensions).
 
-        :return: populated error-correlation matrix
+        :param sli: slice of observation variable to return error-correlation matrix for
+        :return: expanded matrix for use in dot product with error correlation in other dimensions.
         """
-
-        return self.expand_dim_matrix(
-            self.build_matrix(sli), self.get_sliced_dims_errcorr(sli), sli
+        return expand_errcorr_dims(
+            in_corr=self.build_matrix(sli),
+            in_dim=self.get_sliced_dims_errcorr(sli),
+            out_dim=list(self._obj[self._unc_var_name][sli].dims),
+            dim_sizes=self.get_sliced_dim_sizes_uncvar(sli),
         )
 
 
@@ -250,24 +263,14 @@ def build_matrix(self, sli: tuple) -> np.ndarray:
 
         :return: populated error-correlation matrix
         """
-
         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
+        submatrix = self.slice_errcorr_matrix(
+            self._obj[self.params[0]], self.get_varshape_errcorr(), sli_submatrix
         )
 
         return submatrix

obsarray/test/test_err_corr_forms.py

@@ -7,6 +7,7 @@
     ErrCorrForms,
     RandomCorrelation,
     SystematicCorrelation,
+    ErrCorrMatrixCorrelation,
 )
 from obsarray.test.test_unc_accessor import create_ds
 
@@ -65,18 +66,26 @@ class BasicErrCorrForm(BaseErrCorrForm):
             form = "basic"
 
             def build_matrix(self, sli):
-                return None
+                full_matrix = np.arange(144).reshape((12, 12))
+                return self.slice_errcorr_matrix(full_matrix, (2, 2, 3), sli)
 
         self.BasicErrCorrForm = BasicErrCorrForm
 
+    def test_get_varshape_errcorr(self):
+        basicerrcorr = self.BasicErrCorrForm(
+            self.ds, "u_ran_temperature", ["x", "y", "time"], [], []
+        )
+        shape = basicerrcorr.get_varshape_errcorr()
+        np.testing.assert_equal(shape, (2, 2, 3))
+
     def test_slice_full_cov_full(self):
         basicerrcorr = self.BasicErrCorrForm(
             self.ds, "u_ran_temperature", ["x"], [], []
         )
 
         full_matrix = np.arange(144).reshape((12, 12))
-        slice_matrix = basicerrcorr.slice_full_cov(
-            full_matrix, (slice(None), slice(None), slice(None))
+        slice_matrix = basicerrcorr.build_matrix(
+            (slice(None), slice(None), slice(None))
         )
 
         np.testing.assert_equal(full_matrix, slice_matrix)
@@ -118,13 +127,13 @@ def test_get_sliced_shape_errcorr(self):
         shape = basicerrcorr.get_sliced_shape_errcorr((slice(None), 0, slice(0, 2, 1)))
         assert shape == (2, 2)
 
-    def test_slice_flattened_matrix(self):
+    def test_slice_errcorr_matrix(self):
         basicerrcorr = self.BasicErrCorrForm(
-            self.ds, "u_ran_temperature", ["x"], [], []
+            self.ds, "u_ran_temperature", ["x", "y", "z"], [], []
         )
 
         full_matrix = np.arange(144).reshape((12, 12))
-        slice_matrix = basicerrcorr.slice_flattened_matrix(
+        slice_matrix = basicerrcorr.slice_errcorr_matrix(
             full_matrix, (2, 2, 3), (slice(None), slice(None), 0)
         )
 
@@ -136,13 +145,10 @@ def test_slice_flattened_matrix(self):
 
     def test_slice_full_cov_slice(self):
         basicerrcorr = self.BasicErrCorrForm(
-            self.ds, "u_ran_temperature", ["x"], [], []
+            self.ds, "u_ran_temperature", ["x", "y", "z"], [], []
         )
 
-        full_matrix = np.arange(144).reshape((12, 12))
-        slice_matrix = basicerrcorr.slice_full_cov(
-            full_matrix, (slice(None), slice(None), 0)
-        )
+        slice_matrix = basicerrcorr.build_dot_matrix((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]]
@@ -206,5 +212,40 @@ def test_build_dot_matrix(self):
         np.testing.assert_equal((x.dot(y)).dot(time), np.ones((12, 12)))
 
 
+class TestErrCorrMatrixCorrelation(unittest.TestCase):
+    def setUp(self) -> None:
+        self.ds = create_ds()
+
+    def test_build_matrix_full(self):
+        ec = ErrCorrMatrixCorrelation(
+            self.ds,
+            "u_str_temperature",
+            ["x", "time"],
+            ["err_corr_str_temperature"],
+            [],
+        )
+
+        ecrm = ec.build_matrix((slice(None), slice(None), slice(None)))
+        np.testing.assert_equal(ecrm, np.ones((6, 6)))
+
+    def test_build_matrix_sliced(self):
+        ec = ErrCorrMatrixCorrelation(
+            self.ds,
+            "u_str_temperature",
+            ["x", "time"],
+            ["err_corr_str_temperature"],
+            [],
+        )
+
+        ecrm = ec.build_matrix((0, slice(None), slice(None)))
+        np.testing.assert_equal(ecrm, np.ones((3, 3)))
+
+        ecrm = ec.build_matrix((slice(None), 0, slice(None)))
+        np.testing.assert_equal(ecrm, np.ones((6, 6)))
+
+        ecrm = ec.build_matrix((slice(None), slice(None), 0))
+        np.testing.assert_equal(ecrm, np.ones((2, 2)))
+
+
 if __name__ == "main":
     unittest.main()

obsarray/test/test_unc_accessor.py

@@ -98,7 +98,7 @@ def create_ds():
             "units": "K",
             "err_corr": [
                 {
-                    "dim": "x",
+                    "dim": ["x", "time"],
                     "form": "err_corr_matrix",
                     "params": ["err_corr_str_temperature"],
                 },
@@ -107,19 +107,19 @@ def create_ds():
                     "form": "systematic",
                     "params": [],
                 },
-                {
-                    "dim": "time",
-                    "form": "systematic",
-                    "params": [],
-                },
             ],
             "pdf_shape": "gaussian",
         },
     )
 
     ds["err_corr_str_temperature"] = (
-        ["x", "x"],
-        np.eye(temperature.shape[0]),
+        ["x.time", "x.time"],
+        np.ones(
+            (
+                temperature.shape[0] * temperature.shape[2],
+                temperature.shape[0] * temperature.shape[2],
+            )
+        ),
     )
 
     return ds

obsarray/unc_accessor.py

@@ -354,7 +354,7 @@ def err_cov_matrix(self):
         err_cov_matrix = empty_err_corr_matrix(self._obj[self._unc_var_name][self._sli])
 
         err_cov_matrix.values = convert_corr_to_cov(
-            self.err_corr_matrix().values, self.value.values
+            self.err_corr_matrix().values, self.abs_value.values
         )
 
         return err_cov_matrix