HiveBrain v1.2.0
Get Started
← Back to all entries
snippetpythonMinor

Create matrices of a certain format

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
certainformatcreatematrices

Problem

I need to create matrices that follow a specific format: the matrix is a square NxN matrix and the off-diagonals all have the same value (param) and the diagonals have values such that the rows sum to zero (in other words, they are equal to 0 - param * (N-1). The output is a numpy array

Here is my code. Is there any way I can make this more readable/elegant? Speed is also important.

import numpy as np

def makeMatrix(param, N):
    ar = []
    for i in range(N):
        ar.append([ param for st in range(N) ])
        ar[i][i] = 0-(param*(N-1))

    return np.array(ar, dtype=np.double)

Solution

One of the main features of NumPy is that you can work with matrices without explicit looping. To get the diagonal, use np.identity(). To set everything else, use broadcasting.

def make_matrix(value, dim):
    return value - value * dim * np.identity(dim, dtype=np.double)


param and N don't mean anything to me. I suggest value and dim. The function name should preferably be make_matrix, according to PEP 8.

Code Snippets

def make_matrix(value, dim):
    return value - value * dim * np.identity(dim, dtype=np.double)

Context

StackExchange Code Review Q#106986, answer score: 5

Revisions (0)

No revisions yet.