Можете ли вы написать numpy-нарезку в общем виде?

#python #numpy

Вопрос:

Я хочу сделать что-то вроде

 x[i, :, :] = (rhs[i, :, :]-diag[i] * x[i 1, :, :])/diag[i]

где x и rhs — это трехмерные массивы numpy размером (T, L, S). diag — это одномерный массив размером T.

Это будет транслироваться правильно.

Но теперь я хотел бы написать аналогичную функцию для работы с 2D-массивами или некоторым другим количеством измерений. Как я могу написать это в общем виде, чтобы оно работало с любым массивом, имеющим первое измерение размера T. Я не хочу дублировать код с другим количеством двоеточий, поскольку в функции много таких строк.

1. Вы тоже хотите выполнить векторизацию T ? То есть это происходит в for i in range(T) цикле?

2. Вы имели в виду, что количество осей не гарантировано равным 3? Может ли это быть 1, 2, 3, …, N?

Ответ №1:

 x[i] = (rhs[i] - diag[i] * x[i 1])/diag[i]

Эти двоеточия совершенно не нужны.

1. Потрясающе. Спасибо. Не думал, что это так просто.

2. Если вы заинтересованы в том, чтобы сделать что-то подобное, но в последнем измерении, вы можете использовать многоточие, например x[...,i 1] .