Склеить и разделить

Данные часто приходят по частям: продажи за январь и за февраль в разных файлах, каждая метеостанция в своём. Чтобы считать по ним вместе, части нужно склеить в один массив. Бывает и наоборот: месяц нужно разрезать на недели.

Продолжить массив: np.concatenate

np.concatenate([a, b, …]) ставит массивы друг за другом. Массивы передают списком — в квадратных скобках:

import numpy as np

a = np.array([1, 2, 3])
b = np.array([10, 20])
print(np.concatenate([a, b]))
print(np.concatenate([b, a, b]))

[ 1  2  3 10 20]
[10 20  1  2  3 10 20]

jan = np.array([410, 385, 402, 395, 430])
feb = np.array([398, 441, 420, 415])
all_sales = ...
total = ...

jan = np.array([410, 385, 402, 395, 430])
feb = np.array([398, 441, 420, 415])
all_sales = np.concatenate([jan, feb])
total = all_sales.sum()

У двумерных массивов concatenate склеивает по оси axis= — как агрегаты: axis=0 (по умолчанию) ставит таблицы друг под другом, добавляя строки, axis=1 — рядом, добавляя столбцы. Размер по другой оси должен совпадать:

t1 = np.array([[1, 2], [3, 4]])
t2 = np.array([[5, 6], [7, 8]])
print(np.concatenate([t1, t2]))            # axis=0: строки добавились
print(np.concatenate([t1, t2], axis=1))    # axis=1: столбцы добавились

[[1 2]
 [3 4]
 [5 6]
 [7 8]]
[[1 2 5 6]
 [3 4 7 8]]

Новая ось: np.stack

Три метеостанции записали по 28 дней — три одномерных массива формы (28,). Хочется таблицу 3 × 28: строка — станция. concatenate тут не поможет: он продолжит массив до 84 чисел. Нужна новая ось — её добавляет np.stack:

center = np.genfromtxt("data/station_center.csv", delimiter=",", skip_header=1)
airport = np.genfromtxt("data/station_airport.csv", delimiter=",", skip_header=1)
forest = np.genfromtxt("data/station_forest.csv", delimiter=",", skip_header=1)
temps_long = np.concatenate([center[:, 1], airport[:, 1], forest[:, 1]])
temps = np.stack([center[:, 1], airport[:, 1], forest[:, 1]])
print(temps_long.shape)
print(temps.shape)
print(temps[:, :4])     # три станции, первые четыре дня

(84,)
(3, 28)
[[-5.3 -5.3 -5.7 -5.7]
 [-6.4 -6.5 -6.3 -8.1]
 [-6.7 -8.7 -9.2 -9.1]]

np.stack кладёт массивы одинаковой формы друг на друга: из нескольких (n,) получается (k, n), где k — сколько массивов. С axis=1 новой осью станут столбцы: (n, k).

hum = ...
hum_avg = ...

hum = np.stack([center[:, 2], airport[:, 2], forest[:, 2]])
hum_avg = np.nanmean(hum, axis=1)

Короче: np.vstack и np.hstack

Для двух частых случаев есть короткие имена. np.vstack (vertical — вертикально) ставит массивы друг под другом, np.hstack (horizontal — горизонтально) — рядом. Для таблиц это concatenate с axis=0 и axis=1, а одномерные массивы vstack делает строками — как stack:

x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
print(np.vstack([x, y]))
print(np.hstack([x, y]))

[[1 2 3]
 [4 5 6]]
[1 2 3 4 5 6]

Склеивать можно только подходящие куски: у vstack должна совпадать длина строк, у hstack — число строк.

np.vstack([np.array([1, 2, 3]), np.array([4, 5])])

ValueError: all the input array dimensions except for the concatenation axis must match exactly, but along dimension 1, the array at index 0 has size 3 and the array at index 1 has size 2

Частый случай — добавить к таблице столбец. Одномерный массив hstack к таблице не приклеит: у них разное число осей. Сначала превратите его в столбец (n, 1) — через np.newaxis из модуля о транслировании:

scores = np.array([[70, 80], [65, 90], [88, 72]])
bonus = np.array([5, 0, 3])
print(np.hstack([scores, bonus[:, np.newaxis]]))

[[70 80  5]
 [65 90  0]
 [88 72  3]]

order = np.array([[120, 3], [80, 5], [200, 1], [45, 12]])
table = ...

order = np.array([[120, 3], [80, 5], [200, 1], [45, 12]])
table = np.hstack([order, (order[:, 0] * order[:, 1])[:, np.newaxis]])

Разрезать: np.split

np.split(a, n) режет массив на n равных частей и возвращает список массивов. Части можно сразу разложить по переменным. Вместо числа частей можно передать список мест разреза — номера, с которых начинается следующая часть:

days = np.arange(1, 15)
first, second = np.split(days, 2)
print(first)
print(second)
print(np.split(days, [3, 10]))    # новые части начинаются с индексов 3 и 10

[1 2 3 4 5 6 7]
[ 8  9 10 11 12 13 14]
[array([1, 2, 3]), array([ 4,  5,  6,  7,  8,  9, 10]), array([11, 12, 13, 14])]

Список массивов Python печатает как [array([...]), array([...])] — array здесь просто подпись: каждый элемент списка — массив NumPy.

Если массив не делится на равные части, split не угадывает, куда деть остаток, а сообщает об ошибке:

np.split(days, 3)

ValueError: array split does not result in an equal division

Список частей одной длины можно снова собрать np.stack — получится таблица, где строка — часть:

center_temp = center[:, 1]
weeks = ...
week_avg = ...

center_temp = center[:, 1]
weeks = np.stack(np.split(center_temp, 4))
week_avg = np.nanmean(weeks, axis=1)

Ту же таблицу даёт center_temp.reshape(4, 7) из модуля об индексах. split удобнее, когда части разной длины (np.split(a, [7, 14])) или нужны отдельными переменными.

Итоги

• np.concatenate([a, b], axis=…) продолжает массив по существующей оси.
• np.stack([a, b]) добавляет новую ось: из k массивов (n,) — таблица (k, n).
• np.vstack — друг под другом, np.hstack — рядом; размеры по другой оси должны совпадать.
• Столбец к таблице — np.hstack([table, col[:, np.newaxis]]).
• np.split(a, n) — n равных частей, np.split(a, [i, j]) — разрез по местам; результат — список.