Chebyshev series of matrices 7 Jun 2020

cheblinop.py

#!/usr/bin/env python
"""Chebyshev series of matrices

    polyD = Cheblinop( A, [c0 c1 c2 ...] )  # a LinearOperator
        A: a square numpy array
            | scipy.sparse matrix
            | scipy.sparse LinearOperator
        coefs: [c0 c1 ...] arraylike
            | numpy Chebyshev polynomial T( coefs [, domain=] )
            | filename to polyload

    polyD(x): the Chebyshev series of matrix A, dot x
        (c0 I + c1 T1( A ) + c2 T2( A ) ...) dot x
        x: scalar / vec / array / matrix -- anything for which A.dot(x) works

keywords: Chebyshev-polynomial, Chebyshev-series, matrix, sparse-matrix, numpy, scipy, LinearOperator  

https://docs.scipy.org/doc/numpy/reference/routines.polynomials.classes.html
https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html
"""
# Chebyshev polynomials: [Numerical Recipes](https://apps.nrbook.com/empanel/index.html?pg=233) ff.

from __future__ import division, print_function
import re
from six import string_types
import numpy as np
from numpy.polynomial import Chebyshev as T, Polynomial as P 
from scipy import sparse
from scipy.sparse.linalg import LinearOperator  # $scipy/sparse/linalg/interface.py

try:
    from polyutil import polysave, polyload, polystr  # .npz <-> P()
except ImportError:
    pass

__version__ = "2020-06-07 June  denis-bz-py t-online.de"


#...............................................................................
class Cheblinop( LinearOperator ):
    __doc__ = globals()["__doc__"]

    def __init__( s, A, poly, name="", verbose=1 ):
        assert hasattr( A, "dot" )  # or matvec ?
        s.A = A
        if isinstance( poly, T ):  # P ?
            pass
        elif isinstance( poly, string_types ):
            if not name:
                name = _basename( name )
            poly = polyload( poly, verbose=0 )
        else:
            poly = T( poly )

            # shift domain [a b] -> [-1 1], e.g. [0 1]: 2x - 1
            # like numpy T(coef, domain)
            # x -> w = domaintowindow(x) -> T( shifted coefs )( w )
        a, b = poly.domain
        assert np.all( poly.window == [-1, 1] ), poly.window
        if verbose:
            if [a, b] == [-1, 1]:
                dom = ""
            else:
                dom = "domain [%.3g %.3g]" % (a, b)
            print( "Cheblinop: A %s %s  coef %s  %s" % (
                    A.shape, type(A).__name__, poly.coef, dom ))
        if [a, b] != [-1, 1]:
            poly = poly.convert( kind=T, domain=[-1, 1] )  # shift coefs
            s.c1 = 2 / (b - a)
            s.c0 = - (a + b) / (b - a)
        else:
            s.c0 = s.c1 = None

        coef = poly.coef
        s.__str__ = s.__name__ = name or ("Cheblinop-d%d" % (len(coef) - 1))
        if len(coef) < 3:
            coef = np.append( coef, np.zeros( 3 - len(coef) ))
        s.coef = coef


    def chebeval( s, x ):
        """ T( matrix A ) at x
            = (c0 I + c1 T1( matrix A ) + c2 T2(A) ...) dot x
        """
            # wikipedia Clenshaw_algorithm, cheb_matrix
        x = _squeeze( x )
        if s.c0 is not None:
            x = s.c1 * x + s.c0  # e.g. 2x - 1
        A, coef = s.A, s.coef
        c1, c2 = coef[-2:]
        d2 = c2 * x
        d1 = 2 * A.dot( d2 ) + c1 * x
        for c in coef[-3:0:-1]:
            d = 2 * A.dot( d1 ) - d2 + c * x
            d2 = d1
            d1 = d
        d0 = A.dot( d1 ) - d2 + coef[0] * x
        return _squeeze( d0 )


    __call__ = matvec = _matvec = matmat = _matmat = dot = chebeval   # ?
        # __call__ = chebeval / polyeval: no, s.call

    @property
    def shape( s ):
        return s.A.shape

    @property
    def dtype( s ):
        return s.A.dtype


#...............................................................................
# etc --

def _basename( filename ):
    """ a/b/c.d.e -> c.d """
    return re.sub( r"\.[^.]*$", "", 
            re.sub( "^.*/", "", filename ))

def _squeeze( x ):
    return x.squeeze() if hasattr( x, "squeeze" ) \
        else x


#...............................................................................
if __name__ == "__main__":
    import sys

    np.set_printoptions( threshold=20, edgeitems=11, linewidth=140,
            formatter = dict( float = lambda x: "%.2g" % x ))  # float arrays %.2g
    print( 80 * "=" )

    n = 11
    # to change these params, run this.py  D=1  b=None  'c = expr' ... in sh or ipython
    for arg in sys.argv[1:]:
        exec( arg )

    #...........................................................................
    """
    vec d <-> matrix D with diagonal d
        d * x = [di * xi] = D dot x
        d * ones = [di] = D dot ones
        p( d ) <-> diag p( D )
    test:
        T( coef )( d ) == Chebpoly( D, coef )( ones )
    """

    d = np.linspace( 0, 1, n )
    D = np.diag( d )
    ones = np.ones( n )
    print( "d, D.diag:", d )

    for poly in [
        # T( [1], domain=[0, 1] ),
        T( [0, 1], domain=[0, 1] ),  # w = 2x - 1, w
        # T( [0, 2], domain=[0, 1] ),  # w = 2x - 1, 2 w
        T( [1, 2], domain=[0, 1] ),  # w = 2x - 1, 1 + 2w
        T( [1, -1], domain=[0, 1] ),  # 1 - w = 1 - (2x - 1)
        T( [0, 0, 1], domain=[0, 1] ),
            # polyload( "~/.data/T-d10*.npz" )  # domain [0 1]
        ]:
        print( "\n" + 80 * "-" )
        polyd = poly( d )
        print( "poly(d): %s  coef %s  domain %s " % (
                polyd, poly.coef, poly.domain))
        # p = poly.convert( kind=P )
        # print( "P(d):", poly(d) )

        polyD = Cheblinop( D, poly )
        # print( ".coef:", polyD.coef )  # shift
        polyD_ones = polyD( ones )
        print( "(ones):", polyD_ones )
        assert np.allclose( polyd, polyD_ones )