def findMode_Mult_basic(Nvec, Wmat):
  ''' Find modes to multinomial with given parameters
        vectorized to solve many problems simultaneously

      Args
      ------
      Nvec : 1D array, size P
             contains non-negative integers
      Wmat : 2D array, size P x K
             each row is valid probability vector of length K
             Wmat[p] has non-negative entries, sums to one

      Returns
      -------
      Nmat : 2D sparse array, size P x K
          Nmat[p,:] = argmax_{n} \log p_{mult}( Nmat[p] | Nvec[p], Wmat[p] ) 
  '''
  # Unpack and parse input
  Nvec = np.asarray(Nvec, dtype=np.int64)
  if Nvec.ndim < 1:
    Nvec = Nvec[np.newaxis]
  Wmat = np.asarray(Wmat, dtype=np.float64)
  if Wmat.ndim < 2:
    Wmat = Wmat[np.newaxis, :]

  # Allocate and create initial guess, 
  #  which is guaranteed to have rows that sum to less than Nvec
  Nmat = np.zeros(Wmat.shape, dtype=np.int64)
  np.floor(Nvec[:,np.newaxis] * Wmat, out=Nmat)

  # Run Alg. 1 from paper
  f = Nvec[:,np.newaxis] * Wmat - Nmat
  q = (1 - f)/ Wmat
  activeRows = np.flatnonzero(Nmat.sum(axis=1) < Nvec)
  while len(activeRows) > 0:
    minIDs = np.argmin(q[activeRows], axis=1)
    ids = np.ravel_multi_index([activeRows, minIDs], Nmat.shape)
    Nmat.ravel()[ids] += 1

    stillActiveMask = Nmat[activeRows].sum(axis=1) < Nvec[activeRows]
    activeRows = activeRows[stillActiveMask]    
    ids = ids[stillActiveMask]    
    q.ravel()[ids] += 1.0 / Wmat.ravel()[ids]
  # Double-check Nmat satisfies constraints          
  assert np.all(np.abs(Nmat.sum(axis=1) - Nvec) < 0.0001)
  return Nmat
dataset/ETHPy150Open daeilkim/refinery/refinery/bnpy/bnpy-dev/bnpy/util/NumericHardUtil.py/findMode_Mult_basic/original