Как выполнить эту операцию копирования, используя numpy?

Question

Я работал над базовым c моделированием для "диффузии Монте-Карло", чтобы найти энергию основного состояния молекулы водорода. Есть критическая часть алгоритма, которая мучительно замедляет мой код, и я не уверен, как это исправить.

Это то, что делает код. У меня есть массив 6 на N numpy под названием х. Массив представляет собой N случайных «ходоков», которые выбирают 6-мерное фазовое пространство (два электрона на 3 измерения - это 6 измерений). Я предлагаю определенные случайные изменения каждому «ходячему», чтобы получить нового «ходунка», а затем, используя формулу, выкладываю число «м» для каждого нового ходунка.

Число m может быть либо 0,1,2, либо 3. Это то место, где вступает жесткая часть. Если m равно 0, то «ходок», которому он соответствует, удаляется из массива. Если m равен 1, то ходок остается. Если m равно 2, то Уокер остается, и я должен сделать новую копию Уокера в массиве. Если m равно 3, то Уокер остается, и я должен сделать две новые копии Уокера в массиве. После этого код повторяется; случайные изменения предложены моему массиву ходячих, et c.

Итак; Ниже приведен код, который замедляет алгоритм. Это код для последней части, где я должен go через мои буквы и определить, что делать с каждым «ходоком», и создать мой новый массив x для использования в следующей итерации алгоритма.

        j1 = 0
        n1 = len(x[0,:])
        x_n = np.ones((6,2*n1))
        for i in range(0,n1):
            if m[i] == 1:
                x_n[:,j1] = x[:,i]
                j1 = j1 + 1
            if m[i] == 2:
                x_n[:,j1] = x[:,i]
                x_n[:,j1+1] = x[:,i]
                j1 = j1 + 2
            if m[i] == 3:
                x_n[:,j1] = x[:,i]
                x_n[:,j1+1] = x[:,i]
                j1 = j1 + 3
        x = np.ones((6,j1))
        for j in range(0,j1):
            x[:,j] = x_n[:,j]

Мой вопрос заключается в следующем; Есть ли способ сделать то, что я делаю в этом коде, используя numpy сам? Numpy имеет тенденцию быть намного быстрее, чем для циклов в моем опыте. Используя numpy непосредственно в вариационных симуляциях Монте-Карло, я смог добиться 100-кратного улучшения во время выполнения. Если вы хотите, чтобы полный код запускал алгоритм, я могу опубликовать это; это просто довольно долго.

Answer 1

пусть M будет массивом N x 1 значений m для каждого случайного обходчика.

пусть X будет вашим исходным массивом данных 6 x N

# np.where returns a list of indices where the condition is satisfied 
zeros =  np.where(M == 0) # don't actually need this variable, I just did it for completeness
ones =   np.where(M == 1)
twos =   np.where(M == 2)
threes = np.where(M == 3)

# use the lists of indices to access the relevant portions of X
ones_array = X[:,ones]
twos_array = X[:,twos]
threes_array = X[:,threes]

# update X with one copy where m = 1, two copies where m = 2, three copies where m = 3
X = np.concatenate((ones_array,twos_array,twos_array,threes_array,threes_array,threes_array),axis = 1)

Это не сохраняет порядок ходунков, поэтому, если это важно, код будет немного отличаться.