feectools 0.1.0__py3-none-any.whl

feectools/linalg/direct_solvers.py

@@ -0,0 +1,201 @@
+# coding: utf-8
+# Copyright 2018 Jalal Lakhlili, Yaman Güçlü
+
+from abc                 import abstractmethod
+import numpy               as np
+from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
+from scipy.sparse        import spmatrix
+from scipy.sparse.linalg import splu
+
+from feectools.linalg.basic    import LinearSolver
+
+__all__ = ('BandedSolver', 'SparseSolver')
+
+#===============================================================================
+class BandedSolver(LinearSolver):
+    """
+    Solve the equation Ax = b for x, assuming A is banded matrix.
+
+    Parameters
+    ----------
+    u : integer
+        Number of non-zero upper diagonal.
+
+    l : integer
+        Number of non-zero lower diagonal.
+
+    bmat : nd-array
+        Banded matrix.
+
+    """
+    def __init__(self, u, l, bmat, transposed=False):
+
+        self._u    = u
+        self._l    = l
+        self._transposed = transposed
+
+        # ... LU factorization
+        if bmat.dtype == np.float32:
+            self._factor_function = sgbtrf
+            self._solver_function = sgbtrs
+        elif bmat.dtype == np.float64:
+            self._factor_function = dgbtrf
+            self._solver_function = dgbtrs
+        elif bmat.dtype == np.complex64:
+            self._factor_function = cgbtrf
+            self._solver_function = cgbtrs
+        elif bmat.dtype == np.complex128:
+            self._factor_function = zgbtrf
+            self._solver_function = zgbtrs
+        else:
+            msg = f'Cannot create a BandedSolver for bmat.dtype = {bmat.dtype}'
+            raise NotImplementedError(msg)
+
+        self._bmat, self._ipiv, self._finfo = self._factor_function(bmat, l, u)
+
+        self._sinfo = None
+
+        self._space = np.ndarray
+        self._dtype = bmat.dtype
+
+    @property
+    def finfo(self):
+        return self._finfo
+
+    @property
+    def sinfo(self):
+        return self._sinfo
+
+    #--------------------------------------
+    # Abstract interface
+    #--------------------------------------
+    @property
+    def space(self):
+        return self._space
+
+    def transpose(self):
+        cls = type(self)
+        obj = super().__new__(cls)
+
+        obj._u = self._l
+        obj._l = self._u
+        obj._bmat = self._bmat
+        obj._ipiv = self._ipiv
+        obj._finfo = self._finfo
+        obj._factor_function = self._factor_function
+        obj._solver_function = self._solver_function
+        obj._sinfo = None
+        obj._space = self._space
+        obj._dtype = self._dtype
+        obj._transposed = not self._transposed
+
+        return obj
+
+    #...
+    def solve(self, rhs, out=None):
+        """
+        Solves for the given right-hand side.
+
+        Parameters
+        ----------
+        rhs : ndarray
+            The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order,
+            i.e. if multiple right-hand sides are given, then rhs is a two-dimensional array with the 0-th
+            index denoting the number of the right-hand side, and the 1-st index denoting the element inside
+            a right-hand side.
+        
+        out : ndarray | NoneType
+            Output vector. If given, it has to have the same shape and datatype as rhs.
+        """
+        assert rhs.T.shape[0] == self._bmat.shape[1]
+
+        transposed = self._transposed
+
+        if out is None:
+            preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, rhs.T, self._ipiv,
+                                                        trans=transposed)
+            out = preout.T
+
+        else:
+            assert out.shape == rhs.shape
+            assert out.dtype == rhs.dtype
+
+            # support in-place operations
+            if rhs is not out:
+                out[:] = rhs
+
+            # TODO: handle non-contiguous views?
+
+            # we want FORTRAN-contiguous data (default is assumed to be C contiguous)
+            _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True,
+                                                   trans=transposed)
+
+        return out
+
+#===============================================================================
+class SparseSolver (LinearSolver):
+    """
+    Solve the equation Ax = b for x, assuming A is scipy sparse matrix.
+
+    Parameters
+    ----------
+    spmat : scipy.sparse.spmatrix
+        Generic sparse matrix.
+
+    """
+    def __init__(self, spmat, transposed=False):
+
+        assert isinstance(spmat, spmatrix)
+
+        self._space = np.ndarray
+        self._splu  = splu(spmat.tocsc())
+        self._transposed = transposed
+
+    #--------------------------------------
+    # Abstract interface
+    #--------------------------------------
+    @property
+    def space(self):
+        return self._space
+
+    def transpose(self):
+        cls = type(self)
+        obj = super().__new__(cls)
+
+        obj._space = self._space
+        obj._splu = self._splu
+        obj._transposed = not self._transposed
+
+        return obj
+
+    #...
+    def solve(self, rhs, out=None):
+        """
+        Solves for the given right-hand side.
+
+        Parameters
+        ----------
+        rhs : ndarray
+            The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order,
+            i.e. if multiple right-hand sides are given, then rhs is a two-dimensional array with the 0-th
+            index denoting the number of the right-hand side, and the 1-st index denoting the element inside
+            a right-hand side.
+        
+        out : ndarray | NoneType
+            Output vector. If given, it has to have the same shape and datatype as rhs.
+        """
+        
+        assert rhs.T.shape[0] == self._splu.shape[1]
+        transposed = self._transposed
+
+        if out is None:
+            out = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T
+
+        else:
+            assert out.shape == rhs.shape
+            assert out.dtype == rhs.dtype
+
+            # currently no in-place solve exposed
+            out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T
+
+        return out

feectools/linalg/fft.py

@@ -0,0 +1,258 @@
+from feectools.linalg.basic import LinearOperator, LinearSolver
+from feectools.linalg.stencil import StencilVectorSpace
+from feectools.linalg.kron import KroneckerLinearSolver
+
+import numpy as np
+import scipy.fft as scifft
+import os
+
+class DistributedFFTBase(LinearOperator):
+    """
+    A base class for the distributed FFT, DCT and DST.
+    Internally calls a KroneckerLinearSolver on a solver which just applies the FFT or some other function.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The vector space needed for the KroneckerLinearSolver internally.
+    
+    function : callable | list/tuple of callables
+        A list/tuple of callables function, each with one parameter x which applies some function in-place on x.
+        The function at position i is applied to the i-th tensor direction.
+        If only a single callable is given, it is used for all directions.
+    """
+    def toarray(self):
+        raise NotImplementedError('toarray() is not defined for DistributedFFTBase.')
+
+    def tosparse(self):
+        raise NotImplementedError('tosparse() is not defined for DistributedFFTBase.')
+
+    # Possible additions for the future:
+    # * split off the LinearSolver class when used with the space ndarray (as used in the KroneckerLinearSolver),
+    #   and make it state if it works in-place (or if it needs temporary memory), and what its optimal
+    #   size is (FFT might work faster with padding)
+    # * include FFTW support (e.g. pyfftw)
+
+    class OneDimSolver(LinearSolver):
+        """
+        A one-dimensional solver which just applies a given function.
+
+        Parameters
+        ----------
+        function : Callable
+            The given function.
+        """
+        def __init__(self, function):
+            self._function = function
+
+        @property
+        def space(self):
+            return np.ndarray
+
+        def transpose(self):
+            raise NotImplementedError('transpose() is not implemented for OneDimSolvers')
+        
+        def solve(self, rhs, out=None):
+            if out is None:
+                out = np.empty_like(rhs)
+            
+            if out is not rhs:
+                out[:] = rhs
+            
+            self._function(out)
+
+            return out
+
+    def __init__(self, space, functions):
+        assert isinstance(space, StencilVectorSpace)
+        if isinstance(functions, list) or isinstance(functions, tuple):
+            solvers = [DistributedFFTBase.OneDimSolver(function) for function in functions]
+        else:
+            onedimsolver = DistributedFFTBase.OneDimSolver(functions)
+            solvers = [onedimsolver] * space.ndim
+        self._isolver = KroneckerLinearSolver(space, space, solvers)
+
+    # ...
+    @property
+    def domain(self):
+        return self._isolver.space
+
+    # ...
+    @property
+    def codomain(self):
+        return self._isolver.space
+
+    # ...
+    @property
+    def dtype( self ):
+        return self._isolver.dtype
+
+    # ...
+    def dot(self, v, out=None):
+        # just call the KroneckerLinearSolver
+        return self._isolver.solve(v, out=out)
+
+    def transpose(self, conjugate=False):
+        raise NotImplementedError()
+
+# IMPORTANT NOTE: All of these scifft.fft functions currently trust that overwrite_x=True will yield an in-place fft...
+# (this is not completely given to hold forever, so in case these tests fail in some future version, change this)
+
+class DistributedFFT(DistributedFFTBase):
+    """
+    Equals an n-dimensional FFT operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.fft for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional FFT should be run on. Must have a complex data type (i.e. space.dtype.kind == 'c').
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.fft parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.fft parameter.
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None)):
+        # only allow complex data types
+        assert isinstance(space, StencilVectorSpace)
+        assert np.dtype(space.dtype).kind == 'c'
+        workers = int(workers) if workers is not None else None
+
+        super().__init__(space, lambda out: scifft.fft(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm))
+
+
+class DistributedIFFT(DistributedFFTBase):
+    """
+    Equals an n-dimensional IFFT operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.ifft for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional IFFT should be run on. Must have a complex data type (i.e. space.dtype.kind == 'c').
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.ifft parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.ifft parameter.
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None)):
+        # only allow complex data types
+        assert isinstance(space, StencilVectorSpace)
+        assert np.dtype(space.dtype).kind == 'c'
+        workers = int(workers) if workers is not None else None
+        
+        super().__init__(space, lambda out: scifft.ifft(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm))
+
+class DistributedDCT(DistributedFFTBase):
+    """
+    Equals an n-dimensional DCT operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.dct for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional DCT should be run on.
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.dct parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.dct parameter.
+    
+    ttype : int
+        The DCT type to use. (the name of this parameter in the underlying method is actually `type`).
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None), ttype=2):
+        workers = int(workers) if workers is not None else None
+        super().__init__(space, lambda out: scifft.dct(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm, type=ttype))
+
+class DistributedIDCT(DistributedFFTBase):
+    """
+    Equals an n-dimensional IDCT operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.idct for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional IDCT should be run on.
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.idct parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.idct parameter.
+    
+    ttype : int
+        The DCT type to use. (the name of this parameter in the underlying method is actually `type`).
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None), ttype=2):
+        workers = int(workers) if workers is not None else None
+        super().__init__(space, lambda out: scifft.idct(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm, type=ttype)) 
+
+class DistributedDST(DistributedFFTBase):
+    """
+    Equals an n-dimensional DST operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.dst for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional DST should be run on.
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.dst parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.dst parameter.
+    
+    ttype : int
+        The DCT type to use. (the name of this parameter in the underlying method is actually `type`).
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None), ttype=2):
+        workers = int(workers) if workers is not None else None
+        super().__init__(space, lambda out: scifft.dst(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm, type=ttype))
+
+class DistributedIDST(DistributedFFTBase):
+    """
+    Equals an n-dimensional IDST operation, except that it works on a distributed/parallel StencilVector.
+
+    Internally calls scipy.fft.idst for each direction.
+
+    Parameters
+    ----------
+    space : StencilVectorSpace
+        The space the n-dimensional IDST should be run on.
+    
+    norm : str
+        Specifies the normalization factor. See the documentation of the corresponding scipy.fft.idst parameter.
+    
+    workers : Union[int, NoneType]
+        Specifies the number of worker threads. By default set to the number of OpenMP threads, if given.
+        See also the documentation of the corresponding scipy.fft.idst parameter.
+    
+    ttype : int
+        The DCT type to use. (the name of this parameter in the underlying method is actually `type`).
+    """
+    def __init__(self, space, norm=None, workers=os.environ.get('OMP_NUM_THREADS', None), ttype=2):
+        workers = int(workers) if workers is not None else None
+        super().__init__(space, lambda out: scifft.idst(
+                out, axis=1, overwrite_x=True, workers=workers, norm=norm, type=ttype)) 

feectools/linalg/kernels/axpy_kernels.py

@@ -0,0 +1,57 @@
+from typing import TypeVar
+
+T = TypeVar('T', float, complex)
+
+#========================================================================================================
+def axpy_1d(alpha: 'T', x: 'T[:]', y: 'T[:]'):
+    """
+    Kernel for computing y = alpha * x + y.
+
+    Parameters
+    ----------
+    alpha : float | complex
+        Scaling coefficient.
+
+    x, y : 1D Numpy arrays of (float | complex) data
+        Data of the vectors.
+    """
+    n1, = x.shape
+    for i1 in range(n1):
+        y[i1] += alpha * x[i1]
+
+#========================================================================================================
+def axpy_2d(alpha: 'T', x: 'T[:,:]', y: 'T[:,:]'):
+    """
+    Kernel for computing y = alpha * x + y.
+
+    Parameters
+    ----------
+    alpha : float | complex
+        Scaling coefficient.
+
+    x, y : 2D Numpy arrays of (float | complex) data
+        Data of the vectors.
+    """
+    n1, n2 = x.shape
+    for i1 in range(n1):
+        for i2 in range(n2):
+            y[i1, i2] += alpha * x[i1, i2]
+
+#========================================================================================================
+def axpy_3d(alpha: 'T', x: 'T[:,:,:]', y: 'T[:,:,:]'):
+    """
+    Kernel for computing y = alpha * x + y.
+
+    Parameters
+    ----------
+    alpha : float | complex
+        Scaling coefficient.
+
+    x, y : 3D Numpy arrays of (float | complex) data
+        Data of the vectors.
+    """
+    n1, n2, n3 = x.shape
+    for i1 in range(n1):
+        for i2 in range(n2):
+            for i3 in range(n3):
+                y[i1, i2, i3] += alpha * x[i1, i2, i3]

feectools/linalg/kernels/inner_kernels.py

@@ -0,0 +1,100 @@
+
+#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!#
+#!!!!!!!!!!!!!!!!!!! WARNING !!!!!!!!!!!!!!!!!!!#
+#!!!!!!! Conjugate on the first argument !!!!!!!#
+#!!!!!!!!!! This will need an update !!!!!!!!!!!#
+#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!#
+
+from typing import TypeVar
+
+T = TypeVar('T', float, complex)
+
+#==============================================================================
+def inner_1d(v1: 'T[:]', v2: 'T[:]', nghost0: 'int64'):
+    """
+    Kernel for computing the inner product (case of two 1D vectors).
+
+    Parameters
+    ----------
+    v1, v2 : 1D NumPy array
+        Data of the vectors from which we are computing the inner product.
+
+    nghost0 : int
+        Number of ghost cells of the arrays along the index 0.
+
+    Returns
+    -------
+    res : scalar
+        Scalar (real or complex) containing the result of the inner product.
+    """
+    shape0, = v1.shape
+
+    res = v1[0] - v1[0]
+    for i0 in range(nghost0, shape0 - nghost0):
+        res += v1[i0].conjugate() * v2[i0]
+
+    return res
+
+#==============================================================================
+def inner_2d(v1: 'T[:,:]', v2: 'T[:,:]', nghost0: 'int64', nghost1: 'int64'):
+    """
+    Kernel for computing the inner product (case of two 2D vectors).
+
+    Parameters
+    ----------
+    v1, v2 : 2D NumPy array
+        Data of the vectors from which we are computing the inner product.
+
+    nghost0 : int
+        Number of ghost cells of the arrays along the index 0.
+
+    nghost1 : int
+        Number of ghost cells of the arrays along the index 1.
+
+    Returns
+    -------
+    res : scalar
+        Scalar (real or complex) containing the result of the inner product.
+    """
+    shape0, shape1 = v1.shape
+
+    res = v1[0, 0] - v1[0, 0]
+    for i0 in range(nghost0, shape0 - nghost0):
+        for i1 in range(nghost1, shape1 - nghost1):
+            res += v1[i0, i1].conjugate() * v2[i0, i1]
+
+    return res
+
+#==============================================================================
+def inner_3d(v1: 'T[:,:,:]', v2: 'T[:,:,:]', nghost0: 'int64', nghost1: 'int64', nghost2: 'int64'):
+    """
+    Kernel for computing the inner product (case of two 3D vectors).
+
+    Parameters
+    ----------
+    v1, v2 : 3D NumPy array
+        Data of the vectors from which we are computing the inner product.
+
+    nghost0 : int
+        Number of ghost cells of the arrays along the index 0.
+
+    nghost1 : int
+        Number of ghost cells of the arrays along the index 1.
+
+    nghost2 : int
+        Number of ghost cells of the arrays along the index 2.
+
+    Returns
+    -------
+    res : scalar
+        Scalar (real or complex) containing the result of the inner product.
+    """
+    shape0, shape1, shape2 = v1.shape
+
+    res = v1[0, 0, 0] - v1[0, 0, 0]
+    for i0 in range(nghost0, shape0 - nghost0):
+        for i1 in range(nghost1, shape1 - nghost1):
+            for i2 in range(nghost2, shape2 - nghost2):
+                res += v1[i0, i1, i2].conjugate() * v2[i0, i1, i2]
+
+    return res