Data visualization

Lesson 6 — Python

Lesson 6 of 20 Intermediate ~80 min

Learning objectives

  • Build plots with matplotlib’s object-oriented interface
  • Use seaborn for statistical graphics
  • Create interactive plots with plotly
  • Use plotnine if you already think in ggplot2
  • Save publication-quality figures
  • Choose the right library for the job

The landscape

Library Style Best for
matplotlib Imperative, low-level Full control, publication figures
seaborn Declarative, statistical Fast exploratory statistical plots
plotly Interactive Dashboards, exploration, HTML output
plotnine Grammar of graphics ggplot2 users
altair Declarative (Vega-Lite) Interactive, JSON-serialisable specs
bokeh Interactive Web applications, large data

seaborn and plotnine are built on matplotlib, so matplotlib knowledge applies to both.

matplotlib

Use the object-oriented interface. The plt.plot() state-machine style is shorter for one-off plots and becomes unmanageable for anything with subplots.

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(8, 5), dpi=100)

ax.plot(x, y, color="#16355e", linewidth=1.5, label="Drug A")
ax.plot(x, y2, color="#7a8698", linestyle="--", label="Placebo")

ax.set_xlabel("Study week")
ax.set_ylabel("ALT (U/L)")
ax.set_title("Mean ALT over time")
ax.legend(frameon=False)
ax.grid(True, alpha=0.3)

fig.tight_layout()
fig.savefig("output/figures/alt.png", dpi=300, bbox_inches="tight")
plt.close(fig)

fig is the canvas; ax is a set of axes on it. Everything is a method on one or the other.

Plot types

ax.plot(x, y)                       # line
ax.scatter(x, y, s=40, alpha=0.6)   # scatter
ax.bar(categories, values)
ax.barh(categories, values)         # horizontal
ax.hist(values, bins=30)
ax.boxplot([group1, group2], labels=["A", "B"])
ax.errorbar(x, y, yerr=se, capsize=3)
ax.fill_between(x, lower, upper, alpha=0.2)
ax.axhline(0, color="grey", linestyle="--")
ax.axvline(x=0)
ax.step(x, y, where="post")         # for Kaplan-Meier

Subplots

fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharex=True, sharey=True)

for ax, (name, group) in zip(axes.flat, df.groupby("PARAMCD")):
    ax.plot(group["AVISITN"], group["AVAL"])
    ax.set_title(name)

fig.supxlabel("Study week")
fig.supylabel("Value")
fig.tight_layout()

axes.flat flattens the 2-D array of axes, so you can zip it with groups.

Customisation

ax.set_xlim(0, 100)
ax.set_ylim(bottom=0)
ax.set_xticks([0, 4, 8, 12])
ax.set_xticklabels(["Baseline", "Wk 4", "Wk 8", "Wk 12"], rotation=45)
ax.spines[["top", "right"]].set_visible(False)     # remove the box
ax.tick_params(labelsize=9)

from matplotlib.ticker import PercentFormatter, FuncFormatter
ax.yaxis.set_major_formatter(PercentFormatter(xmax=1))
ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:,.0f}"))

Style

plt.style.use("seaborn-v0_8-whitegrid")
plt.style.available

# A house style, set once
plt.rcParams.update({
    "figure.figsize":    (8, 5),
    "figure.dpi":        100,
    "savefig.dpi":       300,
    "savefig.bbox":      "tight",
    "font.family":       "sans-serif",
    "font.size":         11,
    "axes.titlesize":    12,
    "axes.titleweight":  "bold",
    "axes.labelsize":    11,
    "axes.spines.top":   False,
    "axes.spines.right": False,
    "axes.grid":         True,
    "grid.alpha":        0.3,
    "legend.frameon":    False,
})

TRT_COLOURS = {
    "Placebo": "#7a8698",
    "Drug A":  "#16355e",
    "Drug B":  "#0a8f9e",
}
WarningClose your figures
for param in params:
    fig, ax = plt.subplots()
    ...
    fig.savefig(f"{param}.png")
    # figure never closed

matplotlib keeps every unclosed figure in memory. A loop over 200 parameters produces a RuntimeWarning: More than 20 figures have been opened and eventually exhausts memory.

    plt.close(fig)

Or use a context manager:

from contextlib import contextmanager

@contextmanager
def figure(**kwargs):
    fig, ax = plt.subplots(**kwargs)
    try:
        yield fig, ax
    finally:
        plt.close(fig)

seaborn

Statistical plots in one call, working directly from a DataFrame.

import seaborn as sns

sns.set_theme(style="whitegrid", palette="deep", font_scale=1.0)

# Distributions
sns.histplot(data=df, x="AGE", hue="TRT01A", bins=20, kde=True)
sns.kdeplot(data=df, x="AGE", hue="TRT01A", fill=True)
sns.ecdfplot(data=df, x="AGE", hue="TRT01A")

# Categorical
sns.boxplot(data=df, x="TRT01A", y="AGE")
sns.violinplot(data=df, x="TRT01A", y="AGE", inner="quartile")
sns.stripplot(data=df, x="TRT01A", y="AGE", alpha=0.4)
sns.swarmplot(data=df, x="TRT01A", y="AGE")
sns.barplot(data=df, x="TRT01A", y="AGE", errorbar=("ci", 95))
sns.countplot(data=df, x="AGEGR1", hue="SEX")

# Relationships
sns.scatterplot(data=df, x="BASE", y="CHG", hue="TRT01A", size="AGE")
sns.regplot(data=df, x="BASE", y="CHG")
sns.lineplot(data=df, x="AVISITN", y="AVAL", hue="TRT01A", errorbar="se")

# Matrices
sns.heatmap(corr, annot=True, fmt=".2f", cmap="RdBu_r", center=0)
sns.clustermap(corr)

# Figure-level: creates its own figure, supports faceting
sns.relplot(data=df, x="AVISITN", y="AVAL", hue="TRT01A",
            col="PARAMCD", col_wrap=3, kind="line", height=3)
sns.catplot(data=df, x="TRT01A", y="AVAL", col="PARAMCD", kind="box")
sns.displot(data=df, x="AGE", col="TRT01A", kind="hist")
sns.pairplot(df[["AGE", "WEIGHT", "HEIGHT", "BMI"]], hue="SEX")
TipAxes-level versus figure-level

sns.boxplot() draws onto an existing ax — you can combine it with other matplotlib calls. sns.catplot() creates its own figure and supports col= / row= faceting, but you cannot draw onto it as easily.

fig, ax = plt.subplots()
sns.boxplot(data=df, x="TRT01A", y="AGE", ax=ax)     # axes-level
ax.axhline(65, color="red", linestyle="--")           # add to it

g = sns.catplot(data=df, x="TRT01A", y="AGE", col="SEX", kind="box")
g.map(plt.axhline, y=65, color="red", ls="--")        # figure-level

Rule of thumb: use figure-level when you want facets, axes-level otherwise.

plotly

Interactive by default.

import plotly.express as px
import plotly.graph_objects as go

fig = px.scatter(
    df, x="BASE", y="CHG",
    color="TRT01A", size="AGE",
    hover_data=["USUBJID", "AVISIT"],
    facet_col="PARAMCD", facet_col_wrap=2,
    labels={"BASE": "Baseline (U/L)", "CHG": "Change from baseline"},
    title="Change from baseline versus baseline value",
    color_discrete_map=TRT_COLOURS,
)

fig.update_layout(
    template="plotly_white",
    height=700,
    legend_title_text="",
    font=dict(family="Source Sans 3, sans-serif", size=12),
)
fig.add_hline(y=0, line_dash="dash", line_color="grey")

fig.write_html("output/figures/change.html")
fig.write_image("output/figures/change.png", scale=2)   # needs kaleido
fig.show()

plotly.express covers most needs in one call. graph_objects gives full control:

fig = go.Figure()
for arm, group in df.groupby("TRT01A"):
    fig.add_trace(go.Scatter(
        x=group["AVISITN"], y=group["mean"],
        error_y=dict(type="data", array=group["se"] * 1.96),
        mode="lines+markers", name=arm,
        line=dict(color=TRT_COLOURS[arm], width=2),
    ))
fig.update_layout(template="plotly_white",
                  xaxis_title="Study week", yaxis_title="ALT (U/L)")

The trade-off is the same as plotly in R: excellent for a few thousand points, unusable at a hundred thousand, and the HTML output embeds all the data.

plotnine

ggplot2 for Python, closely following the original API.

from plotnine import (
    ggplot, aes, geom_point, geom_line, geom_smooth, geom_hline,
    facet_wrap, labs, theme_minimal, theme, scale_colour_manual,
    element_text,
)

(
    ggplot(df, aes("AVISITN", "AVAL", colour="TRT01A"))
    + geom_line(aes(group="USUBJID"), alpha=0.2)
    + geom_smooth(method="loess", se=True, size=1.2)
    + geom_hline(yintercept=40, linetype="dashed", colour="grey")
    + facet_wrap("~PARAMCD", scales="free_y", ncol=2)
    + scale_colour_manual(values=TRT_COLOURS)
    + labs(x="Study week", y="Value", colour="Treatment",
           title="Laboratory profiles over time")
    + theme_minimal(base_size=11)
    + theme(figure_size=(10, 7), legend_position="bottom")
)

If you already think in ggplot2, this removes the translation cost entirely. The trade-offs: slower than seaborn, a smaller extension ecosystem, and some ggplot2 features are missing. For a team moving between R and Python, the consistency is often worth it.

Saving figures

fig.savefig("figure.png", dpi=300, bbox_inches="tight")
fig.savefig("figure.pdf")                       # vector, for print
fig.savefig("figure.svg")                       # vector, for the web
fig.savefig("figure.png", transparent=True)
Format Use
PNG at 300 dpi Reports, presentations, most purposes
PDF Publication, print — vector, infinitely scalable
SVG Web, further editing
HTML Interactive plotly

bbox_inches="tight" crops whitespace and is almost always what you want.

Figures for clinical outputs

def clinical_figure(data, param, output_path):
    """Mean profile with 95% CI, in the study's house style."""
    summary = (
        data.query("PARAMCD == @param and ANL01FL == 'Y'")
        .groupby(["TRT01A", "AVISITN"], as_index=False)
        .agg(n=("AVAL", "count"), mean=("AVAL", "mean"), sd=("AVAL", "std"))
        .assign(se=lambda d: d["sd"] / np.sqrt(d["n"]))
    )

    fig, ax = plt.subplots(figsize=(9, 6))

    for arm, g in summary.groupby("TRT01A"):
        ax.errorbar(
            g["AVISITN"], g["mean"], yerr=1.96 * g["se"],
            label=arm, color=TRT_COLOURS.get(arm, "#333333"),
            marker="o", markersize=5, capsize=3, linewidth=1.4,
        )

    ax.set_xlabel("Study week")
    ax.set_ylabel(f"{param} (mean ± 95% CI)")
    ax.set_xticks(sorted(summary["AVISITN"].unique()))
    ax.legend(title="", loc="best")
    ax.spines[["top", "right"]].set_visible(False)

    fig.suptitle("Figure 14.2.1", fontweight="bold", x=0.02, ha="left")
    ax.set_title(f"Mean {param} Over Time (Safety Analysis Set)", loc="left")

    fig.text(
        0.02, 0.01,
        f"Program: figures.py    Generated: {pd.Timestamp.now():%d%b%Y %H:%M}",
        fontsize=7, color="grey",
    )

    fig.tight_layout(rect=[0, 0.03, 1, 0.97])
    fig.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close(fig)

    return summary

The provenance footer is the same discipline as in R’s TLF lesson — the figure carries the answer to “what produced this?”.

R and Python side by side

ggplot2 plotnine seaborn matplotlib
ggplot(df, aes(x, y)) identical data=df, x=, y= ax.plot(x, y)
geom_point() identical sns.scatterplot() ax.scatter()
geom_line() identical sns.lineplot() ax.plot()
geom_histogram() identical sns.histplot() ax.hist()
geom_boxplot() identical sns.boxplot() ax.boxplot()
geom_smooth() identical sns.regplot() manual
facet_wrap() identical col= in figure-level plt.subplots()
labs() identical ax.set_*() ax.set_*()
theme_minimal() identical sns.set_theme() plt.style.use()
ggsave() .save() fig.savefig() fig.savefig()

Choosing

Exploratory, statistical, one line?      -> seaborn
Full control, publication figure?         -> matplotlib
Interactive, dashboard, HTML?             -> plotly
Already think in ggplot2?                 -> plotnine
Faceting with statistical layers?         -> seaborn figure-level, or plotnine

Common mistakes

Mistake Consequence Fix
Not closing figures in a loop Memory exhaustion plt.close(fig)
plt.plot() state machine with subplots Draws on the wrong axes Object-oriented fig, ax
Saving before tight_layout() Clipped labels bbox_inches="tight"
Default dpi for print Blurry figures dpi=300
plotly with 100k points Browser hangs Aggregate, or use matplotlib
Mixing axes-level and figure-level seaborn Confusing behaviour Know which you are using
No provenance on a clinical figure Untraceable Footer with program and timestamp

Exercise 6.1 — Faceted profile plot

Build a figure showing mean value over time by treatment arm, faceted by parameter, with error bars and a reference line, in three ways: matplotlib, seaborn and plotnine. Compare them.

Show solution
import numpy as np
import pandas as pd

summary = (
    adlb.query("ANL01FL == 'Y'")
    .groupby(["PARAMCD", "TRT01A", "AVISITN"], as_index=False)
    .agg(n=("AVAL", "count"), mean=("AVAL", "mean"), sd=("AVAL", "std"))
    .assign(
        se=lambda d: d["sd"] / np.sqrt(d["n"]),
        lower=lambda d: d["mean"] - 1.96 * d["se"],
        upper=lambda d: d["mean"] + 1.96 * d["se"],
    )
)

TRT_COLOURS = {"Placebo": "#7a8698", "Drug A": "#16355e", "Drug B": "#0a8f9e"}

matplotlib — most code, most control

import matplotlib.pyplot as plt

params = sorted(summary["PARAMCD"].unique())
ncol = 2
nrow = int(np.ceil(len(params) / ncol))

fig, axes = plt.subplots(nrow, ncol, figsize=(11, 3.2 * nrow), squeeze=False)

for ax, param in zip(axes.flat, params):
    sub = summary[summary["PARAMCD"] == param]
    for arm, g in sub.groupby("TRT01A"):
        ax.errorbar(g["AVISITN"], g["mean"], yerr=1.96 * g["se"],
                    label=arm, color=TRT_COLOURS.get(arm),
                    marker="o", markersize=4, capsize=3, linewidth=1.3)
    ax.axhline(0, color="grey", linestyle="--", linewidth=0.8)
    ax.set_title(param, fontweight="bold", fontsize=10)
    ax.set_xticks(sorted(sub["AVISITN"].unique()))
    ax.spines[["top", "right"]].set_visible(False)

for ax in axes.flat[len(params):]:
    ax.set_visible(False)                    # hide unused panels

handles, labels = axes.flat[0].get_legend_handles_labels()
fig.legend(handles, labels, loc="lower center", ncol=3, frameon=False)
fig.supxlabel("Study week")
fig.supylabel("Mean value (95% CI)")
fig.tight_layout(rect=[0, 0.05, 1, 1])
fig.savefig("output/figures/profiles_mpl.png", dpi=300, bbox_inches="tight")
plt.close(fig)

seaborn — least code

import seaborn as sns

sns.set_theme(style="whitegrid", font_scale=0.95)

g = sns.relplot(
    data=adlb.query("ANL01FL == 'Y'"),
    x="AVISITN", y="AVAL", hue="TRT01A",
    col="PARAMCD", col_wrap=2,
    kind="line", errorbar=("ci", 95),
    palette=TRT_COLOURS, marker="o",
    height=3.2, aspect=1.4,
    facet_kws={"sharey": False},
)
g.set_axis_labels("Study week", "Mean value (95% CI)")
g.set_titles("{col_name}")
g.map(plt.axhline, y=0, color="grey", ls="--", lw=0.8)
g.savefig("output/figures/profiles_sns.png", dpi=300)

seaborn computes the confidence interval itself from the raw data — no separate summary step. That is convenient and also a caution: it bootstraps by default, so the CI will not exactly match a hand-computed normal-approximation interval. For a figure accompanying a table, compute the summary yourself and use errorbar=None with explicit plt.errorbar overlays.

plotnine — most familiar if you know ggplot2

from plotnine import *

p = (
    ggplot(summary, aes("AVISITN", "mean", colour="TRT01A"))
    + geom_hline(yintercept=0, linetype="dashed", colour="grey", size=0.4)
    + geom_line(size=0.8)
    + geom_point(size=2)
    + geom_errorbar(aes(ymin="lower", ymax="upper"), width=0.4, size=0.4)
    + facet_wrap("~PARAMCD", ncol=2, scales="free_y")
    + scale_colour_manual(values=TRT_COLOURS)
    + labs(x="Study week", y="Mean value (95% CI)", colour="")
    + theme_minimal(base_size=11)
    + theme(figure_size=(11, 7), legend_position="bottom",
            strip_text=element_text(weight="bold"))
)
p.save("output/figures/profiles_p9.png", dpi=300)

Comparison

matplotlib seaborn plotnine
Lines of code ~25 ~15 ~18
Faceting Manual loop col=, col_wrap= facet_wrap()
Free y-scales Automatic per axes facet_kws={"sharey": False} scales="free_y"
Error bars Explicit yerr Computed from raw data geom_errorbar
Unused panels Must hide manually Handled Handled
Legend placement Manual Automatic legend_position
Control Total Limited Good
Speed Fast Fast Slower

Which to use

  • seaborn for exploration. Fifteen lines and you can see the data.
  • plotnine if the team also writes R. The cognitive saving of one grammar across both languages is larger than any of the technical differences.
  • matplotlib for a figure going into a CSR, where you need exact control of every element — tick positions, footnote placement, panel sizes — and where the figure will be regenerated identically for years.
The matplotlib version’s for ax in axes.flat[len(params):]: ax.set_visible(False) line is a good illustration of the trade: with 5 parameters in a 2-column grid, the sixth panel is empty, and only the manual version makes you deal with it. seaborn and plotnine handle it; matplotlib makes you notice it.

Recap

  • Use matplotlib’s object-oriented interface: fig, ax = plt.subplots()
  • Always plt.close(fig) in a loop
  • seaborn axes-level draws onto an ax; figure-level creates its own and facets
  • plotly is interactive but does not scale past a few thousand points
  • plotnine removes the translation cost if you already know ggplot2
  • dpi=300 and bbox_inches="tight" for anything going into a document
  • Put a provenance footer on clinical figures

Next: File formats and dataset comparison.

Back to top