Блог

Я работаю над проектом, в котором в какой-то момент мне требуется высокая оптимизация алгоритма, используемого в вычислениях. Я хотел бы знать, каким путем лучше пойти, если это можно эффективно сделать на Python, или, с другой стороны, я должен попытаться реализовать процедуру C , которая делает именно это; вот пример кода узкого места, с которым я сталкиваюсь:

 import numpy as np
from numba import jit

from time import time

# this generates the outer product by rows of arrays 1 and 2
@jit( nopython=True  )
def arrays_product( array_1, array_2 ):
    
    dim_x       = array_1.shape[0]
    dim_y       = array_2.shape[0]
    total_dim   = dim_x*dim_y
    total_depth = array_1.shape[1]
    
    result      = np.empty( ( total_dim, total_depth ), dtype =np.complex128  ) 
    
    counter = 0
    for i in range( dim_x ):
        for j in range( dim_y ):            
            result[ counter, : ] = np.multiply( array_1[ i, : ], array_2[ j, : ] )
            
            counter  = 1
            
    return result


def create_complex_array( dim ):    
    return np.random.rand( dim  )   np.random.rand( dim )*1j


prefactor        = np.pi
total_iterations = 40
tot_points       = 3000
dim_0            = 4000
dim              = 100   # the outer product dimension will be dimm**2
total_combs      = 10000

#create some dictionary keys, that are composed of 12 integers each. To each key, a complex array is associated

combination_keys     = np.random.randint( low=-128, high = 127, size=(total_combs, 12), dtype=np.int8 )
combs_dict           = { tuple( x ): create_complex_array( tot_points )  for x in combination_keys }

comb_keys_array      = np.array( list( combs_dict.keys()  ) ) # put keys in array form

#source array used to make the calculations
the_source_array     = np.random.rand( dim_0, tot_points  )   np.random.rand( dim_0, tot_points )*1j


#create an array of masks, from which different keys of the dictionary will be selected later
comb_masks_p = []


for k in range( total_iterations ):
    
    random_choice  = np.random.randint(  low=0, high = total_combs - 1, size = dim**2   )  
    comb_masks_p.append( random_choice  )
    
comb_masks  = np.array( comb_masks_p )

#this will select two different sets of elements of the_source_array
array_masks = np.random.randint( low=0,    high=dim_0 - 1, size=( total_iterations, 2, dim  )  )


for k in range( total_iterations ):
    
    ts = time()
    
    selection_combs = comb_keys_array[  comb_masks[ k, : ]  ] 
    keys_in_tuples  = list( map( tuple, selection_combs ) ) # convert selected keys for the dictionary in list of tuples
    
    # we iterate over the subset of selected keys and generate the corresponding complex array
    # this will be slow when the total number of keys is high
    array_1         = np.array( [ combs_dict[ key ] for key in keys_in_tuples ] )    
    
    # select left and right components to generate outer product by rows
    
    mask_left       = array_masks[ k, 0, : ]
    mask_right      = array_masks[ k, 1, : ]
    array_left      = the_source_array[  mask_left, : ] 
    array_right     = the_source_array[  mask_right, : ]      
    
    # this part is also slow 
    product_array   = arrays_product( array_left, array_right )
    
    # finally, use array_1 to make a product element-wise and reduce on the first axis
    
    result          = ( prefactor*array_1*product_array ).sum( axis=0 )

    print("Iteration time: ",  time() - ts )
 

ОПИСАНИЕ: У нас есть словарь, хранящий 1-мерные сложные массивы, которые необходимо выбрать на данной итерации, чтобы создать из них 2D-массив. Способ выбора массивов осуществляется с помощью уникальных ключей n-кортежей, получающих доступ к элементам словаря. Созданный массив NumPy из словаря всегда имеет одинаковые размеры ( dim ** 2 , tot_points ).

Кроме того, у нас есть «source_array» (2D), из которого значения будут выбираться по строкам. Способ выбора подходящих значений на данной итерации заключается в использовании причудливой индексации или масок. Этот выбор генерирует два разных массива ( array_left, array_right ), которые позже используются для создания внешнего продукта в смысле строк (см. Функцию «arrays_product» с декоратором Numba ). Как только создается этот внешний «product_array», который имеет размерность (dim ** 2, tot_points ), «then array_1» поэлементно умножается на «product_array», и сумма элементов переносится по оси 0.

Обратите внимание, что во всех случаях критической точкой является увеличение параметра «dim», поскольку это делает массивы из внешнего продукта очень большими, а также общее количество ключей, которые необходимо выбрать из словаря для построения «array_1».

ЧТО ЛУЧШЕ?Я считаю, что приведенный выше код крайне неэффективен для больших наборов данных. Я ищу способ оптимизировать его либо на чистом Python, либо путем создания подпрограммы C , которая будет вызываться для выполнения операций. Вопросы заключаются в следующем:

1. Могу ли я ожидать большой разницы при использовании подхода C ? Мне нужно передать словарь так же, как и все другие объекты, и я не знаю, насколько большой выгоды я могу ожидать здесь. Я также открыт для использования любых других подходов, таких как Cython, но я думаю, что там объект dictionary является большой проблемой, наряду с некоторыми другими вызовами Python, которые могут возникнуть во время одной итерации.
2. Используя чистый Python и, в частности, модуль numpy, есть ли способ выполнить сокращения (последние два шага) за один вызов numpy einsum или аналогичных? Я думаю, что это должно быть что-то прямое из-за типа задействованных операций, но я не могу что-то найти.
3. Во внешнем цикле нет обходного пути. Внешний цикл необходим для существования, потому что в моем основном коде есть дополнительный шаг, не отраженный здесь (что не важно для проблем с производительностью).

Я надеюсь, что кто-то может помочь, спасибо!