This approach allows genes, metabolites, drugs and other vertices to be connected based on shared network topology. However, to date I’ve only discussed edge prediction using a dot-product head, where a vertex-pair’s edge support is a direct readout of their similarity in embedding space (𝐚 · 𝐛). While surprisingly powerful, this head has limitations when vertices are heterogeneous or interact in qualitatively different ways — particularly when we want to distinguish between activation and inhibition.

Here, I explore more expressive approaches for learning mappings between A → B by evaluating both general edge prediction heads (like MLPs) and “relation-aware” heads that can learn distinct mappings for different edge types. The post will cover:

• Data model and training changes enabling relation-specific predictions
• Geometric analysis revealing how relation-aware heads encode regulatory semantics
• PerturbSeq validation demonstrating successful prediction of signed regulatory interactions
• Pre-trained models available on HuggingFace

Edge prediction is a powerful approach for predicting regulatory relationships between molecular species, but not all regulatory relations are equivalent. They vary both in how molecules interact (physically, functionally, mechanistically) and in the consequences of these interactions (activation, inhibition, ambiguous effects, or no effect). While the edge encoder partially captures this information to weight message passing, the ultimate prediction is a single continuous score representing edge likelihood — without distinguishing the type of interaction.

Ideally, I want models that can predict not just whether an interaction occurs, but how it occurs. This led me to relation-aware approaches. Relations are commonly discussed in the context of knowledge graphs, where qualitatively different vertex types are connected by different relationship types. For example, embedding the Open Targets knowledge graph organizes genes and phenotypes in a common manifold while also connecting drugs, chemical probes, and other entity types. Learning relation-aware edges defines specific transformations that map between distinct regions of the embedding space.

However, I will be focusing on learning relation-types within a largely homogeneous vertex set — primarily genes and metabolites that can serve multiple regulatory roles. This presents a greater challenge for relation-aware methods, as vertices cannot be cleanly separated by type, and the same molecule may act as both activator and inhibitor in different contexts.

Workflow updates supporting relation prediction

Building robust relation-aware models required improvements to both the Napistu-Torch training framework and the underlying data model.

General framework improvements:

• Hugging Face integration for reproducible datasets and sharing pre-trained models
• Training enhancements including Weights & Biases sweep support and resumable training
• Transfer learning capabilities for loading pre-trained encoders and fine-tuning models

Relation-specific data model changes:

• Reaction vertex removal to enable direct edge prediction between molecular species
• Relation-type labels derived from source and target Systems Biology Ontology (SBO) role annotations

Restructuring the NapistuGraph - no more reaction vertices

Earlier versions of Napistu included both molecular species (proteins, metabolites, drugs, etc.) and reaction vertices. For complex regulatory mechanisms like enzymatic reactions, this provided a clear functional description anchoring pairwise molecular interactions. This was also useful for network visualization, since reactions from common sources — particularly the many narrowly-scoped Reactome pathways — provide ideal labels for network neighborhoods.

However, including reactions has a major downside — individual edges lose their meaning. For example, an enzyme transforming A → B would be encoded as two separate edges (A → R, and R → B). For many purposes this is fine, but for edge prediction it adds more noise than signal. For the model to learn what an A → R → B reaction represents, it would need to encode B’s embedding within the R embedding. This is both computationally difficult and conceptually unnecessary, so for GNNs I’m moving to direct A → B connections.

Predicting direct connections introduces a wrinkle; I had previously enforced that no more than one edge could connect an A-B pair. Moving to a reaction-less graph, I relaxed this constraint so multiple edges can now connect the same vertex pair. This allows activating, inhibitory, and interaction edges to simultaneously exist — relationships that would have previously been distinguished by their intermediate reaction vertices.

Adding relation_type to NapistuData

When constructing the network graph, I encode each mechanism as a series of pairwise interactions, with each participant assigned a role from the SBO controlled vocabulary. SBO terms — like interactor, stimulator, inhibitor, modifier, and modified — capture the distinct ways molecules participate in regulatory mechanisms. To create relation_type labels for edges, I constructed composite labels by combining each edge’s source and target SBO terms, such as “catalyst → reactant,” “interactor → interactor,” and “stimulator → modified.”

NapistuData (a subclass of PyG’s Data), supports relations through two optional attributes: relation_type and its associated relation_manager (for tracking label metadata). Creating these relation types is elegantly handled as an extension of the existing “edge strata” functionality, which organizes edges based on vertex and/or species type to create hard negative samples.

Fitting relation-(un)aware models

During standard edge prediction, we score a possible edge based on the source and target vertices’ embeddings. For relation prediction, we additionally provide heads with a relation_type integer to distinguish different types of relations. To evaluate different head architectures, I trained a range of relation-aware heads alongside simpler relation-unaware baselines.

To enable fair and efficient comparison across heads, I first trained a 128-dimensional GraphConv message passing encoder with a 32-dimensional edge encoder and a simple dot-product head. I deployed this pre-trained model to Hugging Face, then initialized each head of interest with the pre-trained encoder weights.

To leverage this pretraining, I made three key changes to the training regime:

• Lowered the learning rate substantially from 0.003 (original dot-product head) to 0.0005 (transfer learning experiments)
• Used the one-cycle scheduler to gradually ramp up the learning rate
• Initialized expressive heads with init_as_identity settings (when appropriate) so they started from a similar state as the pre-trained dot-product head

To address the imbalanced distribution of relation types in the training data, I applied relation-weighting to each head’s loss function (binary cross-entropy for most heads and a margin-based loss for TransE and RotatE). Each relation type’s loss contribution is weighted by 1/√(relation-type count), down-weighting abundant relation types (like “interactor → interactor”) while emphasizing rare but biologically important ones (like “inhibitor → modified”). This ensures that the models learn to predict all relation types effectively, rather than primarily optimizing for the most common edges.

I fitted all models using model-specific configs and the Napistu-Torch CLI.

Reproducing this analysis

This analysis is fully reproducible — all code, data, and model configurations are provided so you can run the complete workflow on your own machine.

Environment setup:

1. Install uv (or use pip if preferred).

2. Set up a Python environment:

uv venv --python 3.11
source .venv/bin/activate
# Core dependencies
uv pip install torch==2.8.0
uv pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.8.0+cpu.html
uv pip install napistu==0.8.5
# pin wandb to 0.22.x for compatibility
uv pip install wandb==0.22.3 
uv pip install "napistu-torch[pyg,lightning,analysis]==0.3.6"
# For rendering the notebook
uv pip install ipykernel nbformat nbclient
python -m ipykernel install --user --name=blog-staging

1. Download the relation_prediction.qmd notebook (or copy relevant code blocks).

2. Choose your path:

   • Using pre-trained models (recommended): The notebook will load the models from Hugging Face on-the-fly.
   • Training from scratch: Download the model configs and training shell script to train models yourself.

3. Configure WORKING_DIR in the following code block to point to your working directory.

Configuration and imports

# standard library imports
from itertools import combinations
import os
from pathlib import Path
import textwrap

# 3rd party imports
from torch import abs
import torch.nn.functional as F
from matplotlib import pyplot as plt
from napistu.ingestion.perturbseq import (
    assign_predicted_direction,
    load_harmonizome_perturbseq_datasets,
    _get_distinct_harmonizome_perturbseq_interactions,
)
from napistu.ingestion.constants import (
    SIGNED_PERTURBATION_TYPES,
    STRONG_ORDERED_SIGNED_PERTURBSEQ_DIRECTIONS,
)
import napistu.utils as napistu_utils
import numpy as np
import pandas as pd

# import a couple of functions used by just the posted version of the blog
# pip install git+https://github.com/shackett/shackett-utils.git
from shackett_utils.utils import pd_utils
from shackett_utils.blog.html_utils import display_tabulator

# napistu-torch imports
from napistu_torch.evaluation.manager import RemoteEvaluationManager
from napistu_torch.visualization.basic_metrics import plot_auc_only, _extract_metric
from napistu_torch.visualization.advanced_metrics import plot_combined_grouped_barplot
from napistu_torch.evaluation.relation_prediction import (
    calculate_relation_type_confusion_and_correlation,
    compare_relation_type_predictions_to_perturbseq_truth,
    get_perturbseq_edgelist_tensor,
    summarize_relation_type_aucs,
)
from napistu_torch.models.constants import HEAD_DESCRIPTIONS
from napistu_torch.utils.tensor_utils import (
    compute_correlation_matrix,
    compute_effective_dimensionality,
    compute_spearman_correlation_torch,
)
from napistu_torch.visualization.heatmaps import plot_heatmap

WORKING_DIR = Path(os.path.expanduser("~/Desktop/relation_prediction_experiments"))
PATH_TO_NAPISTU_STORE = WORKING_DIR / ".store"

MODEL_DISPLAY_ORDER = [
    "dot_product",
    "mlp",
    "attention",
    "distmult",
    "rotate",
    "transe",
    "relation_attention",
    "relation_gated_mlp",
    "relation_attention_mlp",
]

MODEL_HF_REPOSITORIES : dict[str, tuple[str, str]] = {
    "dot_product" : ("seanhacks/relation_prediction_dotprod_128e", "20251229"),
    "mlp" : ("seanhacks/relation_prediction_mlp_128e", "20251229"),
    "attention" : ("seanhacks/relation_prediction_attention_128e", "20251229"),
    "distmult" : ("seanhacks/relation_prediction_distmult_128e", "20251229"),
    "rotate" : ("seanhacks/relation_prediction_rotate_128e", "20251229"),
    "transe" : ("seanhacks/relation_prediction_transe_128e", "20251229"),
    "relation_attention" : ("seanhacks/relation_prediction_relationattention_128e", "20251229-2"),
    "relation_gated_mlp" : ("seanhacks/relation_prediction_relationgatedmlp_128e", "20251229"),
    "relation_attention_mlp" : ("seanhacks/relation_prediction_relationattnmlp_128e", "20251229"),
}

RELATION_AWARE_FOCUSED_HEADS = [
    "distmult",
    "transe",
    "relation_gated_mlp",
    "relation_attention_mlp"
]

PERTURBSEQ_RELATION_TYPES = ["inhibitor -> modified", "stimulator -> modified"]

# local caches
LOCAL_HARMONIZOME_DATA_DIR = "/tmp/harmonizome_data"
CROSS_RELATION_PREDICTION_CACHE = "/tmp/cross_relation_prediction_matrices.pkl"

Comparing relation-(un)aware models

To compare the trained models, I will load their checkpoints and evaluation metrics using Napistu-Torch’s RemoteEvaluationManager, which provides a unified interface for accessing model weights, training configs, and Weights & Biases summaries directly from Hugging Face. (If you are working with local models, you can instead use the similar LocalEvaluationManager, which directly interacts with Weights & Biases and local models and data.)

eval_managers = dict()
for model_name, model_info in MODEL_HF_REPOSITORIES.items():
    model_repo, model_version = model_info
    eval_managers[model_name] = RemoteEvaluationManager.from_huggingface(
        model_repo,
        data_store_dir = PATH_TO_NAPISTU_STORE,
        revision = model_version,
    )

# for local evaluation, instead do this:
# from napistu_torch.evaluation.manager import LocalEvaluationManager
# EXPERIMENT
# eval_managers = dict()
# for experiment in MODEL_DISPLAY_ORDER:
#     experiment_path = <<PATH_TO_EXPERIMENT_DIR>>
#     eval_managers[experiment] = LocalEvaluationManager(experiment_path)

# Load pre-calculated WandB summaries directly from HuggingFace
run_summaries = {exp: manager.get_run_summary() for exp, manager in eval_managers.items()}

# Load all of the trained models
models = {k : v.load_model_from_checkpoint() for k, v in eval_managers.items()}

# Count trainable parameters in each head
n_head_params = {k : sum(p.numel() for p in v.task.head.head.parameters()) for k, v in models.items()}

# connect to the NapistuDataStore and load the NapistuData instance which all models were trained on
# all of the experiments have the same value so we just need to pick an arbitrary one

napistu_data_store = eval_managers["distmult"].napistu_data_store
napistu_data = napistu_data_store.load_napistu_data("relation_prediction")
relation_types = list(napistu_data.relation_manager.label_names.values())
species_identifiers = napistu_data_store.load_pandas_df("species_identifiers")
name_to_sid_map = napistu_data_store.load_pandas_df("name_to_sid_map").reset_index()
name_to_sid_map["integer_id"] = range(len(name_to_sid_map))

if not all(name_to_sid_map["name"] == napistu_data.ng_vertex_names):
    raise ValueError("name_to_sid_map does not match napistu_data.ng_vertex_names")

Model architecture overview

I evaluated seven different head architectures spanning three categories:

• Edge prediction (relation-unaware): Simple heads that predict edge existence without distinguishing relation types. These serve as baselines to assess whether relation-aware methods provide meaningful improvements.
• Knowledge graph embedding: Methods originally developed for heterogeneous knowledge graphs (like TransE and DistMult) that learn relation-specific transformations.
• Relation prediction (expressive heads): Custom architectures that combine relation-aware gating or attention mechanisms with MLPs to learn flexible, relation-specific transformations.

# Create a summary table
model_summaries = pd.DataFrame([HEAD_DESCRIPTIONS[k] for k in MODEL_DISPLAY_ORDER])
model_summaries["category"] = model_summaries["category"].str.replace("_", " ").str.capitalize()
model_summaries["N parameters"] = [n_head_params[k] for k in MODEL_DISPLAY_ORDER]

display_tabulator(
    model_summaries.sort_values(by="N parameters", ascending=True),
    caption="Summary of all tested prediction heads",
    wrap_columns=["label", "category", "description"],
    column_widths={"description" : "50%", "N parameters" : "15%"},
    include_index = False
)

Comparing models with standard and relation-weighted AUC

To evaluate model performance, I use two complementary metrics:

Standard AUC treats all edges equally, measuring how well models discriminate real edges from random negatives across the entire graph.

Relation-weighted AUC accounts for the imbalanced distribution of relation types by:

1. Calculating AUC separately for each relation type (comparing real edges to negative samples of the same relation type)
2. Taking a weighted average of these per-relation AUCs, where weights are proportional to √N (the square root of each relation type’s frequency)

This relation-weighted metric is particularly important given the highly imbalanced distribution of relation types. Some (like “interactor → interactor”) are far more common than others (like “inhibitor → modified”). The √N weighting balances standard AUC (which would weight by N) and equal weighting (which would weight by 1), ensuring rare relation types influence the metric without dominating it. I used relation-weighted AUC for early stopping during training, prioritizing models that perform well across all relation types rather than just the most abundant ones.

# extract and reorder based on test relation-weighted AUC
test_aucs = _extract_metric(run_summaries, "test_relation_weighted_auc")
performance_order = [x for _, x in sorted(zip(test_aucs, run_summaries.keys()))]
performance_ordered_summaries = {k: run_summaries[k] for k in performance_order}
ordered_labels = [
    textwrap.fill(HEAD_DESCRIPTIONS[x]["label"], width=20)  # Adjust width as needed
    for x in performance_order
]

# Create figure with two subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

# Plot regular AUC on first axis
plot_auc_only(performance_ordered_summaries, ordered_labels, ax=ax1, title="Standard AUC")

# Plot relation-weighted AUC on second axis
plot_auc_only(
    performance_ordered_summaries,
    ordered_labels,
    test_auc_attribute="test_relation_weighted_auc",
    val_auc_attribute="val_relation_weighted_auc",
    title="Relation-Weighted AUC",
    ax=ax2
)

plt.tight_layout()
plt.show()

Performance across relation types

While the relation-weighted AUC provides an overall performance metric, examining AUC for individual relation types reveals how well each model handles specific regulatory mechanisms. Counterintuitively, “interactor → interactor” edges are among the hardest to predict, possibly because the high density of interaction edges creates competing demands on vertex positioning in the embedding space. In contrast, directed regulatory relation types like “stimulator → reactant” and “catalyst → reactant” achieve higher AUCs across most models.

relation_type_aucs = summarize_relation_type_aucs(run_summaries, relation_types)
ordered_experiments = relation_type_aucs.groupby("experiment").sum().sort_values(by="test_auc", ascending=True).index.tolist()
ordered_relation_types = relation_type_aucs.groupby("relation_type").sum().sort_values(by="test_auc", ascending=True).index.tolist()

fig, ax = plot_combined_grouped_barplot(
    relation_type_aucs,
    category_order=ordered_relation_types,
    attribute_order=ordered_experiments,
    value_vars = ["test_auc"],
    figsize = (6, 8)
)
plt.show()

Performance takeaways

Comparing models’ relation-weighted AUC and relation-level AUCs reveals several patterns:

• Expressive relation-aware MLPs achieve top performance. The relation-gated MLP and relation-attention MLP heads achieve nearly equivalent performance (~0.87 relation-weighted AUC), representing the ceiling for this architecture and training regime. Both combine relation-specific modulation with multi-layer MLPs to learn flexible, relation-specific transformations.
• DistMult achieves remarkable parameter efficiency. Despite using only ~1,400 parameters (1/500th of the top MLP heads), DistMult achieves 0.865 relation-weighted AUC, trailing the top models by less than 0.01 AUC points. DistMult learns relation-specific scalar weights for each embedding dimension—the scoring function $\text{score}(h, r, t) = \sum_i h_i \cdot r_i \cdot t_i$ means each relation type re-weights the embedding space to emphasize dimensions where related vertices show correlated (positive weights) or anti-correlated (negative weights) patterns.
• MLPs enable effective vertex attention. Both lightweight attention heads (attention and relation-attention) substantially underperform (0.833-0.844 AUC), barely exceeding dot-product performance. The key architectural difference in top-performing attention-based heads is the MLP that processes concatenated source-target embeddings before attention, creating learned edge feature representations that attention can then modulate. Raw attention over node embeddings alone lacks sufficient expressivity.
• Knowledge graph embedding methods show variable performance. RotatE underperforms relative to even the simple dot-product baseline, likely because treating the 128-dimensional embedding as 64 complex dimensions reduces the vertex representation’s expressivity. TransE performs moderately better but still lags behind custom relation-aware heads. The margin loss used by both methods may be poorly suited for edge prediction—it enforces pairwise rankings between individual positive-negative pairs rather than learning distributional differences between positive and negative edge populations (as BCE loss does). In contrast, DistMult uses BCE loss, and this likely contributes to its strong performance.
• Relation-unaware models show relation-type variation. Even the simple dot-product and MLP heads show varying performance across relation types. This likely reflects competing demands on vertex positioning—vertices involved in many dense interactions (like “interactor → interactor” edges) face more constraints in the embedding space than vertices in sparser relation types. Retraining the dot-product head with relation-weighted BCE (versus standard BCE) substantially shifts the learned embeddings, demonstrating that relation-type frequency affects vertex positioning even when relation type isn’t used as a model input.

Evaluating relation-aware models

To explore whether relation-aware models are making meaningful signed regulatory predictions, I will evaluate them using three analyses:

1. Using interpretable relation-aware knowledge graph embedding heads, I will examine the geometric representation of activation and inhibition in the learned transformations.
2. Exploring top-performing heads, I will examine the strength of regulatory predictions — are scores similar regardless of the putative relation type, or do models confidently predict relation-type-specific interactions?
3. Leveraging PerturbSeq data, I will assess whether models’ top-scoring relation-type predictions (activation vs. inhibition) align with experimentally observed transcriptional responses to genetic perturbations.

What is the geometry of activation and inhibition?

To understand how relation-aware heads encode regulatory semantics, I will examine the learned relation embeddings from three knowledge graph embedding methods: RotatE (rotation angles), TransE (translation vectors), and DistMult (dimensional scaling weights). Each method learns these transformations as model weights, allowing direct inspection of relation types’ geometric encoding.

RotatE_head = models["rotate"].task.head.head
RotatE_phases = RotatE_head.relation_emb.weight.detach().cpu()

TransE_head = models["transe"].task.head.head
TransE_vectors = TransE_head.relation_emb.weight.detach().cpu()

DistMult_head = models["distmult"].task.head.head
DistMult_scalars = DistMult_head.relation_emb.weight.detach().cpu()
DistMult_deviations = DistMult_scalars - 1.0  # Deviation from identity

# Get relation indices
regulatory_relations = {
    "activation": relation_types.index("stimulator -> modified"),
    "inhibition": relation_types.index("inhibitor -> modified"),
    "interaction": relation_types.index("interactor -> interactor"),
}

# Compute relation-type similarity matrices for each method
similarity_matrices = {}
for method_name, embeddings in [
    ("RotatE", RotatE_phases),
    ("TransE", TransE_vectors),
    ("DistMult", DistMult_deviations)
]:
    embeddings_norm = F.normalize(embeddings, p=2, dim=1)
    sim_matrix = embeddings_norm @ embeddings_norm.T
    similarity_matrices[method_name] = sim_matrix

Rather than examining all relation types globally, I will focus on three biologically meaningful questions.

Do stimulators and inhibitors have opposing transformations?

If activation and inhibition are fundamentally opposite processes, their learned embeddings should anti-correlate:

# Calculate median similarity for context
median_similarities = {}
n_relations = len(relation_types)
mask = np.triu_indices(n_relations, k=1)

for method_name in ["RotatE", "TransE", "DistMult"]:
    sim_flat = similarity_matrices[method_name][mask]
    median_similarities[method_name] = sim_flat.median().item()

stim_vs_inhib = pd.DataFrame({
    "RotatE": [
        similarity_matrices["RotatE"][regulatory_relations["activation"], regulatory_relations["inhibition"]].item(),
        median_similarities["RotatE"]
    ],
    "TransE": [
        similarity_matrices["TransE"][regulatory_relations["activation"], regulatory_relations["inhibition"]].item(),
        median_similarities["TransE"]
    ],
    "DistMult": [
        similarity_matrices["DistMult"][regulatory_relations["activation"], regulatory_relations["inhibition"]].item(),
        median_similarities["DistMult"]
    ]
}, index=["Spearman ρ: activation vs. inhibition", "Spearman ρ: median"])
stim_vs_inhib.index.name = "metric"

pd_utils.format_numeric_columns(stim_vs_inhib, inplace = True)
display_tabulator(
    stim_vs_inhib,
    caption="Reaction-type correlation summaries",
    layout = "fitDataTable",
    include_index=True
)

Activation and inhibition are not geometric opposites. All three methods show weak or positive correlation rather than the expected anti-correlation. The positive correlations suggest that “being a regulator” is more important than the direction of regulation (activation versus inhibition) in structuring these transformations — both relation types emphasize similar regulatory dimensions rather than encoding opposing effects.

How are undirected relation-types encoded?

Protein-protein interactions are undirected—they exist in the training data as both A → B and B → A edges with the same “interactor → interactor” relation type. One might expect this bidirectional structure to push the relation embedding toward identity (zero rotation, zero translation, unit scaling), which would naturally satisfy:

To test whether interactor edges learn identity-like transformations, I will extract the learned relation embeddings from each model and compare each relation type’s deviation from identity to the median deviation across all relation types.

interactor_data = {}
for method_name, embeddings in [
    ("RotatE", RotatE_phases),
    ("TransE", TransE_vectors),
    ("DistMult", DistMult_deviations)
]:
    all_magnitudes = embeddings.abs().sum(dim=1)
    median_magnitude = all_magnitudes.median().item()
    interactor_magnitude = all_magnitudes[regulatory_relations["interaction"]].item()
    interactor_data[method_name] = [interactor_magnitude / median_magnitude]

interactor_df = pd.DataFrame(
    interactor_data, 
    index=["interaction transformation norm\n÷\nmedian transformation norm"]
).round(2)
interactor_df.index.name = "metric"

display_tabulator(
    interactor_df,
    caption = "Interaction transformation magnitudes",
    layout = "fitDataTable",
    wrap_columns = {"metric" : "30%"},
    include_index=True
)

Interactor edges are not near identity. All three methods learn typical or above-median transformations for protein-protein interactions. While this seems to violate the symmetry requirement for undirected edges, the loss functions don’t actually enforce equal scores for A → B and B → A pairs. Instead, they optimize discrimination between real edges and negative samples. For TransE, edges are scored as $\text{score}(h, r, t) = -|h + r - t|$, but the margin-based loss compares these scores to negatives — meaning non-zero transformations can provide discriminative power even without maintaining symmetry.

Do the three methods agree on regulatory semantics?

If all three methods learn similar patterns for how relation types relate to each other, it would suggest they’re converging on shared regulatory semantics.

# Compare similarity matrices pairwise
n_relations = len(relation_types)
mask = np.triu_indices(n_relations, k=1)

agreement_data = {}
for method1, method2 in combinations(similarity_matrices.keys(), 2):
    sim1_flat = similarity_matrices[method1][mask]
    sim2_flat = similarity_matrices[method2][mask]
    rho = compute_spearman_correlation_torch(sim1_flat, sim2_flat, device='cpu')
    agreement_data[f"{method1} vs {method2}"] = [rho]

agreement_df = pd.DataFrame(agreement_data, index=["Spearman ρ"])
agreement_df.index.name = "metric"

pd_utils.format_numeric_columns(agreement_df, inplace = True)
display_tabulator(
    agreement_df,
    caption = "Model-to-model comparison of relation-type correlations",
    layout = "fitDataTable",
    include_index=True
)

The three methods learn different geometric patterns. The weak inter-method correlations show that RotatE, TransE, and DistMult don’t converge on a shared representation of abstract regulatory relationships, instead they learn method-specific solutions to the edge discrimination task.

Geometry summary

The geometric analysis reveals that knowledge graph embedding methods don’t naturally encode biological intuitions about regulatory relationships:

• Activation and inhibition are not geometric opposites, showing weak or positive correlations rather than anti-correlation
• Undirected edges don’t require identity transformations, with interaction edges learning typical or above-median transformation magnitudes despite being bidirectional in the training data
• Methods don’t converge on shared geometric patterns, suggesting they learn different solutions rather than discovering universal principles of abstract biological regulation

The training graph is dense and shaped by experimental ascertainment biases — different relation types are more readily detected for different subsets of vertices. Knowledge graph embedding heads struggle to cleanly separate activators from inhibitors through vertex positioning alone, instead learning transformations that discriminate real edges from negatives without capturing clear regulatory semantics. DistMult’s strong performance despite this limitation is impressive.

Are models predicting relation-specific interactions?

If heads are learning relation-type-specific transformations, edges should score highly for some relation types and poorly for others. If heads rely on vertex embeddings alone, edge scores should remain similar regardless of the assigned relation type.

To evaluate the specificity of relation-type-based scoring, I will score each test set edge under every possible relation type, then calculate the Spearman correlation between relation types’ edge score distributions. High correlations indicate a model assigns similar scores regardless of relation type, while low correlations suggest relation-specific predictions.

if not os.path.exists(CROSS_RELATION_PREDICTION_CACHE):
    cross_relation_prediction_matrices = {"confusion": {}, "correlation": {}}
    for experiment in RELATION_AWARE_FOCUSED_HEADS:
        model = models[experiment]
        cross_relation_prediction_matrices["confusion"][experiment], cross_relation_prediction_matrices["correlation"][experiment] = (
            calculate_relation_type_confusion_and_correlation(model, napistu_data, normalize="true")
        )

    napistu_utils.save_pickle(CROSS_RELATION_PREDICTION_CACHE, cross_relation_prediction_matrices)
else:
    cross_relation_prediction_matrices = napistu_utils.load_pickle(CROSS_RELATION_PREDICTION_CACHE)

fig, axes = plt.subplots(2, 2, figsize=(16, 16))
axes = axes.flatten()

for idx, head in enumerate(RELATION_AWARE_FOCUSED_HEADS):
    correlation_matrix = cross_relation_prediction_matrices["correlation"][head]
    title = textwrap.fill(HEAD_DESCRIPTIONS[head]["label"], width=30)
    plot_heatmap(
        correlation_matrix,
        row_labels=relation_types,
        title=title,
        cmap='magma',
        cbar=False,
        fmt='.2f',
        vmax=1,
        vmin=0,
        cbar_label='Spearman ρ',
        mask_upper_triangle=True,
        square=True,
        cluster='both',
        cluster_method='average',
        cluster_metric='euclidean',
        title_size=22,
        ax=axes[idx],
    )

plt.tight_layout()
plt.show()

From these cross-relation-type score correlations, several patterns emerge:

• Top-performing models show strong relation-type specificity. DistMult, relation-gated MLP, and relation-attention MLP heads all generate predictions that are highly dependent on relation type, with lower cross-relation correlations indicating distinct scoring patterns for different edge types. This relation-type specificity appears essential for achieving top performance (>0.86 relation-weighted AUC).
• High-performing heads share common structural patterns. The top three model’s relation-type score correlation structures are similar (rho ~= 0.75 DistMult-MLPs, and rho = 0.96 MLP-MLP), suggesting they are learning similar patterns. This is particularly apparent for “modifier → modified” edges, which are distinguished from other relation types. This may reflect a source-specific curation quirk that provides a strong training signal: “modifier → modified” edges arise primarily from a single source (Omnipath) when an interaction is annotated as both activation and inhibition, creating a distinctive pattern that these models learn to recognize.
• TransE shows limited relation-type differentiation. TransE assigns similar scores to edges regardless of relation type, as evidenced by uniformly high cross-relation correlations. This indicates that it relies more heavily on source and target vertex embedding similarity than on the learned relation-specific transformations, potentially incorporating relation-agnostic discriminative signals rather than capturing true regulatory semantics.

Validating signed predictions with PerturbSeq

To evaluate whether relation-aware heads can predict not just interaction existence but regulatory direction (activation vs. inhibition), I need datasets where molecular species are systematically perturbed and their directed impacts on other species are measured. While many such experiments exist, most are either already incorporated into the graph through resources like STRING and IntAct, or haven’t been aggregated at sufficient scale for validation.

PerturbSeq experiments — where genes are perturbed using CRISPR and transcriptome-wide impacts are measured — provide an ideal validation source. Large-scale datasets like Replogle et al. (2022, Cell) perturb many genes, while smaller studies investigate targeted hypotheses (e.g., human mutation knock-ins). However, the original Replogle dataset reports only Anderson-Darling q-values, not signed fold-changes. I therefore turned to Harmonizome, an ongoing effort from the Ma’ayan lab at Mount Sinai to generate and compare diverse gene-centric profiles (Diamant et al., 2025, Nucleic Acids Research). Harmonizome provides signed PerturbSeq fold-changes from both the Replogle datasets and PerturbAtlas in consistent data formats.

To evaluate model predictions against PerturbSeq data, I:

1. Mapped PerturbSeq data to graph vertices. Loaded species identifiers (systematic identifier to molecular species ID mappings) and compartmentalized species ID maps to translate PerturbSeq systematic identifiers into graph vertex IDs.
2. Processed Harmonizome PerturbSeq interactions:
   • Mapped source and target genes to vertex IDs using the adapters from step 1
   • Selected strong perturbations by comparing Harmonizome values to their thresholds
   • Inferred regulatory direction from perturbation type and fold-change direction:
     • Overexpression + upregulation → activation
     • Overexpression + downregulation → inhibition\
     • Knockout/knockdown + upregulation → inhibition (de-repression)
     • Knockout/knockdown + downregulation → activation
     • Other perturbation types (e.g., knock-ins) were excluded
3. Generated relation-type predictions. For each PerturbSeq edge (represented as source-target vertex indices), I scored both activating (“stimulator → modified”) and repressive (“inhibitor → modified”) relation types using each model. I assigned the higher-scoring relation type as the predicted regulatory direction.
4. Compared predictions to PerturbSeq ground truth. I constructed 2×2 contingency tables comparing predicted relation types (activation vs. inhibition) to observed PerturbSeq directions, calculated significance using χ² tests, and quantified the agreement between predicted and observed regulatory directions.

# map relation_types to relation_type indices so we can just score select relation_types
relation_type_indices = {k : i for i, k in enumerate(relation_types)}
focused_relation_type_indices = {k : relation_type_indices[k] for k in PERTURBSEQ_RELATION_TYPES}

# load PerturbSeq results from Harmonizome, map to vertices, and  and filter to strong inhibition and activation
distinct_harmonizome_perturbseq_interactions = (
    load_harmonizome_perturbseq_datasets(LOCAL_HARMONIZOME_DATA_DIR, species_identifiers)
    # rollup to 1 entry per dataset, source, target, and perturbation type
    .pipe(_get_distinct_harmonizome_perturbseq_interactions)
    # filter to only perturbation types where predicted direction is clear (e.g., ignore knock-ins and mutations)
    .query("perturbation_type in @SIGNED_PERTURBATION_TYPES")
    # assign predicted direction (strong inhibition, strong activation)
    .assign(perturbseq_prediction=lambda df: assign_predicted_direction(df))
    .query("perturbseq_prediction in @STRONG_ORDERED_SIGNED_PERTURBSEQ_DIRECTIONS")
)

# pull out all unique from-to species id pairs
distinct_perturbseq_pairs = (
    distinct_harmonizome_perturbseq_interactions[["perturbed_species_id", "target_species_id"]]
    .drop_duplicates()
    .reset_index(drop=True)
)

# Map from species_id to integer_ids and convert to tensor
perturbseq_edgelist_tensor = get_perturbseq_edgelist_tensor(
    distinct_perturbseq_pairs,
    name_to_sid_map
)

# create predictions for the models of interest
predicted_relation_type_vs_perturbseq_truth_predictions = {}
predicted_relation_type_vs_perturbseq_truth_pvalues = {}
for experiment in RELATION_AWARE_FOCUSED_HEADS:
    
    model = models[experiment]

    summaries = compare_relation_type_predictions_to_perturbseq_truth(
        model,
        focused_relation_type_indices,
        perturbseq_edgelist_tensor,
        napistu_data,
        distinct_perturbseq_pairs,
        distinct_harmonizome_perturbseq_interactions,
        STRONG_ORDERED_SIGNED_PERTURBSEQ_DIRECTIONS,
        PERTURBSEQ_RELATION_TYPES,
    )

    predicted_relation_type_vs_perturbseq_truth_predictions[experiment], predicted_relation_type_vs_perturbseq_truth_pvalues[experiment] = summaries

# create the heatmaps

fig, axes = plt.subplots(2, 2, figsize=(10, 10))
axes = axes.flatten()

for idx, head in enumerate(RELATION_AWARE_FOCUSED_HEADS):
    
    dat = predicted_relation_type_vs_perturbseq_truth_predictions[head]
    title = textwrap.fill(HEAD_DESCRIPTIONS[head]["label"], width=25) + "\n" + f"log10p = {predicted_relation_type_vs_perturbseq_truth_pvalues[head].round(2)}"
    row_labels = [textwrap.fill(x, width=20) for x in PERTURBSEQ_RELATION_TYPES]
    col_labels = [textwrap.fill(x, width=10) for x in STRONG_ORDERED_SIGNED_PERTURBSEQ_DIRECTIONS]
    x_label = "PerturbSeq Prediction" if idx >= 2 else None
    y_label = "Top Scoring Relation-Type" if idx % 2 == 0 else None
    
    plot_heatmap(
        dat,
        row_labels=row_labels,
        column_labels=col_labels,
        title=title,
        xlabel=x_label,
        ylabel=y_label,
        cmap='magma',
        fmt='.2f',
        vmin=0.3,
        vmax=0.7,
        cbar_label='Proportion',
        square=True,
        cbar = False,
        title_size=16,
        label_size=16,
        axis_title_size=12,
        annot_size=16,
        ax=axes[idx],
    )

plt.tight_layout()
plt.show()

Several patterns emerge from these contingency tables:

• Relation-type scores lack cross-calibration. Average scores can vary substantially between relation types within a model. This is most apparent for TransE, where inhibitory edges score systematically higher than activating edges, resulting in predominantly inhibitory predictions regardless of ground truth. Since the loss function compares real edges to negative samples within the same relation-type stratum, models have no incentive to calibrate scores across relation types.
• TransE predictions are independent of PerturbSeq ground truth. TransE’s top-scoring relation types are only weakly associated with observed regulatory direction (log10p = -2.98). There is actually a slight enrichment in the off-diagonal quadrants (top-right and bottom-left), where predictions disagree with CRISPR results.
• Top-performing relation-aware heads achieve strong agreement with ground truth. All three of the top-performing models (DistMult, the relation-gated MLP and relation-aware attention heads) show striking enrichment along the diagonal (top-left and bottom-right quadrants), indicating correct prediction of regulatory direction. The statistical significance is overwhelming (log10p < -300), with meaningful visual enrichment patterns where predicted activation/inhibition aligns with observed PerturbSeq responses.

The agreement between top-scoring relation types and CRISPR ground truth is particularly impressive given several important caveats:

• Regulatory ground truth is inherently muddy. Harmonizome’s predicted regulatory calls show limited alignment with the Anderson-Darling q-values reported in the original Replogle supplement, highlighting the fundamental difficulty of establishing definitive in vivo regulatory ground truth from experimental data.
• PerturbSeq captures both direct and indirect effects. CRISPRi perturbations measure transcriptome-wide changes that include both immediate regulatory targets and downstream cascade effects. While I’d expect models to predict direct interactions more accurately, the training data itself contains a mixture of direct and indirect interactions, and relation types may provide a means for expressing this distinction.
• Signed regulatory edges are rare in the training data. Activation (“stimulator → modified”) and inhibition (“inhibitor → modified”) edges comprise a small fraction of the graph compared to undirected physical interactions (~3% of edges). The fact that models can accurately distinguish these regulatory directions despite their relative scarcity demonstrates that these concepts are effectively encoded into the vertex and relation-type embeddings.

Published models

All of the data and models used in this analysis are available on Hugging Face.

The best-performing relation-aware models may be of particular interest to others. These are 128-dim GraphConv encoders trained for relation-stratified edge prediction with the following heads:

• DistMult head
• Relation-gated MLP head
• Relation-attention MLP head

Summary

Relation-aware graph neural networks offer a promising path forward for predicting signed regulatory interactions — a major blind spot in current virtual cell modeling efforts. While large-scale single-cell RNA-seq atlases have enabled unprecedented molecular profiling, translating these observations into predictive models of cellular regulation requires distinguishing activation from inhibition, not just identifying that interactions exist.

The results here validate both the expressive, expansive Napistu graphs and the power of mining them with graph neural networks:

• Appropriate architectures matter more than parameter count. Top-performing heads (DistMult, relation-gated MLP, relation-attention MLP) all achieve strong relation-type specificity, but through different mechanisms. DistMult accomplishes this with minimal parameters (~1,400) through dimensional weighting, while MLP-based heads use 60-90K parameters for gating or attention mechanisms. Critically, raw attention heads substantially underperform despite having 20× more parameters than DistMult, demonstrating that architectural choices trump raw model size.
• Learned relation embeddings prioritize discrimination over semantic meaning. Activation and inhibition — biologically opposing processes—produce similar rather than anti-correlated geometric transformations. Undirected edges are not encoded with symmetric transformations that would score A→B and B→A equally. The geometric patterns learned by these methods reflect statistical structure useful for edge discrimination, rather than interpretable regulatory semantics.
• PerturbSeq validation demonstrates biological grounding. Top-performing models show impressive agreement with CRISPR perturbation ground truth, correctly distinguishing activation from inhibition with overwhelming statistical significance (log10p < -300). This validation against orthogonal experimental data confirms the models have learned biologically meaningful representations of regulation.

This work opens several opportunities for refinement. Adapting loss functions to calibrate scores across relation types would enable more interpretable cross-relation comparisons and better support for predicting novel interaction types. Continued architectural innovation in encoder–decoder designs — particularly in how relation information gates or modulates vertex representations — could further improve the semantic encoding of regulatory concepts. These advances would strengthen the foundation for computational models capable of predicting not just molecular associations, but also their functional consequences in cellular systems.

This puts the Octopus in uncharted territory: large enough to capture genome-scale complexity, yet structured enough to preserve the biological interpretability that makes network analysis valuable. GNNs scale well beyond genome-scale requirements (100M+ nodes in social networks), but remain unexplored for comprehensive biological networks that integrate regulatory, metabolic, and interaction data. Bridging this gap requires infrastructure that handles both the biological complexity of multi-source networks and the engineering complexity of training GNNs at scale.

In this post, I’ll introduce Napistu-Torch — the infrastructure that finally makes this space navigable. Available from PyPI and indexed by the Napistu MCP server, Napistu-Torch provides a modular, reproducible framework for training GNNs on comprehensive biological networks. I’ll demonstrate that it’s feasible to train graph convolutional networks on the complete Octopus network using just a laptop (albeit with 2 days of training time for the full suite of models). But the real contribution is the ecosystem: the data structures, pipelines, and evaluation strategies that unlock far more sophisticated analyses.

Specifically, I’ll walk through the key components of the Napistu-Torch ecosystem:

• Data engineering: Converting NapistuGraph objects into PyTorch Geometric Data objects while preserving the Octopus network’s rich vertex and edge metadata. The NapistuDataStore manages caching and lazy loading of derived artifacts, eliminating the overhead of rebuilding datasets — you can immediately start training or evaluating models.

• Model components: Breaking down the anatomy of a GNN into its core building blocks — encoders that learn vertex representations via message passing, heads that make predictions from embeddings, and optional edge encoders that weight edges based on metadata. I’ll compare several architectures (GCN, GraphSAGE, GraphConv) with and without edge encoding.

• Training infrastructure: Leveraging PyTorch Lightning to orchestrate model training with minimal boilerplate. Configuration files define entire experiments, making it easy to reproduce results or modify architectures without touching code. The CLI supports the full train-test workflow with automatic experiment tracking via Weights & Biases.

• Self-supervised learning: Training without ground truth labels by framing the task as edge prediction. The key challenge is forcing the model to learn real biological patterns through careful negative sampling — ensuring that negative examples aren’t trivially distinguishable from true edges while remaining computationally tractable at the scale of millions of edges.

• Model interpretation: Evaluating what the models learn through three lenses: (1) vertex embeddings that capture molecular similarity and pathway membership, (2) learned edge weights that reveal what makes a high-confidence interaction, and (3) edge prediction patterns that assess whether the model is learning biological structure versus discovering topological constraints.

From Napistu to PyTorch Geometric

Graph neural networks learn representations of nodes and edges by iteratively aggregating information from local neighborhoods — a process called message passing. Unlike traditional neural networks that operate on fixed-size inputs, GNNs can handle graphs of arbitrary size and structure, making them well-suited for biological networks where connectivity patterns encode meaningful relationships. Through multiple rounds of message passing, GNNs capture increasingly complex structural patterns, from immediate neighbors to broader network motifs.

Graph neural networks in Python typically use PyTorch Geometric (PyG), a library that extends PyTorch with data structures and operations optimized for graph-structured data. PyG represents graphs using the Data class, which stores node features, edge connectivity, and optional edge attributes as PyTorch tensors — the fundamental format needed for GPU-accelerated training.

Napistu networks, however, live in a different ecosystem. A NapistuGraph (subclass of igraph.Graph) stores biological networks with rich vertex and edge metadata—species types, reaction mechanisms, database provenance, confidence scores. Training GNNs on these networks requires bridging these two worlds: preserving Napistu’s biological metadata while converting graphs into PyG’s tensor-based format.

This is where Napistu-Torch comes in. Following the same design philosophy as Napistu-Py — extending established frameworks rather than reinventing them — Napistu-Torch builds on PyG with biology-aware data structures and methods. The goal is to lean on well-established frameworks like PyG, PyTorch Lightning, and Weights & Biases so the codebase can focus on domain-specific challenges: encoding biological signals, integrating diverse metadata, and evaluating models with biologically meaningful metrics

NapistuGraph → NapistuData

A key data structure in Napistu-Torch is NapistuData, which extends PyG’s Data class to handle biological network metadata. At its core, it contains the same PyTorch tensor components that any PyG model expects:

• x: vertex attributes [# vertices × # of vertex features]
• edge_index: graph connectivity [2 × # of edges]
• edge_attr (optional): edge attributes [# edges × # of edge features]
• edge_weight (optional): edge weights [# of edges × 1]
• y (optional): node labels for supervised tasks [# of vertices × 1]

But NapistuData also tracks Napistu-specific metadata—feature encoders, vertex and edge masks for train/val/test splits, and mappings back to the original NapistuGraph identifiers.

Creating a NapistuData

Constructing a NapistuData instance involves three conceptual steps:

1. Load the network: Start with a NapistuGraph and its associated SBML_dfs database — here, the 8-source Octopus consensus network downloaded from Google Cloud Storage
2. Augment with attributes: Add relevant vertex and edge metadata as described in the Octopus network post
3. Encode as tensors: Convert attributes to torch.Tensors using sklearn-based encoders with automatic type detection (binary→passthrough, categorical→one-hot, continuous→standardization) and train/val/test splitting

In practice, you rarely construct NapistuData objects manually. Instead, the NapistuDataStore handles this process automatically — loading raw data, applying transformations, caching results, and managing related artifacts. This is what enables immediate model training without rebuild overhead. I’ll demonstrate the store-based workflow after covering environment setup.

Following along

This analysis is fully reproducible — all code, data, and model configurations are provided so you can run the complete workflow on your own machine. This section covers environment setup and file locations.

Environment setup

To reproduce this notebook:

1. Install uv (or use pip if preferred).

2. Set up a Python environment:

uv venv --python 3.11
source .venv/bin/activate
# Core dependencies
uv pip install torch==2.8.0
uv pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.8.0+cpu.html
uv pip install napistu==0.7.5
uv pip install "napistu-torch[pyg,lightning]==0.2.6"
# if you'd like to render the notebook, you'll need to install these additional dependencies
uv pip install seaborn ipykernel nbformat nbclient umap-learn
python -m ipykernel install --user --name=blog-staging

1. Download the napistu_nets.qmd notebook (or copy and paste the relevant code blocks).

2. Choose your path:

   • 4a. Using pre-trained models (recommended): Download pre-trained models and configs (~50MB) and extract to your experiments directory.
   • 4b. Training from scratch: Download model configs only to train models yourself (requires ~2 days on an M4 Max MacBook Pro, 8-12 hours per model).

3. Configure EXPERIMENTS_DIR and other paths in the env_setup code block to point to your local directories.

Configuration and imports

# imports

import logging
import os
import re
from pathlib import Path

import numpy as np
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import pandas as pd
import torch
import seaborn as sns

from napistu_torch.evaluation.edge_prediction import summarize_edge_predictions_by_strata, plot_edge_predictions_by_strata
from napistu_torch.evaluation.edge_weights import compute_edge_feature_sensitivity, format_edge_feature_sensitivity, plot_edge_feature_sensitivity
from napistu_torch.evaluation.evaluation_manager import EvaluationManager
from napistu_torch.evaluation.model_comparison import compare_embeddings
from napistu_torch.evaluation.pathways import calculate_pathway_similarities
from napistu_torch.lightning.tasks import get_edge_encoder
from napistu_torch.lightning.workflows import predict
from napistu_torch.load.gcs import gcs_model_to_store
from napistu_torch.utils.torch_utils import select_device
from napistu_torch.visualization.basic_metrics import plot_model_comparison
from napistu_torch.visualization.embeddings import layout_umap, plot_coordinates_with_masks

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("napistu_torch")

# globals

OVERWRITE = False

EXPERIMENTS_DIR = Path(os.path.expanduser("~/Desktop/EXPERIMENTS/20251106_edge_prediction"))
NAPISTU_DATA_DIR = os.path.join(EXPERIMENTS_DIR, ".napistu_data")
STORE_DIR = os.path.join(EXPERIMENTS_DIR, ".store")
CACHE_DIR = EXPERIMENTS_DIR

EXPERIMENTS = [
    # leave this as a list so it defines plot order
    "20251106_gcn_baseline",
    "20251106_gcn_edge_encoding",
    "20251106_sage_baseline",
    "20251106_graphconv_baseline",
    "20251106_graphconv_edge_encoding"
]

EXPERIMENT_LABELS = {
    "20251106_gcn_baseline" : "GCN",
    "20251106_gcn_edge_encoding" : "GCN + Edge Encoding",
    "20251106_sage_baseline" : "SAGE",
    "20251106_graphconv_baseline" : "GraphConv",
    "20251106_graphconv_edge_encoding" : "GraphConv + Edge Encoding"
}

ordered_labels = [EXPERIMENT_LABELS[exp] for exp in EXPERIMENTS]

TOP_MODEL_NAME = "GraphConv + Edge Encoding"
EMBEDDING_COMPARISONS_PATH = CACHE_DIR / "embedding_comparisons.tsv"

# validation
if not os.path.isdir(EXPERIMENTS_DIR):
    raise FileNotFoundError(f"Experiments directory not found: {EXPERIMENTS_DIR}")
if not os.path.isdir(CACHE_DIR):
    raise FileNotFoundError(f"Cache directory not found: {CACHE_DIR}")

Managing artifacts with NapistuDataStore

Training and evaluating GNN models requires more than just the graph structure — we need encoded features, train/val/test splits, pathway metadata for evaluation, and edge stratification data for negative sampling. Building these artifacts from scratch involves loading the full SBML_dfs database (several minutes) and running various preprocessing steps. Doing this repeatedly during development would be painfully slow.

The NapistuDataStore solves this by managing a registry of cached artifacts. Once built, artifacts load in seconds rather than minutes. Named artifact definitions in the napistu_torch.load.artifacts module support common workflows and integrate seamlessly with config-driven training.

Importantly, the store provides a clean abstraction layer over Napistu-Py. All the logic for loading SBML databases, decorating graphs with metadata, and extracting biological annotations is baked into Napistu-Torch objects. Users can work entirely with NapistuData, encoders, and dataloaders without ever touching Napistu-Py code — though if you’re curious about the underlying biological data model, Napistu-Py is pretty cool.

Initializing the store

A store can be initialized directly from one of the bundled Napistu networks on Google Cloud Storage using gcs_model_to_store. This creates and manages two local directories:

• napistu_data_dir: Raw data including the NapistuGraph and SBML_dfs
• store_dir: Cached artifacts and the registry file tracking what’s been built

napistu_data_store = gcs_model_to_store(
    napistu_data_dir = NAPISTU_DATA_DIR,
    store_dir = STORE_DIR,
    asset_name = "human_consensus",
    # pin to a stable version of the dataset for reproducibility
    asset_version = "20250923" 
)

Building and caching artifacts

The ensure_artifacts method checks whether requested artifacts exist and builds any that are missing. For this analysis, we need four artifacts:

napistu_data_store.ensure_artifacts([
    "edge_prediction",
    "comprehensive_pathway_memberships",
    "edge_strata_by_node_type",
    "edge_strata_by_node_species_type"
])

These artifacts are:

• edge_prediction: A NapistuData instance with train/val/test edge masks, used for self-supervised learning
• comprehensive_pathway_memberships: Detailed pathway associations for all vertices (including fine-grained Reactome pathways), used for evaluating whether embeddings capture biological organization
• edge_strata_by_node_type: Edge categories based on source/target node types (species→species, species→reaction, etc.), used for stratified negative sampling
• edge_strata_by_node_species_type: Finer-grained edge categories including species types (protein, metabolite, RNA), used for assessing prediction biases

Once built, these artifacts load almost instantly:

napistu_data = napistu_data_store.load_napistu_data("edge_prediction")
comprehensive_pathway_memberships = napistu_data_store.load_vertex_tensor("comprehensive_pathway_memberships")
edge_strata_by_node_type = napistu_data_store.load_pandas_df("edge_strata_by_node_type")
edge_strata_by_node_species_type = napistu_data_store.load_pandas_df("edge_strata_by_node_species_type")

    INFO:napistu_torch.napistu_data_store:Loading NapistuData from /Users/sean/Desktop/EXPERIMENTS/20251106_edge_prediction/.store/napistu_data/edge_prediction.pt
    INFO:napistu_torch.napistu_data_store:Loading VertexTensor from /Users/sean/Desktop/EXPERIMENTS/20251106_edge_prediction/.store/vertex_tensors/comprehensive_pathway_memberships.pt
    INFO:napistu_torch.napistu_data_store:Loading pandas DataFrame from /Users/sean/Desktop/EXPERIMENTS/20251106_edge_prediction/.store/pandas_dfs/edge_strata_by_node_type.parquet
    INFO:napistu_torch.napistu_data_store:Loading pandas DataFrame from /Users/sean/Desktop/EXPERIMENTS/20251106_edge_prediction/.store/pandas_dfs/edge_strata_by_node_species_type.parquet

The store abstraction means that downstream code (training scripts, evaluation notebooks) can simply request artifacts by name without worrying about paths, versions, or rebuild logic.

Anatomy of a GNN

Training a GNN requires coordinating several components: the model architecture itself, the task definition, data management, and training infrastructure. This section breaks down each component — starting with a conceptual overview of what it does, then showing how it’s implemented in Napistu-Torch. We’ll begin with the high-level system architecture, then work through the task definition, model components (encoder, head, edge encoder), and finally the training infrastructure that orchestrates everything.

System architecture

Training deep learning models involves coordinating several standard components: the Task defines what we’re learning (loss function, metrics), the Model implements the neural network architecture, the DataModule handles data loading and batching, and the Trainer orchestrates the optimization loop. Each component is configured independently, providing modularity and clear separation of concerns.

The following subsections examine how these components work in the context of biological network analysis, covering the task definition (edge prediction), model architecture (GNN encoders and heads), and training infrastructure. Later, in the Workflow Management section, I’ll show how Napistu-Torch packages these component-level configs into a single ExperimentConfig for experiment reproducibility.

Task: what are we trying to predict?

At the highest level, a GNN task defines the learning objective — what predictions we want the model to make and how we’ll evaluate its performance. Different tasks require different architectural components and training strategies. Common GNN tasks include node classification (predicting properties of individual nodes), graph classification (predicting properties of entire graphs), and edge prediction (predicting whether edges should exist between nodes).

Edge prediction in Napistu-Torch. For this analysis, we’re using the edge prediction task (also called link prediction). The goal is to predict whether an edge should exist between two nodes in a biological network. This is particularly valuable for discovering potential protein-protein interactions, metabolic relationships, or regulatory connections that may be missing from current databases. Crucially, edge prediction is self-supervised — it doesn’t require vertex or edge labels, which are often difficult to obtain or overly contrived for biological networks.

The training process works by teaching the model to discriminate between:

• Positive edges: Real edges that exist in the training set
• Negative edges: Node pairs sampled as non-edges (chosen to maintain biological plausibility)

The EdgePredictionTask class in Napistu-Torch orchestrates this process, handling negative sampling, computing the loss function (binary cross-entropy), and evaluating performance using metrics like AUC and average precision. The task operates in a transductive setting — the model generates node embeddings using only training edges for message passing; validation and test edges are excluded from neighborhood aggregation but used as supervision to evaluate the decoder’s edge predictions.

With the task defined, let’s examine the three core model components: the encoder that learns node embeddings, the head that produces predictions, and the optional edge encoder that can weight message passing.

Encoder: learning vertex representations

The encoder is the core of a GNN — it transforms raw node features into learned embeddings by aggregating information from each node’s local neighborhood. Through multiple layers of message passing, encoders capture increasingly complex patterns of connectivity and feature similarity. The encoder’s architecture determines how information flows through the graph and what kinds of structural patterns the model can learn.

Message passing encoders in Napistu-Torch. Napistu-Torch leverages PyG’s library of encoder architectures, wrapping them through the MessagePassingEncoder class for easy configuration:

• GraphSAGE (SAGE)
  • Samples and aggregates features from neighbors using various aggregation functions (mean, max, add)
  • Efficient and scalable, well-suited for large biological networks
  • Does not support edge weighting — treats all edges uniformly
• GraphConv
  • Similar to SAGE with a simplified message passing scheme
  • Supports optional edge weighting
• Graph Convolutional Networks (GCN)
  • Uses symmetric normalization that accounts for node degrees
  • Supports edge weighting

All encoders follow a multi-layer architecture, progressively refining node embeddings through repeated neighborhood aggregation. The framework provides a unified interface, allowing you to swap architectures through configuration files without changing code.

Head: making predictions from embeddings

Once the encoder has produced node embeddings, the head (or decoder) transforms these embeddings into task-specific predictions. The head’s role is to adapt the general-purpose embeddings produced by the encoder to the specific prediction task — edge prediction, node classification, graph classification, etc.

Heads in Napistu-Torch. For edge prediction, the most commonly used head is the dot product, which computes the inner product of source and target node embeddings — assuming that nodes with similar embeddings should be connected. This is the simplest and most efficient option, serving as a strong baseline in most GNN edge prediction work. Napistu-Torch also implements more expressive alternatives (MLP, bilinear) that can learn non-linear relationships between node pairs, though these come with increased computational cost.

For this analysis, all models use the dot product head due to its efficiency on large biological networks and strong empirical performance. Napistu-Torch also provides heads for other tasks like node classification, all accessible through configuration files.

Edge encoder: weighting message passing

While many GNN implementations either ignore edge attributes entirely or only accept a single edge weight, biological networks often have rich edge metadata. The challenge is that most message passing architectures (like SAGE) don’t support edge attributes at all, while others (like GCN and GraphConv) only accept a single scalar weight per edge.

Learned edge weighting in Napistu-Torch. Napistu-Torch addresses this through the EdgeEncoder class, which compresses multi-dimensional edge attributes into a single learned weight that modulates message passing strength. The edge encoder is a lightweight MLP that takes edge features as input and outputs a scalar weight in [0, 1] via sigmoid activation. These learned weights control how strongly each edge contributes during neighborhood aggregation.

This approach filters noisy edges and amplifies reliable ones, focusing message passing on the most informative connections — crucial in biological networks where edge quality varies widely. The encoder trains end-to-end with the rest of the model, learning which edge attributes are most predictive for the task while remaining compatible with standard GNN architectures that expect scalar edge weights.

Model overview

The diagram shows how model components connect during a forward pass. Node features and edge connectivity flow through the message passing encoder to produce node embeddings. The head then transforms these embeddings into task-specific predictions — edge scores for edge prediction, class probabilities for node classification.

The optional edge encoder (shown in dashed lines) learns to weight edges based on edge attributes, modulating how strongly each edge contributes during message passing. This is particularly useful when edge reliability varies across data sources, as in the Octopus network.

These model components (encoder, head, edge encoder) define the core prediction logic, but training a GNN requires additional infrastructure to manage data loading, optimization, and evaluation. Napistu-Torch uses PyTorch Lightning to orchestrate this training workflow.

Training infrastructure: PyTorch Lightning

While the core GNN components (encoder, head, task) are pure PyTorch, actually training a model requires substantial boilerplate: optimizer setup, learning rate scheduling, checkpoint saving, logging metrics, handling different hardware accelerators, and coordinating training/validation loops.

Lightning integration in Napistu-Torch. Napistu-Torch uses PyTorch Lightning to handle this training infrastructure automatically. Lightning separates scientific code (model architecture, loss functions) from engineering code (training loops, GPU management), making experiments more reproducible and less error-prone.

The key Lightning components are:

• LightningModule: Wraps the core task (encoder + head) and defines training/validation steps, metrics computation, and optimizer configuration. Napistu-Torch provides task-specific adapters like EdgePredictionLightning that bridge pure PyTorch implementations with Lightning’s training infrastructure.

• Trainer: Orchestrates the training loop, handles checkpointing, manages device placement (CPU/GPU/MPS), integrates with experiment tracking tools like Weights & Biases, and implements callbacks like early stopping.

With this architecture, you can concentrate on defining model and task logic, while Lightning takes care of all training mechanics.

Data management: batching strategies

The NapistuDataModule is Lightning’s interface for data loading. It can be initialized directly from an ExperimentConfig, automatically handling artifact loading from the NapistuDataStore, data validation, and dataloader creation.

Full-batch vs. mini-batch training. Napistu-Torch provides two DataModule implementations with fundamentally different training strategies:

• FullGraphDataModule: Returns the complete graph in each batch, processing all training edges simultaneously for a single gradient update per epoch. With only one update per epoch, the model can converge prematurely before exploring the optimization landscape effectively.

• EdgeBatchDataModule: Splits training edges into mini-batches while still using all training edges for message passing. Each batch computes loss and gradients on a subset of training edges but uses the full graph structure for neighborhood aggregation. This enables multiple gradient updates per epoch by subdividing the supervision signal — effectively trading fewer epochs for more updates per epoch, allowing more thorough optimization.

For the models in this post, I used EdgeBatchDataModule with 20 batches per epoch, meaning the model updates its weights 20 times per epoch rather than once.

Workflow management

For this post, I’m comparing 5 models across different encoder architectures and edge encoding strategies:

• GraphConv (+/- edge encoding)
• GCN (+/- edge encoding)
• SAGE (edge encoding not supported)

These models use identical hyperparameters (200 epochs, same batch configuration, same hidden dimensions) for a fair comparison. These models are deliberately unoptimized, with no hyperparameter tuning and simple dot-product heads, as the focus is on feasibility and infrastructure rather than peak performance.

Training configuration in Napistu-Torch. The ExperimentConfig composes lower-level Pydantic configs (DataConfig, ModelConfig, TaskConfig, TrainerConfig, WandBConfig) with validation and sensible defaults. Define an experiment in a minimal YAML file, inherit defaults automatically.

The Napistu-Torch CLI supports training and testing directly from the command line:

napistu-torch train graphconv_baseline.yaml --out-dir 20251106_graphconv_baseline
napistu-torch test 20251106_graphconv_baseline

Training/validation/test metrics log to Weights & Biases for easy comparison across experiments. Each run saves a RunManifest containing the Weights & Biases run ID and complete ExperimentConfig (with all defaults expanded), making experiments fully reproducible.

The configs and training script for these 5 models are available here. On an M4 Max MacBook Pro with 48GB of RAM, training the full suite takes ~2 days (~8-12 hours per model).

Now let’s load these trained models and compare their performance.

Model comparison

Model evaluation in Napistu-Torch. The EvaluationManager class loads a run’s RunManifest and provides methods for accessing checkpoints, the NapistuDataStore, and Weights & Biases metrics — eliminating the need for manual path management or API queries. Here, I’ll load all five trained models and compare their performance metrics and learned representations.

eval_managers = {
    EXPERIMENT_LABELS[out_dir]: EvaluationManager(EXPERIMENTS_DIR / out_dir) for out_dir in EXPERIMENTS
}

# Extract model summaries directly from Weights & Biases using their API
run_summaries = {exp: manager.get_run_summary() for exp, manager in eval_managers.items()}

# visualize model comparison
fig, (ax1, ax2) = plot_model_comparison(run_summaries, ordered_labels)

The training loss shown is the final epoch’s binary cross-entropy, aggregated across all mini-batches, computed on equal numbers of real edges (70% of the network) and negative samples. Validation AUC measures how well the model ranks held-out real edges (15% of the network, excluded from message passing) above an equal number of negative samples. This metric is evaluated after each epoch for checkpoint selection and early stopping. Test AUC evaluates the same ranking task on the final 15% of edges.

Performance differences across models are modest but consistent. Encoder architecture matters: SAGE > GraphConv >> GCN. Edge encoding provides a clear improvement across architectures that support it. While these models haven’t been optimized — no hyperparameter tuning, simple dot product heads—the consistent trends across architectures validate the training infrastructure and provide a baseline for future work.

Next, I’ll compare the learned representations across models. Specifically: Do different encoder architectures produce similar vertex embeddings? Do models with edge encoders learn comparable edge weights? These comparisons reveal whether the biological signal is robust to architectural choices or whether different models capture fundamentally different patterns.

top_model_manager = eval_managers[TOP_MODEL_NAME]
napistu_data = top_model_manager.load_napistu_data()
napistu_data_store = top_model_manager.get_store()
napistu_graph = napistu_data_store.load_napistu_graph()

# pull out the node types and create a mask to distinguish species and reactions
node_types = napistu_graph.get_vertex_series("node_type")
is_species_mask = (node_types == "species").values

## Extract model embeddings
edge_encodings = {}
species_embeddings = {}
for exp, evaluation_manager in eval_managers.items():
    # load the model and data
    model = evaluation_manager.load_model_from_checkpoint()
    napistu_data = evaluation_manager.load_napistu_data()
    
    # pull out learned edge weights (if an edge encoder is present)
    if model.task.encoder.edge_weighting_type == "learned_encoder":
        edge_encodings[exp] = model.get_learned_edge_weights(napistu_data)
    
    # pull out the vertex embeddings
    embeddings = model.get_embeddings(napistu_data)
    species_embeddings[exp] = embeddings[is_species_mask]

    # cleanup
    evaluation_manager.experiment_dict = None

Comparing vertex embeddings

To compare vertex embeddings across models, I compute species-species cosine similarities within each model’s embedding space (# species × hidden dimension), then calculate the Spearman correlation between these similarity matrices across model pairs. This approach works regardless of embedding dimension and avoids the need for explicit alignment (e.g., Procrustes rotation).

def create_correlation_heatmap(
    embedding_comparisons: pd.DataFrame,
    model_order: list[str] | None = None
) -> pd.DataFrame:
    
    # Get all unique models
    all_models = (
        set(embedding_comparisons['model1'].unique()) | \
        set(embedding_comparisons['model2'].unique())
    )
    
    # Use provided order or default to sorted
    if model_order is None:
        models = sorted(all_models)
    else:
        # Validate that all models in data are in the provided order
        missing_models = all_models - set(model_order)
        if missing_models:
            raise ValueError(f"Models in data but not in model_order: {missing_models}")
        # Use only models that exist in the data, in the specified order
        models = [m for m in model_order if m in all_models]
    
    # Initialize matrix with 1s on diagonal
    corr_matrix = pd.DataFrame(
        np.eye(len(models)),
        index=models, 
        columns=models
    )
    
    # Fill in the correlations (both upper and lower triangles)
    for _, row in embedding_comparisons.iterrows():
        corr_matrix.loc[row['model1'], row['model2']] = row['spearman_rho']
        corr_matrix.loc[row['model2'], row['model1']] = row['spearman_rho']
    
    return corr_matrix

# compute the embedding comparisons (cached since this takes a few minutes to run)

if os.path.isfile(EMBEDDING_COMPARISONS_PATH) and not OVERWRITE:
    embedding_comparisons = pd.read_csv(EMBEDDING_COMPARISONS_PATH, sep="\t", index_col=0)
else:
    device = select_device(mps_valid = True)
    embedding_comparisons = compare_embeddings(species_embeddings, device)
    embedding_comparisons.to_csv(EMBEDDING_COMPARISONS_PATH, sep="\t", index_col=False)

# visualize the embedding comparisons

corr_matrix = create_correlation_heatmap(embedding_comparisons, model_order=ordered_labels)

# Display as heatmap
mask = np.triu(np.ones_like(corr_matrix, dtype=bool), k=1)

plt.figure(figsize=(10, 8))
sns.heatmap(
    corr_matrix,
    annot=True,
    fmt='.3f',
    cmap='RdYlBu_r',
    mask=mask,
    vmin=0,
    vmax=1,
    square=True,
    cbar_kws={'label': 'Spearman ρ'},
)
plt.title(
    'Vertex embedding similarity across models',
    fontsize=15,
    fontweight='bold',
    pad=20,
    loc='left'
)
plt.tight_layout()
plt.show()

All models produce highly correlated embeddings (ρ > 0.6 across all pairs), indicating they’ve captured similar biological structure despite architectural differences. However, encoder choice does matter: GCN embeddings correlate less strongly with GraphConv/SAGE (ρ ≈ 0.6-0.7) than GraphConv and SAGE correlate with each other (ρ ≈ 0.9). This suggests that while all models learn similar biological signals, GCN’s symmetric normalization produces somewhat different vertex representations than the mean aggregation used by GraphConv and SAGE.

Comparing learned edge weights

As I’ve previously discussed, confidence in regulatory interactions varies greatly across data sources, and edge weights should capture this uncertainty for downstream network analysis. However, determining appropriate edge weights is challenging when multiple data source attributes each capture different aspects of reliability.

The edge encoder provides a way to learn what makes a high-confidence edge empirically. While I’ll examine the learned edge features in detail later, here I’ll assess whether edge weights are consistent across model architectures by directly comparing the ~10M learned weights from GCN + edge encoding and GraphConv + edge encoding.

Since visualizing 10M points requires aggregation, I’ll use a hexbin plot (bivariate histogram) with logit-transformed weights to map the sigmoid outputs back to ℝ:

def safe_logit(p: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
    p = torch.clamp(p, eps, 1 - eps)
    return torch.log(p / (1 - p))

def plot_edge_encoding_hexbin(
    tensor1: torch.Tensor,
    tensor2: torch.Tensor,
    label1: str = "Model 1",
    label2: str = "Model 2",
    transform_to_logit: bool = True,
    gridsize: int = 50,
    cmap: str = 'viridis',
    figsize: tuple = (10, 8),
) -> tuple:
    # Apply logit transformation if requested
    if transform_to_logit:
        tensor1 = safe_logit(tensor1)
        tensor2 = safe_logit(tensor2)
    
    # Convert to numpy
    if torch.is_tensor(tensor1):
        x = tensor1.cpu().numpy()
    else:
        x = np.array(tensor1)
    
    if torch.is_tensor(tensor2):
        y = tensor2.cpu().numpy()
    else:
        y = np.array(tensor2)
    
    # Create figure
    fig, ax = plt.subplots(figsize=figsize)
    
    # Create hexbin plot with log-scaled color
    hexbin = ax.hexbin(
        x, y,
        gridsize=gridsize,
        cmap=cmap,
        mincnt=1,
        norm=LogNorm()
    )
    
    # Add colorbar
    cb = plt.colorbar(hexbin, ax=ax)
    cb.set_label('Count (log scale)', fontsize=12, fontweight='bold')
    
    # Labels and title
    ax.set_xlabel(label1, fontsize=13, fontweight='bold')
    ax.set_ylabel(label2, fontsize=13, fontweight='bold')
    ax.set_title(
        'Learned edge weight similarity across models',
        fontsize=15,
        fontweight='bold',
        pad=20,
        loc='left'
    )
    
    # Add diagonal reference line (y=x)
    min_val = min(x.min(), y.min())
    max_val = max(x.max(), y.max())
    ax.plot(
        [min_val, max_val],
        [min_val, max_val],
        'r--',
        alpha=0.5,
        linewidth=2,
        label='y=x'
    )
    ax.legend(loc='upper left', fontsize=11)
    
    # Equal aspect ratio
    ax.set_aspect('equal', adjustable='box')
    
    plt.tight_layout()
    
    return fig, ax

fig, ax = plot_edge_encoding_hexbin(
    edge_encodings["GCN + Edge Encoding"],
    edge_encodings["GraphConv + Edge Encoding"],
    label1="GCN + Edge Encoding",
    label2="GraphConv + Edge Encoding",
    transform_to_logit=True,
    gridsize=50,
    cmap='viridis'
)
plt.show()

The edge encoder paired with GraphConv uses a wider dynamic range (logit -5 to 7, corresponding to sigmoid weights 0.007-0.999) compared to GCN + edge encoder (logit -3 to 0, corresponding to 0.05-0.5). This means GraphConv’s edge encoder more thoroughly distinguishes between low-, medium-, and high-confidence edges. Both model-encoder combinations agree on which edges to effectively ignore (lower-left quadrant, sigmoid weights < 0.05), but differ in how strongly they upweight reliable edges.

Having established these cross-model comparisons, I’ll now examine the top-performing model (GraphConv + edge encoding) in detail to understand what biological patterns it has captured and where its limitations lie.

Evaluating the top model

Having compared models on performance and learned representations, I’ll now examine what the best-performing model (GraphConv + edge encoding) has actually learned. GNN-based edge prediction offers three potential contributions beyond predicting missing edges:

1. Vertex embeddings capture molecular similarity - Embeddings group similar vertices across entity types, enabling community detection and direct similarity queries. Community detection could identify functional modules — sets of proteins, metabolites, and reactions that cluster together in the embedding space. Similarity queries could assess how closely any two entities resemble each other, even across different entity types.

2. Learned edge weights for network analysis - The edge encoder learns weights that reflect edge reliability, potentially replacing hand-crafted heuristics in downstream analyses like network layouts, shortest paths, propagation algorithms, and shallow embedding methods. For multi-source networks like the Octopus, this is particularly valuable, rather than manually deciding how to weight STRING coexpression scores versus IntAct citation counts, the model learns what combinations of edge attributes indicate reliability.

3. Edge predictions for hypothesis generation - While self-supervised training is the primary motivation for edge prediction, the predictions themselves may identify plausible regulatory connections absent from current databases. This becomes more promising with expressive heads that can model directional regulation rather than the symmetric similarity assumed by the dot product.

I’ll explore each of these potential contributions using the GraphConv + edge encoding model.

Molecular similarity

Embedding structure

To assess the structure of the vertex embeddings, I’ll use UMAP to project the 128-dimensional embeddings into 2D, then overlay vertex attributes to explore what determines similarity in the embedding space.

embeddings = species_embeddings[TOP_MODEL_NAME]
umap_layout_species = layout_umap(embeddings, n_neighbors=20)

mask = [bool(re.search("__species_type", x)) for x in napistu_data.get_vertex_feature_names()]
indices = [i for i, m in enumerate(mask) if m]
masks = napistu_data.x[:, indices]

# only look at the vertices with embedding values
masks = masks[is_species_mask]
mask_names = [x for x, m in zip(napistu_data.get_vertex_feature_names(), mask) if m]

# drop empty masks
empty_masks = masks.sum(axis=0) == 0
masks = masks[:, ~empty_masks]
mask_names = [x for x, m in zip(mask_names, ~empty_masks) if m]

fig, axes = plot_coordinates_with_masks(
    coordinates=umap_layout_species,
    masks=masks,
    mask_names=mask_names,
    figsize=(10, 15),
    ncols=2,
    cmap_bg='lightblue',
    cmap_fg='darkred',
    alpha=0.5,
    s=5
)
plt.show()

The UMAP visualization shows clear clustering by entity type: proteins cluster with proteins, and metabolites with metabolites. This isn’t surprising given the network’s strong homophily — entities of the same type preferentially connect. STRING alone contributes >80% of the edges in the network, and STRING edges are exclusively protein-protein interactions.

However, entity types don’t completely segregate. Proteins and metabolites intermix at cluster boundaries, and the embedding shows finer-grained structure within each entity type. This suggests the GNN captures more than just entity type — it’s learning the biological organization within these categories. To reveal what additional information the embedding encodes, I will analyze how pathway membership and data source annotations are reflected in the representations

Pathway similarity

To assess pathway organization in the embeddings, I’ll use the comprehensive pathway membership artifact created earlier. This binary tensor encodes both coarse-grained data sources (8 sources) and fine-grained pathway annotations (2800+ Reactome pathways) for each vertex.

For each source and pathway, I’ll calculate the average cosine similarity between all vertex pairs that belong to that category. The fine-grained Reactome pathways are then aggregated to assess whether individual pathways produce tighter clusters than Reactome as a whole.

# not really needed, but we can check artifacts alignments to the canonical vertex and edge feature names from the NapistuData instance
pathway_assignments = comprehensive_pathway_memberships.align_to_napistu_data(napistu_data, inplace=False).data

pathway_similarities = calculate_pathway_similarities(
    embedding_matrix = embeddings,
    pathway_assignments = pathway_assignments[is_species_mask],
    pathway_names = comprehensive_pathway_memberships.feature_names,
)
# rename categories for clarity
pathway_similarities["Reactome (by pathway)"] = pathway_similarities.pop("other")
pathway_similarities["Reactome (overall)"] = pathway_similarities.pop("Reactome")
del pathway_similarities["overall"]

# Sort by value
sorted_items = sorted(pathway_similarities.items(), key=lambda x: x[1])
categories = [item[0] for item in sorted_items]
values = [item[1] for item in sorted_items]

# Create figure and axis
fig, ax = plt.subplots(figsize=(10, 6))

# Create barplot
bars = ax.barh(categories, values, color='steelblue', edgecolor='black', linewidth=0.5)

# Customize the plot
ax.set_xlabel('Within-category cosine similarity', fontsize=12, fontweight='bold')
ax.set_ylabel('Data source', fontsize=12, fontweight='bold')
ax.set_title(
    'Within-category cosine similarity by data source', 
    fontsize=14,
    fontweight='bold',
    pad=20,
    loc='left'
)

# Add value labels on the bars
for i, (cat, val) in enumerate(zip(categories, values)):
    ax.text(val + 0.01, i, f'{val:.3f}', 
            va='center', fontsize=9)

# Add grid for easier reading
ax.grid(axis='x', alpha=0.3, linestyle='--')
ax.set_axisbelow(True)
plt.tight_layout()
plt.show()

The embedding structure reflects the network’s edge composition. STRING contributes >80% of edges—all protein-protein interactions — which means the training objective is dominated by getting protein relationships right. This pushes the model to spread proteins across the embedding space to capture their diverse interaction patterns, resulting in low within-source similarity for protein-rich databases: STRING (0.061), IntAct (0.046), OmniPath (0.044).

In contrast, specialized sources with fewer, lower-degree entities get pushed into tighter regions of the embedding space. Reactome (0.550 overall, 0.586 by pathway) and Recon3D (0.531) show much higher within-source similarity. These sources contribute distinctive entity types — complexes and proteoforms for Reactome, and detailed metabolic species for Recon3D. The model learns to distinguish these entities from standard proteins. However, because they contribute few edges to the training signal, the model clusters them together rather than resolving fine-grained structure within them.

This explains why Reactome pathways show only modest additional cohesion (0.586) compared to Reactome overall (0.550), the model learns “this is a Reactome entity” but doesn’t strongly differentiate between specific Reactome pathways.

Learned edge weights

Next, I’ll explore what makes a high-confidence edge using sensitivity analysis on the edge encoder.

The edge encoder maps 69 edge attributes to a single weight in (0, 1). To evaluate each attribute’s importance, I calculate its average gradient with respect to the learned edge weight across 1M randomly sampled edges. This sensitivity score reveals which features most strongly influence the model’s confidence in an edge.

device = select_device(mps_valid = True)

top_model = top_model_manager.load_model_from_checkpoint(top_model_manager.best_checkpoint_path)
edge_encoder = get_edge_encoder(top_model)
feature_sensitivities = compute_edge_feature_sensitivity(edge_encoder, napistu_data.edge_attr, 1000000, device)

formatted_feature_sensitivities = format_edge_feature_sensitivity(feature_sensitivities, napistu_data)

fig, ax = plot_edge_feature_sensitivity(formatted_feature_sensitivities, top_n=20, figsize=(16, 8), truncate_names=50)
plt.show()

The model shows clear preferences: literature-derived evidence (OmniPath primary sources, STRING text mining) increases edge confidence, while indirect functional evidence (STRING coexpression transfer, experimental transfer) decreases it. This suggests the model may be learning to construct mechanistic regulatory relationships — combinations of physical interactions and functional associations — rather than simply upweighting physical interactions or functional signals independently.

Interpreting individual features is complicated by the edge encoder’s nonlinear combinations. For example, OmniPath primary sources show strong positive sensitivity while total OmniPath sources show negative sensitivity, suggesting the model values concentrated evidence from specific high-quality sources over diffuse evidence from many sources.

Notably, none of the hand-crafted confidence scores from the original databases — STRING combined score, IntAct MI score, Reactome FI score — appear among the most sensitive features. This suggests the edge encoder is learning data quality signals that differ from expert-designed heuristics, reinforcing the value of end-to-end training for edge weighting.

Finally, I’ll explore the types of edges being predicted by the GraphConv GNN.

Edge predictions

As previously discussed, when negative samples differ too greatly from real edges, the model can exploit vertex attributes that indicate implausible connections, rather than capturing the underlying biological network structure. To address this shortcut, negative samples were generated using the observed node_type strata (i.e., sampling an equal number of species→species, species→reaction, and reaction→species edges as in the real set of edges). Yet, certain vertex features could trivially separate real and negative edges; for instance, regulatory RNAs never interact with metabolites in the dataset

To evaluate whether the model exploits potential misalignment between real edges and negative samples, I’ll compare predicted edge probabilities for each edge class to the probabilities expected from their relative frequencies in real and negative edges

edge_predictions = predict(
    top_model_manager.get_experiment_dict(),
    checkpoint=top_model_manager.best_checkpoint_path
)

species_strata_recovery = summarize_edge_predictions_by_strata(edge_predictions, edge_strata_by_node_species_type)

# Filter to categories with >= 100 edges
species_strata_recovery_filtered = species_strata_recovery[species_strata_recovery['count'] >= 100]

fig, ax = plot_edge_predictions_by_strata(species_strata_recovery_filtered)
plt.show()

The plot reveals two distinct patterns. For common edge types — particularly protein→protein interactions (bright yellow, log₂ O/E ≈ 0) — real edges and negative samples occur at similar frequencies, yet the model’s predictions span a wide range (0.3-1.0). This indicates the model is learning patterns of vertex similarity that go beyond trivial features like entity type. The model differentiates among protein pairs based on their learned embeddings, not just their shared protein identity.

For rare edge types involving specialized entities (complexes, metabolites from Recon3D), the pattern changes. These categories often show extreme enrichment or depletion (|log₂ O/E| > 2), yet the model assigns them consistently high prediction probabilities regardless of whether they’re enriched or depleted. This likely reflects the tight embedding clusters observed earlier for Reactome and Recon3D entities — specialized molecular species cluster strongly in the embedding space, leading the dot product head to predict high edge probabilities between them even when such edges are rare in the training data.

This analysis suggests the model has learned meaningful biological structure within major edge types while potentially overgeneralizing for rare, specialized entities. The wide prediction spread for protein→protein edges is encouraging for future work, with more expressive heads and validation datasets, these learned similarity patterns could identify novel regulatory connections within well-represented entity types.

Summary

This post introduced Napistu-Torch and demonstrated that training GNNs on genome-scale biological networks is feasible with standard hardware — the complete suite of models trained in ~2 days on a laptop. More importantly, I’ve established the foundational infrastructure needed to explore this space systematically:

• The NapistuDataStore handles conversion from biological networks to PyTorch Geometric format with caching that eliminates rebuild overhead.
• Modular encoders, heads, and edge encoders enable architectural exploration through configuration files rather than code changes.
• PyTorch Lightning integration and CLI-driven workflows make experiments reproducible and trackable via Weights & Biases.
• Edge prediction provides self-supervised training without ground-truth labels while stratified negative sampling maintains computational tractability.
• Vertex embeddings, learned edge weights, and prediction patterns reveal what biological structure the models capture.

The most impactful findings are that different GNN architectures converge on similar biological representations — suggesting the signal is robust—and not trivially recovered from vertex attributes like its data source or molecule type.

This work creates a low-activation-energy foundation for exploring how GNNs can tap the potential of comprehensive biological networks like the Octopus. The hard infrastructure work is done — what remains is the interesting part: using these tools to build accurate genome-scale representations and discover novel biology.

• Provide an overview of the Octopus model and its construction
• Show side-by-side summaries of individual data sources highlighting their complementarity
• Demonstrate that the model successfully merges results, creating a dense network covering the complete cellular repertoire of genes, metabolites, drugs, and complexes
• Illustrate how source-level information can be carried forward to the Octopus’s graphical network to augment its vertex and edge features

The model is distributed as a set of related Napistu assets bundled together. The core components are two major Napistu data structures:

• SBML_dfs: An in-memory relational database organizing molecular species (genes, metabolites, complexes, drugs) and their relationships (reactions, interactions). I’ll provide a thorough review of this format below.
• NapistuGraph: A directed graph representation of the same network, translating molecular species and reactions into a network optimized for downstream analysis.

Building the 🐙

I built the Octopus using Napistu’s CLI, which processes individual pathway sources, merges them into consensus models, and translates them into genome-scale molecular networks. The build process runs as a cached Quarto notebook — sufficient for present needs, but a dedicated workflow manager like NextFlow or Airflow would be better suited for broader adoption within the research community.

The Octopus build process follows seven sequential steps:

1. Ingest data source-specific content and format as SBML_dfs objects
2. Standardize compartmentalization — the Octopus uses an uncompartmentalized approach for simplicity
3. Merge SBML_dfs objects into a single consensus model
4. Filter cofactors to prevent molecules like water from appearing as hub regulators
5. Convert the SBML_dfs into a NapistuGraph network representation
6. Generate derived summaries including precomputed molecular distances
7. Package all components into a single artifact and deploy to Google Cloud Storage

Follow Along!

Environment setup

To follow along with the code in this post, you’ll need a Python environment with the napistu package installed. Here’s a simple setup using venv:

1. Install uv (or use pip if preferred)

2. Setup a Python environment:

uv venv --python 3.11
source .venv/bin/activate

# Core dependencies
uv pip install "napistu==0.7.1"
# if you'd like to render the notebook, you'll need to install these additional dependencies
uv pip install seaborn ipykernel nbformat nbclient
python -m ipykernel install --user --name=blog-staging

1. Download the octopus_network.qmd notebook (or copy and paste the relevant code blocks)

2. Configure data_dir in the setup code to a path where you’re comfortable saving the consensus SBML_dfs model

Configuring the Python notebook

import os
from pathlib import Path
import requests
from math import pi

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from napistu import utils as napistu_utils
from napistu.gcs import downloads
from napistu import sbml_dfs_utils
from napistu.network import ng_utils
from napistu.sbml_dfs_core import SBML_dfs
from napistu.network.ng_core import NapistuGraph
from napistu.ontologies.constants import SPECIES_TYPE_PLURAL

from shackett_utils.blog.html_utils import display_tabulator

# globals
DATA_DIR = "/tmp/napistu_data"
ASSET = "human_consensus"
VERSION_TAG = "20250923"
INPUT_SBML_DFS_SUMMARIES_URL = "https://raw.githubusercontent.com/shackett/shackett/main/assets/data/octopus_input_sbml_dfs_summaries.json"

# utils
def cooccurrence_to_conditional_prob(cooccur_df):
    set_sizes = np.diag(cooccur_df.values)
    intersection = cooccur_df.values
    conditional_prob = intersection / set_sizes  # P(A|B) = |A ∩ B| / |B|
    return pd.DataFrame(conditional_prob, index=cooccur_df.index, columns=cooccur_df.columns)

def simple_pd_heatmap(df, plot_title, colorbar_label="Counts", fmt="d"):
    # Set up the figure size and style
    plt.rcParams.update({'font.size': 15})  # Base font size
    
    # Create clustermap with proper sizing
    g = sns.clustermap(
        df, 
        annot=True,  # Show values in cells
        cmap='Blues', 
        fmt=fmt,
        cbar_kws={'label': colorbar_label},
        figsize=(12, 10),  # Larger figure
        annot_kws={'size': 12},  # Annotation font size
        cbar_pos=(0.02, 0.83, 0.03, 0.15),  # Colorbar position (left, bottom, width, height)
    )
    
    # Increase font sizes for axis labels
    g.ax_heatmap.set_xlabel(g.ax_heatmap.get_xlabel(), fontsize=14, fontweight='bold')
    g.ax_heatmap.set_ylabel(g.ax_heatmap.get_ylabel(), fontsize=14, fontweight='bold')
    
    # Rotate and adjust tick labels for better readability
    g.ax_heatmap.tick_params(axis='x', labelsize=11, rotation=45)
    g.ax_heatmap.tick_params(axis='y', labelsize=11, rotation=0)
    
    # Add title with proper positioning (left-aligned)
    g.fig.suptitle(plot_title, fontsize=16, fontweight='bold', y=0.98, horizontalalignment='left', x=0.05)
    
    # Adjust layout to prevent clipping
    plt.tight_layout()
    
    # Return the clustermap object for further customization if needed
    return g

def create_pathway_radar_plot(df, figsize=(8, 7), title='Pathway Analysis Radar Plot'):
    
    # Get categories (columns) and pathways (rows)
    categories = list(df.columns)
    pathways = list(df.index)
    
    # Number of variables
    num_vars = len(categories)
    
    # Compute angle for each axis
    angles = [n / float(num_vars) * 2 * pi for n in range(num_vars)]
    angles += angles[:1]  # Complete the circle
    
    # Initialize the plot
    fig, ax = plt.subplots(figsize=figsize, subplot_kw=dict(projection='polar'))
    
    # Color palette for different pathways
    colors = plt.cm.tab10(np.linspace(0, 1, len(pathways)))
    
    # Plot each pathway
    for idx, pathway in enumerate(pathways):
        values = df.loc[pathway].values.astype(float)
        
        # Log10 transform (add 1 to avoid log(0))
        log_values = np.log10(values + 1)
        
        # Complete the circle
        plot_values = list(log_values) + [log_values[0]]
        
        # Plot
        ax.plot(angles, plot_values, 'o-', linewidth=2, 
                label=pathway, color=colors[idx], alpha=0.7)
        ax.fill(angles, plot_values, alpha=0.15, color=colors[idx])
    
    # Fix axis to go in the right order and start at 12 o'clock
    ax.set_theta_offset(pi / 2)
    ax.set_theta_direction(-1)
    
    # Set category labels - use built-in matplotlib positioning
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(categories, size=11)
    
    # Adjust label padding to move them outside the plot
    ax.tick_params(axis='x', pad=20)
    
    # Set y-axis labels to show original values at powers of 10
    # Determine the max value to set appropriate range
    max_log_value = np.max(np.log10(df.values.astype(float) + 1))
    
    # Create ticks at powers of 10: 10, 100, 1000, 10000, etc.
    max_power = int(np.ceil(max_log_value))
    ytick_values = [10**i for i in range(1, max_power + 1)]
    ytick_positions = [np.log10(v + 1) for v in ytick_values]
    
    ax.set_yticks(ytick_positions)
    ax.set_yticklabels([f'{v:,}' for v in ytick_values], size=9)
    ax.set_ylim(0, np.log10(10**max_power + 1))
    
    # Add grid
    ax.grid(True, linestyle='--', alpha=0.7)
    
    # Add legend
    plt.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1), fontsize=10)
    
    # Add title
    plt.title(title, size=16, pad=20)
    
    plt.tight_layout()
    
    return fig, ax

Data sources

The Octopus’s integration success stems from Napistu’s flexible SBML_dfs data structure, which standardizes diverse pathway sources while preserving their unique molecular and mechanistic contributions.

Overview of the SBML_dfs pathway representation

The core SBML_dfs data representation involves five tables linked by primary key-foreign key relationships:

• Compartments: Define distinct cellular locations (e.g., cytosol, nucleoplasm). Uncompartmentalized models contain only “cellular component” by convention.
• Species: Catalog distinct molecular entities including proteins, metabolites, complexes, and drugs.
• Compartmentalized Species: Map each species to its specific compartmental locations.
• Reactions: Represent distinct biochemical events including metabolic reactions, complex formation, and physical/functional interactions.
• Reaction Species: Define each compartmentalized species’ role in specific reactions (substrate, catalyst, inhibitor, etc.).

Additional optional tables store quantitative annotations beyond the core schema:

• Species Data: Contains tables with molecular species-specific quantitative attributes.
• Reactions Data: Contains tables with reaction-specific quantitative attributes.

Source descriptions

Each data source is formatted as a separate SBML_dfs object encapsulating its molecular species, their interactions, and any quantitative data.

• Reactome is the human gold-standard pathway database, employing rigorous expert curation with multi-tier review by over 820 scientists to produce reaction-centric models of cellular processes.
• Recon3D is a comprehensive human metabolic model that enables quantitative flux balance analysis and phenotype prediction.
• STRING is a comprehensive protein interaction database that integrates evidence from seven distinct channels with probabilistic scoring to capture functional associations rather than directional causality. Its strength is broad multi-organism coverage with confidence scores calibrated to known pathway relationships.
• IntAct is a manually curated database of experimentally verified molecular interactions with unprecedented annotation depth, making it the gold standard for high-confidence molecular interaction data.
• Reactome-FI transforms Reactome’s detailed biochemical reactions into simplified functional interaction networks using machine learning.
• OmniPath is a comprehensive integration database that combines data from over 100 resources into unified directed signaling networks with sophisticated consensus-building mechanisms. It specializes in literature-curated activity flow interactions with effect signs (activation/inhibition).
• TRRUST uses sentence-based text mining to identify transcription factor-target regulatory relationsh1ips from Medline abstracts.
• Dogma is a Napistu-specific resource that contributes gene annotations to help merge species across different primary ontologies without adding reactions to the consensus model.

Source comparisons

To understand how these sources complement each other, I’ll provide four side-by-side analyses examining the scale and characteristics of each database:

• Scale: How many species and reactions does each source contain?
• Molecular diversity: What types of entities exist (proteins, metabolites, complexes, drugs)?
• Interaction mechanisms: How do molecules connect (undirected associations, directed regulation, metabolic transformations)?
• Quantitative data: What additional measurements do sources provide (confidence scores, expression levels, binding affinities)?

These comparisons use summary statistics extracted from each source’s SBML_dfs during the Octopus build process. The summaries were saved to a public GitHub repository for reproducibility and transparency. For reference, the source summaries were generated using this non-runnable code block:

from napistu import consensus
from napistu.sbml_dfs_core import SBML_dfs
from napistu import utils

sbml_dfs_uris = [
    # mechanisms
    "napistu_data/human_consensus/cache/reactome/reactome.pkl",
    "napistu_data/human_consensus/cache/bigg.pkl",
    # consensus interactions
    "napistu_data/human_consensus/cache/hpa_filtered_string.pkl",
    # PPIs
    "napistu_data/human_consensus/cache/intact.pkl",
    # regulatory mechanisms
    "napistu_data/human_consensus/cache/omnipath.pkl",
    "napistu_data/human_consensus/cache/reactome_fi.pkl",
    "napistu_data/human_consensus/cache/trrust.pkl",
    # gene annotations
    "napistu_data/human_consensus/cache/dogma_sbml_dfs.pkl",
]

sbml_dfs_list = [SBML_dfs.from_pickle(uri) for uri in sbml_dfs_uris]

# reorganize as a list and table containing model-level metadata from the individual SBML_dfs
sbml_dfs_dict, pw_index = consensus.prepare_consensus_model(sbml_dfs_list)
sbml_dfs_dict_summaries = {k: v.get_summary() for k, v in sbml_dfs_dict.items()}
utils.save_json("<<MY_LOCAL_PATH>>/input_sbml_dfs_summaries.json", sbml_dfs_dict_summaries)

I can load these pre-computed summaries directly from GitHub:

sbml_dfs_dict_summaries = requests.get(INPUT_SBML_DFS_SUMMARIES_URL).json()

These summaries enable direct side-by-side comparison of each source’s unique characteristics and contributions to the consensus model.

Scale

I’ll count entities in each source to assess their relative sizes.

# load the SBML_dfs' summaries from GitHub
entity_type_counts = (
    pd.DataFrame({k: v["n_entity_types"] for k, v in sbml_dfs_dict_summaries.items()})
    .T
    .sort_index(axis = 1)
)

display_tabulator(entity_type_counts, layout = "fitDataTable", caption = "Entity counts per source")

Sources contain similar numbers of molecular species but reaction counts vary dramatically — from TRRUST’s ~8K reactions to STRING’s 4.2M. Dogma contains only one placeholder reaction since it contributes gene annotations rather than interactions, helping merge species across different ontologies (Ensembl, UniProt, Entrez).

Molecular diversity

Each source specializes in different molecular entity types based on their ontological annotations.

species_type_counts = (
    pd.DataFrame({k: v["n_species_per_type"] for k, v in sbml_dfs_dict_summaries.items()})
    .astype('Int64')
    .fillna(0)
    .T
    .sort_index(axis = 1)
    .rename(columns = SPECIES_TYPE_PLURAL)
)
display_tabulator(
  species_type_counts,
  layout = "fitDataTable",
  caption = "Counts of molecular species types in each source"
)

RADAR_ORDER = ["proteins", "metabolites", "drugs", "unknowns", "other", "regulatory RNAs", "complexes"]
fig, ax = create_pathway_radar_plot(
    species_type_counts[RADAR_ORDER],
    title = "Species types by pathway source",
)
plt.show()

Clear specialization patterns emerge: gene-centric sources (STRING, Dogma, Reactome-FI, TRRUST), metabolite-focused databases (BiGG), and comprehensive resources covering diverse molecular species (Reactome, IntAct, OmniPath).

Interaction mechanisms

I’ll examine interaction types using Systems Biology Ontology (SBO) terms that define molecular roles:

• Interactor: Undirected associations
• Stimulator/Inhibitor/Modifier: Regulators of expression or activity
• Modified: Targets of regulation
• Catalyst: Enzymes and transporters
• Substrate/Product: Consumed or produced molecules

sbo_term_counts = (
    pd.DataFrame({k: v["sbo_name_counts"] for k, v in sbml_dfs_dict_summaries.items()})
    .astype('Int64')
    .fillna(0)
    .T
    .sort_index(axis = 1)
)
display_tabulator(
    sbo_term_counts,
    width="auto",
    layout="fitDataStretch",
    caption = "Counts of reaction participant roles in each source"
)

RADAR_ORDER = ["catalyst", "reactant", "product", "stimulator", "inhibitor", "modifier", "modified", "interactor"]
fig, ax = create_pathway_radar_plot(
    sbo_term_counts[RADAR_ORDER],
    title = "SBO terms by pathway source",
    )
plt.show()

Broad sources like STRING favor generic “interactor” classifications, while specialized databases like Recon3D and Reactome capture specific mechanistic detail more faithfully.

Quantitative data

Beyond structural information, sources provide additional annotations and metadata for both molecular species and their interactions.

data_summaries = {k: v["data_summary"] for k, v in sbml_dfs_dict_summaries.items() if len(v["data_summary"]["reactions"]) > 0}

data_summaries_list = []
for k, v in data_summaries.items():
    for entity_type, entity_data in v.items():
        if len(entity_data) > 0:
            for table_name, table_data in entity_data.items():
                table_summary ={
                    "table_name" : table_name,
                    "entity_type" : entity_type,
                    "n_rows" : table_data["n_rows"],
                    "columns" : ", ".join(table_data["columns"]),
                }
                data_summaries_list.append(table_summary)
data_summaries_df = pd.DataFrame(data_summaries_list)

display_tabulator(
    data_summaries_df,
    wrap_columns = "columns",
    column_widths = {"columns" : "65%"},
    caption = "Additional species- and/or reactions-data in each source"
)

This additional data falls into two key categories: confidence scoring systems (STRING interaction scores, IntAct experimental evidence) and mechanistic granularity (OmniPath activation/inhibition breakdowns, IntAct interaction types). Both provide crucial context for assessing interaction reliability and biological mechanisms.

Source compatibility

Many data sources used by Napistu, like STRING and OmniPath, already aim to integrate multiple upstream data sources into a consistent consensus. Napistu builds on these resources to merge what would otherwise be incompatible data sources into a single, well-mixed model. Without proper integration, sources would separate like oil and water — each forming disconnected subnetworks with minimal overlap. Instead, we need to gel them together by establishing a unified molecular vocabulary that enables seamless integration of source-specific interactions.

Napistu accomplishes this integration through:

• Data standardization: Systematic identifiers and SBO ontology terms create a common vocabulary for describing molecules and their interactions across diverse sources
• Algorithmic merging: A consensus procedure that identifies equivalent entities and reconciles overlapping information into a single integrated model

Merging SBML_dfs objects into a consensus SBML_dfs

Merging multiple SBML_dfs objects into a consensus model requires resolving entities by determining which compartments, species, and reactions are shared across sources. This process works through tables in logical order (compartments & species $\rightarrow$ compartmentalized species $\rightarrow$ reactions & reaction species), aggregating a single table drawn from all models to construct:

• Consensus tables: New unified tables with standardized structure
• Key mapping tables: Lookup tables linking old source-specific primary keys to new consensus keys for updating foreign key relationships

Two variables are critical for successful merging:

• Identifiers: Determine what can be merged by organizing systematic identifiers as curated lists, each defined by ontology, identifier, and bioqualifier (e.g., “BQB_IS” or “BQB_HAS_PART”)
• Sources: Track what has been merged by associating each consensus entity with all contributing data sources

The consensus algorithm proceeds through four steps:

1. Resolve foundational entities: Use greedy network-based matching to identify species and compartments through shared identifiers, connecting entities that share BQB-coded systematic identifiers
2. Define compartmentalized species: Map resolved species to their appropriate compartmental locations
3. Merge reactions: Update reaction species annotations and identify redundant reactions based on participants and mechanisms
4. Harmonize data tables: Update species_data and reactions_data with consensus primary keys, aggregating results to ensure one row per consensus entity

Loading the 🐙

With the consensus algorithm framework established, let’s examine the actual eight-source Octopus model to see how well these theoretical merging principles work in practice. To do this, I’ll download the pre-built model from Google Cloud Storage and provide some quick summaries of its core properties.

The Octopus network is available through GCS and gets updated periodically as sources are added and Napistu data structures evolve. To ensure reproducibility for this post and others like Network Biology with Napistu, Part 2: Translating Statistical Associations into Biological Mechanisms, tagged versions are preserved for reliable future access. Here, I’ll load a tagged version compatible with Napistu 0.7.1.

# ~3 min load
# download and cache the Octopus sbml_dfs and the other assets its bundled with
sbml_dfs_path = downloads.load_public_napistu_asset(
    asset = ASSET,
    subasset = "sbml_dfs",
    data_dir = DATA_DIR,
    # download the tagged version for reproducibility and Python env compatibility
    version = VERSION_TAG,
)

sbml_dfs = SBML_dfs.from_pickle(sbml_dfs_path)

🐙 summary

With the core SBML_dfs object loaded, I’ll examine its high-level properties using the show_summary method.

summary_stats = sbml_dfs.get_summary()
summary_table = sbml_dfs_utils.format_sbml_dfs_summary(summary_stats)
display_tabulator(
    summary_table,
    width="auto",
    layout="fitDataStretch",
    caption = "Consensus SBML_dfs summaries"
)

The consensus model contains genes/proteins, metabolites, complexes, drugs, and regulatory RNAs within a single compartment — cellular component (the root term of GO’s cellular component category). The model encompasses approximately 4.5M reactions spanning undirected interactions, directed regulation, and complex multi-participant regulatory mechanisms.

While the earlier source comparisons demonstrated each database’s potential contributions, the key question remains: did the sources actually merge into a single well-mixed model? Successful integration requires extensive molecular species sharing across sources and meaningful reaction overlap. Rather than separate, highly connected subnetworks with minimal inter-source connections, we want a unified network where sources are genuinely integrated.

The model’s Source objects provide the answer — they track which data sources contributed to each species, compartment, and reaction, enabling direct assessment of integration success.

Shared molecular vocabulary

To assess integration success, I’ll examine which data sources contributed to each molecular species through contingency tables of species-source occurrences. Values reflect how many of a source’s molecular species merged into each consensus species — typically 0 (not present) or 1 (exact match), though occasionally higher when multiple protein annotations roll up to a single gene.

species_source_occurrence = sbml_dfs.get_source_occurrence("species")
display_tabulator(
    species_source_occurrence.head(),
    layout = "fitDataTable",
    caption = "Example molecular species and the sources they were originally found in"
)

I can visualize species sharing patterns by converting the occurrence matrix ($X$) into a cooccurrence matrix ($C$) using:

Where $B = \mathbf{1}(X \neq 0)$ is the binary matrix obtained by converting non-zero entries of $X$ to 1.

species_source_cooccurrence = (
    sbml_dfs.get_source_cooccurrence("species")
    .rename_axis('Database', axis=0)
    .rename_axis('Database', axis=1)
)

simple_pd_heatmap(species_source_cooccurrence, "Species Source Co-Occurrence")

The heatmap reveals that gene-centric, dense sources (STRING, Dogma, Reactome-FI) cluster together with similar molecular coverage, while Reactome, Recon3D, and TRRUST remain more isolated. This reflects TRRUST’s smaller size and the molecular specialization of Reactome and Recon3D compared to comprehensive protein databases.

To quantify this specialization, I’ll identify species unique to individual sources.

private_species = species_source_occurrence.loc[(species_source_occurrence != 0).sum(axis=1) == 1]

private_species_source_counts = (
    (private_species != 0)
    .sum()
    .sort_values(ascending=False)
    .rename("Private species")
    .to_frame()
    .T
)

display_tabulator(
    private_species_source_counts,
    width="auto",
    layout="fitDataStretch",
    include_index=False,
    caption = "Private molecular species from each source"
)

Several sources contribute substantial numbers of private species, each for logical reasons:

• Reactome: Detailed complex mechanisms with fine-grained complex definitions
• OmniPath: Extensive drug collections (PubChem) and microRNAs (MirBase)
• IntAct: Small molecules and microRNAs alongside core protein interactions
• Recon3D: Extensive coverage of metabolites and lipids

The Octopus successfully integrates molecular species, with proteins shared across multiple sources while specialized molecular types arise from domain-specific resources.

Reaction overlap reveals data source specialization

To understand what individual sources contribute, I’ll analyze reaction source occurrences and cooccurrences using a similar approach to species analysis.

To interpret this analysis, readers should understand two important points:

• Strict merging criteria: Reactions merge only with identical participants and SBO terms. A reaction between genes A and B won’t merge if one source labels them inhibitor $\rightarrow$ modified while another uses modifier $\rightarrow$ modified, explaining the low overlap we’ll observe.
• Analysis scope: This analysis excludes interactor-interactor interactions because the existing tooling is designed to surface information relevant for graph construction, where these interactions become direct edges between molecular species with no reaction vertices added.

reactions_source_occurrence = sbml_dfs.get_source_occurrence("reactions")

display_tabulator(
    reactions_source_occurrence.head(),
    width="auto",
    layout="fitDataStretch",
    caption = "Example reactions and the sources they were originally found in"
)

The reaction occurrence data is notably sparse, and the order-of-magnitude differences in reaction counts between sources complicate direct cooccurrence visualization.

To better assess source dependencies, I’ll calculate conditional probabilities $\Pr(A|B)$ from the cooccurrence matrix, showing how likely a reaction from source $B$ also appears in source $A$.

reactions_source_cooccurrence = (
    sbml_dfs.get_source_cooccurrence("reactions")
    .rename_axis('Database', axis=0)
    .rename_axis('Database', axis=1)
)

reactions_source_conditional_prob = cooccurrence_to_conditional_prob(reactions_source_cooccurrence)

simple_pd_heatmap(reactions_source_conditional_prob, "Conditional probability of reaction found in column source\n given that it occurs in the row source", fmt=".3f", colorbar_label="Probability")

The conditional probability analysis reveals distinct patterns: Reactome and Recon3D reactions remain largely unique, while meaningful overlap exists between Reactome-FI, OmniPath, and TRRUST. The strongest overlap occurs between TRRUST and OmniPath (50% of TRRUST interactions also appear in OmniPath) — an expected result since TRRUST is one of the resources incorporated into OmniPath.

These patterns demonstrate successful species integration alongside preserved source-specific reaction diversity, with each database contributing substantial unique mechanistic content to the consensus model.

Decorating the 🐙 graph with species and reaction data

While SBML_dfs comprehensively organizes pathway data, network analyses like personalized PageRank require graph representations. NapistuGraphs convert this tabular data into networks where compartmentalized species and reactions become vertices connected by information flow edges. Built on igraph’s foundation, they combine versatile graph operations with biological annotations, data provenance, and specialized network biology methods.

A NapistuGraph was bundled with the SBML_dfs downloaded above, enabling direct loading and analysis:

napistu_graph_path = downloads.load_public_napistu_asset(
    asset = ASSET,
    subasset = "napistu_graph",
    data_dir = DATA_DIR,
    version = VERSION_TAG
)

napistu_graph = NapistuGraph.from_pickle(napistu_graph_path)

summary_stats = napistu_graph.get_summary()
summary_table = ng_utils.format_napistu_graph_summary(summary_stats)
display_tabulator(
    summary_table,
    wrap_columns = True,
    column_widths = {"Value" : "80%"},
    caption = "Summaries of the NapistuGraph network"
)

Many vertex and edge attributes mirror those from the SBML_dfs summaries.

The critical quantitative attribute is edge weight, representing each interaction’s plausibility and strength. Edge weights drive most graph algorithms — from shortest path calculations and network layouts to propagation methods. However, capturing this complexity in a single attribute becomes increasingly challenging as more sources contribute quantitative information relevant to regulatory plausibility.

The default NapistuGraph contains a limited array of vertex and edge attributes, but there’s actually a wealth of quantitative information highlighted throughout this post that can be integrated directly into the Octopus network. The power of Napistu’s design becomes apparent when we start layering in this additional context. I’ll demonstrate by augmenting the graph with two particularly valuable information types:

• add_sbml_dfs_summaries: Generates source and ontology occurrence data for all vertices, revealing which databases contributed to each node and what biological categories they represent
• add_all_entity_data: Transfers comprehensive quantitative measurements from reactions_data and species_data tables directly onto their corresponding edges and vertices

# augment the graph
# add ontology and source data to vertices
napistu_graph.add_sbml_dfs_summaries(sbml_dfs, stratify_by_bqb = False)

# add reactions_data to edges
napistu_graph.add_all_entity_data(sbml_dfs, "reactions", overwrite=True)
napistu_graph.add_all_entity_data(sbml_dfs, "species", mode = "extend")

summary_stats = napistu_graph.get_summary()
summary_table = ng_utils.format_napistu_graph_summary(summary_stats)
display_tabulator(
    summary_table,
    wrap_columns = True,
    column_widths = {"Value" : "80%"},
    caption = "Post-augmentation summaries of the `NapistuGraph` network"
)

Now, we’ve gone from a relatively spartan set of vertex and edge attributes to a comprehensive graph of human cellular physiology enriched with detailed biological annotations that describe what each vertex and edge represents. This is a robust foundation for training expressive network-based methods like graph neural networks.

Summary

The Octopus model integrates eight diverse pathway databases into a unified, genome-scale network of human cellular physiology. Through systematic merging of complementary data sources, the model establishes a shared molecular vocabulary while preserving each source’s specialized contributions:

• Molecular species integration: Proteins and other species are effectively shared across sources, creating a common parts list that specialized databases can extend with domain-specific molecules (metabolites from Recon3D, complexes from Reactome).
• Reaction specialization: Sources show modest but meaningful overlap in reactions, with each database contributing unique mechanisms that reflect its individual curation focus.

The resulting NapistuGraph provides a framework for layering extensive biological information onto network structures. Source provenance, confidence scores, ontological classifications, and mechanistic annotations can be systematically integrated as vertex and edge attributes, enabling sophisticated analyses from network propagation to machine learning approaches.

The Octopus model is now ready for use. I’m excited to build on this foundation and to see how the community engages with this new resource.

This tool empowers new users, advanced contributors, and AI agents alike to quickly access relevant project knowledge.

Before MCP, I fed Claude a mix of README files, wikis, and raw code hoping for useful answers. Tools like Cursor struggled with the tangled structure, sparking the idea for the Napistu MCP server.

I’ll cover:

• Why I built the Napistu MCP server and the problems it solves
• How I deployed it using GitHub Actions and Google Cloud Run
• Case studies showing how AI agents perform with — and without — MCP context

The AI development paradox

We’re at an interesting inflection point in software development. AI both accelerates and hinders the creation of high-quality code.

✅ Acceleration

AI speeds up development by:

• Handling repetitive tasks
• Lowering the barrier to entry
• Simplifying debugging

Sometimes I hand an agent a stack of failing pytest errors and say, “You handle this.” And it does. It’s pure magic.

❌ Friction

But AI also introduces chaos:

• Repeats patterns instead of reusing code
• Misses domain-specific idioms
• Adds unnecessary abstractions
• Produces brittle, poorly-structured code

This goes beyond simple messiness — it can rapidly escalate into a technical debt time bomb.

🎯 Key Insight

AI isn’t inherently good or bad; its performance is task-dependent. Most AI failures can be traced back to missing context: the model simply wasn’t given the information it needed to succeed. If an AI agent understands your domain, your design patterns, and your project structure, it can generate excellent code. Without that context, it’s flying blind — and that’s where most of the frustration comes from.

Many of us are seeking the optimal balance where AI maximizes productivity, while simultaneously working to expand that potential by enhancing its performance on critical tasks.

Information is everything

Context is the central challenge in any domain-specific codebase — be it a financial trading system, a game engine, or a scientific library. Until AI agents understand the context of the codebase, they will struggle to follow existing patterns and conventions, and misuse domain-specific approaches.

Off-the-shelf approaches to AI integration are improving, either by seamlessly integrating with external services (like Claude talking with GitHub, Google Docs, etc.), or by directly interfacing with the codebase itself (as tools like Cursor and Copilot do). But, as you’ll see, while these options help, they still leave something to be desired.

Case study: the value of context

To highlight the value of context, I’ll use a real-world example. Say I ask an agent to help me with the following question:

How do I create a consensus network from multiple pathway databases in Napistu? Please create a single artifact with your initial thoughts.

Let’s explore a few scenarios to see how this plays out.

Scenario 1: no context

Without any context, Claude has no idea what I’m talking about and starts Googling:

I’m not familiar with Napistu as a specific software tool or platform for pathway database analysis. … Since I cannot locate specific documentation for Napistu, I’ll create an artifact with general guidance on creating consensus networks from multiple pathway databases

Claude suggests some helpful ideas — but says nothing about Napistu.

Scenario 2: some context

With relevant code context, by pointing Claude to relevant .py files from GitHub:

Looking at the Napistu codebase, I can see this is a systems biology toolkit for working with pathway models. Let me create a comprehensive guide on how to create a consensus network from multiple pathway databases.

The resulting artifact highlights key classes and functions, organizing them into a clear, orderly progression. Nonetheless, the response feels disjointed, as if it pulled snippets from many sources without adequately synthesizing them. Moreover, producing this response required me to provide specific .py files to Claude because only ~15% of the napistu.py codebase could fit into Claude’s context window.

Scenario 3: expert knowledge

If you asked an expert (me, 🙃) this question, the guidance would draw from multiple information sources:

Start with the merging_models_into_a_consensus tutorial — it provides a step-by-step walkthrough of this exact workflow. Building consensus involves calling consensus.construct_consensus_model() with multiple SBML_dfs objects and a pathway index, which organizes the objects’ metadata. This is currently being reworked in Issue 169 to remove the pathway index requirement. Finally, review the consensus Napistu wiki page to gain a high-level understanding of the key algorithms.

This response demonstrates true understanding — it connects theory (the algorithm), practice (the tutorial), current development (the GitHub issue), and conceptual framework (the wiki) into actionable guidance.

Providing AI agents with expert knowledge

For an AI agent to match the expert response, we would need to do more than just expand its context window; we need to address two key challenges:

1. information fragmentation - Relevant information is scattered across multiple sources such as code repositories, wikis, issue trackers, tutorials, README files, and more. This dispersion makes providing relevant information to agents a cumbersome and often manual process.
2. signal vs. noise - Critical context can easily be obscured by large volumes of irrelevant or low-priority information, making it challenging for AI agents to identify what truly matters.

Solution preview: What if an AI could retrieve domain-specific information on demand?

Before diving into how to provide agent-friendly information, let me first show you the results of that effort in Napistu.

First, I’ll install Napistu with MCP dependencies enabled.

pip install 'napistu[mcp]'

Then, I can configure the remote documentation server’s URL and port with production_client_config.

from IPython.display import HTML, display
from napistu.mcp.config import production_client_config
config = production_client_config()

display(HTML(f"""
    <div>
    <b>Client config:</b><br>
    <b>Host:</b> {config.host}<br>
    <b>Port:</b> {config.port}<br><br>
    </div>
    """))

Since this is a remote server, I can now start interacting with it directly. I’ll pose the consensus modeling question again and then reformat the AI-friendly JSON output as human-readable tables.

import html
import re
import pandas as pd
from napistu.mcp.client import search_component
from shackett_utils.utils import pd_utils
from shackett_utils.blog.html_utils import display_tabulator

def sanitize_content(text):
    if pd.isna(text):
        return ""
    # Remove/replace problematic characters
    text = str(text)
    text = re.sub(r'[^\w\s\-.,!?():]', ' ', text)  # Keep only basic chars
    text = re.sub(r'\s+', ' ', text)  # Normalize whitespace
    return html.escape(text)  # HTML escape

QUERY = "How do I create a consensus network from multiple pathway databases in Napistu?"
COMPONENTS = ["codebase", "documentation", "tutorials"]

# Returns actual Napistu function signatures, docs, and usage examples

combined_results = list()
for component in COMPONENTS:
    results = await search_component(
        component,
        QUERY,
        config=config
    )
    results_df = pd.DataFrame(results["results"]).assign(component=component)
    
    combined_results.append(results_df)

combined_results_df = pd.concat(combined_results).sort_values(by="similarity_score", ascending=False)[["component",  "similarity_score", "source", "content"]]

display_combined_results_df = combined_results_df.copy()
pd_utils.format_numeric_columns(display_combined_results_df, inplace = True)
pd_utils.format_character_columns(display_combined_results_df, inplace = True)
display_combined_results_df['content'] = display_combined_results_df['content'].apply(sanitize_content)
display_combined_results_df['source'] = display_combined_results_df['source'].apply(sanitize_content)

display_tabulator(
    display_combined_results_df,
    caption="Top search results by cosine similarity",
    wrap_columns=["source", "content"],
    column_widths={"source" : "25%", "content" : "50%"},
    include_index = False
)

Because the top result is Markdown, I’ll display it as a blockquote.

import re
from IPython.display import Markdown

def quote_markdown(markdown_content):
    
    suppressed_headings = re.sub(r"^#+ (.*)$", r"**\1**", markdown_content, flags=re.MULTILINE)
    blockquoted = "\n".join(f"> {line}" for line in suppressed_headings.split("\n"))

    return blockquoted

# Add blockquote formatting
markdown_content = combined_results_df["content"].iloc[0]
display(Markdown(quote_markdown(markdown_content)))

Tutorials

These tutorials are intended as stand-alone demonstrations of Napistu’s core functionality. Most examples will focus on small pathways so that results can easily be reproduced by users.

• Downloading pathway data
• Understanding the sbml_dfs format
• Merging networks with the consensus module
• Using the CPR Command Line Interface (CLI)
• Formatting sbml_dfs as napistu_graph networks
• Suggesting mechanisms with network approaches
• Adding molecule- and reaction-level information to graphs
• R-based network visualization

Much of the information that the expert provided is returned in this initial query. However, the goal is not to deliver all relevant information at once because this would inevitably include a significant amount of irrelevant data. Rather, we can use tools like search_component to give agents agency, putting information at the tips of their virtual fingers. This allows agents to nimbly explore a problem drawing upon relevant resources on-demand. As a result, rather than generating generic or hallucinated responses, agents can uncover actual patterns, locate pertinent tutorials, and gain a deeper understanding of our domain-specific approaches.

Enter the Model Context Protocol (MCP)

MCP provides a standardized way for AI models to access external information sources. Think of MCP as giving AI agents a research assistant who knows your project inside and out — someone who can instantly locate relevant documentation, code examples, and domain-specific implementation patterns specific to your domain.

AI agents can interact with MCP servers through two primary mechanisms: tools and resources. Think of resources as reference materials agents can read (like a library catalog), and tools as actions they can execute (like asking a librarian to retrieve specific materials).

Let’s look at how this works in practice with the Napistu MCP server:

Tools enable agents to perform actions and searches:

• search_documentation() - Find relevant project docs and issues
• search_codebase() - Discover functions, classes, and methods
• search_tutorials() - Locate implementation examples

The search_component function used in the solution preview above indirectly uses a tool by calling the lower-level function call_server_tools. call_server_tools in turn calls the actual FastMCP Client method call_tool. This method accepts a tool name and arguments, returning structured results.

Resources in Napistu MCP provide read-only access to structured information:

• napistu://health - Server status and component health
• napistu://documentation/summary - Overview of available documentation
• napistu://tutorials/index - Available tutorial content

To call a resource endpoint like napistu://documentation/summary in Python, we can use the read_server_resource function, which calls the FastMCP Client method read_resource. This method takes a resource URI and returns the contents of the resource.

from napistu.mcp.client import read_server_resource

content = await read_server_resource("napistu://documentation/summary", config)
print(content)

    {
      "readme_files": [
        "napistu",
        "napistu-py",
        "napistu-r",
        "napistu/tutorials"
      ],
      "issues": [
        "napistu",
        "napistu-py",
        "napistu-r"
      ],
      "prs": [
        "napistu",
        "napistu-py",
        "napistu-r"
      ],
      "wiki_pages": [
        "Environment-Setup",
        "Data-Sources",
        "Napistu-Graphs",
        "Model-Context-Protocol-(MCP)-server",
        "SBML-DFs",
        "SBML",
        "Dev-Zone",
        "Exploring-Molecular-Relationships-as-Networks",
        "Precomputed-distances",
        "GitHub-Actions-napistu‐py",
        "Consensus",
        "History"
      ],
      "packagedown_sections": []
    }

This architecture solves our earlier problem; instead of manually curating context, AI agents can dynamically discover and retrieve exactly the information they need.

Anatomy of the Napistu MCP server

Before diving into the technical implementation, it’s worth understanding why I built this system. The Napistu MCP server serves three key purposes:

1. Dramatically lowers the barrier to entry for new users who struggle with the “cold start” problem
2. Democratizes domain expertise to encourage broader community contributions
3. Gives core developers’ AI agents comprehensive project knowledge to efficiently extend the codebase

These objectives directly address the information fragmentation and context limitations we identified earlier. With this motivation in mind, I’ll provide an overview of the server’s architecture.

FastMCP foundation

The Model Context Protocol provides a standard way for AI models to access external information. FastMCP provides a Flask-like Python implementation of the MCP protocol.

from fastmcp import FastMCP

mcp = FastMCP("napistu-server")

@mcp.resource("napistu://health")
async def health_check():
    return {"status": "healthy", "components": [...]}

@mcp.tool()
async def search_documentation(query: str) -> dict:
    return {"results": [...]}

FastMCP manages the protocol details, while we focus on exposing Napistu’s knowledge.

The server.py module orchestrates the entire lifecycle through a simple three-step process:

1. Create a FastMCP server instance with the validated host, port, and server name from a standard, or manually defined, configuration.
2. Based on the selected profile (execution, docs, or full), register the enabled components with the server; each component adds its own resources and tools to the endpoint registry.
3. Asynchronously initialize all registered components, loading their data sources and setting up semantic search indexing in parallel.

Once this process completes, the server starts listening for incoming MCP requests.

Components

Napistu employs a component-based architecture that ensures separation of concerns — each component manages its own data sources and search logic. This design supports graceful degradation; for instance, a failed GitHub API call won’t disrupt tutorial searches. It also enables flexible deployment, allowing the activation of only the necessary components. This modularity lets me create servers tailored to specific use cases—for example, a local server capable of executing Napistu code or a remote server focused solely on documentation.

The current components are:

• Documentation: READMEs, wiki pages, GitHub issues and PRs\
• Codebase: API documentation and function signatures sourced from Read The Docs\
• Tutorials: Jupyter notebooks converted into searchable Markdown\
• Execution (in development): Interaction with a live Python environment
• Health: Server monitoring and diagnostics for system components

Each component follows a consistent pattern: load data, register endpoints, and handle search.

class DocumentationComponent(MCPComponent):
    async def initialize(self, semantic_search: SemanticSearch = None) -> bool:
        """Load READMEs, wiki pages, GitHub issues"""
        # Load external data and populate component state
        return success
    
    def register(self, mcp: FastMCP) -> None:
        """Register resources and tools with MCP server"""
        @mcp.tool()
        async def search_documentation(query: str):
            return self.state.semantic_search.search(query, "documentation")

**USE THIS WHEN:**
- Looking for specific Napistu functions, classes, or modules
- Finding API documentation for Napistu features

**DO NOT USE FOR:**
- General programming concepts not specific to Napistu
- Documentation for other libraries or frameworks

Smart search: semantic + vector embeddings

We support two search methods: exact keyword search (e.g., “create_consensus”) and semantic search (e.g., “How do I merge pathway data?”). Semantic search is powered by a SemanticSearch object used across components.

The pipeline includes:

• Content Processing: Load content and chunk long documents at natural boundaries
• Embedding Generation: Convert chunks to 384-dimensional vectors using all-MiniLM-L6-v2 sentence transformer, selected for its ease of implementation and effectiveness with general text
• Vector Storage: Store embeddings in ChromaDB, along with metadata to support fast similarity search
• Query Processing: Embed user queries and find nearest neighbors using cosine similarity

class SemanticSearch:
    def __init__(self, persist_directory: str = "./chroma_db"):
        self.client = chromadb.PersistentClient(path=persist_directory)
        self.embedding_function = SentenceTransformerEmbeddingFunction(
            model_name="all-MiniLM-L6-v2"
        )
    
    def search(self, query: str, collection_name: str):
        # Convert query to vector, find similar content by cosine similarity
        return similarity_results_with_scores

The agent experience

With the server architecture in place, let’s explore how this translates to the actual user experience for both humans and AI agents. The MCP protocol uses structured JSON messages over HTTP, with one key advantage: both humans and AI agents interact through the same unified interface.

Human developers engage directly using the Napistu client utilities or the MCP command line interface.

from napistu.mcp.client import search_component
from napistu.mcp.config import production_client_config

config = production_client_config()
results = await search_component("documentation", "how to install Napistu", config=config)

AI agents (such as Claude, Cursor, or any MCP-compatible tool) send equivalent requests via their respective MCP client implementations.

For example, when an agent asks “How do I create consensus networks in Napistu?””, it automatically:

AI agents (like Claude, Cursor, or any MCP-compatible tool) send the same underlying requests but through their MCP client implementations. When an agent asks “How do I create consensus networks in Napistu?”, it automatically:

1. Calls the search_tutorials tool with the query
2. Receives structured results with similarity scores and content snippets
3. May follow up with search_documentation or search_codebase for additional context
4. Uses all this information to provide comprehensive, accurate guidance.

The key insight is that agents receive the same rich, structured responses as human developers but can instantly process and integrate information across multiple sources. This transforms a simple Q&A interaction into an expert-level consultation.

From local to global: deployment story

Local development

It’s easy to set up a local MCP server that digests relevant documents and interacts with local agents.

# Install Napistu with MCP dependencies
pip install 'napistu[mcp]'

# Start full development server (all components)
python -m napistu.mcp server full

# Health check shows component loading
python -m napistu.mcp health --local

🏥 Napistu MCP Server Health Check

Server URL: http://127.0.0.1:8765/mcp

Server Status: healthy

Components:
  ✅ documentation: healthy
  ✅ codebase: healthy  
  ✅ tutorials: healthy
  ✅ semantic_search: healthy

This local approach works well for individual developers, but creates barriers for broader adoption. It requires installing Napistu, maintaining a background process, and keeping it running — imposing a significant burden on users who simply want to explore the project or collaborate.

The always-up solution

Instead, I aimed to create an always-available service that I and others access easily, without any local setup. This meant deploying the server to the cloud with automatic updates triggered by changes in codebase, integrated seamlessly into my GitHub Actions CI/CD workflows.

Every tagged release triggers deployment to Google Cloud Run.

# Deploy workflow - simplified view
on:
  workflow_run:
    workflows: ["Release"]  # Auto-deploy after successful release
    types: [completed]
  schedule:
    - cron: '0 10 * * *'  # Daily content refresh at 2 AM PST (10 AM UTC)

jobs:
  deploy:
    steps:
      - name: Deploy to Cloud Run
        run: |
          gcloud run deploy napistu-mcp-server \
            --image="us-west1-docker.pkg.dev/.../napistu-mcp-server:latest" \
            --cpu=1 --memory=2Gi \
            --set-env-vars="MCP_PROFILE=docs"

The production setup runs the “docs” profile (documentation + codebase + tutorials, no execution component since these are meant to operate in a user environment) with 1 CPU and 2Gi memory, costing less than $1 per day. Content is refreshed both upon new release of napistu-py and nightly to ensure the latest documentation changes are captured.

This deployment strategy creates a powerful feedback loop between development and documentation.

The payoff

Now any AI tool can access the Napistu knowledge base instantly at https://napistu-mcp-server-844820030839.us-west1.run.app. Users don’t need to install software, run local processes, or handle maintenance; they can simply configure their AI tools to connect to the shared knowledge base. The service automatically updates with the latest documentation and code changes, while Google Cloud Run handles scaling, health checks, and automatic restarts to ensure high availability.

// Claude Desktop / Cursor configuration
// Add this to your MCP settings to access Napistu knowledge
{
  "mcpServers": {
    "napistu": {
      "command": "npx",
      "args": ["mcp-remote", "https://napistu-mcp-server-844820030839.us-west1.run.app/mcp/"]
    }
  }
}

The result is clear: Napistu’s entire knowledge base becomes instantly searchable by AI agents worldwide, dramatically lowering the barrier to contribution and collaboration.

🛡️ Security and privacy

The remote Napistu MCP server is intended for narrowly-scoped information retrieval, so security and privacy issues should be minimal:

• Scoped to Napistu-specific resources only
• No user data stored or processed. Standard cloud platform logging may occur because this runs on Google Cloud Run
• Auditable: all resources and tools are openly available in the public GitHub repository

Case studies: AI agents in action

As previously noted, the server’s core goals are to make the codebase more accessible to new users and collaborators, while also enhancing the capabilities of AI agents used by core developers. In this section, I’ll present two case studies that A/B test the impact of the MCP server on both onboarding and development experiences. In both cases, even when MCP was not enabled, I provided substantial contextual information through standard channels — Claude accessed files via GitHub, and Cursor had access to the full codebase. This makes MCP’s impact less binary and instead highlights its marginal contributions in more realistic, real-world scenarios.

Case Study 1: Learning with Claude

To illustrate MCP’s value for training, I provided Claude the same prompt both with and without MCP enabled for comparison:

I’m new to Napistu. Can you provide a high-level overview of the structure, creation, and usage of SBML_dfs? Please think deeply and incorporate your response into a Markdown file.

Without MCP

The resulting artifact is a mixed bag:

• ✅ Provides a good overview of the core and optional tables and their relationships.
• ✅ Public methods are grouped logically, with some light-weight explanations.
• ✅ Most code appears syntactically correct.
• ❌ Logistically, I had to manually add select files from GitHub which requires prior knowledge of the codebase — something a new user would likely lack.
• ❌ The “Creating SBML_dfs Objects” section mentions a subset of the approaches and includes the consensus logic, which feels out of place.
• ❌ Advanced usage and “integration with the Napistu ecosystem” consists of random functionality inferred from the CLI.

With MCP

Armed with the MCP, the artifact is well rounded, but far from perfect.

• ✅ Good overview of the core and optional tables and their relationships.
• ✅ Good high-level overview of how to create SBML_dfs and its public methods.
• ✅ Advanced usage, best practices, and general use cases are solid
• ❌ model_source = source.Source("MyDatabase", "v1.0"). This line captures the gist of the model_source object, but it won’t actually run. Since this was a new addition to the codebase, this is probably a case where the documentation has lagged behind the code. This serves as a helpful reminder that you’ll only get relevant information when your sources are up-to-date.

I definitely prefer the artifact generated with MCP, although it was generated from a single initial prompt. In the absence of MCP, follow-up questions by a user would quickly devolve into hallucinations. With access to the Napistu MCP, Claude can continue to guide the user through the complex codebase while maintaining a high-level perspective.

The Result: From “intimidating research codebase” to “approachable, guided experience”

Case Study 2: Building with Cursor

There are several areas where access to the MCP server would significantly benefit Cursor:

• 🚀 When NOT working on the actual Napistu codebase, Cursor could still look up classes and functions.
• 🚀 For general usage questions and training prompts, as demonstrated in Case Study 1, having access to related content and leveraging semantic search to handle synonyms would be especially valuable. However, this is not a major Cursor use case, and users would likely receive better responses from a general-purpose LLM like Claude in such scenarios.

And, situations where having access to the Napistu MCP would be entirely unnecessary:

• 🤷 When working directly on the Napistu codebase, Cursor can efficiently look up function signatures and search for functionality using its native methods.

Rather than setting up a strawman by denying Cursor access to the codebase, I wanted to explore scenarios when MCP could help Cursor in more nuanced situations. In exploring this question, I asked Cursor directly — and I found its response quite insightful.

MCP tools in Cursor are about bridging the gap between “what the code can do” and “how it’s meant to be used”. They’re not replacing code navigation; rather, they’re adding the intent and context that lives in documentation and tutorials.

Intent and context become particularly relevant when applying a framework like Napistu, rather than directly extending it. In this sense, the Napistu MCP is particularly valuable when using Napistu to explore scientific questions within a notebook environment. Given that Cursor recently added support for Jupyter notebooks (albeit still in an early and somewhat rough state), this represents a particularly compelling use case. To make the tasks straightforward, I asked Cursor to extend one of the Napistu tutorials because tutorials should be a mix of code and explanatory prose, just like a good biological analysis:

Can you help me extend the understanding_sbml_dfs.ipynb tutorial to flesh out the “from_edgelist” workflow and to include any recent updates to the core data structure? THINK DEEPLY AND (DO NOT USE THE NAPISTU MCP / USE THE NAPISTU MCP AS NEEDED). Since you’ll have trouble directly editing the ipynb, please suggest what I should incorporate in a separate Markdown file. Edit for readability and to prioritize high value content. Limit the total content to less than 30 new sentences.

Without MCP

The artifact has its high and low points:

• ✅ The summary of the edgelist format with running code is quite good.
• ❌ It has no idea what I meant by “recent updates” and just provides code snippets for random public functions.

With MCP

Third-party integrations with Cursor are in a far rawer state than the major models, so you really have to twist its arm to use it. In practice, the experience is underwhelming — Cursor tends to rely solely on the search_codebase tool, which surfaces information it already has access to. As a result, the actual output is fairly poor:

• ❌ Only pseudocode describing the edgelist format
• ❌ Misinterprets “recent updates,” instead listing all public methods not already covered in the tutorial

Between these two scenarios, I would choose the output generated without MCP access. Much of this comes down to how Cursor used the MCP server. It tends to follow a one-track mindset — fixated on “code, code, code” — even when equipped with tools that could broaden its scope. The key takeaway is that the agentic coding space still has significant room to mature, particularly in fields like computational biology, where effective work demands both technical execution and deep domain insight.

Agents for science

In this post, I’ve shared how I’ve improved my AI-based code development experience by creating a remote Model Context Protocol (MCP) server for my scientific codebase, Napistu. The server delivers on-demand, contextually relevant information to AI agents, enabling them to efficiently surface relevant heterogeneous information and synthesize this information into actionable guidance. By deploying the server to Google Cloud Run via GitHub Actions, it can easily be used by both new and advanced users, at little cost to me.

To clarify the benefits of on-demand context, I provided case studies comparing agent behavior with and without access to the MCP server. These examples highlight the server’s impact on both onboarding efficiency and the overall development experience.

What’s next

More content

• External library documentation: Because the codebase is formatted by directly scraping the Read the docs, it would be straightforward to ingest documentation for non-Napistu libraries like igraph.
• More Napistu docs: Include the Napistu CLI, and READMEs, and the napistu.r’s pkgdown site.
• Supporting multiple Napistu versions: By preparing the core data as part of Napistu’s CI/CD workflow and saving the results to GCS, the server can download and cache a local version based on the user’s request.

More power

• Cross component semantic search: This would allow agents to search across multiple components, providing a more comprehensive understanding of the codebase.
• Execution components running Napistu functions: The execution components enable agents to register Python objects and apply transformations using Napistu functions. While still experimental, these components would allow agents to execute multiple steps in a live Python environment (like looking up two genes and finding the shortest path between them, entirely within the execution context).

More science

• Planning features and updating documentation with Claude
• Efficiently implementing features and squashing bugs in Cursor

🔧 Getting Started

Want to explore or contribute?

• Configure Claude Desktop or your favorite LLM with the MCP server.
• Ask questions about Napistu’s internals and architecture.
• Start contributing to open issues with AI-assisted development.
• Join our community discussions to collaborate, share ideas, and help shape the project.

From statistical associations to biological mechanisms

Modern genomics excels at identifying disease-associated genes and proteins through statistical analysis. Methods like Gene Set Enrichment Analysis (GSEA) group these genes into functional categories, offering useful biological context. However, we aim to go beyond simply identifying which genes and gene sets change. Our goal is to understand why these genes change together, uncovering the mechanistic depth typically seen in Figure 1 of a Cell paper. To achieve this, we must identify key molecular components, summarize their interactions, and characterize the dynamic cascades that drive emergent biological behavior.

In this post, I’ll demonstrate how to gain this insight by mapping statistical disease signatures onto genome-scale biological networks. Then, using personalized PageRank, I’ll trace signals from dysregulated genes back to their shared regulatory origins. This transforms lists of differentially expressed genes into interconnected modules that reveal upstream mechanisms driving coordinated molecular changes.

Napistu’s implementation

Napistu makes network biology practical through three core capabilities:

1. Pathway representation with SBML_dfs

   Napistu uses a custom format, SBML_dfs, to faithfully capture regulatory mechanisms. It tracks genes, metabolites, protein complexes, and drugs as molecular species, connecting them through regulatory interactions and biochemical transformations.

2. Translation to NapistuGraphs

   It provides tools to convert these pathway representations into NapistuGraph, where vertices represent molecular species and edges represent direct regulatory relationships.

3. Biological query capabilities

   Napistu enables users to ask general-purpose questions of these networks, such as:

   • What is the relationship between two genes?
   • What are the direct and indirect regulators of a molecular target?
   • And — as we’ll explore in this post — what shared mechanisms unite a set of disease-associated genes?

Throughout this post, I’ll use two types of asides to add context without interrupting the main flow:

• 🟩 Green boxes offer biological background and systems biology “inside baseball.”
• 🟦 Blue boxes share reflections on building scientific software in the age of AI.

Series overview

Part 1: Creating Multimodal Disease Profiles established the foundation for this post by systematically extracting disease-relevant molecular signatures from the Forny et al. methylmalonic acidemia dataset. Through careful batch effect correction and both supervised and unsupervised analyses, I uncovered coordinated gene and protein expression programs linked to key disease phenotypes. The result? Clean, quantitative profiles — ready for network-level mechanistic exploration. Most profiles were generated using Generalized Additive Models (GAMs), each combining regression summaries (effect size, statistic, p/q-value) with phenotypes — such as case vs. control, MMA urine levels reflecting metabolic burden, or OHCblPlus as a proxy for enzyme activity.

In this post, I’ll decode these statistical signals by mapping them onto genome-scale biological networks with Napistu. The goal is to trace disease signals from dysregulated genes and proteins upstream to their common regulatory drivers. I’ll begin by mapping statistical results onto genes within the pathway model, then transfer these signals to nodes in a regulatory graph. Finally, using personalized PageRank with empirical null models, I’ll identify subgraphs enriched for disease signals — revealing the upstream regulatory mechanisms driving MMA pathophysiology.

Integrating genome-scale networks and genome-wide data

Environment setup

To reproduce this analysis:

1. Follow the setup instructions and run the first notebook in the series creating_multimodal_profiles.qmd. This will set up both the Python environment and the input data required for this analysis.

2. Download the napistu_network_propagation.qmd notebook

3. Modify the following code block in your copy of the notebook to set appropriate paths:

   a. CACHE_DIR should match the value used in creating_multimodal_profiles.qmd b. INPUT_DATA_DIR should be a suitable location for saving the network representations (~4 GB in size)

4. Run the notebook and render an HTML output by executing:

   quarto render napistu_network_propagation.qmd
   

First, I’ll load the necessary Python modules, configure file paths, set global parameters, and define utility functions.

import os
import re
from pathlib import Path
from types import SimpleNamespace

import mudata as md
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from IPython.display import display, HTML

from napistu import utils as napistu_utils
from napistu.sbml_dfs_core import SBML_dfs
from napistu.network.ng_core import NapistuGraph
from napistu.source import unnest_sources
from napistu.gcs import downloads
from napistu.matching import mount
from napistu.network import net_propagation
from napistu.network import data_handling
from napistu.network import ng_utils
from napistu.scverse.loading import prepare_anndata_results_df
from napistu.scverse.loading import prepare_mudata_results_df
from napistu.constants import ONTOLOGIES, MINI_SBO_TO_NAME

from shackett_utils.statistics import hypothesis_testing
from shackett_utils.statistics import multi_model_fitting
from shackett_utils.blog.html_utils import display_tabulator
from shackett_utils.utils.pd_utils import format_numeric_columns

# setup logging
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)
logging.getLogger('matplotlib.pyplot').setLevel(logging.WARNING) 

# File paths and data organization
# All input data should be placed in the SUPPLEMENTAL_DATA_DIR
# Cached results and models will be stored in CACHE_DIR

# paths
PROJECT_DIR = os.path.expanduser("~/napistu_mma_posts")
INPUT_DATA_DIR = os.path.join(PROJECT_DIR, "input")
CACHE_DIR = os.path.join(PROJECT_DIR, "cache")

# inputs
# model to download from GCS and store in NAPISTU_DATA_DIR
NAPISTU_ASSET = "human_consensus"
NAPISTU_ASSET_VERSION = "20250901"
# H5Mu file containing the optimal model from MOFA+ and regression summaries
OPTIMAL_MODEL_H5MU_OUTFILE = "mofa_optimal_model.h5mu"

# intermediate files
PPR_NULL_CACHE_OUTFILE = "ppr_null_cache.tsv"

# outputs
PPR_RESULTS_OUTFILE = "ppr_results.tsv"
SBML_DFS_W_DATA_OUTFILE = "sbml_dfs_w_data.pkl"
NAPISTU_GRAPH_W_DATA_OUTFILE = "napistu_graph_w_data.pkl"

# Paths to input/output files
OPTIMAL_MODEL_H5MU_PATH = os.path.join(CACHE_DIR, OPTIMAL_MODEL_H5MU_OUTFILE)
PPR_NULL_TMP_PATH = os.path.join(CACHE_DIR, PPR_NULL_CACHE_OUTFILE)
PPR_RESULTS_PATH = os.path.join(PROJECT_DIR, PPR_RESULTS_OUTFILE)
SBML_DFS_W_DATA_PATH = os.path.join(PROJECT_DIR, SBML_DFS_W_DATA_OUTFILE)
NAPISTU_GRAPH_W_DATA_PATH = os.path.join(PROJECT_DIR, NAPISTU_GRAPH_W_DATA_OUTFILE)

# dataset metadata
FORNY_MODALITIES = SimpleNamespace(
    TRANSCRIPTOMICS = "transcriptomics",
    PROTEOMICS = "proteomics"
)

MODALITIES = list(FORNY_MODALITIES.__dict__.values())

# Napistu controlled vocabulary
FORNY_ONTOLOGIES = SimpleNamespace(
    ENSEMBL_GENE = ONTOLOGIES.ENSEMBL_GENE,
    UNIPROT = ONTOLOGIES.UNIPROT
)

FORNY_DEFS = SimpleNamespace(
    # varm table names set in part
    LFS = "LFs",
    # table names used to add data sources to `sbml_dfs`
    MOFA_LFS = "mofa_lfs",
    VAR_LEVEL_RESULTS = "var_level_results",
    # template for is_X variables will be used to restrict vertex permutation
    # to measured proteins/transcripts
    INDICATOR_STR = 'is_{modality}',
    MODALITY_VAR_LEVEL_RESULTS_STR = "{modality}_var_level_results" 
)

MUDATA_ONTOLOGIES = {
    # these dicts indicate the ontology that we want to match against for each modality
    # and indicate that this ontology's identifiers are present in the .var table's index
    FORNY_MODALITIES.TRANSCRIPTOMICS :
        {
            "ontologies" : [FORNY_ONTOLOGIES.ENSEMBL_GENE],
            "index_which_ontology" : FORNY_ONTOLOGIES.ENSEMBL_GENE
        },
    FORNY_MODALITIES.PROTEOMICS :
        {
            "ontologies" : [FORNY_ONTOLOGIES.UNIPROT],
            "index_which_ontology" : FORNY_ONTOLOGIES.UNIPROT
        }
}

# attributes to use for network propagation
LFS_OF_INTEREST = ["LF1", "LF2", "LF3", "LF4", "LF5"]

# regression terms to add from var table
PPR_LINEAR_PHENOTYPES = {"MMA_urine", "OHCblPlus", "case", "responsive_to_acute_treatment"}
PPR_SMOOTH_PHENOTYPES = {"date_freezing", "proteomics_runorder"}

VAR_VARS = list()
for phenotype in PPR_LINEAR_PHENOTYPES:
    VAR_VARS.append(f"est_{phenotype}")
    VAR_VARS.append(f"stat_{phenotype}")
    # using -log10p since normal p- and q-values will underflow
    VAR_VARS.append(f"log10p_{phenotype}")
    VAR_VARS.append(f"q_{phenotype}")
for phenotype in PPR_SMOOTH_PHENOTYPES:
    VAR_VARS.append(f"q_{phenotype}")

STAT_PREFIXES = ["est", "log10p", "q", "stat"]
VAR_METADATA = pd.DataFrame([
  *[{ "phenotype" : "latent factors", "summary" : x, "variable" : x} for x in LFS_OF_INTEREST],
  *[{ "phenotype" : x, "summary" : y, "variable" : f"{y}_{x}"} for x in (PPR_LINEAR_PHENOTYPES | PPR_SMOOTH_PHENOTYPES) for y in STAT_PREFIXES if f"{y}_{x}" in VAR_VARS]
])
VAR_METADATA["summary"] = VAR_METADATA["summary"].str.replace('_', ' ')

# defining variables to add as vertex attributes and how to transform them so
# they are appropriate for personalized pagerank reset probability
ATTRIBUTES_TO_GRAPH_SPEC = [
    {
        "attribute_names": "LF",
        "table_name": FORNY_DEFS.MOFA_LFS,
        "transformation": "square"
    },
    {
        "attribute_names": "^est_",
        "table_name": FORNY_DEFS.VAR_LEVEL_RESULTS,
        "transformation": "square"
    },
    {
        "attribute_names": "^stat_",
        "table_name": FORNY_DEFS.VAR_LEVEL_RESULTS,
        "transformation": "abs"
    },
    {
        "attribute_names": "^log10p_",
        "table_name": FORNY_DEFS.VAR_LEVEL_RESULTS,
        "transformation": "negate"
    },
    {
        "attribute_names": "^q_",
        "table_name": FORNY_DEFS.VAR_LEVEL_RESULTS,
        "transformation": "underflow_guarded_nlog10"
    },
]

ATTRIBUTES_TO_GRAPH_SPEC = ATTRIBUTES_TO_GRAPH_SPEC + [{
    "attribute_names": FORNY_DEFS.INDICATOR_STR.format(modality = m),
    "table_name": FORNY_DEFS.MODALITY_VAR_LEVEL_RESULTS_STR.format(modality = m),
    "transformation": "identity"
} for m in MODALITIES]

# masks from vertex attribute name to modality to use during vertex permutation 
REGEXES_TO_MASKS = { x: FORNY_DEFS.INDICATOR_STR.format(modality = x) for x in MODALITIES }

# utility functions

def underflow_guarded_nlog10(x):
    if x < 1e-12:
        return 1e-12 # underflow guard
    else:
        return -np.log10(x)

CUSTOM_TRANSFORMATIONS = {
    # take the absolute value
    "abs" : lambda x: abs(x),
    "negate" : lambda x: -x,
    # -log10[pvalue]
    "underflow_guarded_nlog10" : underflow_guarded_nlog10,
    "square" : lambda x: x**2
}

def floor_pvalue_by_resolution(p_value, n_samples):
    """
    Floor p-values by resolution.
    """
    
    return (p_value + 1 / n_samples) * (n_samples / (n_samples + 1))

def create_stacked_barplot_seaborn(df):
    """
    Alternative version using seaborn styling
    """
    # Set seaborn style
    sns.set_style("whitegrid")
    
    # Group by measure and sum counts across modalities
    total_counts = df.groupby('variable')['count'].sum().sort_values(ascending=False)
    
    # Create pivot table
    pivot_df = df.pivot_table(index='variable', columns='modality', values='count', fill_value=0)
    pivot_df = pivot_df.reindex(total_counts.index)
    
    # Create the plot
    fig, ax = plt.subplots(figsize=(16, 8))
    
    # Use seaborn color palette
    colors = sns.color_palette("husl", len(pivot_df.columns))
    
    # Plot stacked bars
    pivot_df.plot(kind='bar', stacked=True, ax=ax, color=colors, alpha=0.8)
    
    # Customize
    ax.set_title('Stacked Barplot by Attributes and Modality', fontsize=16, fontweight='bold')
    ax.set_xlabel('Attributes (ordered by total count)', fontsize=12)
    ax.set_ylabel('Count', fontsize=12)
    ax.legend(title='Modality', bbox_to_anchor=(1.05, 1), loc='upper left')
    
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.show()
    
    return fig, ax

def plot_ppr_enrichment_histograms(fdr_controlled_results):
    fig, axes = plt.subplots(1, 2, figsize=(8, 4), sharey=True)

    fdr_controlled_results[fdr_controlled_results["is_enriched"] == False]["p_value"].hist(
        bins=50, ax=axes[0]
    )
    axes[0].set_title("Depleted (False)")
    axes[0].set_xlabel("P-value")
    axes[0].set_ylabel("Count")

    fdr_controlled_results[fdr_controlled_results["is_enriched"] == True]["p_value"].hist(
        bins=50, ax=axes[1]
    )
    axes[1].set_title("Enriched (True)")
    axes[1].set_xlabel("P-value")

    plt.tight_layout()
    plt.show()

def reorder_by_rank_sum(df):
    """Reorder rows by sum of ranks (lower sum = better overall rank)"""
    df_num = df.replace('.', pd.NA).apply(pd.to_numeric, errors='coerce')
    max_val = df_num.max().max()
    df_filled = df_num.fillna(max_val + 1)
    rank_sums = df_filled.sum(axis=1)
    return df.loc[rank_sums.sort_values().index]

# constants affecting behavior
N_NULL_SAMPLES = 500
OVERWRITE = False