CODEWITHSUNDEEP
pythonnumpymatrix
user3440639
user3440639

Reputation: 185

Python build numpy matrix

I am given this matrix and am trying to write a function to build this matrix for any size of n. I am told the height of the matrix is n, but not sure the width.

enter image description here

Below is my code and output, is this correct? I am slightly confused by the notation of the matrix itself.

def buildMatrix(n, a):
    matrix = np.zeros([n, n], dtype=float)
    x_diag, y_diag = np.diag_indices_from(matrix)
    for (x,y) in zip(x_diag, y_diag):
        if x > (n / 2):
            matrix[x][y] = -2*a
        elif x == (n / 2):
            matrix[x][y] = -(1 + a)
        else:
            matrix[x][y] = -2

        if x != n - 1:
            matrix[x + 1][y] = a if x >= (n / 2) else 1
            matrix[x][y + 1] = a if x >= (n / 2) else 1

    return matrix

Output with buildMatrix(5, 2)

[[-2.  1.  0.  0.  0.]
 [ 1. -2.  1.  0.  0.]
 [ 0.  1. -3.  2.  0.]
 [ 0.  0.  2. -4.  2.]
 [ 0.  0.  0.  2. -4.]]

Can anyone help me out?

Upvotes: 2

Views: 90

Answers (1)

alexdor
alexdor

Reputation: 866

To answer your first question, the matrix has to have a width of n in order for the matrix-vector product to be compatible.

The picture of the matrix is ambiguous on where the switch from -2 to -(1-a) to -2a occurs. In your code, you check if x==n/2 to set the switch. This is fine in python2 but will cause problems in python3 since x/2 returns 2.5. Using safer x==n//2 since n//2 return an integer in python2 as well as python3.

For generality, I'm going to assume that the switch happens at row m. The matrix can be built easier using slicing and the np.diag command. 

def buildmat(n, m, a):
    diag = np.zeros(n)
    offdiag = np.zeros(n-1)

    offdiag[0:m] = 1
    offdiag[m:n-1] = a

    diag[0:m] = -2
    diag[m] = -(1+a)
    diag[m+1:n] = -2*a

    matrix =  np.diag(diag) + np.diag(offdiag, 1) + np.diag(offdiag, -1)
    return matrix

Running

buildmat(5, 2, 3)

produces

[[-2.  1.  0.  0.  0.]
 [ 1. -2.  1.  0.  0.]
 [ 0.  1. -3.  2.  0.]
 [ 0.  0.  2. -4.  2.]
 [ 0.  0.  0.  2. -4.]]

Upvotes: 4

Related Questions