Как векторизовать цикл через матрицу NumPy

Question

Предположим, у меня есть матрица 100000 x 100

.

import numpy as np

mat = np.random.randint(2, size=(100000,100))

Я хочу пройти через эту матрицу, и если каждый row содержит полностью либо 1, либо 0, я хочу изменить переменную state на это значение. Если состояние не изменяется, я хочу установить row значение state. Начальное значение state равно 0.

Наивно в цикле for это можно сделать следующим образом

state = 0

for row in mat:
    if set(row) == {1}:
        state = 1
    elif set(row) == {0}:
        state = 0
    else:
        row[:] = state

Однако, когда размер матрицы увеличивается, это занимает непрактичное количество времени. Может ли кто-нибудь указать мне, как использовать numpy, чтобы векторизовать этот цикл и ускорить его?

Так для примера ввода

array([[0, 1, 0],
       [0, 0, 1],
       [1, 1, 1],
       [0, 0, 1],
       [0, 0, 1]])

Ожидаемый результат в этом случае будет

array([[0, 0, 0],
       [0, 0, 0],
       [1, 1, 1],
       [1, 1, 1],
       [1, 1, 1]])

Answer 1

Подход № 1: NumPy-векторизация

Вот векторизация -

def check_all(a, state): # a is input matrix/array
    # Get zeros and ones all masks
    zm = (a==0).all(1)
    om = (a==1).all(1)

    # "Attach" boundaries with False values at the start of these masks.
    # These will be used to detect rising edges (as indices) on these masks.
    zma = np.r_[False,zm]
    oma = np.r_[False,om]

    omi = np.flatnonzero(oma[:-1] < oma[1:])
    zmi = np.flatnonzero(zma[:-1] < zma[1:])

    # Group the indices and the signatures (values as 1s and -1s)
    ai = np.r_[omi,zmi]
    av = np.r_[np.ones(len(omi),dtype=int),-np.ones(len(zmi),dtype=int)]

    # Sort the grouped-indices, thus we would know the positions
    # of these group starts. Then index into the signatures/values
    # and indices with those, giving us the information on how these signatures
    # occur through the length of the input
    sidx = ai.argsort()
    val,aidx = av[sidx],ai[sidx]

    # The identical consecutive signatures are to be removed
    mask = np.r_[True,val[:-1]!=val[1:]]
    v,i = val[mask],aidx[mask]

    # Also, note that we are assigning all 1s as +1 signature and all 0s as -1
    # So, in case the starting signature is a 0, assign a value of 0
    if v[0]==-1:
        v[0] = 0

    # Initialize 1D o/p array, which stores the signatures as +1s and -1s.
    # The bigger level idea is that performing cumsum at the end would give us the
    # desired 1D output
    out1d = np.zeros(len(a),dtype=a.dtype)

    # Assign the values at i positions
    out1d[i] = v

    # Finally cumsum to get desired output
    out1dc = out1d.cumsum()

    # Correct the starting positions based on starting state value
    out1dc[:i[0]] = state

    # Convert to 2D view for mem. and perf. efficiency
    out = np.broadcast_to(out1dc[:,None],a.shape)
    return out

Подход № 2: На основе Numba

Вот еще одна основанная на нумбе память и, следовательно, перф.КПД -

@njit(parallel=True)
def func1(zm, om, out, start_state, cur_state):
    # This outputs 1D version of required output.

    # Start off with the starting given state
    newval = start_state

    # Loop through zipped zeros-all and ones-all masks and in essence do :
    # Switch between zeros and ones based on whether the other ones
    # are occuring through or not, prior to the current state
    for i,(z,o) in enumerate(zip(zm,om)):
        if z and cur_state:
            cur_state = ~cur_state
            newval = 0
        if o and ~cur_state:
            cur_state = ~cur_state
            newval = 1
        out[i] = newval
    return out

def check_all_numba(a, state):
    # Get zeros and ones all masks
    zm = (a==0).all(1)
    om = (a==1).all(1)

    # Decide the starting state
    cur_state = zm.argmax() < om.argmax()

    # Initialize 1D o/p array with given state values
    out1d = np.full(len(a), fill_value=state)
    func1(zm, om, out1d, state, cur_state)

    # Broadcast into the 2D view for memory and perf. efficiency
    return np.broadcast_to(out1d[:,None],a.shape)

Answer 2

Вы можете сделать это без каких-либо циклов, используя np.accumulate :

R = 5 # 100000
C = 3 # 100

mat   = np.random.randint(2, size=(R,C))
print(mat) # original matrix

state    = np.zeros((1,C))                        # or np.ones((1,C))
mat      = np.concatenate([state,mat])            # insert state row
zRows    = np.isin(np.sum(mat,1),[0,C])           # all zeroes or all ones
iRows    = np.arange(R+1) * zRows.astype(np.int)  # base indexes
mat      = mat[np.maximum.accumulate(iRows)][1:]  # indirection, remove state
print(mat) # modified

#original
[[0 0 1]
 [1 1 1]
 [1 0 1]
 [0 0 0]
 [1 0 1]]

# modified
[[0 0 0]
 [1 1 1]
 [1 1 1]
 [0 0 0]
 [0 0 0]]

Это работает путем подготовки массива косвенных ссылок для строк, которые необходимо изменить.Это делается из np.arange индексов строк, в которых мы устанавливаем в ноль индексы, которые будут нуждаться в замене.При накоплении максимального индекса каждая замененная строка будет отображаться в строку со всеми нулями или со всеми единицами перед ней.

Например:

  [ 0, 1, 2, 3, 4, 5 ] # row indexes
  [ 0, 1, 0, 0, 1, 0 ] # rows that are all zeroes or all ones (zRows)
  [ 0, 1, 0, 0, 4, 0 ] # multiplied (iRows)
  [ 0, 1, 1, 1, 4, 4 ] # np.maximum.accumulate

Это дает нам список индексов, где содержимое строкидолжно быть взято из.

Состояние представлено дополнительной строкой, вставленной в начало матрицы перед выполнением операции и удаленной впоследствии.

Это решение будет незначительно медленнее для очень маленьких матриц(5x3), но он может дать вам 20-кратное увеличение скорости для более крупных (100000x100: 0,7 секунды против 14 секунд).

Answer 3

Вот простой и быстрый метод NumPy:

import numpy as np

def pp():
    m,n = a.shape
    A = a.sum(axis=1)    
    A = np.where((A==0)|(A==n))[0]
    if not A.size:
        return np.ones_like(a) if state else np.zeros_like(a)
    st = np.concatenate([np.arange(A[0]!=0), A, [m]])
    v = a[st[:-1],0]
    if A[0]:
        v[0] = state
    return np.broadcast_to(v.repeat(st[1:]-st[:-1])[:,None],(m,n))

Я сделал несколько таймингов, используя этот

state=0
a = (np.random.random((100000,100))<np.random.random((100000,1))).astype(int)

простой тестовый пример:

0.8655898020006134   # me
4.089095343002555    # Alain T.
2.2958932030014694   # Divakar 1
2.2178015549980046   # & 2