across equivalent in pandas python

Question

I would like to perform function down bellow in python using preferably pandas as a part of my data aggregation process. In tidyverse there is nice function across that, combined with starts_with function, is a very powerfull combination.

library(tidyverse)

df <- tibble(mtcars)

aggregate <- function(...) {
  df %>% 
    group_by(...) %>% 
    summarise(across(ends_with("p"), ~ weighted.mean(.x, w = am))) %>% 
    ungroup()
}

res1 <- aggregate(gear)
res2 <- aggregate(cyl)

Answer 1

Checkout datar package in python.

With it, you will feel smooth to manipulate your data in tidyverse way:

from datar.all import f, group_by, summarise, ungroup, across, ends_with, weighted_mean
from datar.datasets import mtcars as df

def aggregate (*gvars):
  return df >> group_by(*gvars) >> summarise(
    across(ends_with("p"), weighted_mean, w=f.am)
  ) >> ungroup()

res1 = aggregate(f.gear)
#    gear      disp       hp
# 0     3       NaN      NaN
# 1     4  106.6875   83.875
# 2     5  202.4800  195.600

res2 = aggregate(f.cyl)
#   cyl      disp          hp
# 0    4   93.6125   81.875000
# 1    6  155.0000  131.666667
# 2    8  326.0000  299.500000

I am the author of the package. Please feel free to submit issues if you have any questions in using it.