def findMode_Mult_massToAdd(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
  activeErrors = Nvec - Nmat.sum(axis=1)
  activeRows = np.flatnonzero(activeErrors)

  if len(activeRows) > 0:
    f = Nvec[:,np.newaxis] * Wmat - Nmat
    q = (1 - f)/ Wmat
    q = np.take(q, activeRows, axis=0)
    activeErrors = np.take(activeErrors, activeRows)    

  while len(activeRows) > 0:
    qcopy = q.copy()
    minIDs = np.argmin(q, axis=1)
    ids = np.ravel_multi_index([activeRows, minIDs], Nmat.shape)

    qids = np.ravel_multi_index([np.arange(len(minIDs)), minIDs], q.shape)
    qcopy.ravel()[qids] = np.inf
    rank2IDs = np.argmin(qcopy, axis=1)
    qids2 = np.ravel_multi_index([np.arange(len(minIDs)), rank2IDs], q.shape)

    gap = q.ravel()[qids2] - q.ravel()[qids]
    massToAdd = np.floor(gap * Wmat.ravel()[ids])
    massToAdd = np.maximum(1.0, np.minimum(massToAdd, activeErrors) )   

    activeErrors -= massToAdd
    Nmat.ravel()[ids] += massToAdd

    stillActiveMask = activeErrors > 0
    activeErrors = activeErrors[stillActiveMask]
    activeRows = activeRows[stillActiveMask]    
    ids = ids[stillActiveMask]
    minIDs = minIDs[stillActiveMask]
    massToAdd = massToAdd[stillActiveMask]

    q = np.take(q, np.flatnonzero(stillActiveMask), axis=0)
    qids = np.ravel_multi_index([np.arange(len(minIDs)), minIDs], q.shape)
    q.ravel()[qids] += massToAdd / Wmat.ravel()[ids]
  # Double-check Nmat satisfies constraints          
  assert np.all(np.abs(Nmat.sum(axis=1) - Nvec) < 0.0001)
  return Nmat

########################################################### colwisesumLogFact
###########################################################
dataset/ETHPy150Open daeilkim/refinery/refinery/bnpy/bnpy-dev/bnpy/util/NumericHardUtil.py/findMode_Mult_massToAdd/original