def findMode_Mult_fastest_skipsingles(Nvec, Wmat):
  ''' Find modes to multinomial with given parameters
        vectorized to solve many problems simultaneously
        tries to handle "single" case (Nvec=1) very fast

      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, :]
  
  Nmat = toHardAssignmentMatrix(Wmat, dtype=np.int64)

  bigRows = np.flatnonzero(Nvec > 1)
  Nmatorig = np.take(Wmat, bigRows, axis=0)
  Nmatorig *= Nvec[bigRows][:, np.newaxis]
  np.floor(Nmatorig, out=np.take(Nmat, bigRows, axis=0))

  # Run Alg. 1 from paper
  activeErrors = Nvec - Nmat.sum(axis=1)
  activeMask = activeErrors > 0
  activeRows = np.flatnonzero(activeErrors)

  if len(activeRows) > 0:
    q = 1.0 - np.take(Nmatorig, np.flatnonzero(activeMask[bigRows]), axis=0) \
            + np.take(Nmat, activeRows, axis=0)
    q /= np.take(Wmat, activeRows, axis=0)
    activeErrors = np.take(activeErrors, activeRows)

  while len(activeRows) > 0:
    minIDs = np.argmin(q, axis=1)
    ids = np.ravel_multi_index([activeRows, minIDs], Nmat.shape)
    Nmat.ravel()[ids] += 1
    activeErrors -= 1

    stillActiveMask = activeErrors > 0
    activeRows = activeRows[stillActiveMask]   

    # Prepare for next round...
    if len(activeRows) > 0:
      activeErrors = activeErrors[stillActiveMask]
      ids = ids[stillActiveMask]
      minIDs = minIDs[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] += 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_fastest_skipsingles/original