square_matrix.py

from flame_math.matrix import Matrix
from flame_math.vector import Vector
from flame_math import vector

from itertools import izip

class SquareMatrix(Matrix):
    """
    This matrix class represent a squared matrix 
    """
    def __init__(self, size=4, data=None,  identity = True):
        Matrix.__init__(self,size, size)
        """
        The constructor
        Args:

        :size: integer, the size of row and column matrix
        :data: Vector[], an array of vectors used to initialize the columns of the mat
        :identity: bool, wheter or not to initialize the matrix as idenitiy
        """

        if identity:
            self.set_identity()
        if not identity and data:
            self.data = data
    
    @classmethod
    def from_list(cls, size=4 , data=None):
        """
        Alternative constuctor used to build a square matrix from a list of elements
        Args:

        :size: integer, the size of the NxN matrix
        :data: float[], the data used for thea matrix columns
        """

        vectors = []
        for c in range(size):
            vec = Vector([ data[c*size + r] for r in range(size)])
            vectors.append(vec)
    
        inst = cls(size,vectors,False)
        return inst
    
    @classmethod
    def createMatrix44(cls, data=None, identity = True):
        """
        Alternative constructor to crate a 4x4 matrix

        :data: Vector[], an array of vectors used to initialize the columns of the mat
        :identity: bool, wheter or not to initialize the matrix as idenitiy
        """ 
        
        return cls(4, data, identity)
    
    @classmethod
    def createMatrix33(cls, data=None, identity = True):
        """
        Alternative constructor to crate a 3x3 matrix

        :data: Vector[], an array of vectors used to initialize the columns of the mat
        :identity: bool, wheter or not to initialize the matrix as idenitiy
        """ 
        
        return cls(3,data,  identity)

    @classmethod
    def copy(cls, mat):
        """
        Alternative constructor to generate a matrix identical to the given one
        """
        return cls(mat.columns(),[Vector.copy(v) for v in mat.column_iterator()],False)
    
    def set_identity(self):
        """
        This function reset the matrix to the identity matrix
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            TL.zero()
            value=1
            BL.zero()

            dt.merge(TL, value, BL)
    
    def set_to_diagonal(self, x):
        """
        Set the matrix to be a diagonal matrix and set the diagonal to 
        the provided values

        :x: floa[], the values of the diagonal
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            TL.zero()
            value= x[dt.counter-1] 
            BL.zero()

            dt.merge(TL, value, BL)

    def set_upper_triangular(self):
        """
        Set the given matrix to upper diagonal matrix
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            BL.zero()
            dt.merge(TL, value, BL)
    
    def set_lower_triangular(self):
        """
        Set the matrix to lower diagonal matrix
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            TL.zero()
            dt.merge(TL, value, BL)


    def set_strictly_upper_triangular(self):
        """
        Set to strictly upper diagonal matrix, means that
        also the diagonal is zeroed
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            BL.zero()
            value= 0
            dt.merge(TL, value, BL)
    
    def set_strictly_lower_triangular(self):
        """
        Set to strictly lower diagonal matrix, means that
        also the diagonal is zeroed
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            TL.zero()
            value=0
            dt.merge(TL, value, BL)

    def set_unit_upper_triangular(self):
        """
        Setting the matrix to unit upper diagonal meaning the
        diagonal will be ones
        """

        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            BL.zero()
            value= 1
            dt.merge(TL, value, BL)
    
    def set_unit_lower_triangular(self):
        """
        Setting the matrix to unit lower diagonal meaning the
        diagonal will be ones
        """
        dt = self.diagonal_iterator()
        for TL, value, BL in dt:
            TL.zero()
            value=1
            dt.merge(TL, value, BL)

    def set_transpose(self):
        """
        In place transposition of the matrix, and usues a full iterator
        explotiing the squareness to to the transpose
        """
        it = self.full_iterator()
        for data in it:
            temp = data["a01"]
            
            data["a01"] = data["at10"]
            data["at10"] = temp

            it.merge(data)

    def set_upper_simmetry(self):
        """
        Sets the upper matrix to be symmetrix to the lower
        part
        """
        it = self.full_iterator()
        for data in it:
            data["a01"] = data["at10"]
            it.merge(data)
    
    def set_lower_simmetry(self):
        """
        Sets the lower part of the matrix to be symmetrix to the
        uppper part of the matrix
        """
        it = self.full_iterator()
        for data in it:
            data["at10"] = data["a01"]
            it.merge(data)
    
    #####################
    ## DOT OPERATIONS
    ####################
    
    def mult_vec_upper_triangular_by_dot(self,vec):
        """
        This function is multiplying a vector by an upper triangular matrix
        using a parted matrix and the dot product approach, the optimmization comes
        from the fact we don't need to perform the computation of the row before the
        current diagonal index
        
        :vec: Vector, the value to multiply
        :returns: Vector
        """
        
        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = vec.vector_iterator()

        for i,data in enumerate(it):
            #we slice the vector calling manually the iterator next
            tl, value, bl = vit.next() 
            #we accumualte the result in the proper index of the result vector
            #the result is composed of the result of diagonal value being multiplied
            #plus the row after the diagonal 
            result[i] = ( data["a11"] * value +
                      bl.dot(data["at12"])) 

        return result 
 
    def mult_vec_lower_triangular_by_dot(self,vec):
        """
        This function is multiplying a vector by an upper triangular matrix
        using a parted matrix and the dot product approach, the optimmization comes
        from the fact we don't need to perform the computation of the row after the
        current diagonal index
        
        :vec: Vector, the value to multiply
        :returns: Vector
        """

        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = vec.vector_iterator()

        for i,data in enumerate(it):
            #we slice the vector calling manually the iterator next
            tl, value, bl = vit.next() 
            #we accumualte the result in the proper index of the result vector
            #the result is composed of the result of diagonal value being multiplied
            #plus the row before the diagonal 
            result[i] = ( data["a11"] * value +
                      tl.dot(data["at10"])) 

        return result 

    def mult_vec_upper_symmetric_by_dot(self,vec):
        """
        This function multiplies the given vector by an upper simmetric matrix, this will
        allow to perform less computation by exploiting the simmetry

        :vec: Vector, the vector to be transformed
        :returns: Vector
        """
        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = vec.vector_iterator()

        for i,data in enumerate(it):
            #we slice the vector calling manually the iterator next
            tl, value, bl = vit.next() 
            
            #to nothe here instead the first part of the row (before the diagonal
            #we use the column before the diagona, exploiting the simmetry
            result[i] = (data["a01"].dot(tl)  + 
                      data["a11"] * value +
                      data["at12"].dot(bl)) 
        return result 


    def mult_vec_lower_symmetric_by_dot(self,vec):
        """
        This funtion transforms a vector using a lower simmetric matrix, means 
        we don't neet to use  the upper traingular matrix which might old other
        data, we use dot product under the hood to perform the computation

        :vec: Vector, the vector to be transformed
        """
        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = vec.vector_iterator()

        for i,data in enumerate(it):
            #we slice the vector calling manually the iterator next
            tl, value, bl = vit.next() 
            
            #here instead of the upper column before the diagonal we use the 
            #row before the diagonal
            result[i] = (data["at10"].dot(tl)  + 
                      data["a11"] * value +
                      data["a21"].dot(bl)) 
        return result 

    
    
    
    #####################
    # AXPY OPERATIONS
    ####################
    def mult_vec_upper_triangular_by_axpy(self, vec):
        """
        This function implements an upper triangular matrix vector mult based on AXPY 

        :vec: Vector, the vector to work on
        :returns: Vector
        """
        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = result.vector_iterator()
        
        for i,data in enumerate(it):
            #slicing the vector
            tl, value, bl = vit.next() 
            
            #computing the diagonal value
            value = data["a11"] * vec[i] + value
            #computing the upper diagonal value, using an axpy
            #and accumulating in the up vector
            tl = data["a01"].axpy(tl,vec[i])    
            
            #merging the result back in
            vit.merge(tl,value,bl)
        return result 


    def mult_vec_lower_triangular_by_axpy(self, vec):
        """
        This function implements a lower triangular matrix vector mult based on AXPY 

        :vec: Vector, the vector to work on
        :returns: Vector
        """


        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = result.vector_iterator()
        
        for i,data in enumerate(it):
            #slicing the vector
            tl, value, bl = vit.next() 
            
            #computing the diagonal value
            value = data["a11"] * vec[i] + value
            #computing the upper diagonal value, using an axpy
            #and accumulating in the up vector
            bl = data["a21"].axpy(bl,vec[i])    
            
            #merging the result back in
            vit.merge(tl,value,bl)
        return result 


    def mult_vec_upper_symmetric_by_axpy(self, vec):
        """
        This function implements a matrix vector mult based on AXPY using an upper simmetrix matric 

        :vec: Vector, the vector to work on
        :returns: Vector
        """

        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = result.vector_iterator()
        
        for i,data in enumerate(it):
            #slicing the vector
            tl, value, bl = vit.next() 
            
            #computing the diagonal value
            value = data["a11"] * vec[i] + value
            
            #we perform two different axpy and we make up for the missing part
            #by using the equivalent transposed of the column below the diagonal
            #aka the row after the diagonal
            tl = data["a01"].axpy(tl,vec[i])    
            bl = data["at12"].axpy(bl,vec[i])    
            
            #merging the result back in
            vit.merge(tl,value,bl)
        
        return result 

    def mult_vec_lower_symmetric_by_axpy(self, vec):
        """
        This function implements a matrix vector mult based on AXPY using an upper simmetrix matric 

        :vec: Vector, the vector to work on
        :returns: Vector
        """

        result = Vector([0]* self.rows())
        it = self.full_iterator()
        vit = result.vector_iterator()
        
        for i,data in enumerate(it):
            #slicing the vector
            tl, value, bl = vit.next() 
            
            #computing the diagonal value
            value = data["a11"] * vec[i] + value
            
            tl = data["at10"].axpy(tl,vec[i])    
            bl = data["a21"].axpy(bl,vec[i])    
            
            #merging the result back in
            vit.merge(tl,value,bl)
        
        return result 


    def gaussian_reduction_upper_triangular(self, store_coeff=False, vec=None):
        """
        This function used gaussian reduction to reduce the linear system represented
        by this linear matrix to an upper triangular matrix
        
        :store_coeff: boolean, instead of zeroing the lower part here we will store the coefficents
                      of reduction
        :vec: Vector, this is the appended column of the system = the column after the equals in the 
                        equations system

        :returns: Matrix
        """
        mat = SquareMatrix.copy(self);
        it = mat.full_iterator()

        if vec:
            vit = vec.vector_iterator()

        for data in it:
            
            if data["a21"].size() == 0:
                break 
            coeff = (1.0 / data["a11"]) * data["a21"]
            data["A22"].set_add(-(coeff.outer_product(data["at12"])))
            if vec:
                TL, value, BL = vit.next()
                
                #here we could use += but is not implemented vector side yet
                BL =  BL + (-(value*coeff)) 
                vit.merge(TL,value,BL)
            
            if store_coeff:
                data["a21"] = coeff
            else:
                data["a21"].zero()

        
            
            it.merge(data)

        return mat


    def gaussian_appended_reduction_upper_triangular(self, vec):
        """
        This function uses a gaussian matrix that has been previously reduced 
        and has the reduction coefficent stored in the lower triangular matrix
        otherwise you will get crap

        :vec: Vector, the vector representing the appended part of the system
        """

        it = self.full_iterator()
        cpy = vector.Vector.copy(vec)
        vit = cpy.vector_iterator() 
        
        for data in it:
            #we need to break the iteration in step earlier 
            #because make no sense to work with a column below
            #the diagonal martix of size 0
            if data["a21"].size() ==0:
                break
            TL, value, BL = vit.next()
            BL = BL + (-(value*data["a21"]))
            vit.merge(TL,value,BL)

        return cpy 

    def lu_factorization_compact(self):
        """
        This function retunrs the LU factorization for the matrix,
        the result is stored in a single matrix, means that the diagonal
        for the lower unit matrix wont be written but we know it s unit lower
        diagonal

        :returns: SquareMatrix
        """

        mat = SquareMatrix.copy(self)
        it = mat.full_iterator()

        for data in it:
            
            #we need to break the iteration in step earlier 
            #because make no sense to work with a column below
            #the diagonal martix of size 0
            if data["a21"].size() ==0:
                break
            #we accumulate directly on the matrix itself
            data["a21"] =   (1.0/ data["a11"]) * data["a21"]
            data["A22"].set_add(-( data["a21"].outer_product(data["at12"])))

            it.merge(data)

        return mat

    def lu_factorization(self):
        """
        This function retunrs the LU factorization for the matrix,

        :returns: SquareMatrix L, SquareMatrix U
        """

        it = self.full_iterator()
        L = SquareMatrix.copy(self)
        U = SquareMatrix.copy(self)
        
        L.set_identity()
        #U.set_upper_triangular()
        
        lit = L.full_iterator()
        uit = U.full_iterator()

        for data , ldata, udata in izip(it, lit, uit):
            
            if data["a21"].size() ==0:
                break
            ldata["a21"] =   (1.0/ udata["a11"]) * udata["a21"]
            udata["A22"].set_add(-( ldata["a21"].outer_product(udata["at12"])))
            udata["a21"].zero()

            lit.merge(ldata)
            uit.merge(udata)

        return L, U

    def solve_lower_triangular_system(self, b):
        """
        This function is going to solve for 
        L*z = b 
        where we solve for z, b is given and L is a lower unit triangular matrix
        (most likely the result of a LU operation)

        :b: Vector, the vector needed to solve the system
        :returns: Vector
        """

        res = Vector.copy(b)
        bit = res.vector_iterator()
        
        it = self.full_iterator()

        for TL,value,BL in bit:
            data = it.next()

            BL = BL + (-(value * data["a21"]))
            bit.merge(TL,value,BL)

        return res


    def solve_upper_triangular_system(self,b):
        """
        This function is the complementary one for solving a liner equation system
        based on LU,
        with solve_lower_triangular_system we solved L* z = b
        where now we can solve U* x = b where z = U*x and U is the upper triangular 
        matrix from the factoriazition, basically we are performing a backward sostiuition 

        :b: Vector, the vector needed to solve the system
        :returns: Vector
        """ 
        res = Vector.copy(b)
        bit = res.vector_iterator(True)
        it = self.full_iterator(True)

        for TL,value,BL in bit:
            data = it.next()
            
            value -= data["at12"].dot(BL)
            value /= data["a11"]

            bit.merge(TL,value,BL)

        return res

    def solve_linear_system(self,b):
        """
        This function solves a linear system composed by the equation at 
        left of the equal signe composed by the matrix (self), and the right
        part (after the equal) composed by the vector b
        the function will fist compute the L, U factorization then solve L(Ux) =b

        :b: Vector
        :returns: Vector
        """
        L,U = self.lu_factorization()

        res = L.solve_lower_triangular_system(b)
        res = U.solve_upper_triangular_system(res)

        return res