Python для сбора данных

Алла Тамбовцева, НИУ ВШЭ

Работа с датафреймами pandas: часть 3

Часть 3 включает:

  • добавление новых столбцов: продолжение;
  • метод .apply() и lambda-функции;
  • группировка и агрегирование;
  • результаты агрегирования и мультииндексы.

Добавление новых столбцов: продолжение

Продолжим работать с файлом firtree.csv с вымышленными результатами опроса посетителей елочного базара. Импортируем pandas и загрузим файл по ссылке:

In [1]:
import pandas as pd
tree = pd.read_csv("https://allatambov.github.io/pydj/seminars/firtree.csv")

Вспомним, как выглядит датафрейм tree:

In [2]:
tree.head() 
Out[2]:
Unnamed: 0 gender ftype height score expenses wish
0 1 female пихта Нобилис 190 3 1051 да
1 2 male пихта Нобилис 174 3 2378 нет
2 3 female сосна Крым 248 4 655 да
3 4 female сосна Крым 191 1 2934 да
4 5 female сосна Крым 147 3 1198 нет

Добавим в датафрейм бинарный столбец female, где 1 соответствует респондентам женского пола, а 0 – мужского. Сделать это можно разными способами. Мы пойдем по простому пути – создадим столбец из True и False, а потом превратим его в целочисленный (True превратятся в 1, а False – в 0).

Получить набор из True и False легко, достаточно сформулировать условие с помощью операторов к столбцу датафрейму (в этом массивы и последовательности pandas Series похожи):

In [3]:
tree["gender"] == "female"
Out[3]:
0        True
1       False
2        True
3        True
4        True
        ...  
1195    False
1196     True
1197    False
1198    False
1199    False
Name: gender, Length: 1200, dtype: bool

Теперь добавим полученный столбец в датафрейм и изменим его тип на integer с помощью метода .astype:

In [4]:
tree["female"] = tree["gender"] == "female"
tree["female"] = tree["female"].astype(int)
tree.head()
Out[4]:
Unnamed: 0 gender ftype height score expenses wish female
0 1 female пихта Нобилис 190 3 1051 да 1
1 2 male пихта Нобилис 174 3 2378 нет 0
2 3 female сосна Крым 248 4 655 да 1
3 4 female сосна Крым 191 1 2934 да 1
4 5 female сосна Крым 147 3 1198 нет 1

Ранее мы уже обсуждали, что новый столбец на основе старого можно создать с помощью метода .apply(), в который можно вписать свою функцию для преобразований. Мы создавали функцию trans_comm() и применяли ее к столбцу score, чтобы получить текстовый комментарий для каждой оценки. Часто вместе с .apply() используют lambda-функции, которые позволяют компактно определить функцию прямо внутри метода. Сделаем небольшое отступление и обсудим lambda-функции.

Lambda-функции и метод .apply()

Ранее мы создавали классические функции с помощью оператора def. Для примера возьмем простую функцию square(), которая принимает на вход число и возвращает его квадрат.

In [5]:
def square(x):
    return x ** 2

Как создать lambda-функцию в одну строчку, которая будет выполнять то же самое? Вот так:

In [6]:
square = lambda x: x ** 2

Сначала создаем функцию square, потом через = присваиваем ей значение. Пишем оператор lambda, который указывает на начало lambda-функции, добавляем аргумент функции x, описывая, что функция принимает на вход, а после двоеточия ставим то, что функция должна возвращать. Проверим, как она работает:

In [7]:
square(6)
Out[7]:
36

Lambda-функции удобны тем, что они более компактные, и тем, что они могут существовать без названия (их еще называют анонимными). Рассмотрим пример с фильтрацией элементов списка. Для фильтрации элементов списка иногда используют функцию filter(). Работает она так: внутри filter() указываем функцию для отбора элементов, которая возвращает True и False в зависимости от выполнения условия, и filter() выбирает из списка только те элементы, где было возвращено True.

Отберем из списка L неотрицательные элементы: внутри filter() поместим lambda-функцию, которая будет возвращать True, если число больше или равно 0, и False – в противном случае.

In [8]:
L = [0, 8, 4, -4, -5, 6]
In [9]:
# преобразуем в list(), чтобы получить список

list(filter(lambda x: x >= 0, L)) 
Out[9]:
[0, 8, 4, 6]

Что интересного в примере выше? То, что нам совсем не понадобилось как-то называть функцию и сохранять ее отдельно. Мы ее использовали один раз, применили и забыли. Для простых функций, которые создаются для одной маленькой задачи, это может быть актуально.

Вернемся к датафреймам. Создадим столбец yes, который будет состоять из 1 (ответ "да" в столбце wish) и 0 (ответ "нет" в столбце wish). Применим метод .apply() к столбцу wish и внутри него напишем lambda-функцию:

In [10]:
tree["yes"] = tree["wish"].apply(lambda x: int(x == "да"))
tree.head() 
Out[10]:
Unnamed: 0 gender ftype height score expenses wish female yes
0 1 female пихта Нобилис 190 3 1051 да 1 1
1 2 male пихта Нобилис 174 3 2378 нет 0 0
2 3 female сосна Крым 248 4 655 да 1 1
3 4 female сосна Крым 191 1 2934 да 1 1
4 5 female сосна Крым 147 3 1198 нет 1 0

Внутри lambda-функции можно использовать и более сложные конструкции, например, if-else. Выглядеть это будет так:

In [11]:
# сначала результат для if, потом само условие
# потом else и результат для else

check = lambda n: "Ч" if n % 2 == 0 else "Н"
In [12]:
check(9)
Out[12]:
'Н'
In [13]:
check(10)
Out[13]:
'Ч'

Если у lambda-функции несколько аргументов, они указываются через запятую:

In [14]:
my_sum = lambda x, y: x + y
In [15]:
my_sum(8, 4)
Out[15]:
12

Группировка и агрегирование

Для начала сгруппируем данные по типу дерева (ftype). Группировка осуществляется с помощью метода .groupby().

In [16]:
tree.groupby('ftype')
Out[16]:
<pandas.core.groupby.generic.DataFrameGroupBy object at 0x10931ea90>

Результат группировки от нас скрыт, он хранится в объекте особого типа DataFrameGroupBy. Чтобы посмотреть, что внутри, воспользуемся циклом:

In [18]:
for g in tree.groupby('ftype'):
    print(g)

Цикл выше выдает нам кортежи, в которых заключены пары значений: название группы и маленький датафрейм со строками, соответствующими этой группе. Для чего это можно использовать? Например, для сохранения данных по каждой группе в отдельный файл. Сделаем перебор в цикле сразу по элементам внутри пары (вспомните словари и перебор по .items()):

In [19]:
# на первом месте название, на втором – датафрейм

for name, dat in tree.groupby('ftype'):
    dat.to_csv(name + ".csv")

Теперь в рабочей папке появились четыре файла, один для каждого типа дерева.

Перейдем к агрегированию. Тут на помощь придет метод .agg(), который выполняет агрегирование по группам. Сгруппируем данные по столбцу ftype и посчитаем для каждого столбца среднее значение.

In [20]:
tree.groupby('ftype').agg('mean')
Out[20]:
Unnamed: 0 height score expenses female yes
ftype
ель обыкновенная 575.895349 155.503876 2.980620 1603.813953 0.542636 0.480620
пихта Нобилис 623.076687 160.101227 3.082822 1634.463190 0.555215 0.527607
сосна Крым 596.599388 160.657492 2.987768 1572.076453 0.495413 0.492355
сосна датская 601.411765 159.280277 2.958478 1709.916955 0.446367 0.532872

Столбцы текстового типа были исключены из расчетов, так как метод для вычисления среднего к ним неприменим. Если нам нужно сразу несколько характеристик сразу, их нужно перечислить в виде списка. Посчитаем среднее, медиану и число заполненных ячеек:

In [24]:
tree.groupby('ftype').agg(['mean', 'median', 'count']) 
Out[24]:
Unnamed: 0 height score expenses female yes
mean median count mean median count mean median count mean median count mean median count mean median count
ftype
ель обыкновенная 575.895349 570.5 258 155.503876 149.0 258 2.980620 3 258 1603.813953 1580 258 0.542636 1 258 0.480620 0 258
пихта Нобилис 623.076687 639.5 326 160.101227 159.5 326 3.082822 3 326 1634.463190 1634 326 0.555215 1 326 0.527607 1 326
сосна Крым 596.599388 614.0 327 160.657492 157.0 327 2.987768 3 327 1572.076453 1552 327 0.495413 0 327 0.492355 0 327
сосна датская 601.411765 557.0 289 159.280277 159.0 289 2.958478 3 289 1709.916955 1795 289 0.446367 0 289 0.532872 1 289

В результате получили датафрейм с более сложной структурой, где внутри одного столбца содержится несколько маленьких.

Как быть, если мы хотим для одного столбца посчитать одну характеристику по группам, а для другого – другую? Задать наши пожелания в виде словаря внутри .agg(). Посчитаем для столбца height среднее и стандартное отклонение, а для столбца score – медиану:

In [21]:
# ключи в словаре – названия столбцов
# значения в словаре – нужные характеристики

gr = tree.groupby('ftype').agg({'height' : ['mean', 'std'], 
                          'score': 'median'}) 
gr
Out[21]:
height score
mean std median
ftype
ель обыкновенная 155.503876 53.209976 3
пихта Нобилис 160.101227 52.472283 3
сосна Крым 160.657492 49.903024 3
сосна датская 159.280277 51.564492 3

В заключение отметим, что внутрь метода .agg() можно помещать название функции, написанной самостоятельно (не встроенные mean, std и прочие), только тогда ее название должно указываться без кавычек. Более сложный пример с группировкой и агрегированием можно посмотреть здесь.

Результаты агрегирования и мультииндексы

Результат группировки и агрегирования выше мы сохранили в переменную gr. Давайте посмотрим на структуру этого датафрейма. Запросим названия столбцов:

In [22]:
gr.columns
Out[22]:
MultiIndex([('height',   'mean'),
            ('height',    'std'),
            ( 'score', 'median')],
           )

Объект выше имеет особый тип MultiIndex, при этом его элементами являются кортежи – пары строк название столбца - название показателя. Это объяснимо: ранее мы заметили, что столбцы в gr имеют вложенную структуру, поэтому, чтобы дойти, например, до столбца со стандартным отклонением высоты деревьев, сначала придется зайти внутрь столбца height.

Строки здесь у нас пока обычные, поэтому внутри .index чего-то совсем нового нет:

In [23]:
gr.index
Out[23]:
Index(['ель обыкновенная', 'пихта Нобилис', 'сосна Крым', 'сосна датская'], dtype='object', name='ftype')

Если нам понадобится извлечь из gr данные по пихте Нобилис, мы сможем воспользоваться методом .loc:

In [24]:
# строка пихта Нобилис, все столбцы

gr.loc["пихта Нобилис", :]
Out[24]:
height  mean      160.101227
        std        52.472283
score   median      3.000000
Name: пихта Нобилис, dtype: float64

Если мы захотим извлечь данные по высоте пихты Нобилис, укажем в loc обе координаты, название строки и название столбца:

In [25]:
gr.loc["пихта Нобилис", "height"]
Out[25]:
mean    160.101227
std      52.472283
Name: пихта Нобилис, dtype: float64

А вот уже из таблицы выше мы сможем извлечь отдельное значение – среднее:

In [26]:
gr.loc["пихта Нобилис", "height"]["mean"]
Out[26]:
160.10122699386503

Иногда методов loc и iloc недостаточно. Особенно, если мы имеем дело с более сложным датафреймом, где мультииндексы присутствуют как в строках, так и в столбцах. Такой датафрейм может получиться, если группировка производится более, чем по одному показателю. Сгруппируем данные по типу дерева и по полу респондента и посмотрим, на сколько, в среднем, женщины и мужчины оценили деревья разных видов:

In [27]:
gr2 = tree.groupby(['ftype', 'gender']).agg({'score':
                                             ['mean', 'median']})
gr2
Out[27]:
score
mean median
ftype gender
ель обыкновенная female 3.050000 3
male 2.898305 3
пихта Нобилис female 3.088398 3
male 3.075862 3
сосна Крым female 2.981481 3
male 2.993939 3
сосна датская female 2.891473 3
male 3.012500 3

Как из такого датафрейма извлечь данные, соответствующие только женщинам, то есть извлечь только строки, где в gender указано "female"? Для этого существует метод .xs(), который позволяет извлечь данные с учетом определенного уровня группировки. Здесь у нас два уровня группировки – ftype и gender. Мы же хотим получить все данные, где уровень gender равен "female".

In [28]:
gr2.xs('female', level = 'gender')
Out[28]:
score
mean median
ftype
ель обыкновенная 3.050000 3
пихта Нобилис 3.088398 3
сосна Крым 2.981481 3
сосна датская 2.891473 3

Если бы нас интересовал уровень ftype, в данном случае результат, полученный с помощью .xs() не отличался бы от результата, полученного с помощью обычного .loc. Сравним:

In [29]:
gr2.xs('сосна Крым', level = 'ftype')
Out[29]:
score
mean median
gender
female 2.981481 3
male 2.993939 3
In [30]:
gr2.loc['сосна Крым', :]
Out[30]:
score
mean median
gender
female 2.981481 3
male 2.993939 3