pythonforbio.
[WASM idle]
pandas18/21

groupby

Starting sandbox…
Beginnerlesson

groupby

In one line: pandas' implementation of Split-Apply-Combine — and the method you will use more than any other.

The four apply modes

The mode determines the output shape. This is the distinction to internalise.

g = df.groupby("gene")

g["count"].agg("mean")          # one row per GROUP
g["count"].transform("mean")    # one row per ORIGINAL ROW (broadcast back)
g.filter(lambda x: len(x) > 5)  # SUBSET of original rows
g.apply(func)                    # anything; slowest; pandas guesses the shape

agg

df.groupby("gene")["count"].mean()
df.groupby("gene")["count"].agg(["count", "mean", "std", "min", "max"])
df.groupby(["gene", "sample"])["count"].sum()

# named aggregation — the clearest form, and it names your output columns
df.groupby("gene").agg(
    n_variants=("pos", "size"),
    mean_af=("af", "mean"),
    max_impact=("impact", "max"),
    n_samples=("sample", "nunique"),
)

# different functions per column
df.groupby("gene").agg({"af": ["mean", "max"], "depth": "median"})

Named aggregation (new_name=("col", "func")) is the modern idiom. It produces flat, well-named columns instead of a MultiIndex you then have to flatten.

Use the string names ("mean", "sum", "size", "nunique", "first", "idxmax") rather than lambdas or np.mean. The strings dispatch to compiled Cython implementations; a lambda runs Python once per group. On 20,000 gene groups that is a 50–100× difference.

size vs count

g.size()            # rows per group, INCLUDING NaN
g.count()           # non-NaN values per column
g["col"].count()    # non-NaN in that column
g["col"].nunique()  # distinct values

size counts rows; count counts values. They differ exactly where you have missing data — which is where you most want to know.

transform

The one people underuse. It returns a result the same length as the input, broadcast back to every row of its group.

# centre expression within each gene
df["centred"] = df["count"] - df.groupby("gene")["count"].transform("mean")

# z-score within batch
df["z"] = df.groupby("batch")["value"].transform(lambda s: (s - s.mean()) / s.std())

# proportion of each sample's total
df["frac"] = df["count"] / df.groupby("sample")["count"].transform("sum")

# rank within group
df["rank_in_gene"] = df.groupby("gene")["af"].rank(ascending=False)

The alternative — aggregate, then merge back — is three lines and slower:

means = df.groupby("gene")["count"].mean().rename("gene_mean")
df = df.merge(means, on="gene")
df["centred"] = df["count"] - df["gene_mean"]

Any time you find yourself writing that pattern, transform is the answer. Per-batch normalization, per-sample scaling, within-gene centring — all transform.

filter

df.groupby("gene").filter(lambda g: len(g) > 10)
df.groupby("gene").filter(lambda g: g["af"].max() > 0.01)

Keeps whole groups based on a group-level predicate. It runs a Python lambda per group, so for simple count-based filtering the value_counts + isin version is much faster:

counts = df["gene"].value_counts()
df.loc[df["gene"].isin(counts[counts > 10].index)]

apply — last resort

df.groupby("gene").apply(lambda g: g.nlargest(3, "af"))       # slow

apply is flexible and slow. Before reaching for it, check whether the operation is really:

Want Instead of apply
top N per group df.sort_values(...).groupby(...).head(n)
best row per group df.loc[df.groupby("g")["x"].idxmin()]
several aggregations named agg
per-group normalization transform
per-group cumulative df.groupby("g")["x"].cumsum()

apply also has awkward semantics — pandas inspects the return value to decide the output shape, which changes if your data changes. Pass include_groups=False in modern pandas to avoid the deprecation warning about the grouping column appearing in the passed frame.

Grouping keys

df.groupby("gene")                       # a column
df.groupby(["chrom", "gene"])            # several
df.groupby(level="chrom")                # an index level
df.groupby(df["pos"] // 1_000_000)       # a derived Series — 1 Mb bins
df.groupby(pd.cut(df["af"], bins=[0, .01, .05, 1]))     # binned
df.groupby("chrom", observed=True)       # categoricals: only present values
df.groupby("gene", dropna=False)         # keep NaN keys as a group
df.groupby("gene", sort=False)           # skip sorting — faster
df.groupby("gene", as_index=False)       # keys as columns, not the index

sort=False is free speed when you do not need ordered output. as_index=False saves a reset_index().

dropna=False matters: by default, rows with a missing group key are silently dropped. If 5% of your variants have no gene annotation, they vanish from every per-gene summary without a word.

Iterating

for gene, group in df.groupby("gene"):
    ...                                  # group is a DataFrame

groups = dict(list(df.groupby("gene")))
df.groupby("gene").get_group("TP53")

Fine for plotting per group or writing per-group files. Not fine as a substitute for agg/transform.

Bioinformatics examples

# variants per gene, with a breakdown
summary = variants.groupby("gene").agg(
    n=("pos", "size"),
    n_rare=("gnomad_af", lambda s: (s.fillna(0) < 0.001).sum()),
    max_impact=("impact", "max"),          # works because impact is an ORDERED category
    n_samples=("sample_id", "nunique"),
).sort_values("n", ascending=False)

# mean coverage per exon
cov.groupby(["gene", "exon"])["depth"].agg(["mean", "min"])

# per-sample QC
reads.groupby("sample").agg(
    total=("read_id", "size"),
    mapped=("mapped", "sum"),
    dup_rate=("is_dup", "mean"),
    median_mapq=("mapq", "median"),
)

# normalise counts within each sample (CPM)
long["cpm"] = long["count"] / long.groupby("sample")["count"].transform("sum") * 1e6

# most significant variant per gene
best = variants.loc[variants.groupby("gene")["pvalue"].idxmin()]

# variants per 1 Mb bin, for a Manhattan-style plot
variants.groupby(["chrom", variants["pos"] // 1_000_000]).size()

The max_impact=("impact", "max") line only produces the right answer because impact is an ordered categorical — otherwise max is alphabetical and returns MODIFIER. See Categorical dtype.

Performance

df.groupby("gene", sort=False)["x"].mean()          # skip sorting
df["gene"] = df["gene"].astype("category")           # group on int codes
df.groupby("gene")["x"].agg("mean")                  # NOT .agg(lambda s: s.mean())

Rough ordering, fastest first: built-in string aggregations → transform with a string → agg with a lambda → filterapply.

Common mistakes

  • Lambdas where a string name exists. 50–100× slower.
  • agg-then-merge instead of transform.
  • NaN group keys silently dropped. Pass dropna=False.
  • size vs count confusion when NaNs are present.
  • apply for everything.
  • Forgetting observed=True/False with categorical keys.
  • Grouping on a float column. Floating-point equality makes the groups unreliable. Round or bin first.
  • Not checking group sizes. A "mean" over a group of size 1 is a data point, not an estimate. Always aggregate size alongside.

See also

Split-Apply-Combine · Window Functions · Merging and Joining · Categorical dtype · Axes and Reductions · pandas Performance · Tidy Data

scratch

No output yet — run the code to populate this drawer.