Update pyfaust.poly function signatures and remove special case ChebyshevPhi.

wrapper/python/pyfaust/poly.py

@@ -4,44 +4,24 @@
 
 import scipy.sparse as sp
 from scipy.sparse import csr_matrix
-import pyfaust as pf
+from pyfaust import Faust, isFaust
 from scipy.sparse.linalg import eigsh
 
 
-def ChebyshevPhi(L, K, phi=None, dev='cpu'):
+def Chebyshev(L, K, ret_gen=False, dev='cpu', T0=None):
     """
-    Builds the Faust of the Chebyshev polynomials defined on the symmetric matrix L using its greatest eigenvalue phi.
-
-    Args:
-        L: the symmetric matrix.
-        phi: the greatest eigen value of the L matrix (if None, it is computed
-        automatically using scipy.sparse.linalg.eigsh).
-        K: the degree of the last polynomial, i.e. the K+1 first polynomials are built.
-        dev: the destination device of the polynomial Faust.
-
-    Returns:
-        The Faust of the K+1 Chebyshev polynomials.
-
-    """
-    if phi is None:
-        phi = eigsh(L, k=1, return_eigenvectors=False)[0] / 2
-    d = L.shape[0]
-    if not isinstance(L, csr_matrix):
-        L = csr_matrix(L)
-    L_phi = (1/phi)*L
-    Id = T0 = sp.eye(d, format="csr")
-    T1 = sp.vstack((Id, -Id+L_phi))
-    rR = sp.hstack((-Id, 2*(L_phi-Id)), format="csr")
-    return _chebyshev(L, K, T0, T1, rR, dev)
-
-def Chebyshev(L, K, dev='cpu'):
-    """
-    Builds the Faust of the Chebyshev polynomials defined on the symmetric matrix L.
+    Builds the Faust of the Chebyshev polynomial basis defined on the symmetric matrix L.
 
     Args:
         L: the symmetric matrix.
         K: the degree of the last polynomial, i.e. the K+1 first polynomials are built.
         dev: the destination device of the polynomial Faust.
+        ret_gen: to return a generator of polynomials in addition to the
+        polynomial itself (the generator starts from the the
+        K+1-degree polynomial, and allows this way to compute the next
+        polynomial simply with the instruction: next(generator)).
+        T0: to define the 0-degree polynomial as something else than the
+        identity.
 
     Returns:
         The Faust of the K+1 Chebyshev polynomials.
@@ -50,49 +30,94 @@ def Chebyshev(L, K, dev='cpu'):
         L = csr_matrix(L)
     twoL = 2*L
     d = L.shape[0]
-    Id = T0 = sp.eye(d, format="csr")
+    Id = sp.eye(d, format="csr")
+    if isinstance(T0, type(None)):
+        T0 = Id
     T1 = sp.vstack((Id, L))
     rR = sp.hstack((-Id, twoL), format="csr")
-    return _chebyshev(L, K, T0, T1, rR, dev)
+    if ret_gen:
+        g = _chebyshev_gen(L, T0, T1, rR, dev)
+        for i in range(0, K):
+            next(g)
+        return next(g), g
+    else:
+        return _chebyshev(L, K, T0, T1, rR, dev)
 
 
 def _chebyshev(L, K, T0, T1, rR, dev='cpu'):
-    Id = T0
     d = L.shape[0]
-    factors = [Id]
+    factors = [T0]
     if(K > 0):
         factors.insert(0, T1)
         for i in range(2, K + 1):
-            if i <= 2:
-                R = rR
-            else:
-                zero = csr_matrix((d, (i-2)*d), dtype=float)
-                R = sp.hstack((zero, rR), format="csr")
-            Ti = sp.vstack((sp.eye(d*i, format="csr"), R),
-                           format="csr")
+            Ti = _chebyshev_Ti_matrix(rR, L, i)
             factors.insert(0, Ti)
-    T = pf.Faust(factors, dev=dev)
+    T = Faust(factors, dev=dev)
     return T  # K-th poly is T[K*L.shape[0]:,:]
 
 
-def polyFaust(coeffs, L=None, poly=Chebyshev, dev='cpu'):
+def _chebyshev_gen(L, T0, T1, rR, dev='cpu'):
+    T = Faust(T0)
+    yield T
+    T = Faust(T1) @ T
+    yield T
+    i = 2
+    while True:
+        Ti = _chebyshev_Ti_matrix(rR, L, i)
+        T = Faust(Ti) @ T
+        yield T
+        i += 1
+
+
+def _chebyshev_Ti_matrix(rR, L, i):
+    d = L.shape[0]
+    if i <= 2:
+        R = rR
+    else:
+        zero = csr_matrix((d, (i-2)*d), dtype=float)
+        R = sp.hstack((zero, rR), format="csr")
+    Ti = sp.vstack((sp.eye(d*i, format="csr"), R),
+                   format="csr")
+    return Ti
+
+def basis(L, K, basis_name, ret_gen=False, dev='cpu', T0=None):
+    """
+    Builds the Faust of the polynomial basis defined on the symmetric matrix L.
+
+    Args:
+        L: the symmetric matrix.
+        K: the degree of the last polynomial, i.e. the K+1 first polynomials are built.
+        basis_name: 'chebyshev', and others yet to come.
+        dev: the destination device of the polynomial Faust.
+        ret_gen: to return a generator of polynomials in addition to the
+        polynomial itself (the generator starts from the the
+        K+1-degree polynomial, and allows this way to compute the next
+        polynomial simply with the instruction: next(generator)).
+
+    Returns:
+        The Faust of the K+1 Chebyshev polynomials.
+    """
+    if basis_name.lower() == 'chebyshev':
+        return Chebyshev(L, K, ret_gen=ret_gen, dev=dev, T0=None)
+
+def poly(coeffs, L=None, basis=Chebyshev, dev='cpu'):
     """
-        Returns the linear combination of the polynomials defined by poly.
+        Returns the linear combination of the polynomials defined by basis.
 
         Args:
             coeffs: the linear combination coefficients (numpy.array).
-            poly: either the function to build the polynomials on L or the Faust of
-            polynomials if already built.
+            basis: either the function to build the polynomials basis on L or the Faust of
+            polynomials if already built externally.
             L: the symmetric matrix on which the polynomials are built, can't be None
-            if poly is a function (not a Faust).
+            if basis is a function (not a Faust).
             dev: the device to instantiate the returned Faust ('cpu' or 'gpu').
 
         Returns:
             The linear combination Faust.
     """
     K = coeffs.size-1
-    if pf.isFaust(poly):
-        F = poly
+    if isFaust(basis):
+        F = basis
         if F.device != dev:
             F = F.clone(dev=dev)
     else:
@@ -100,9 +125,9 @@ def polyFaust(coeffs, L=None, poly=Chebyshev, dev='cpu'):
             raise ValueError('The L matrix must be set to build the'
                              ' polynomials.')
         F = poly(L, K, dev=dev)
-    Id = sp.eye(F.shape[1], format="csr")
+    Id = sp.eye(L.shape[1], format="csr")
     scoeffs = sp.hstack(tuple(Id*coeffs[i] for i in range(0, K+1)),
                         format="csr")
-    Fc = pf.Faust(scoeffs, dev=dev) @ F
+    Fc = Faust(scoeffs, dev=dev) @ F
     return Fc
 # experimental block end