From 8a0cc0dfe9f666ae85f37dc60b65b9f12c553f34 Mon Sep 17 00:00:00 2001
From: Jonathan Wenger <jonathan.wenger@uni-tuebingen.de>
Date: Fri, 9 Oct 2020 10:09:59 +0200
Subject: [PATCH] Package restructuring and naming (#223)

* moved linops up in the namespace

* isort

* imports unified

* tests moved to correct dir and changed name back to linops

* renamed back to linops

* doctest problems fixed

* refactored aslinop into own file

* pylint fix

* doctest bug fixed in aslinop

* api doc structure unified

* unified import in tests
---
 benchmarks/random_variables.py                |  2 +-
 docs/source/index.rst                         |  7 ++--
 docs/source/public_api/linalg.rst             |  7 ----
 docs/source/public_api/linops.rst             |  5 +++
 .../{probnum.rst => random_variables.rst}     |  8 ++--
 src/probnum/__init__.py                       |  2 +-
 .../linalg/linearsolvers/linearsolvers.py     |  3 +-
 .../linalg/linearsolvers/matrixbased.py       |  5 ++-
 src/probnum/{linalg => }/linops/__init__.py   | 24 ++++++++---
 .../kronecker.py => linops/_kronecker.py}     | 20 ++++-----
 .../_linear_operator.py}                      | 40 +-----------------
 src/probnum/linops/_utils.py                  | 41 +++++++++++++++++++
 src/probnum/random_variables/_arithmetic.py   |  6 +--
 src/probnum/random_variables/_normal.py       |  2 +-
 .../test_linearsolvers/test_linearsolvers.py  |  6 +--
 .../{test_linalg => }/test_linops/__init__.py |  0
 .../test_linops/test_kronecker.py             |  2 +-
 .../test_linops/test_linearoperators.py       |  2 +-
 tests/test_random_variables/test_normal.py    |  6 ++-
 .../test_random_variable.py                   |  2 +-
 20 files changed, 102 insertions(+), 88 deletions(-)
 create mode 100644 docs/source/public_api/linops.rst
 rename docs/source/public_api/{probnum.rst => random_variables.rst} (54%)
 rename src/probnum/{linalg => }/linops/__init__.py (50%)
 rename src/probnum/{linalg/linops/kronecker.py => linops/_kronecker.py} (96%)
 rename src/probnum/{linalg/linops/linearoperators.py => linops/_linear_operator.py} (93%)
 create mode 100644 src/probnum/linops/_utils.py
 rename tests/{test_linalg => }/test_linops/__init__.py (100%)
 rename tests/{test_linalg => }/test_linops/test_kronecker.py (99%)
 rename tests/{test_linalg => }/test_linops/test_linearoperators.py (99%)

diff --git a/benchmarks/random_variables.py b/benchmarks/random_variables.py
index 0601a197f..074578a65 100644
--- a/benchmarks/random_variables.py
+++ b/benchmarks/random_variables.py
@@ -4,8 +4,8 @@
 
 import numpy as np
 
-import probnum.linalg.linops as linops
 from benchmarks.benchmark_utils import SPD_MATRIX_5x5
+from probnum import linops
 from probnum import random_variables as rvs
 
 # Module level variables
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5a84ca048..0ac745bfb 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -48,11 +48,12 @@ To learn how to use ProbNum check out the `quickstart guide <introduction/quicks
    :maxdepth: 1
    :caption: API Documentation
 
-   public_api/probnum
-   public_api/linalg
-   public_api/quad
    public_api/diffeq
    public_api/filtsmooth
+   public_api/linalg
+   public_api/linops
+   public_api/quad
+   public_api/random_variables
    public_api/utils
 
 .. toctree::
diff --git a/docs/source/public_api/linalg.rst b/docs/source/public_api/linalg.rst
index 13d1fc3fa..8af33a09e 100644
--- a/docs/source/public_api/linalg.rst
+++ b/docs/source/public_api/linalg.rst
@@ -1,12 +1,5 @@
 probnum.linalg
 ==============
 
-probnum.linalg
-**************
 .. automodapi:: probnum.linalg
-	:no-heading:
-
-probnum.linalg.linops
-*********************
-.. automodapi::probnum.linalg.linops
 	:no-heading:
\ No newline at end of file
diff --git a/docs/source/public_api/linops.rst b/docs/source/public_api/linops.rst
new file mode 100644
index 000000000..a064f6180
--- /dev/null
+++ b/docs/source/public_api/linops.rst
@@ -0,0 +1,5 @@
+probnum.linops
+==============
+
+.. automodapi:: probnum.linops
+	:no-heading:
\ No newline at end of file
diff --git a/docs/source/public_api/probnum.rst b/docs/source/public_api/random_variables.rst
similarity index 54%
rename from docs/source/public_api/probnum.rst
rename to docs/source/public_api/random_variables.rst
index 83b4fe020..b5d5be5f4 100644
--- a/docs/source/public_api/probnum.rst
+++ b/docs/source/public_api/random_variables.rst
@@ -1,7 +1,5 @@
-probnum
-=======
-
 probnum.random_variables
-************************
+========================
+
 .. automodapi:: probnum.random_variables
-	:no-heading:
+	:no-heading:
\ No newline at end of file
diff --git a/src/probnum/__init__.py b/src/probnum/__init__.py
index 87ff82777..3ece2da91 100644
--- a/src/probnum/__init__.py
+++ b/src/probnum/__init__.py
@@ -25,7 +25,7 @@
 # -*- coding: utf-8 -*-
 from pkg_resources import DistributionNotFound, get_distribution
 
-from . import diffeq, filtsmooth, linalg, quad, random_variables, utils
+from . import diffeq, filtsmooth, linalg, linops, quad, random_variables, utils
 from .random_variables import RandomVariable, asrandvar
 
 try:
diff --git a/src/probnum/linalg/linearsolvers/linearsolvers.py b/src/probnum/linalg/linearsolvers/linearsolvers.py
index 5ba840b45..fccd507df 100644
--- a/src/probnum/linalg/linearsolvers/linearsolvers.py
+++ b/src/probnum/linalg/linearsolvers/linearsolvers.py
@@ -13,8 +13,7 @@
 import scipy.sparse
 
 import probnum
-from probnum import utils
-from probnum.linalg import linops
+from probnum import linops, utils
 from probnum.linalg.linearsolvers.matrixbased import (
     AsymmetricMatrixBasedSolver,
     NoisySymmetricMatrixBasedSolver,
diff --git a/src/probnum/linalg/linearsolvers/matrixbased.py b/src/probnum/linalg/linearsolvers/matrixbased.py
index d9e4d8388..9f02427ce 100644
--- a/src/probnum/linalg/linearsolvers/matrixbased.py
+++ b/src/probnum/linalg/linearsolvers/matrixbased.py
@@ -11,8 +11,8 @@
 import numpy as np
 
 import probnum
+from probnum import linops
 from probnum import random_variables as rvs
-from probnum.linalg import linops
 
 
 class ProbabilisticLinearSolver(abc.ABC):
@@ -760,7 +760,8 @@ def _matvec(x):
         A = rvs.Normal(mean=self.A_mean, cov=linops.SymmetricKronecker(A=_A_covfactor))
 
         Ainv = rvs.Normal(
-            mean=self.Ainv_mean, cov=linops.SymmetricKronecker(A=_Ainv_covfactor)
+            mean=self.Ainv_mean,
+            cov=linops.SymmetricKronecker(A=_Ainv_covfactor),
         )
         # Induced distribution on x via Ainv
         # Exp(x) = Ainv b, Cov(x) = 1/2 (W b'Wb + Wbb'W)
diff --git a/src/probnum/linalg/linops/__init__.py b/src/probnum/linops/__init__.py
similarity index 50%
rename from src/probnum/linalg/linops/__init__.py
rename to src/probnum/linops/__init__.py
index ec102903e..b722e7678 100644
--- a/src/probnum/linalg/linops/__init__.py
+++ b/src/probnum/linops/__init__.py
@@ -6,23 +6,35 @@
 linear operator as a matrix in memory.
 """
 
-from probnum.linalg.linops.kronecker import *
-from probnum.linalg.linops.linearoperators import *
+from ._kronecker import Kronecker, Svec, SymmetricKronecker, Symmetrize, Vec
+from ._linear_operator import Identity, LinearOperator, MatrixMult, ScalarMult
+from ._utils import aslinop
 
 # Public classes and functions. Order is reflected in documentation.
 __all__ = [
+    "aslinop",
     "LinearOperator",
     "Identity",
     "ScalarMult",
     "MatrixMult",
     "Kronecker",
     "SymmetricKronecker",
+    "Symmetrize",
     "Vec",
     "Svec",
-    "Symmetrize",
-    "aslinop",
 ]
 
 # Set correct module paths. Corrects links and module paths in documentation.
-LinearOperator.__module__ = "probnum.linalg.linops"
-ScalarMult.__module__ = "probnum.linalg.linops"
+LinearOperator.__module__ = "probnum.linops"
+
+Identity.__module__ = "probnum.linops"
+ScalarMult.__module__ = "probnum.linops"
+MatrixMult.__module__ = "probnum.linops"
+
+Kronecker.__module__ = "probnum.linops"
+Svec.__module__ = "probnum.linops"
+SymmetricKronecker.__module__ = "probnum.linops"
+Symmetrize.__module__ = "probnum.linops"
+Vec.__module__ = "probnum.linops"
+
+aslinop.__module__ = "probnum.linops"
diff --git a/src/probnum/linalg/linops/kronecker.py b/src/probnum/linops/_kronecker.py
similarity index 96%
rename from src/probnum/linalg/linops/kronecker.py
rename to src/probnum/linops/_kronecker.py
index 849e0baf3..d475738a4 100644
--- a/src/probnum/linalg/linops/kronecker.py
+++ b/src/probnum/linops/_kronecker.py
@@ -5,10 +5,10 @@
 """
 import numpy as np
 
-from probnum.linalg.linops.linearoperators import LinearOperator, aslinop
+from . import _linear_operator, _utils
 
 
-class Symmetrize(LinearOperator):
+class Symmetrize(_linear_operator.LinearOperator):
     """
     Symmetrizes a vector in its matrix representation.
 
@@ -32,7 +32,7 @@ def _matvec(self, x):
         return Y.reshape(-1, 1)
 
 
-class Vec(LinearOperator):
+class Vec(_linear_operator.LinearOperator):
     """
     Vectorization operator.
 
@@ -68,7 +68,7 @@ def _matmat(self, X):
             return X.ravel(order="F")
 
 
-class Svec(LinearOperator):
+class Svec(_linear_operator.LinearOperator):
     """
     Symmetric vectorization operator.
 
@@ -148,7 +148,7 @@ def _matmat(self, X):
             )
 
 
-class Kronecker(LinearOperator):
+class Kronecker(_linear_operator.LinearOperator):
     """
     Kronecker product of two linear operators.
 
@@ -191,8 +191,8 @@ class Kronecker(LinearOperator):
 
     # todo: extend this to list of operators
     def __init__(self, A, B, dtype=None):
-        self.A = aslinop(A)
-        self.B = aslinop(B)
+        self.A = _utils.aslinop(A)
+        self.B = _utils.aslinop(B)
         super().__init__(
             dtype=dtype,
             shape=(
@@ -258,7 +258,7 @@ def trace(self):
             raise NotImplementedError
 
 
-class SymmetricKronecker(LinearOperator):
+class SymmetricKronecker(_linear_operator.LinearOperator):
     """
     Symmetric Kronecker product of two linear operators.
 
@@ -297,13 +297,13 @@ class SymmetricKronecker(LinearOperator):
 
     def __init__(self, A, B=None, dtype=None):
         # Set parameters
-        self.A = aslinop(A)
+        self.A = _utils.aslinop(A)
         self._ABequal = False
         if B is None:
             self.B = self.A
             self._ABequal = True
         else:
-            self.B = aslinop(B)
+            self.B = _utils.aslinop(B)
         self._n = self.A.shape[0]
         if self.A.shape != self.B.shape or self.A.shape[1] != self._n:
             raise ValueError(
diff --git a/src/probnum/linalg/linops/linearoperators.py b/src/probnum/linops/_linear_operator.py
similarity index 93%
rename from src/probnum/linalg/linops/linearoperators.py
rename to src/probnum/linops/_linear_operator.py
index 1e971629e..a0f62b457 100644
--- a/src/probnum/linalg/linops/linearoperators.py
+++ b/src/probnum/linops/_linear_operator.py
@@ -45,8 +45,6 @@ class LinearOperator(scipy.sparse.linalg.LinearOperator):
     Implementing ``_adjoint`` is preferable; ``_rmatvec`` is mostly there for backwards
     compatibility.
 
-    This class inherits from :class:`scipy.sparse.linalg.LinearOperator`.
-
     Parameters
     ----------
     shape : tuple
@@ -70,7 +68,7 @@ class LinearOperator(scipy.sparse.linalg.LinearOperator):
     Examples
     --------
     >>> import numpy as np
-    >>> from probnum.linalg.linops import LinearOperator
+    >>> from probnum.linops import LinearOperator
     >>> def mv(v):
     ...     return np.array([2 * v[0] - v[1], 3 * v[1]])
     ...
@@ -575,39 +573,3 @@ def trace(self):
             raise ValueError("The trace is only defined for square linear operators.")
         else:
             return np.trace(self.A)
-
-
-def aslinop(A):
-    """
-    Return `A` as a :class:`LinearOperator`.
-
-    Parameters
-    ----------
-    A : array-like or LinearOperator or RandomVariable or object
-        Argument to be represented as a linear operator. When `A` is an object it needs
-        to have the attributes `.shape` and `.matvec`.
-
-    Notes
-    -----
-    If `A` has no `.dtype` attribute, the data type is determined by calling
-    :func:`LinearOperator.matvec()` - set the `.dtype` attribute to prevent this
-    call upon the linear operator creation.
-
-    See Also
-    --------
-    LinearOperator : Class representing linear operators.
-
-    Examples
-    --------
-    >>> from probnum.linalg.linops import aslinop
-    >>> M = np.array([[1,2,3],[4,5,6]], dtype=np.int32)
-    >>> aslinop(M)
-    <2x3 MatrixMult with dtype=int32>
-    """
-    if isinstance(A, scipy.sparse.linalg.LinearOperator):
-        return A
-    elif isinstance(A, (np.ndarray, scipy.sparse.spmatrix)):
-        return MatrixMult(A=A)
-    else:
-        op = scipy.sparse.linalg.aslinearoperator(A)
-        return LinearOperator(op)
diff --git a/src/probnum/linops/_utils.py b/src/probnum/linops/_utils.py
new file mode 100644
index 000000000..df38f2474
--- /dev/null
+++ b/src/probnum/linops/_utils.py
@@ -0,0 +1,41 @@
+import numpy as np
+import scipy.sparse
+
+from . import _linear_operator
+
+
+def aslinop(A) -> _linear_operator.LinearOperator:
+    """
+    Return ``A`` as a :class:`LinearOperator`.
+
+    Parameters
+    ----------
+    A : array-like or LinearOperator or RandomVariable or object
+        Argument to be represented as a linear operator. When `A` is an object it needs
+        to have the attributes `.shape` and `.matvec`.
+
+    See Also
+    --------
+    LinearOperator : Class representing linear operators.
+
+    Notes
+    -----
+    If `A` has no `.dtype` attribute, the data type is determined by calling
+    :func:`LinearOperator.matvec()` - set the `.dtype` attribute to prevent this
+    call upon the linear operator creation.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from probnum.linops import aslinop
+    >>> M = np.array([[1,2,3],[4,5,6]], dtype=np.int32)
+    >>> aslinop(M)
+    <2x3 MatrixMult with dtype=int32>
+    """
+    if isinstance(A, scipy.sparse.linalg.LinearOperator):
+        return A
+    elif isinstance(A, (np.ndarray, scipy.sparse.spmatrix)):
+        return _linear_operator.MatrixMult(A=A)
+    else:
+        op = scipy.sparse.linalg.aslinearoperator(A)
+        return _linear_operator.LinearOperator(op)
diff --git a/src/probnum/random_variables/_arithmetic.py b/src/probnum/random_variables/_arithmetic.py
index ebf54395a..869ed8631 100644
--- a/src/probnum/random_variables/_arithmetic.py
+++ b/src/probnum/random_variables/_arithmetic.py
@@ -6,8 +6,8 @@
 
 import numpy as np
 
+import probnum.linops as _linear_operators
 from probnum import utils as _utils
-from probnum.linalg import linops as _linops
 
 from ._dirac import Dirac as _Dirac
 from ._normal import Normal as _Normal
@@ -280,8 +280,8 @@ def _matmul_normal_dirac(norm_rv: _Normal, dirac_rv: _Dirac) -> _Normal:
             ),
         )
     elif norm_rv.ndim == 2 and norm_rv.shape[0] > 1:
-        cov_update = _linops.Kronecker(
-            _linops.Identity(dirac_rv.shape[0]), dirac_rv.support
+        cov_update = _linear_operators.Kronecker(
+            _linear_operators.Identity(dirac_rv.shape[0]), dirac_rv.support
         )
 
         return _Normal(
diff --git a/src/probnum/random_variables/_normal.py b/src/probnum/random_variables/_normal.py
index eae08cda4..2c9568a09 100644
--- a/src/probnum/random_variables/_normal.py
+++ b/src/probnum/random_variables/_normal.py
@@ -6,8 +6,8 @@
 import scipy.linalg
 import scipy.stats
 
+from probnum import linops
 from probnum import utils as _utils
-from probnum.linalg import linops
 from probnum.type import (
     ArrayLikeGetitemArgType,
     FloatArgType,
diff --git a/tests/test_linalg/test_linearsolvers/test_linearsolvers.py b/tests/test_linalg/test_linearsolvers/test_linearsolvers.py
index f4c49a29c..e9d20cadd 100644
--- a/tests/test_linalg/test_linearsolvers/test_linearsolvers.py
+++ b/tests/test_linalg/test_linearsolvers/test_linearsolvers.py
@@ -7,9 +7,8 @@
 import scipy.sparse.linalg
 
 import probnum
-from probnum import linalg
+from probnum import linalg, linops
 from probnum import random_variables as rvs
-from probnum.linalg import linops
 from tests.testing import NumpyAssertions
 
 
@@ -383,7 +382,8 @@ def callback_iterates_CG(xk):
             cov=linops.SymmetricKronecker(A=linops.Identity(A.shape[1])),
         )
         A0 = rvs.Normal(
-            mean=linops.Identity(A.shape[1]), cov=linops.SymmetricKronecker(A)
+            mean=linops.Identity(A.shape[1]),
+            cov=linops.SymmetricKronecker(A),
         )
         for kwargs in [{"assume_A": "sympos", "rtol": 10 ** -6}]:
             with self.subTest():
diff --git a/tests/test_linalg/test_linops/__init__.py b/tests/test_linops/__init__.py
similarity index 100%
rename from tests/test_linalg/test_linops/__init__.py
rename to tests/test_linops/__init__.py
diff --git a/tests/test_linalg/test_linops/test_kronecker.py b/tests/test_linops/test_kronecker.py
similarity index 99%
rename from tests/test_linalg/test_linops/test_kronecker.py
rename to tests/test_linops/test_kronecker.py
index d32232d79..87ffbaf65 100644
--- a/tests/test_linalg/test_linops/test_kronecker.py
+++ b/tests/test_linops/test_kronecker.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from probnum.linalg import linops
+from probnum import linops
 from tests.testing import NumpyAssertions
 
 
diff --git a/tests/test_linalg/test_linops/test_linearoperators.py b/tests/test_linops/test_linearoperators.py
similarity index 99%
rename from tests/test_linalg/test_linops/test_linearoperators.py
rename to tests/test_linops/test_linearoperators.py
index b49b504c6..8933a5472 100644
--- a/tests/test_linalg/test_linops/test_linearoperators.py
+++ b/tests/test_linops/test_linearoperators.py
@@ -5,7 +5,7 @@
 import numpy as np
 import scipy.sparse
 
-from probnum.linalg import linops
+from probnum import linops
 from tests.testing import NumpyAssertions
 
 
diff --git a/tests/test_random_variables/test_normal.py b/tests/test_random_variables/test_normal.py
index 7e499d59a..8a9741fca 100644
--- a/tests/test_random_variables/test_normal.py
+++ b/tests/test_random_variables/test_normal.py
@@ -7,8 +7,8 @@
 import scipy.stats
 
 import probnum
+from probnum import linops
 from probnum import random_variables as rvs
-from probnum.linalg import linops
 from tests.testing import NumpyAssertions
 
 
@@ -223,7 +223,9 @@ def test_symmetric_samples(self):
         A = np.random.uniform(size=(n, n))
         A = 0.5 * (A + A.T) + n * np.eye(n)
         rv = rvs.Normal(
-            mean=np.eye(A.shape[0]), cov=linops.SymmetricKronecker(A=A), random_state=1
+            mean=np.eye(A.shape[0]),
+            cov=linops.SymmetricKronecker(A=A),
+            random_state=1,
         )
         rv = rv.sample(size=10)
         for i, B in enumerate(rv):
diff --git a/tests/test_random_variables/test_random_variable.py b/tests/test_random_variables/test_random_variable.py
index d416b1419..5fba46f91 100644
--- a/tests/test_random_variables/test_random_variable.py
+++ b/tests/test_random_variables/test_random_variable.py
@@ -7,8 +7,8 @@
 import scipy.stats
 
 import probnum
+from probnum import linops
 from probnum import random_variables as rvs
-from probnum.linalg import linops
 from tests.testing import NumpyAssertions