How to apply window case-when functions on pandas groups

Question

I have a table with two identifier columns w and x. The columns y brings order to the column z, which is the destination variable. Every sequence identified by w and x has at least two rows and all distinct values in z. I want to count for each of the zs how often they occur in the sequences in which position. Position-wise I'm interested in whether it's the first, last or any other. My approach in R's dplyris the following (for those unfamiliar with %>%, it's the pipe operator that takes the left hand output and puts it as first argument into the right hand function, you can read it as "and then"):

library(tidyverse)
library(reticulate)

data <- tribble(
  ~w,   ~x,  ~y,   ~z,
   1,    1,    1,  "a",
   1,    1,    2,  "b",
   1,    2,    1,  "b",
   1,    2,    2,  "a",
   1,    2,    3,  "c",
   1,    2,    4,  "d",
   2,    1,    1,  "b",
   2,    1,    2,  "a",
   2,    1,    3,  "c",
   2,    1,    4,  "d"
)

First I add a sequence index to each ordered w and x groups and a marker that determine the categorical position of z in the sequence from it (the position_in_sequence).

data_with_markers <- 
  data %>%
  group_by(w, x) %>%
  arrange(y, .by_group = TRUE) %>%
  mutate(
    sequence_id = row_number(),
    position_in_sequence = case_when(
      sequence_id == first(sequence_id) ~ "first",
      sequence_id == last(sequence_id) ~ "last",
      TRUE ~ "other" # this is the "else"
    )
  ) %>%
  ungroup()


>data_with_markers
# A tibble: 10 x 6
       w     x     y z     sequence_id position_in_sequence 
              
 1     1     1     1 a               1 first
 2     1     1     2 b               2 last 
 3     1     2     1 b               1 first
 4     1     2     2 a               2 other
 5     1     2     3 c               3 other
 6     1     2     4 d               4 last 
 7     2     1     1 b               1 first
 8     2     1     2 a               2 other
 9     2     1     3 c               3 other
10     2     1     4 d               4 last

Then I do a simple count by the position_in_sequence and z.

data_summary <- data_with_markers %>% 
  group_by(position_in_sequence, z) %>% 
  count() %>% 
  ungroup()

> data_summary
# A tibble: 6 x 3
  position_in_sequence  z         n
    
1 first a         1
2 first b         2
3 last  b         1
4 last  d         2
5 other a         2
6 other c         2

For pandas, I'm stuck with getting the position_in_sequence variable (I'm using the reticulate package here):

import pandas as pd
data = r.data
data

     w    x    y  z
0  1.0  1.0  1.0  a
1  1.0  1.0  2.0  b
2  1.0  2.0  1.0  b
3  1.0  2.0  2.0  a
4  1.0  2.0  3.0  c
5  1.0  2.0  4.0  d
6  2.0  1.0  1.0  b
7  2.0  1.0  2.0  a
8  2.0  1.0  3.0  c
9  2.0  1.0  4.0  d

data_sorted = data.sort_values(['w', 'x', 'y'])

data_sorted['sequence_id'] = data_sorted.groupby(['w', 'x']).cumcount() + 1
data_sorted

     w    x    y  z  sequence_id
0  1.0  1.0  1.0  a            1
1  1.0  1.0  2.0  b            2
2  1.0  2.0  1.0  b            1
3  1.0  2.0  2.0  a            2
4  1.0  2.0  3.0  c            3
5  1.0  2.0  4.0  d            4
6  2.0  1.0  1.0  b            1
7  2.0  1.0  2.0  a            2
8  2.0  1.0  3.0  c            3
9  2.0  1.0  4.0  d            4

I fiddled around with .apply but I would need to access the column sequence_id at a certain row and the column's min and max at the same time to compare, but I didn't get it to work.

user_5088165 · Answer

You can do it exactly as you did in R with datar in python:

>>> from datar.all import (
...     f, tribble, group_by, arrange, row_number, 
...     case_when, first, last, ungroup, mutate, count
... )
>>> 
>>> data = tribble(
...    f.w,  f.x,  f.y, f.z,
...    1,    1,    1,   "a",
...    1,    1,    2,   "b",
...    1,    2,    1,   "b",
...    1,    2,    2,   "a",
...    1,    2,    3,   "c",
...    1,    2,    4,   "d",
...    2,    1,    1,   "b",
...    2,    1,    2,   "a",
...    2,    1,    3,   "c",
...    2,    1,    4,   "d"
... )
>>> 
>>> data_with_markers = (
...     data >>
...     group_by(f.w, f.x) >>
...     arrange(f.y, _by_group=True) >>
...     mutate(
...         sequence_id=row_number(),
...         position_in_sequence=case_when(
...             f.sequence_id == first(f.sequence_id), "first",
...             f.sequence_id == last(f.sequence_id),  "last",
...             True, "other"
...         )
...     ) >>
...     ungroup()
... )
>>> 
>>> data_with_markers
        w       x       y        z  sequence_id position_in_sequence

user_13227036 · Answer

im not sure if this will help but i see you have used a variable called "type", this code might be giving you error because "type" is a keyword and probably not the best choice for variable/column name.. try "type1" or "typ"? keyword assignment as variable usually doesn't give any problems initially, but later leaves you with inexplicable errors (im just brainstorming here, i havent used reticulate)

user_7175713 · Answer

One option is a combination of groupby and case_when:

grp = df.groupby(['w', 'x'])

From here, we can run through with chaining and some help from pyjanitor case_when:

# pip install pyjanitor
import pandas as pd
import janitor

out = (df
      .assign(sequence_id = grp.cumcount().add(1))
      .case_when(
          lambda df: df.sequence_id == 1, 'first', # condition, result
          lambda df: df.sequence_id == grp.y.transform('size'), 'last',
          'other', # default if conditions fail
          column_name = 'position_in_sequence')
      )


out

     w    x    y  z  sequence_id position_in_sequence
0  1.0  1.0  1.0  a            1                first
1  1.0  1.0  2.0  b            2                 last
2  1.0  2.0  1.0  b            1                first
3  1.0  2.0  2.0  a            2                other
4  1.0  2.0  3.0  c            3                other
5  1.0  2.0  4.0  d            4                 last
6  2.0  1.0  1.0  b            1                first
7  2.0  1.0  2.0  a            2                other
8  2.0  1.0  3.0  c            3                other
9  2.0  1.0  4.0  d            4                 last

You can subsequently get your count:

out.groupby(['position_in_sequence', 'z']).size()

position_in_sequence  z
first                 a    1
                      b    2
last                  b    1
                      d    2
other                 a    2
                      c    2
dtype: int64

Another option, with just pandas and numpy:


grp = df.groupby(['w', 'x'])

out = df.assign(sequence_id = grp.cumcount().add(1))

conditions = [out.sequence_id == 1, 
              out.sequence_id == grp.y.transform('size')]
results = ['first', 'last']

out = out.assign(position_in_sequence = np.select(conditions, results, 'other'))

out

     w    x    y  z  sequence_id position_in_sequence
0  1.0  1.0  1.0  a            1                first
1  1.0  1.0  2.0  b            2                 last
2  1.0  2.0  1.0  b            1                first
3  1.0  2.0  2.0  a            2                other
4  1.0  2.0  3.0  c            3                other
5  1.0  2.0  4.0  d            4                 last
6  2.0  1.0  1.0  b            1                first
7  2.0  1.0  2.0  a            2                other
8  2.0  1.0  3.0  c            3                other
9  2.0  1.0  4.0  d            4                 last

How to apply window case-when functions on pandas groups

3 Answers