Однотипные срезы массивов в python

Известно, что в python есть такая замечательная вещь, как срезы массивов (списков, кортежей и т.д.). По сути, в качестве начала и конца среза передаются положения элементов в массиве. Например,

x = range(10)
print x
> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
print x[1:3]
> [1, 2]
print x[0:5]
> [0, 1, 2, 3, 4]

Видно, что по сути отсчёт ведётся не по элементам, а по промежуткам между элементами, при этом на выходе даются все элементы, попадающие в заданный интервал промежутков. Особый интерес представляют отрицательные индексы, являющиеся по сути обратным отсчётом промежутков с конца массива.

print x[1:-1]
> [1, 2, 3, 4, 5, 6, 7, 8]
print x[0:-2]
> [0, 1, 2, 3, 4, 5, 6, 7]

Однако, что делать, если, скажем, имеется набор экспериментальных данных, которые мы хотим обрезать, основываясь на значениях, нежели на положении элементов в массивах? Допустим, есть список x,

from math import *
import numpy as np
x = np.arange(0, 2*pi, 0.1)

из которого надо вырезать 1.2<x<2.7.

Традиционный способ

В принципе, это можно осуществить перебором по всем элементам списка:

begin,end = 0,0
x0, x1 = 1.2, 2.7
for i in range(len(x)):
    if begin == 0:
        if x[i] > x0:
            begin = i
        else: pass
    else: pass
    if int(end) == 0:
        if x[i] > x1:
            end = i
        else: pass
    else: pass
print begin, end
print x[begin:end]

В результате получим:

> 12 28
[ 1.2  1.3  1.4  1.5  1.6  1.7  1.8  1.9  2.   2.1  2.2  2.3  2.4  2.5  2.6 2.7]

В принципе. то, чего и ожидали. Теперь можно этот цикл for оформить внутри функции, которая будет принимать в качестве аргументов список и границы.

def conventional_cut(x, (x0,x1)):
    begin,end = 0,0
    for i in range(len(x)):
        if begin == 0:
            if x[i] > x0:
                begin = i
            else: pass
        else: pass
        if end == 0:
            if x[i] > x1:
                end = i
            else: pass
        else: pass
    return x[begin:end]
print conventional_cut(x, (0.9, 1.4))

> [ 1.   1.1  1.2  1.3]

Несколько массивов

С практической точки зрения удобно резать сразу два списка: агрумент и функцию. Допустим, отснимали большой спектр y от x, из которого интерес представляет только часть. Таким образом, необходимо откусить часть от x, а также соответствующую часть y.

Сделать это легко, слегка модернизировав функцию:

def cut_xy((x,y), (x0,x1)):
    begin,end = 0,0
    for i in range(len(x)):
        if begin == 0:
            if x[i] > x0:
                begin = i
            else: pass
        else: pass
        if end == 0:
            if x[i] > x1:
                end = i
            else: pass
        else: pass
    return x[begin:end], y[begin:end]

То есть мы будем использовать те же индексы для списка y, что и для x. Приведу пример, для большей наглядности, графический:

from math import *
import numpy as np
import pylab as pl

x = np.arange(0, 4*pi, 0.01)
y = np.sin(x)/x
y = y + 0.1*(np.random.rand(len(y))-0.5)
# представим, что x,y - экспериментальные данные

xnew, ynew = cut_xy((x,y), (0.75*pi, 1.5*pi))
pl.plot(x, y, 'r-')
pl.plot(xnew, ynew, 'b-')

pl.xlabel("x")
pl.ylabel("Signal")
pl.show()

В результате получим то, что и следовало ожидать.

Выборочное соответствующее вырезание двух массивов

Разумеется, никто не запрещает резать не два, а три и более массивов одновременно, немного видоизменив функцию cut_xy.

Использование numpy.nonzero

Традиционный метод работает, однако может показаться корявым. По-видимому, более элегантное решение может быть получено с использованием функции nonzero из библиотеки numpy.

Пример работы numpy.nonzero:

from math import *
import numpy as np
x = np.arange(0, 2*pi, 0.1)
print x[np.nonzero(x<2.7)]
print x[np.nonzero(x>1.4)]

> [ 0.   0.1  0.2  0.3  0.4  0.5  0.6  0.7  0.8  0.9  1.   1.1  1.2  1.3  1.4
  1.5  1.6  1.7  1.8  1.9  2.   2.1  2.2  2.3  2.4  2.5  2.6]
> [ 1.4  1.5  1.6  1.7  1.8  1.9  2.   2.1  2.2  2.3  2.4  2.5  2.6  2.7  2.8
  2.9  3.   3.1  3.2  3.3  3.4  3.5  3.6  3.7  3.8  3.9  4.   4.1  4.2  4.3
  4.4  4.5  4.6  4.7  4.8  4.9  5.   5.1  5.2  5.3  5.4  5.5  5.6  5.7  5.8
  5.9  6.   6.1  6.2]

В скобках np.nonzero() указывается булево выражение, при этом nonzero возвращает индексы ненулевых элементов. Выражение x<2.7 является булевым и имеет следующий вид:

[ True  True  True  True  True  True  True  True  True  True  True  True
  True  True  True  True  True  True  True  True  True  True  True  True
  True  True  True False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False]

То есть, если x[i]>2.7, то результат True, в противном случае False. Стало быть, передавая x<2.7 в качестве аргумента функции nonzero, и скармливая результат исходному массиву x в качестве индексов среза, получаем обрезанный x, где каждый элемент x не превышает 2.7.

Трудность заключается в том, что не удаётся в качестве аргумента nonzero затолкать сразу несколько условий. Таким образом, если надо откусить массив с двух краёв, то необходимо сначала откусить массив слева, а потом результат — справа:

X = x[np.nonzero(x>1.4)]
X = X[np.nonzero(X<2.7)]
print X
> [ 1.4  1.5  1.6  1.7  1.8  1.9  2.   2.1  2.2  2.3  2.4  2.5  2.6]

Замечу, что при использовании nonzero левая граница промежутка включается в искомый массив.

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

X = x[np.nonzero(x>1.4)]
X = X[np.nonzero(X<2.7)]
Y = y[np.nonzero(x>1.4)]
Y = Y[np.nonzero(X<2.7)]

print "X =", X
print "Y =", Y

На выходе:

X = [ 1.4  1.5  1.6  1.7  1.8  1.9  2.   2.1  2.2  2.3  2.4  2.5  2.6]
Y = [ 0.66167563  0.62029339  0.59024345  0.57468639  0.5094329   0.5281753
  0.4542492   0.39952115  0.37289637  0.36049412  0.266617    0.19706189
  0.247443  ]

Можно аналогичным образом затолкать это в функцию:

def cut((x,y),(x0,x1)):
    after = np.nonzero(x>x0)
    X = x[after]
    before = np.nonzero(X<x1)
    X = X[before]
    Y = y[after]
    Y = Y[before]
    return X,Y

которая принимает в качестве аргументов исходные (x,y) и интервалы x.

Графический пример:

from math import *
import numpy as np
import pylab as pl

x = np.arange(0, 2*pi, 0.1)
y = np.sin(x)/x
y = y + 0.1*(np.random.rand(len(y))-0.5)
# представим, что x,y - экспериментальные данные

pl.plot(x, y, 'r-')
X,Y = cut((x,y),(4,5))
pl.plot(X,Y,'g-')

pl.xlabel("x [channels]")
pl.ylabel("Signal [counts]")
pl.show()

Результат аналогичной процедуры при использовании numpy.nonzero

Копируете статью - поставьте ссылку на оригинал!