Блог

Обновление: Еще немного упростило проблему: у меня есть 2d — матрица m и предикат p . Как извлечь последний элемент в каждой строке m , который удовлетворяет предикату, p без зацикливания на строках и столбцах?

 import numpy as np

m = np.array([[0.1, 0.2, 1.1, 1.2], [0.1, 1.5, 0.9, 1.6], [2, 3, 4, 5]])

idxs = np.full((3,), np.nan)
values = np.full((3,), np.nan)
for i in range(3):
  for j in reversed(range(4)):
    if 0 <= m[i, j] <= 1:
      idxs[i] = j
      values[i] = m[i, j]
      break

print(f'idxs: {idxs}')
print(f'values: {values}')
 

выходы

 idxs: [ 1. 2. nan]
values: [0.2 0.9 nan]
 

Оригинальный вопрос: У меня есть код GPU, выполняющий несколько оптимизаций вдоль многомерной поверхности. Это возвращает мне два массива Numpy:

• xs формы (optimization_steps, run_count, x_dimension) , содержащей траектории аргумента функции
• ys формы (optimization_steps, run_count) , содержащей траектории соответствующих значений функций

Мне нужно ограничиться xs некоторой областью и игнорировать точки за ее пределами; из каждой траектории выберите последнюю точку, которая xs все еще находится в пределах области; затем, наконец, получите максимальное y значение среди этих последних точек и соответствующее x . Я написал этот код, как упрощено ниже, но, надеюсь, должен быть какой-то способ его векторизации.

Есть какие-нибудь предложения?

 # Demo of non-vectorized approach
import numpy as np

linears = np.linspace(0, 10, num=11)
sins = np.sin(linears)
exps = np.exp(linears)
xs = np.stack([sins, 0.5 * exps, 0.5   0.4 * sins], axis=1)
ys = np.stack([linears] * 3, axis = 1)
print(f'xs: {xs}')
print(f'ys: {ys}')

last_x_within = np.full((3,), np.nan)
last_y_within = np.full((3,), np.nan)
for run_idx in range(3):
  for traj_idx in reversed(range(11)):
    x = xs[traj_idx, run_idx]
    if x >= 0 and x <= 1:
      last_x_within[run_idx] = x
      last_y_within[run_idx] = ys[traj_idx, run_idx]
      break
best_run = np.argmax(last_y_within)
best_x = last_x_within[best_run]

print(f'last_x_within: {last_x_within}')
print(f'last_y_within: {last_y_within}')
print(f'Best: x {best_x}, run {best_run}')