Улучшение производительности python-матрицы

Question

У меня есть большая матрица (форма: 2e6, 6), содержащая геофизические данные. У меня есть 3 for циклы, прежде чем искать значения в матрице для назначения переменных.

Мое первое решение есть с np.where. Это слишком медленно! Я прочитал it would be better, чтобы использовать другой цикл for для повышения производительности. Однако код, который я придумал, даже немного медленнее.

Есть ли у кого-нибудь идеи, как улучшить производительность, пожалуйста?

Первое решение (np.where)

for lat in LATS: 
    for lon in LONS: 
     for depth in range(1,401,1): 

      node_point_line = matrix[np.where((matrix[:,0]==lat) * (matrix[:,1]==lon) * (matrix[:,2]==depth))][0] 

      var1 = node_point_line[3] 
      var2 = node_point_line[4] 
      var3 = node_point_line[5] 
      ... 

Второе решение (дополнительный for петля)

for lat in LATS: 
    for lon in LONS: 
     for depth in range(1,401,1): 

      matrix_flat = matrix.flatten() 
      for i in range(len(matrix_flat)): 
       if matrix_flat[i]==lat and matrix_flat[i+1]==lon and matrix_flat[i+2]==depth: 
        var1 = matrix_flat[i+3] 
        var2 = matrix_flat[i+4] 
        var3 = matrix_flat[i+5] 
        ... 

Опять же, оба решения слишком медленно. Я избегаю Fortran или C++ (я знаю, что это быстрее). Какие-либо предложения?

Answer 1

Насколько велики ваши LATS и LONS массивы? Даже если они являются одним элементом длинным, тогда ваша программа должна выполнить приблизительно 1*1*400*1e6*3 ~ 1.2e9 операций, что слишком много. Даже если он правильно реализован на C++, может потребоваться второй или более.

Я думаю, вы должны оптимизировать свой алгоритм. Не совсем понятно, что вы пытаетесь сделать, но я предполагаю, что вы пытаетесь найти любое место в матрице с широтой в вашем списке LATS, долготе в вашем списке LONS и глубине от 1 до 400. Места обмена циклов: бегите вдоль матрицы во внешнем контуре, а затем проверьте, что все три условия выполнены. Впоследствии вы сможете заменить список широт на set широт, а поиск set намного быстрее, чем поиск по списку.

Answer 2

Учитывая цифры, он должен быть достаточно эффективным только для lexsort вашего массива данных. Это 2 секунды вложили в начало удаление всех операторов where.

фрагмент кода предполагает, что LATS и LONS сортируют:

order = np.lexsort([matrix.T][2::-1] # takes < 2 sec 
sorted = matrix[order, :] 

data_per_lat = np.split(sorted, sorted[:, 0].searchsorted(LATS[1:], axis=0) 
for lat, dpla in zip(LATS, data_per_lat): 
    data_per_lon = np.split(dpla, dpla[:, 1].searchsorted[LONS[1:], axis=0) 
    for lon, dplo in zip(LONS, data_per_lon): 
     depth, var1, var2, var3 = dplo.T[2:] 
     # the variables are now vectors, containing all data matching lon and lat sorted by depth 

Answer 3

Ну вы можете сжать исходную матрицу с векторизованного процесса, который сделает ваш цикл быстрее.

DPTH=np.arange(1,401,1) 
mask=np.in1d(matrix[:,0],LATS) * np.in1d(matrix[:,1],LONS) * np.in1d(matrix[:,2],DPTH) 
matrix_masked=matrix[mask] 

Тогда просто for перебирает matrix_masked либо nditer (сложного) или

for i in range(matrix_masked.shape[0]): 
    var1 = matrix_masked[i,3] 
    . . .