Search code examples
pythonnumpysympy

copying only upper triangular values of sympy Matrix into array from numpy.triu()?


I have a square matrix A (could be any size) and I want to take the upper triangular part and place those values in an array without the values below the center diagonal (k=0).

A = sympy.Matrix([[ 4,  0,  3],
                  [ 2,  4, -2],
                  [-2, -3,  7]])

using A_upper = numpy.triu(A) gets me to

A_Upper = sympy.Matrix([[ 4,  0,  3],
                        [ 0,  4, -2],
                        [ 0,  0,  7]])

but from here how would I copy only the upper triangular elements into a simply array? Such as:

[4, 0, 3, 4, -2, 7]

I was going to just iterate though and copy all non-zero elements, however zero's in the upper triangular are allowed.


Solution

  • Give a numpy array, this is an easy operation using numpy.triu_indices(N, k=0) where N is the size of the square array:

    In [28]: B = np.array([[4, 0, 3], [2, 4, -2], [-2, -3, 7]])
    
    In [29]: B[np.triu_indices(B.shape[0])]
    Out[29]: array([ 4,  0,  3,  4, -2,  7])
    

    The B.shape[0] is just there in case you don't want to hard code the size of the array (3).


    Given a sympy Matrix, this isn't as easy, but close enough. Just convert to a numpy array and make sure you change the dtype from object. This should work well if your matrices are reasonably sized. If they get really big, you might want to rethink this.

    In [36]: A = sp.Matrix([[4, 0, 3], [2, 4, -2], [-2, -3, 7]])
    
    # you can change the dtype of the new array to match the first array
    # e.g., .astype(int), .astype(sp.Symbol)
    # or you can just leave the default (dtype=object)
    In [37]: C = np.array(A) #.astype(new_dtype) 
    
    In [38]: C[np.triu_indices(C.shape[0])]
    Out[38]: array([ 4,  0,  3,  4, -2,  7])
    

    To get them into just a plain list, do

    In [39]: C[np.triu_indices(C.shape[0])].tolist()
    Out[39]: [4, 0, 3, 4, -2, 7]