Задача: выбрать из массива данных только данные удовлетворяющие условию и сделать с ними какую ни-будь гадость
Решение: numpy.where
Слухи о тормознутости Python сильно преувеличены, просто нужно уметь его готовить. Конечно скорости C или Fortran вы на нем не добьетесь, но и его вполне можно заставить быстро обрабатывать огромные массивы информации. Если вы хотите считать быстро, то ваш враг номер один в Python (также как и в MATLAB, кстати) это циклы, заданные в явном виде (оператор for). От большинства из них можно избавиться, применяя нехитрые приемы. Об одном таком приеме, позволившем увеличить скорость обработки массива размером более 100 гигабайт в 46 раз, я очень коротко расскажу в этом посте.
Дано: четерехмерный массив с полями температуры (x,y,глубина, время). Поскольку сетка у нас не регулярная, имеются два двумерных массива широт (lat) и долгот (lon).
Задача: вырезать из этого массива хитрую область, и получить среднее по этой области для каждой глубины, за каждый отсчет времени.
Область выглядит следующим образом:
Одним из возможных решений данной задачи является такой код:
for ttime in range(0,487):
for lev in range(0,51):
THO_new = []
for ff in range(285):
for zz in range(3602):
if 100 < lon[ff,zz] < 140:
if lat[ff,zz] > 66:
THO_new.append(ffile.variables["tho"][ttime,lev,ff,zz])
elif 100 > lon[ff,zz] > 0:
if lat[ff,zz] > 80:
THO_new.append(ffile.variables["tho"][ttime,lev,ff,zz])
elif 300 < lon[ff,zz] < 360:
if lat[ff,zz] > 80:
THO_new.append(ffile.variables["tho"][ttime,lev,ff,zz])
У нас имеется внешний цикл по времени (ttime), внутри него цикл по глубинам (lev), затем цикл по x (ff) и по y (zz). Вот в этом последнем цикле мы и проверяем каждую точку на соответствие нашим условиям по широте и долготе и если она им удовлетворяет, добавляем ее значение в массив (THO_new), с которым будем дальше совершать все необходимые непотребства.
Количество точек в одном шаге по времени, которые небходимо прогнать через циклы и обработать условными операторами составляет в данном случае 52355070 штук, что занимает 2 минуты 18 секунд. Шагов по времени у нас 487, так что ждать окончания работы программы мы будем дооолго.
Способ второй:
i1,j1 = numpy.where((100<lon)&(lon<140)&(lat>66))
i2,j2 = numpy.where((100>lon)&(lon>0)&(lat>80))
i3,j3 = numpy.where((300<lon)&(lon<360)&(lat>80))
for ttime in range(0,487):
for lev in range(0,51):
THO_new = []
a = ffile.variables["tho"][ttime,lev,:,:]
THO_new = numpy.concatenate((THO_new, a[i1,j1]), 0)
THO_new = numpy.concatenate((THO_new, a[i2,j2]), 0)
THO_new = numpy.concatenate((THO_new, a[i3,j3]), 0)
Понятное дело, что точки с необходимыми нам координатами будут занимать одинаковое положение во всех полях. Поэтому проще сразу определиться с индексками тех точек что нас интересуют. Для этого используем оператор numpy.where, которому в качестве условия передаем наши пожелания по ограничению региона. К сожалению, конструкции вида 100 <> он не воспринимает, так что небоходимо пользоваться логическими операторами (естественно можно использовать не только AND (&) но и, напримерр OR (|)). В качестве вывода, мы получаем индексы тех точек, что удовлетворяют нашим условиям.
Циклы по времени и глубинам мы оставляем, а вот перебирать каждую точку двумерного поля нам уже не нужно, мы точно знаем какие точки нам нужны. Сначала создаем переменную a, в которую заливаем двумерное поле (это сделано исключительно из-за того что 100 гигабайт в память не влезало :)). Затем добавляем в переменную THO_new те значения, что вписываются в заданную нами область. Дальше уже происходит наложение маски (пропущенных значений) и осреднение THO_new, что для нас в данном контексте не особо важно, поэтому эту часть я не привожу. Время работы нового кода для одного шага по времени составляет 3 секунды.
Блог python на хабрахабре
NumPy — это расширение языка Python, добавляющее поддержку больших многомерных массивов и матриц, вместе с большой библиотекой высокоуровневых математических функций для операций с этими массивами.
