pandas
Lesson 5 — Python
Learning objectives
- Work with Series and DataFrame, and understand the index
- Select rows and columns with
.locand.iloccorrectly - Filter, transform, group and aggregate
- Merge and reshape data
- Handle missing values with the nullable dtypes
- Avoid
SettingWithCopyWarningand the chained-assignment trap
Series and DataFrame
import pandas as pd
import numpy as np
s = pd.Series([1, 2, 3], name="age")
s = pd.Series([1, 2, 3], index=["a", "b", "c"])
df = pd.DataFrame({
"usubjid": ["001", "002", "003"],
"age": [45, 72, 38],
"sex": ["F", "M", "F"],
"arm": ["Placebo", "Drug A", "Placebo"],
})
df.shape # (3, 4)
df.columns
df.dtypes
df.info()
df.head(); df.tail()
df.describe()
df.describe(include="all")A DataFrame is a dict of Series sharing an index. That is the mental model — and the index is the part with no R equivalent.
The index
df.index # RangeIndex(start=0, stop=3, step=1)
df = df.set_index("usubjid")
df.loc["001"] # select by index label
df = df.reset_index() # index becomes a column againThe index is a labelled axis used for alignment. It is powerful and it is the source of most pandas confusion.
a = pd.Series([1, 2, 3], index=["x", "y", "z"])
b = pd.Series([10, 20, 30], index=["z", "y", "x"])
a + b
#> x 31
#> y 22
#> z 13The addition aligned on the index labels, not position. This is usually what you want and occasionally a nasty surprise — particularly after a groupby or a filter, where the index carries the original row positions.
reset_index(drop=True) after filtering avoids most of the trouble.
Selection
The two accessors, and using them correctly:
df.loc[row_label, column_label] # label-based, stop INCLUSIVE
df.iloc[row_pos, column_pos] # position-based, stop exclusive# Columns
df["age"] # a Series
df[["age", "sex"]] # a DataFrame
df.age # works, but fails on names with spaces or that
# collide with methods — avoid in code
# Rows
df.loc[0] # by index label
df.iloc[0] # by position
df.loc[0:2] # labels 0, 1 AND 2 — inclusive!
df.iloc[0:2] # positions 0, 1 — exclusive
# Both
df.loc[0, "age"]
df.loc[0:2, ["age", "sex"]]
df.iloc[0:2, 1:3]
df.loc[df["age"] > 40, "usubjid"]
# Scalar access, fastest
df.at[0, "age"]
df.iat[0, 1].loc slices are inclusive
df.loc[0:2] # three rows: 0, 1, 2
df.iloc[0:2] # two rows: 0, 1.loc slices on labels, and label slicing includes both endpoints — unlike every other slice in Python. This is deliberate (a label slice "2026-01":"2026-03" should include March) and it catches everyone once.
Filtering
df[df["age"] > 40]
df[(df["age"] > 40) & (df["sex"] == "F")]
df[df["arm"].isin(["Placebo", "Drug A"])]
df[df["age"].between(40, 70)]
df[df["age"].isna()]
df[df["usubjid"].str.startswith("00")]
df[~df["arm"].isin(["Screen Failure"])]
# query() — string-based, often more readable
df.query("age > 40 and sex == 'F'")
df.query("arm in ['Placebo', 'Drug A']")
threshold = 40
df.query("age > @threshold") # @ references a Python variableSame rules as NumPy: &, |, ~, with parentheses. and/or raise.
Adding and modifying columns
df["bmi"] = df["weight"] / (df["height"] / 100) ** 2
# assign() returns a new DataFrame — chainable, no mutation
df = df.assign(
bmi=lambda d: d["weight"] / (d["height"] / 100) ** 2,
bmi_cat=lambda d: pd.cut(
d["bmi"],
bins=[-np.inf, 18.5, 25, 30, np.inf],
labels=["Underweight", "Normal", "Overweight", "Obese"],
),
)assign with lambdas is the pandas equivalent of dplyr::mutate() — later expressions can reference earlier ones via d, and the whole thing composes into a pipeline.
Conditional logic:
# Two branches
df["flag"] = np.where(df["age"] >= 65, "Y", "N")
# Many branches
df["agegr"] = np.select(
[df["age"] < 18, df["age"] < 65],
["<18", "18-64"],
default=">=65",
)
# Binning
df["agegr"] = pd.cut(
df["age"],
bins=[0, 18, 65, np.inf],
labels=["<18", "18-64", ">=65"],
right=False, # [0,18), [18,65), [65,inf)
)
# Mapping
df["armn"] = df["arm"].map({"Placebo": 0, "Drug A": 1, "Drug B": 2})SettingWithCopyWarning
The most confusing pandas message.
subset = df[df["age"] > 40]
subset["flag"] = "Y"
# SettingWithCopyWarning: A value is trying to be set on a copy of a sliceThe problem: pandas cannot tell whether subset is a view or a copy, so the assignment may or may not affect df. The fixes:
# 1. Be explicit that you want a copy
subset = df[df["age"] > 40].copy()
subset["flag"] = "Y"
# 2. Assign into the original with .loc
df.loc[df["age"] > 40, "flag"] = "Y"
# 3. Use assign in a chain (no mutation at all)
result = df.query("age > 40").assign(flag="Y")pandas 3.0 makes Copy-on-Write the default, which removes this ambiguity entirely — chained assignment simply never propagates, and the warning disappears. Enable it now:
pd.options.mode.copy_on_write = TrueTurning it on early surfaces any code that was relying on the ambiguous behaviour, which is code you want to find.
Grouping and aggregation
df.groupby("arm")["age"].mean()
df.groupby("arm")["age"].agg(["count", "mean", "std", "min", "max"])
df.groupby(["arm", "sex"])["age"].mean()
# Named aggregations — the clearest form
summary = df.groupby("arm").agg(
n=("usubjid", "count"),
n_distinct=("usubjid", "nunique"),
mean_age=("age", "mean"),
sd_age=("age", "std"),
n_female=("sex", lambda s: (s == "F").sum()),
).reset_index()
# Multiple columns, multiple functions
df.groupby("arm").agg({"age": ["mean", "std"], "weight": "mean"})
# transform — returns the original shape, broadcast back
df["arm_mean_age"] = df.groupby("arm")["age"].transform("mean")
df["centred"] = df["age"] - df.groupby("arm")["age"].transform("mean")
# filter — keep whole groups
df.groupby("arm").filter(lambda g: len(g) >= 10)
# apply — arbitrary function per group (slowest; use only when needed)
df.groupby("arm").apply(lambda g: g.nlargest(3, "age"))groupby(...).transform(...) is the equivalent of mutate(..., .by = ...) in dplyr, and agg is summarise. Note that groupby puts the grouping keys in the index by default — .reset_index() or as_index=False brings them back as columns.
df.groupby("arm", as_index=False).agg(n=("usubjid", "count"))
df.groupby("arm", dropna=False) # keep NaN groups — default drops them!dropna=False matters: by default groupby silently discards rows where the grouping key is missing, which can quietly remove subjects from a summary.
Merging
pd.merge(adsl, adae, on="usubjid", how="left")
pd.merge(adsl, adae, on="usubjid", how="inner")
pd.merge(adsl, adae, on="usubjid", how="outer")
pd.merge(adsl, sites, left_on="siteid", right_on="site_number")
pd.merge(a, b, on=["usubjid", "paramcd"])
pd.merge(a, b, on="usubjid", suffixes=("_adsl", "_adae"))
df.join(other) # joins on the index
pd.concat([df1, df2]) # stack rows
pd.concat([df1, df2], axis=1) # side by sidevalidate= is the equivalent of dplyr’s relationship=
pd.merge(adsl, sites, on="siteid", how="left", validate="many_to_one")
# MergeError: Merge keys are not unique in right dataset;
# not a many-to-one mergeOptions: "one_to_one", "one_to_many", "many_to_one", "many_to_many".
Always declare it. A merge on a non-unique key silently multiplies rows, and that is the single most common way an analysis produces wrong numbers. See the same argument in R’s dplyr lesson.
Also useful:
pd.merge(a, b, on="usubjid", how="left", indicator=True)["_merge"].value_counts()
#> both 1847
#> left_only 12 <- subjects with no eventsReshaping
# Wide -> long
long = df.melt(
id_vars=["usubjid"],
value_vars=["week0", "week4", "week8"],
var_name="visit",
value_name="value",
)
# Long -> wide
wide = long.pivot(index="usubjid", columns="visit", values="value")
# With aggregation
pd.pivot_table(
df,
index="arm",
columns="sex",
values="age",
aggfunc=["mean", "count"],
margins=True,
)
# Cross-tabulation
pd.crosstab(df["arm"], df["sex"], margins=True, normalize="index")
# Stack / unstack — move between index and columns
df.stack()
df.unstack()melt is pivot_longer; pivot is pivot_wider. pivot raises on duplicate index/column pairs, where pivot_table aggregates them — the pandas equivalent of tidyr’s list-column warning, and the same advice applies: investigate the duplicates rather than aggregating them away.
Missing values
df.isna(); df.notna()
df.isna().sum() # missing count per column
df.isna().sum().sum() # total
df.dropna() # any NaN
df.dropna(subset=["age"])
df.dropna(how="all")
df.fillna(0)
df.fillna({"age": 0, "sex": "Unknown"})
df["age"].fillna(df["age"].median())
df.ffill(); df.bfill() # forward / backward fillNullable dtypes
The traditional problem: an integer column with a missing value becomes float.
pd.Series([1, 2, None])
#> 0 1.0
#> 1 2.0
#> 2 NaN
#> dtype: float64 — became float!
pd.Series([1, 2, None], dtype="Int64") # capital I
#> 0 1
#> 1 2
#> 2 <NA>
#> dtype: Int64 — stays integerThe nullable dtypes:
"Int64", "Int32" # nullable integers
"Float64" # nullable float
"boolean" # nullable boolean
"string" # nullable string (better than object)df = df.convert_dtypes() # convert everything to the nullable equivalentsThese behave much more like R’s typed NA values, and they are worth adopting for clinical data where the distinction between an integer count and a float matters for the output.
The Arrow backend goes further:
df = pd.read_csv("data.csv", dtype_backend="pyarrow")Faster, lower memory, better string handling, and NA semantics throughout.
Strings and dates
df["usubjid"].str.upper()
df["usubjid"].str.strip()
df["usubjid"].str.len()
df["usubjid"].str.contains("001")
df["usubjid"].str.startswith("STUDY")
df["usubjid"].str.replace("-", "_", regex=False)
df["usubjid"].str.split("-", expand=True) # -> a DataFrame
df["usubjid"].str.extract(r"(\d{3})-(\d{4})")
df["usubjid"].str[0:3]
pd.to_datetime(df["dtc"])
pd.to_datetime(df["dtc"], format="%Y-%m-%d", errors="coerce")
df["date"].dt.year; df["date"].dt.month; df["date"].dt.day
df["date"].dt.dayofweek
df["date"].dt.strftime("%d%b%Y")
(df["end"] - df["start"]).dt.days
df["date"] + pd.Timedelta(days=30)
df["date"] + pd.DateOffset(months=1)errors="coerce" turns unparseable values into NaT rather than raising — essential for partial dates, and it means you must then check how many became NaT.
Method chaining
The pandas equivalent of the pipe:
result = (
adsl
.query("SAFFL == 'Y'")
.assign(
agegr=lambda d: pd.cut(d["AGE"], [0, 65, 80, np.inf],
labels=["<65", "65-80", ">80"], right=False)
)
.groupby(["TRT01A", "agegr"], observed=True, dropna=False)
.agg(n=("USUBJID", "nunique"), mean_age=("AGE", "mean"))
.reset_index()
.sort_values(["TRT01A", "agegr"])
)Wrap the whole chain in parentheses so each method can go on its own line. Use .pipe() for custom functions:
def add_flags(df, threshold):
return df.assign(high=df["AVAL"] > threshold)
result = df.query("PARAMCD == 'ALT'").pipe(add_flags, threshold=40).head().pipe(f, args) is f(df, args) — exactly R’s pipe.
R and Python side by side
| Task | dplyr | pandas |
|---|---|---|
| Filter rows | filter(df, age > 40) |
df[df["age"] > 40] or df.query("age > 40") |
| Select columns | select(df, a, b) |
df[["a", "b"]] |
| Drop columns | select(df, -a) |
df.drop(columns="a") |
| Add column | mutate(df, x = y * 2) |
df.assign(x=lambda d: d["y"] * 2) |
| Rename | rename(df, new = old) |
df.rename(columns={"old": "new"}) |
| Sort | arrange(df, x) |
df.sort_values("x") |
| Sort descending | arrange(df, desc(x)) |
df.sort_values("x", ascending=False) |
| Distinct | distinct(df, x) |
df.drop_duplicates("x") |
| Group + summarise | summarise(..., .by = g) |
df.groupby("g").agg(...) |
| Group + mutate | mutate(..., .by = g) |
df.groupby("g")[c].transform(...) |
| Count | count(df, x) |
df["x"].value_counts() |
| Left join | left_join(a, b, "id") |
pd.merge(a, b, on="id", how="left") |
| Bind rows | bind_rows(a, b) |
pd.concat([a, b]) |
| Long | pivot_longer() |
df.melt() |
| Wide | pivot_wider() |
df.pivot() |
| Rows | nrow(df) |
len(df) |
| Columns | names(df) |
df.columns |
| Structure | glimpse(df) |
df.info() |
| Pipe | \|> |
Method chain or .pipe() |
Common mistakes
| Mistake | Consequence | Fix |
|---|---|---|
and/or in a filter |
ValueError |
&, | with parentheses |
SettingWithCopyWarning ignored |
Assignment may not take effect | .copy() or .loc[] |
.loc slice treated as exclusive |
Off by one | .loc includes both endpoints |
Merge without validate= |
Silent row multiplication | Always declare the cardinality |
groupby dropping NaN keys |
Subjects silently excluded | dropna=False |
| Integer column becomes float | Type changes on missing | Nullable "Int64" |
df.column attribute access |
Breaks on names with spaces or method collisions | df["column"] |
apply where a vectorised method exists |
100× slower | Use the built-in method |
Forgetting reset_index() after groupby |
Keys stuck in the index | as_index=False or .reset_index() |
Exercise 5.1 — Demographics summary
From an ADSL DataFrame, produce a summary by treatment arm with N, mean and SD of age, count and percentage female, and count and percentage aged 65+. Restrict to SAFFL == "Y".
Show solution
import pandas as pd
import numpy as np
pd.options.mode.copy_on_write = True
def demographics_summary(adsl: pd.DataFrame) -> pd.DataFrame:
"""Summarise demographics by treatment arm for the safety population."""
return (
adsl
.query("SAFFL == 'Y'")
.groupby("TRT01A", as_index=False, dropna=False)
.agg(
N=("USUBJID", "nunique"),
age_mean=("AGE", "mean"),
age_sd=("AGE", "std"),
age_n=("AGE", "count"), # non-missing — the small n
n_female=("SEX", lambda s: (s == "F").sum()),
n_elderly=("AGE", lambda a: (a >= 65).sum()),
)
.assign(
pct_female=lambda d: 100 * d["n_female"] / d["N"],
pct_elderly=lambda d: 100 * d["n_elderly"] / d["N"],
)
.sort_values("TRT01A")
.reset_index(drop=True)
)
summary = demographics_summary(adsl)
print(summary.to_string(index=False)) TRT01A N age_mean age_sd age_n n_female n_elderly pct_female pct_elderly
Placebo 86 75.21 8.590 86 53 73 61.63 84.88
Xanomeline High Dose 84 74.38 7.886 84 40 71 47.62 84.52
Xanomeline Low Dose 84 75.67 8.286 84 40 73 47.62 86.90Points that matter:
NusesnuniqueonUSUBJID, notcount. If ADSL is correct these are identical, butnuniquefails safe: a duplicated subject inflatescountand every percentage silently, whilenuniquegives the right denominator. It also makes the intent explicit.age_nis separate fromN.Nis the population size and the percentage denominator;age_nis the number of non-missing ages. They differ when age is missing, and conflating them is the most common finding in a demographics table review — see TLF generation.dropna=Falseon the groupby. A subject with a missingTRT01Awould otherwise vanish from the summary entirely and the totals would not add up.Percentages in a separate
assignafter the aggregation, so they use the aggregatedNrather than being recomputed.
Formatting for display
formatted = summary.assign(
Age=lambda d: d.apply(lambda r: f"{r['age_mean']:.1f} ({r['age_sd']:.2f})", axis=1),
Female=lambda d: d.apply(lambda r: f"{r['n_female']:.0f} ({r['pct_female']:.1f})", axis=1),
Elderly=lambda d: d.apply(lambda r: f"{r['n_elderly']:.0f} ({r['pct_elderly']:.1f})", axis=1),
)[["TRT01A", "N", "Age", "Female", "Elderly"]]
print(formatted.to_string(index=False)) TRT01A N Age Female Elderly
Placebo 86 75.2 (8.59) 53 (61.6) 73 (84.9)Exercise 5.2 — Baseline and change
From a long-format lab DataFrame with USUBJID, PARAMCD, AVISITN, ADT and AVAL, derive BASE (last value at or before visit 0), CHG and PCHG, and a flag for the last record per subject/parameter/visit. Verify the result.
Show solution
import pandas as pd
import numpy as np
pd.options.mode.copy_on_write = True
def derive_baseline_change(adlb: pd.DataFrame) -> pd.DataFrame:
"""Derive BASE, CHG, PCHG and the analysis record flag.
Baseline is the last non-missing value on or before visit 0.
"""
df = adlb.sort_values(["USUBJID", "PARAMCD", "ADT", "AVISITN"]).copy()
key = ["USUBJID", "PARAMCD"]
# --- Baseline flag: last non-missing record at or before visit 0 --------
pre = df["AVAL"].notna() & (df["AVISITN"] <= 0)
# Position of the last qualifying record within each subject/parameter
last_pre_idx = (
df[pre]
.groupby(key, dropna=False)
.tail(1)
.index
)
df["ABLFL"] = np.where(df.index.isin(last_pre_idx), "Y", None)
# --- Baseline value broadcast to every record --------------------------
baseline = (
df.loc[df["ABLFL"] == "Y", key + ["AVAL"]]
.rename(columns={"AVAL": "BASE"})
)
df = df.merge(baseline, on=key, how="left", validate="many_to_one")
# --- Change ------------------------------------------------------------
df["CHG"] = df["AVAL"] - df["BASE"]
df["PCHG"] = np.where(
df["BASE"].notna() & (df["BASE"] != 0),
100 * (df["AVAL"] - df["BASE"]) / df["BASE"],
np.nan,
)
# --- Analysis record flag: last record per subject/parameter/visit -----
anl_idx = (
df[df["AVAL"].notna()]
.sort_values(["USUBJID", "PARAMCD", "AVISITN", "ADT"])
.groupby(["USUBJID", "PARAMCD", "AVISITN"], dropna=False)
.tail(1)
.index
)
df["ANL01FL"] = np.where(df.index.isin(anl_idx), "Y", None)
return df.reset_index(drop=True)Verification
result = derive_baseline_change(adlb)
# 1. Exactly one baseline per subject and parameter
dupe_base = (
result[result["ABLFL"] == "Y"]
.groupby(["USUBJID", "PARAMCD"])
.size()
)
assert (dupe_base == 1).all(), f"Multiple baselines: {dupe_base[dupe_base > 1]}"
# 2. Baseline is never after visit 0
assert (result.loc[result["ABLFL"] == "Y", "AVISITN"] <= 0).all()
# 3. CHG is 0 on the baseline record itself
base_chg = result.loc[result["ABLFL"] == "Y", "CHG"]
assert (base_chg.fillna(0) == 0).all(), "CHG must be 0 at baseline"
# 4. Row count unchanged
assert len(result) == len(adlb), "Derivation changed the row count"
# 5. Subjects with no baseline — expected, but report them
no_base = (
result.groupby(["USUBJID", "PARAMCD"])["ABLFL"]
.apply(lambda s: (s == "Y").any())
.pipe(lambda s: s[~s])
)
print(f"{len(no_base)} subject/parameter combinations with no baseline")
# 6. PCHG must be NaN where BASE is zero, not inf
assert not np.isinf(result["PCHG"]).any(), "PCHG contains infinity"Four design decisions worth explaining:
validate="many_to_one" on the baseline merge. If the baseline extraction somehow produced two rows for a subject/parameter, this merge would silently double every record. Declaring the cardinality turns that into an immediate, located error. This is the pandas equivalent of dplyr’s relationship= argument and it should be on every merge.
groupby(...).tail(1) rather than a boolean condition. Getting “the last qualifying record per group” with a filter requires computing a max date and then matching it, which breaks on ties. tail(1) after an explicit sort gives exactly one row per group by construction.
The PCHG guard against BASE == 0. Without it, division by zero produces inf rather than NaN, and inf propagates through means and standard deviations to produce a summary table full of inf. Assertion 6 catches this.
Sorting includes AVISITN as a tiebreaker. Two records on the same date with different visit numbers must order deterministically, or the same code produces different baselines on different runs. This is the same requirement as deterministic --SEQ derivation.
CHG throughout, and how they are handled in the analysis is a statistical decision, not a programming one.
Recap
- A DataFrame is a dict of Series sharing an index; the index aligns operations automatically
.locis label-based and slice-inclusive;.ilocis positional and exclusive&,|,~with parentheses;query()for readable filters.assign()with lambdas ismutate();.pipe()is the pipe- Always pass
validate=tomerge()— silent row multiplication is the main data bug groupby(dropna=False)or missing keys silently drop rows- Nullable dtypes (
"Int64","string") behave like R’s typedNA - Enable Copy-on-Write to eliminate
SettingWithCopyWarning
Next: Data visualization.