"""
RMSF-per-residue plot.
Assumes RMSF has already been computed elsewhere (e.g. via MDAnalysis'
`rms.RMSF`) and lives as one row per residue (optionally per
system/replica), not one row per frame -- RMSF is already a
time-averaged-per-residue quantity by the time it reaches this
function. See the docstring for the exact expected shape.
"""
try:
import matplotlib.pyplot as plt
import seaborn as sns
except ImportError: # pragma: no cover
plt = None
sns = None
try:
from ._shading import generate_group_shades
except ImportError: # pragma: no cover - fallback for direct source usage
from _shading import generate_group_shades
[docs]
def plot_rmsf(
df,
residue_col,
rmsf_col,
sim_name_col=None,
replica_col=None,
kind="line",
hue="sim_name",
palette="tab10",
ax=None,
**kwargs,
):
"""
Plot RMSF as a function of residue number/id.
Parameters
----------
df : pandas.DataFrame
One row PER RESIDUE (per system/replica if applicable) --
NOT per frame. Expected columns:
- `residue_col`: residue number or id (x-axis). Should be
numeric or at least sortable; the plot is sorted by this
column before drawing.
- `rmsf_col`: precomputed RMSF value for that residue
(y-axis), typically in Angstrom or nm.
- `sim_name_col` / `replica_col` (optional): identify system/
replica, used for coloring/averaging depending on `hue`.
residue_col : str
Column with residue number/id.
rmsf_col : str
Column with the precomputed RMSF value.
sim_name_col : str or None, default None
Column identifying the system; used for coloring. If None, all
rows are drawn as a single series.
replica_col : str or None, default None
Column identifying the replica.
kind : {"line", "bar"}, default "line"
"line" draws a curve per group. "bar" draws grouped bars --
more readable for a small number of residues, cluttered for a
full protein. `hue="replica_shaded"` (individual per-replica
lines) is only supported for `kind="line"`.
hue : {"sim_name", "replica_shaded", None}, default "sim_name"
What to color by:
- "sim_name" (default): one curve per system. If `replica_col`
is also given, replicas are AVERAGED per (system, residue)
into that one curve, with the across-replica std shown as a
shaded band (`kind="line"` only) -- this is a summary view,
individual replicas are not separately visible.
- "replica_shaded": one line per individual (system, replica)
pair, NOT averaged -- a per-system base color shaded
lighter/darker per replica (same idea as `plot_vonmises`'s
per-replica mode and `plot_rmsd`'s `hue="replica_shaded"`),
so replicas of the same system are visually grouped by hue
family while remaining individually distinguishable. This is
what you want for "RMSF per replica, same color family per
system". Requires both `sim_name_col` and `replica_col`.
- None: no hue, all rows drawn as a single series (with
`sim_name_col`/`replica_col` ignored for coloring).
palette : str or list, optional
Seaborn palette (name) or explicit color list. For
`hue="replica_shaded"` it's the *base* palette (one color per
system, then shaded per replica). Default "tab10".
ax : matplotlib.axes.Axes or None, default None
Axes to draw on. A new figure/axes is created if not given.
**kwargs
Forwarded to `seaborn.lineplot` / `seaborn.barplot`.
Returns
-------
matplotlib.axes.Axes
Always returned -- grab it to customize the plot further, e.g.
`ax.set_ylim(...)`, `ax.figure.savefig(...)`.
Examples
--------
RMSF per system, averaged over replicas with a std band::
ax = plot_rmsf(df, "residue", "rmsf", sim_name_col="sim_name",
replica_col="replica") # hue="sim_name" default
RMSF per individual replica, same color family per system::
ax = plot_rmsf(df, "residue", "rmsf", sim_name_col="sim_name",
replica_col="replica", hue="replica_shaded")
"""
if plt is None or sns is None:
raise ImportError("plot_rmsf requires matplotlib and seaborn to be installed.")
if ax is None:
_, ax = plt.subplots(figsize=kwargs.pop("figsize", (10, 4)))
plot_df = df.sort_values(residue_col)
hue_col = None
errorbar = None
if hue == "sim_name" and sim_name_col is not None:
hue_col = sim_name_col
errorbar = "sd" if replica_col is not None else None
elif hue == "replica_shaded":
if kind != "line":
raise ValueError("hue='replica_shaded' is only supported for kind='line'.")
if sim_name_col is None or replica_col is None:
raise ValueError(
"hue='replica_shaded' requires both sim_name_col and replica_col."
)
plot_df = plot_df.copy()
hue_col = "_sim_replica"
plot_df[hue_col] = (
plot_df[sim_name_col].astype(str) + " - " + plot_df[replica_col].astype(str)
)
sim_of_replica = list(
dict.fromkeys(zip(plot_df[sim_name_col], plot_df[replica_col]))
)
_, replica_colors = generate_group_shades(sim_of_replica, base_palette=palette)
palette = {f"{s} - {r}": color for (s, r), color in replica_colors.items()}
if kind == "line":
# `errorbar="sd"` gives the across-replica std band automatically
# when averaging (hue="sim_name" with replica_col given); no band
# for "replica_shaded" (each line is already an individual
# replica, nothing to average).
sns.lineplot(
data=plot_df,
x=residue_col,
y=rmsf_col,
hue=hue_col,
palette=palette if hue_col is not None else None,
errorbar=errorbar,
ax=ax,
**kwargs,
)
elif kind == "bar":
sns.barplot(
data=plot_df,
x=residue_col,
y=rmsf_col,
hue=hue_col,
palette=palette if hue_col is not None else None,
errorbar=errorbar,
ax=ax,
**kwargs,
)
# Residue-number x-axes get unreadable fast with bar ticks;
# thin them out.
xticklabels = ax.get_xticklabels()
step = max(1, len(xticklabels) // 30)
for i, label in enumerate(xticklabels):
if i % step != 0:
label.set_visible(False)
else:
raise ValueError(f"kind must be 'line' or 'bar', got {kind!r}")
ax.set_xlabel(residue_col)
ax.set_ylabel(rmsf_col)
return ax