Search code examples
pythonnumpyreadability

Using Numpy without getting extra dimensions in arrays


I'm getting started with Numpy, and I'm attempting to use some of it's functionality to fill in a 2 dimensional square matrix where all the rows sum to 1. This accomplishes what I've set out to do so far, however, there were several previous iterations of this code that were giving me variations of a list of 1 dimensional arrays. This solution doesn't feel very elegant- What would be the best practices to write this code legibly?

import numpy as np


    class Markov:

        def __init__(self, states=5):
            self.prob_dist = []
            print(self.prob_dist)
            for i in range(states):
                change_prob = [1 for x in range(states)]
                print(change_prob)
                self.prob_dist.append(list(np.random.dirichlet(change_prob, size=1)[0]))
            print(self.prob_dist)


def main():
    output = Markov()


if __name__ == '__main__':
    main()

Solution

  • First create a random matrix of the desired size. In your case 5x5 since you have 5 states:

    A = np.random.random((5, 5))
    

    Then normalize it so that each column sums to 1:

    B = A / A.sum(axis = 0)
    

    Then transpose it so that each row sums to 1:

    B = B.T