title: Efficiently applying transition to state matrix given transition matrix
body: <p>I'm looking to apply a state change to a large categorical matrix (M) with k categories where I know the transition probabilities of every category to each other category in k (T)</p>
  
  <p>Essentially I want to be able to efficiently take each element in M, simulate a state change given the probabilities in T, and replace the element with the calculated change. </p>
  
  <p>I've tried a few solutions :</p>
  
  <ul>
  <li>brute force nested for loops with indexing (much much too long)</li>
  <li>numba assisted nested for loops (~500ms which is too long for my purposes)</li>
  <li>precalculated draws for each category and replacement (~400ms)</li>
  </ul>
  
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  
  
  def categorical_transition(mat, t_mat, k=4):
  
      transformed_mat = mat.copy()
      cat_counts = np.bincount(mat.reshape(-1,))
  
      for i in range(k):
          rand_vec = np.random.multinomial(1, t_mat[i], cat_counts[i])
  
          choice = np.where(rand_vec)[1]
  
          transformed_mat[mat == i] = choice
  
      return transformed_mat
  
  
  # load data
  mat = np.random.choice(4, (16000, 256))
  t_mat = np.random.random((4, 4))
  
  # normalize transition matrix
  for i in range(t_mat.shape[0]):
      t_mat[i] = t_mat[i] / t_mat[i].sum()
  
  transformed_mat = categorical_transition(mat, t_mat)
  
  </code></pre>
  
  <p>This method works, but it is slow and I would appreciate any advice on more efficient ways to implement it</p>