refactor randvars with pylint (#684)

* refactor pylint * python black formatting * Update src/probnum/randvars/_categorical.py Co-authored-by: Marvin Pförtner <[email protected]> Co-authored-by: Anshul Kapoor <[email protected]> Co-authored-by: Jonathan Wenger <[email protected]> Co-authored-by: Marvin Pförtner <[email protected]>
probabilistic-numerics · Apr 3, 2022 · ae017dd · ae017dd
1 parent 0864fcf
commit ae017dd
Show file tree

Hide file tree

Showing 8 changed files with 133 additions and 128 deletions.
diff --git a/src/probnum/randvars/_arithmetic.py b/src/probnum/randvars/_arithmetic.py
@@ -255,16 +255,17 @@ def _mul_normal_constant(
             return _Constant(
                 support=np.zeros_like(norm_rv.mean),
             )
+
+        if norm_rv.cov_cholesky_is_precomputed:
+            cov_cholesky = constant_rv.support * norm_rv.cov_cholesky
         else:
-            if norm_rv.cov_cholesky_is_precomputed:
-                cov_cholesky = constant_rv.support * norm_rv.cov_cholesky
-            else:
-                cov_cholesky = None
-            return _Normal(
-                mean=constant_rv.support * norm_rv.mean,
-                cov=(constant_rv.support**2) * norm_rv.cov,
-                cov_cholesky=cov_cholesky,
-            )
+            cov_cholesky = None
+
+        return _Normal(
+            mean=constant_rv.support * norm_rv.mean,
+            cov=(constant_rv.support**2) * norm_rv.cov,
+            cov_cholesky=cov_cholesky,
+        )
 
     return NotImplemented
 
@@ -276,7 +277,8 @@ def _mul_normal_constant(
 def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal:
     """Normal random variable multiplied with a vector or matrix.
 
-    Computes the distribution of the random variable :math:`Y = XA`, where :math:`X` is a matrix- or multi-variate normal random variable and :math:`A` a constant.
+    Computes the distribution of the random variable :math:`Y = XA`, where :math:`X`
+    is a matrix- or multi-variate normal random variable and :math:`A` a constant.
     """
     if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[0] == 1):
         if norm_rv.cov_cholesky_is_precomputed:
@@ -293,25 +295,25 @@ def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal
             cov = cov.reshape((1, 1))
 
         return _Normal(mean=mean, cov=cov, cov_cholesky=cov_cholesky)
+
+    # This part does not do the Cholesky update,
+    # because of performance configurations: currently, there is no way of switching
+    # the Cholesky updates off, which might affect (large, potentially sparse)
+    # covariance matrices of matrix-variate Normal RVs. See Issue #335.
+    if constant_rv.support.ndim == 1:
+        constant_rv_support = constant_rv.support[:, None]
     else:
-        # This part does not do the Cholesky update,
-        # because of performance configurations: currently, there is no way of switching
-        # the Cholesky updates off, which might affect (large, potentially sparse) covariance matrices
-        # of matrix-variate Normal RVs. See Issue #335.
-        if constant_rv.support.ndim == 1:
-            constant_rv_support = constant_rv.support[:, None]
-        else:
-            constant_rv_support = constant_rv.support
+        constant_rv_support = constant_rv.support
 
-        cov_update = _linear_operators.Kronecker(
-            _linear_operators.Identity(norm_rv.shape[0]), constant_rv_support.T
-        )
+    cov_update = _linear_operators.Kronecker(
+        _linear_operators.Identity(norm_rv.shape[0]), constant_rv_support.T
+    )
 
-        # Cov(rvec(XA)) = Cov((I (x) A.T)rvec(X)) = (I (x) A.T)Cov(rvec(X))(I (x) A.T).T
-        return _Normal(
-            mean=norm_rv.mean @ constant_rv.support,
-            cov=cov_update @ (norm_rv.cov @ cov_update.T),
-        )
+    # Cov(rvec(XA)) = Cov((I (x) A.T)rvec(X)) = (I (x) A.T)Cov(rvec(X))(I (x) A.T).T
+    return _Normal(
+        mean=norm_rv.mean @ constant_rv.support,
+        cov=cov_update @ (norm_rv.cov @ cov_update.T),
+    )
 
 
 _matmul_fns[(_Normal, _Constant)] = _matmul_normal_constant
@@ -320,7 +322,8 @@ def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal
 def _matmul_constant_normal(constant_rv: _Constant, norm_rv: _Normal) -> _Normal:
     """Matrix-multiplication with a normal random variable.
 
-    Computes the distribution of the random variable :math:`Y = AX`, where :math:`X` is a matrix- or multi-variate normal random variable and :math:`A` a constant.
+    Computes the distribution of the random variable :math:`Y = AX`, where :math:`X` is
+    a matrix- or multi-variate normal random variable and :math:`A` a constant.
     """
     if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[1] == 1):
         if norm_rv.cov_cholesky_is_precomputed:
@@ -334,26 +337,26 @@ def _matmul_constant_normal(constant_rv: _Constant, norm_rv: _Normal) -> _Normal
             cov=constant_rv.support @ (norm_rv.cov @ constant_rv.support.T),
             cov_cholesky=cov_cholesky,
         )
+
+    # This part does not do the Cholesky update,
+    # because of performance configurations: currently, there is no way of switching
+    # the Cholesky updates off, which might affect (large, potentially sparse)
+    # covariance matrices of matrix-variate Normal RVs. See Issue #335.
+    if constant_rv.support.ndim == 1:
+        constant_rv_support = constant_rv.support[None, :]
     else:
-        # This part does not do the Cholesky update,
-        # because of performance configurations: currently, there is no way of switching
-        # the Cholesky updates off, which might affect (large, potentially sparse) covariance matrices
-        # of matrix-variate Normal RVs. See Issue #335.
-        if constant_rv.support.ndim == 1:
-            constant_rv_support = constant_rv.support[None, :]
-        else:
-            constant_rv_support = constant_rv.support
+        constant_rv_support = constant_rv.support
 
-        cov_update = _linear_operators.Kronecker(
-            constant_rv_support,
-            _linear_operators.Identity(norm_rv.shape[1]),
-        )
+    cov_update = _linear_operators.Kronecker(
+        constant_rv_support,
+        _linear_operators.Identity(norm_rv.shape[1]),
+    )
 
-        # Cov(rvec(AX)) = Cov((A (x) I)rvec(X)) = (A (x) I)Cov(rvec(X))(A (x) I).T
-        return _Normal(
-            mean=constant_rv.support @ norm_rv.mean,
-            cov=cov_update @ (norm_rv.cov @ cov_update.T),
-        )
+    # Cov(rvec(AX)) = Cov((A (x) I)rvec(X)) = (A (x) I)Cov(rvec(X))(A (x) I).T
+    return _Normal(
+        mean=constant_rv.support @ norm_rv.mean,
+        cov=cov_update @ (norm_rv.cov @ cov_update.T),
+    )
 
 
 _matmul_fns[(_Constant, _Normal)] = _matmul_constant_normal

diff --git a/src/probnum/randvars/_categorical.py b/src/probnum/randvars/_categorical.py
@@ -56,15 +56,15 @@ def _sample_categorical(rng, size=()):
             return self.support[indices]
 
         def _pmf_categorical(x):
-            """PMF of a categorical distribution.
+            """PMF of a categorical distribution."""
 
-            This implementation is defense against cryptic warnings such as:
+            # This implementation is defense against cryptic warnings such as:
             # https://stackoverflow.com/questions/45020217/numpy-where-function-throws-a-futurewarning-returns-scalar-instead-of-list
-            """
             x = np.asarray(x)
             if x.dtype != self.dtype:
                 raise ValueError(
-                    "The data type of x does not match with the data type of the support."
+                    "The data type of x does not match with the data type of the "
+                    "support."
                 )
 
             mask = (x == self.support).nonzero()[0]
@@ -109,7 +109,8 @@ def resample(self, rng: np.random.Generator) -> "Categorical":
         Returns
         -------
         Categorical
-            Categorical random variable with resampled support (according to self.probabilities).
+            Categorical random variable with resampled support
+            (according to self.probabilities).
         """
         num_events = len(self.support)
         new_support = self.sample(rng=rng, size=num_events)

diff --git a/src/probnum/randvars/_constant.py b/src/probnum/randvars/_constant.py
@@ -146,8 +146,8 @@ def _sample(self, rng: np.random.Generator, size: ShapeLike = ()) -> ValueType:
 
         if size == ():
             return self._support.copy()
-        else:
-            return np.tile(self._support, reps=size + (1,) * self.ndim)
+
+        return np.tile(self._support, reps=size + (1,) * self.ndim)
 
     # Unary arithmetic operations
 

diff --git a/src/probnum/randvars/_normal.py b/src/probnum/randvars/_normal.py
@@ -38,11 +38,14 @@ class Normal(_random_variable.ContinuousRandomVariable[ValueType]):
     cov :
         (Co-)variance of the random variable.
     cov_cholesky :
-        (Lower triangular) Cholesky factor of the covariance matrix. If None, then the Cholesky factor of the covariance matrix
-        is computed when :attr:`Normal.cov_cholesky` is called and then cached. If specified, the value is returned by :attr:`Normal.cov_cholesky`.
-        In this case, its type and data type are compared to the type and data type of the covariance.
-        If the types do not match, an exception is thrown. If the data types do not match,
-        the data type of the Cholesky factor is promoted to the data type of the covariance matrix.
+        (Lower triangular) Cholesky factor of the covariance matrix. If None, then the
+        Cholesky factor of the covariance matrix is computed when
+        :attr:`Normal.cov_cholesky` is called and then cached. If specified, the value
+        is returned by :attr:`Normal.cov_cholesky`. In this case, its type and data type
+        are compared to the type and data type of the covariance. If the types do not
+        match, an exception is thrown. If the data types do not match,
+        the data type of the Cholesky factor is promoted to the data type of the
+        covariance matrix.
 
     See Also
     --------
@@ -142,7 +145,8 @@ def __init__(
             compute_cov_cholesky = self.dense_cov_cholesky
 
             # Ensure that the Cholesky factor has the same type as the covariance,
-            # and, if necessary, promote data types. Check for (in this order): type, shape, dtype.
+            # and, if necessary, promote data types. Check for (in this order): type,
+            # shape, dtype.
             if cov_cholesky is not None:
 
                 if not isinstance(cov_cholesky, type(cov)):
@@ -171,8 +175,8 @@ def __init__(
                     if m != n or n != cov.A.shape[0] or n != cov.B.shape[1]:
                         raise ValueError(
                             "Normal distributions with symmetric Kronecker structured "
-                            "kernels must have square mean and square kernels factors with "
-                            "matching dimensions."
+                            "kernels must have square mean and square kernels factors "
+                            "with matching dimensions."
                         )
 
                     if cov.identical_factors:
@@ -193,8 +197,8 @@ def __init__(
                             or n != cov.B.shape[1]
                         ):
                             raise ValueError(
-                                "Kronecker structured kernels must have factors with the same "
-                                "shape as the mean."
+                                "Kronecker structured kernels must have factors with "
+                                "the same shape as the mean."
                             )
 
                 else:
@@ -269,16 +273,16 @@ def dense_mean(self) -> Union[np.floating, np.ndarray]:
         """Dense representation of the mean."""
         if isinstance(self.mean, linops.LinearOperator):
             return self.mean.todense()
-        else:
-            return self.mean
+
+        return self.mean
 
     @cached_property
     def dense_cov(self) -> Union[np.floating, np.ndarray]:
         """Dense representation of the covariance."""
         if isinstance(self.cov, linops.LinearOperator):
             return self.cov.todense()
-        else:
-            return self.cov
+
+        return self.cov
 
     def __getitem__(self, key: ArrayIndicesLike) -> "Normal":
         """Marginalization in multi- and matrixvariate normal random variables,
@@ -486,10 +490,11 @@ def _dense_sample(
     def _arg_todense(x: Union[np.ndarray, linops.LinearOperator]) -> np.ndarray:
         if isinstance(x, linops.LinearOperator):
             return x.todense()
-        elif isinstance(x, np.ndarray):
+
+        if isinstance(x, np.ndarray):
             return x
-        else:
-            raise ValueError(f"Unsupported argument type {type(x)}")
+
+        raise ValueError(f"Unsupported argument type {type(x)}")
 
     @staticmethod
     def _dense_in_support(x: ValueType) -> bool:

diff --git a/src/probnum/randvars/_random_variable.py b/src/probnum/randvars/_random_variable.py
@@ -439,17 +439,16 @@ def cdf(self, x: ValueType) -> np.float_:
             return RandomVariable._ensure_numpy_float(
                 "cdf", self.__cdf(self._as_value_type(x))
             )
-        elif self.__logcdf is not None:
-            cdf = np.exp(self.logcdf(self._as_value_type(x)))
 
+        if self.__logcdf is not None:
+            cdf = np.exp(self.logcdf(self._as_value_type(x)))
             assert isinstance(cdf, np.float_)
-
             return cdf
-        else:
-            raise NotImplementedError(
-                f"Neither the `cdf` nor the `logcdf` of the random variable object "
-                f"with type `{type(self).__name__}` is implemented."
-            )
+
+        raise NotImplementedError(
+            f"Neither the `cdf` nor the `logcdf` of the random variable object "
+            f"with type `{type(self).__name__}` is implemented."
+        )
 
     def logcdf(self, x: ValueType) -> np.float_:
         """Log-cumulative distribution function.
@@ -466,17 +465,16 @@ def logcdf(self, x: ValueType) -> np.float_:
             return RandomVariable._ensure_numpy_float(
                 "logcdf", self.__logcdf(self._as_value_type(x))
             )
-        elif self.__cdf is not None:
-            logcdf = np.log(self.__cdf(x))
 
+        if self.__cdf is not None:
+            logcdf = np.log(self.__cdf(x))
             assert isinstance(logcdf, np.float_)
-
             return logcdf
-        else:
-            raise NotImplementedError(
-                f"Neither the `logcdf` nor the `cdf` of the random variable object "
-                f"with type `{type(self).__name__}` is implemented."
-            )
+
+        raise NotImplementedError(
+            f"Neither the `logcdf` nor the `cdf` of the random variable object "
+            f"with type `{type(self).__name__}` is implemented."
+        )
 
     def quantile(self, p: FloatLike) -> ValueType:
         """Quantile function.
@@ -1023,17 +1021,16 @@ def pmf(self, x: ValueType) -> np.float_:
         """
         if self.__pmf is not None:
             return DiscreteRandomVariable._ensure_numpy_float("pmf", self.__pmf(x))
-        elif self.__logpmf is not None:
-            pmf = np.exp(self.__logpmf(x))
 
+        if self.__logpmf is not None:
+            pmf = np.exp(self.__logpmf(x))
             assert isinstance(pmf, np.float_)
-
             return pmf
-        else:
-            raise NotImplementedError(
-                f"Neither the `pmf` nor the `logpmf` of the discrete random variable "
-                f"object with type `{type(self).__name__}` is implemented."
-            )
+
+        raise NotImplementedError(
+            f"Neither the `pmf` nor the `logpmf` of the discrete random variable "
+            f"object with type `{type(self).__name__}` is implemented."
+        )
 
     def logpmf(self, x: ValueType) -> np.float_:
         """Natural logarithm of the probability mass function.
@@ -1050,17 +1047,16 @@ def logpmf(self, x: ValueType) -> np.float_:
             return DiscreteRandomVariable._ensure_numpy_float(
                 "logpmf", self.__logpmf(self._as_value_type(x))
             )
-        elif self.__pmf is not None:
-            logpmf = np.log(self.__pmf(self._as_value_type(x)))
 
+        if self.__pmf is not None:
+            logpmf = np.log(self.__pmf(self._as_value_type(x)))
             assert isinstance(logpmf, np.float_)
-
             return logpmf
-        else:
-            raise NotImplementedError(
-                f"Neither the `logpmf` nor the `pmf` of the discrete random variable "
-                f"object with type `{type(self).__name__}` is implemented."
-            )
+
+        raise NotImplementedError(
+            f"Neither the `logpmf` nor the `pmf` of the discrete random variable "
+            f"object with type `{type(self).__name__}` is implemented."
+        )
 
 
 class ContinuousRandomVariable(RandomVariable[ValueType]):
@@ -1267,15 +1263,13 @@ def logpdf(self, x: ValueType) -> np.float_:
             return ContinuousRandomVariable._ensure_numpy_float(
                 "logpdf", self.__logpdf(self._as_value_type(x))
             )
-        elif self.__pdf is not None:
 
+        if self.__pdf is not None:
             logpdf = np.log(self.__pdf(self._as_value_type(x)))
-
             assert isinstance(logpdf, np.float_)
-
             return logpdf
-        else:
-            raise NotImplementedError(
-                f"Neither the `logpdf` nor the `pdf` of the continuous random variable "
-                f"object with type `{type(self).__name__}` is implemented."
-            )
+
+        raise NotImplementedError(
+            f"Neither the `logpdf` nor the `pdf` of the continuous random variable "
+            f"object with type `{type(self).__name__}` is implemented."
+        )