From 95a12eef159422c77919aecf8e8dd517ea64db5b Mon Sep 17 00:00:00 2001 From: Louisa Cornelis Date: Fri, 11 Sep 2026 10:48:24 -0700 Subject: [PATCH] Add optional learnable node identity features. Append or add a trainable embedding for each node before feature encoding and message passing. Keep the feature independent of k-fold splitting and external GenePT data, with focused tests for graph and flat encoders. Co-authored-by: Cursor --- README.md | 16 ++ configs/gene_identity/disabled.yaml | 2 + configs/gene_identity/learnable.yaml | 11 ++ configs/train.yaml | 1 + ogbench/nn/encoders/__init__.py | 2 + ogbench/nn/encoders/gene_identity.py | 221 ++++++++++++++++++++++++ ogbench/run.py | 8 + tests/nn/encoders/test_gene_identity.py | 178 +++++++++++++++++++ 8 files changed, 439 insertions(+) create mode 100644 configs/gene_identity/disabled.yaml create mode 100644 configs/gene_identity/learnable.yaml create mode 100644 ogbench/nn/encoders/gene_identity.py create mode 100644 tests/nn/encoders/test_gene_identity.py diff --git a/README.md b/README.md index 3ed9f116..16e11e88 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,7 @@ python ogbench/run.py dataset=addneuromed model=graph_sage trainer=ddp OGBench uses [Hydra](https://hydra.cc/) for configuration management. Key config groups: - `configs/dataset/` — dataset-specific settings (features, classes, splits, baselines) +- `configs/gene_identity/` — optional learnable node identity before message passing - `configs/model/` — model architectures and hyperparameters - `configs/trainer/` — training backend (`cpu`, `gpu`, `mps`, `ddp`, `ddp_sim`) - `configs/logger/` — logging backends (WandB, TensorBoard, CSV, MLflow, etc.) @@ -130,6 +131,21 @@ python ogbench/run.py dataset=brca model=gin \ seed=123 ``` +### Learnable node identity + +Graphs share a fixed node order, but expression alone does not identify which +gene or marker each row represents. Append a trainable embedding for every node +before the feature encoder and message-passing layers: + +```bash +python -m ogbench dataset=brca model=gcn gene_identity=learnable +``` + +The default `combine=concat` appends 32 identity channels and updates the +configured encoder dimensions. Set `gene_identity.embed_dim=64` to change the +embedding size, or use `gene_identity.combine=add` to project identity into the +existing feature dimension. + ## Baselines — GNN-Features Pipeline OGBench supports a hybrid baseline approach: train a GNN to learn node embeddings, then use those embeddings as features for sklearn classifiers. This isolates the value of the graph structure from the classifier head. diff --git a/configs/gene_identity/disabled.yaml b/configs/gene_identity/disabled.yaml new file mode 100644 index 00000000..ffabba97 --- /dev/null +++ b/configs/gene_identity/disabled.yaml @@ -0,0 +1,2 @@ +# No gene / marker identity features (default). +mode: disabled diff --git a/configs/gene_identity/learnable.yaml b/configs/gene_identity/learnable.yaml new file mode 100644 index 00000000..e29623fb --- /dev/null +++ b/configs/gene_identity/learnable.yaml @@ -0,0 +1,11 @@ +# Learnable per-node identity embedding for shared, fixed-order graphs. +# +# Usage: +# python -m ogbench gene_identity=learnable +# python -m ogbench gene_identity=learnable gene_identity.embed_dim=64 + +mode: learnable +embed_dim: 32 + +# concat appends identity channels; add projects identity to the feature dimension. +combine: concat diff --git a/configs/train.yaml b/configs/train.yaml index 27d63983..eb2c2ba1 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -17,6 +17,7 @@ defaults: - extras: default - hydra: default - hf: default + - gene_identity: disabled # experiment configs allow for version control of specific hyperparameters # e.g. best hyperparameters for given model and datamodule diff --git a/ogbench/nn/encoders/__init__.py b/ogbench/nn/encoders/__init__.py index e88c5ca9..b2f7192a 100644 --- a/ogbench/nn/encoders/__init__.py +++ b/ogbench/nn/encoders/__init__.py @@ -3,6 +3,7 @@ from ogbench.nn.encoders.all_cell_encoder import AllCellFeatureEncoder from ogbench.nn.encoders.dgm_encoder import DGMStructureFeatureEncoder from ogbench.nn.encoders.flat_encoder import FlatEncoder +from ogbench.nn.encoders.gene_identity import GeneIdentityFeatureEncoder # Create dictionary of all feature encoders FEATURE_ENCODERS: dict[str, type] = { @@ -20,4 +21,5 @@ 'AllCellFeatureEncoder', 'DGMStructureFeatureEncoder', 'FlatEncoder', + 'GeneIdentityFeatureEncoder', ] diff --git a/ogbench/nn/encoders/gene_identity.py b/ogbench/nn/encoders/gene_identity.py new file mode 100644 index 00000000..8989e99a --- /dev/null +++ b/ogbench/nn/encoders/gene_identity.py @@ -0,0 +1,221 @@ +"""Learnable node-identity features injected before message passing.""" + +from __future__ import annotations + +import logging +from typing import Any + +import torch +import torch.nn as nn +import torch_geometric +from omegaconf import DictConfig, OmegaConf, open_dict + +from ogbench.nn.encoders.base import AbstractFeatureEncoder + +logger = logging.getLogger(__name__) + + +def is_gene_identity_enabled(gene_cfg: DictConfig | dict | None) -> bool: + """Return whether learnable node identity is enabled.""" + if gene_cfg is None: + return False + mode = str(gene_cfg.get('mode', 'disabled')).lower() + return mode not in {'disabled', 'none', 'null', ''} + + +class LearnableGeneIdentityBank(nn.Module): + """Learnable identity embedding indexed by shared node order.""" + + def __init__(self, num_nodes: int, embed_dim: int) -> None: + super().__init__() + if num_nodes <= 0: + raise ValueError(f'num_nodes must be positive, got {num_nodes}') + if embed_dim <= 0: + raise ValueError(f'embed_dim must be positive, got {embed_dim}') + + self.emb = nn.Embedding(num_nodes, embed_dim) + self.out_dim = embed_dim + nn.init.normal_(self.emb.weight, mean=0.0, std=0.02) + + def forward(self) -> torch.Tensor: + """Return the identity table with shape ``[num_nodes, embed_dim]``.""" + return self.emb.weight + + +class GeneIdentityFeatureEncoder(AbstractFeatureEncoder): + """Wrap a feature encoder and inject node identity into ``data.x`` first.""" + + def __init__( + self, + base_encoder: nn.Module, + bank: LearnableGeneIdentityBank, + *, + combine: str = 'concat', + add_in_channels: int | None = None, + ) -> None: + super().__init__() + self.base_encoder = base_encoder + self.bank = bank + self.combine = combine.lower() + if self.combine not in {'concat', 'add'}: + raise ValueError(f"combine must be 'concat' or 'add', got {combine!r}") + + self.add_proj: nn.Linear | None = None + if self.combine == 'add': + if add_in_channels is None: + raise ValueError("combine='add' requires the input feature dimension") + self.add_proj = nn.Linear(bank.out_dim, add_in_channels, bias=False) + + # Mirror attributes consumed by model/readout configuration. + self.in_channels = getattr(base_encoder, 'in_channels', None) + self.out_channels = getattr(base_encoder, 'out_channels', None) + + @staticmethod + def _expand_bank(embeddings: torch.Tensor, num_rows: int) -> torch.Tensor: + """Tile one shared identity table across a batch of fixed-order graphs.""" + num_nodes = embeddings.size(0) + if num_rows == num_nodes: + return embeddings + if num_rows % num_nodes != 0: + raise ValueError( + f'Cannot align gene-identity bank with {num_nodes} nodes to ' + f'{num_rows} feature rows. Graphs must share a fixed node order.' + ) + batch_size = num_rows // num_nodes + return embeddings.unsqueeze(0).expand(batch_size, -1, -1).reshape(num_rows, -1) + + def forward(self, data: torch_geometric.data.Data) -> torch_geometric.data.Data: + """Inject identity features, then apply the configured feature encoder.""" + if not hasattr(data, 'x') or data.x is None: + raise AttributeError('GeneIdentityFeatureEncoder expects data.x') + + embeddings = self.bank() + embeddings = embeddings.to(device=data.x.device, dtype=data.x.dtype) + embeddings = self._expand_bank(embeddings, data.x.size(0)) + + if self.combine == 'concat': + data.x = torch.cat([data.x, embeddings], dim=-1) + else: + assert self.add_proj is not None + data.x = data.x + self.add_proj(embeddings) + + return self.base_encoder(data) + + +def _dataset_num_nodes(dataset: Any) -> int: + """Read the actual shared node count from a loaded graph.""" + if len(dataset) == 0: + raise ValueError('Gene identity requires a non-empty dataset') + + sample = dataset[0] + num_nodes = getattr(sample, 'num_nodes', None) + if num_nodes is None and getattr(sample, 'x', None) is not None: + num_nodes = sample.x.size(0) + if num_nodes is None: + raise ValueError('Could not determine num_nodes from the loaded dataset') + return int(num_nodes) + + +def _resolve_encoder_name(model_cfg: DictConfig) -> str: + name = OmegaConf.select(model_cfg, 'feature_encoder.encoder_name') + if name: + return str(name) + target = str(OmegaConf.select(model_cfg, 'feature_encoder._target_') or '') + return target.rsplit('.', 1)[-1] + + +def _bump_channels_for_concat(cfg: DictConfig, embed_dim: int, num_nodes: int) -> None: + """Adjust feature-encoder dimensions for appended identity channels.""" + model = OmegaConf.to_container(cfg.model, resolve=True) + assert isinstance(model, dict) + feature_encoder = model.get('feature_encoder') or {} + encoder_name = _resolve_encoder_name(cfg.model) + + with open_dict(cfg.model.feature_encoder): + if encoder_name == 'FlatEncoder': + in_channels = int(feature_encoder['in_channels']) + cfg.model.feature_encoder.out_channels = num_nodes * (in_channels + embed_dim) + return + + in_channels = feature_encoder.get('in_channels') + if isinstance(in_channels, list): + updated = list(in_channels) + updated[0] = int(updated[0]) + embed_dim + cfg.model.feature_encoder.in_channels = updated + elif in_channels is not None: + cfg.model.feature_encoder.in_channels = int(in_channels) + embed_dim + + +def setup_gene_identity( + cfg: DictConfig, + dataset: Any, +) -> LearnableGeneIdentityBank | None: + """Build the configured identity bank and adjust model channel dimensions.""" + gene_cfg = cfg.get('gene_identity') + if not is_gene_identity_enabled(gene_cfg): + return None + + mode = str(gene_cfg.get('mode')).lower() + if mode != 'learnable': + raise ValueError(f"Unsupported gene_identity.mode={mode!r}; expected 'learnable'") + + num_nodes = _dataset_num_nodes(dataset) + configured_nodes = OmegaConf.select(cfg, 'dataset.parameters.num_nodes') + if configured_nodes is not None and int(configured_nodes) != num_nodes: + raise ValueError( + 'Gene identity requires dataset.parameters.num_nodes to match the loaded graph: ' + f'configured={configured_nodes}, actual={num_nodes}' + ) + + embed_dim = int(gene_cfg.get('embed_dim', 32)) + bank = LearnableGeneIdentityBank(num_nodes, embed_dim) + combine = str(gene_cfg.get('combine', 'concat')).lower() + if combine == 'concat': + _bump_channels_for_concat(cfg, embed_dim, num_nodes) + elif combine != 'add': + raise ValueError(f"combine must be 'concat' or 'add', got {combine!r}") + + with open_dict(cfg.gene_identity): + cfg.gene_identity._runtime = { + 'mode': mode, + 'combine': combine, + 'embed_dim': embed_dim, + 'num_nodes': num_nodes, + } + logger.info( + 'Gene identity [%s]: embed_dim=%d, combine=%s, num_nodes=%d', + mode, + embed_dim, + combine, + num_nodes, + ) + return bank + + +def apply_gene_identity_to_model( + model: nn.Module, + bank: LearnableGeneIdentityBank | None, + cfg: DictConfig, +) -> nn.Module: + """Wrap ``model.feature_encoder`` when learnable identity is enabled.""" + if bank is None: + return model + + combine = str(cfg.gene_identity.get('combine', 'concat')).lower() + add_in_channels = None + if combine == 'add': + in_channels = getattr(model.feature_encoder, 'in_channels', None) + if isinstance(in_channels, list | tuple) and in_channels: + add_in_channels = int(in_channels[0]) + elif isinstance(in_channels, int): + add_in_channels = in_channels + else: + add_in_channels = int(OmegaConf.select(cfg, 'dataset.parameters.num_features') or 1) + + model.feature_encoder = GeneIdentityFeatureEncoder( + model.feature_encoder, + bank, + combine=combine, + add_in_channels=add_in_channels, + ) + return model diff --git a/ogbench/run.py b/ogbench/run.py index 16400d82..c0ccbe97 100644 --- a/ogbench/run.py +++ b/ogbench/run.py @@ -12,6 +12,10 @@ from ogbench.data.preprocessor import PreProcessor from ogbench.dataloader import TBDataloader +from ogbench.nn.encoders.gene_identity import ( + apply_gene_identity_to_model, + setup_gene_identity, +) from ogbench.utils import ( RankedLogger, extras, @@ -124,6 +128,9 @@ def run(cfg: DictConfig) -> tuple[dict[str, Any], dict[str, Any]]: else: raise ValueError('Invalid task_level') + # Inject optional learnable node identity before the feature encoder. + gene_identity_bank = setup_gene_identity(cfg, dataset) + # Model for us is Network + logic: inputs backbone, readout, losses log.info(f'Instantiating model <{cfg.model._target_}>') @@ -133,6 +140,7 @@ def run(cfg: DictConfig) -> tuple[dict[str, Any], dict[str, Any]]: optimizer=cfg.optimizer, loss=cfg.loss, ) + model = apply_gene_identity_to_model(model, gene_identity_bank, cfg) log.info('Instantiating callbacks...') callbacks: list[Callback] = instantiate_callbacks(cfg.get('callbacks')) diff --git a/tests/nn/encoders/test_gene_identity.py b/tests/nn/encoders/test_gene_identity.py new file mode 100644 index 00000000..6b7262d2 --- /dev/null +++ b/tests/nn/encoders/test_gene_identity.py @@ -0,0 +1,178 @@ +"""Tests for learnable node-identity injection.""" + +from types import SimpleNamespace + +import pytest +import torch +from omegaconf import OmegaConf +from torch_geometric.data import Data + +from ogbench.nn.encoders.all_cell_encoder import AllCellFeatureEncoder +from ogbench.nn.encoders.flat_encoder import FlatEncoder +from ogbench.nn.encoders.gene_identity import ( + GeneIdentityFeatureEncoder, + LearnableGeneIdentityBank, + apply_gene_identity_to_model, + is_gene_identity_enabled, + setup_gene_identity, +) + + +def test_is_gene_identity_enabled(): + assert not is_gene_identity_enabled(None) + assert not is_gene_identity_enabled(OmegaConf.create({'mode': 'disabled'})) + assert is_gene_identity_enabled(OmegaConf.create({'mode': 'learnable'})) + + +def test_learnable_bank_shape_and_validation(): + bank = LearnableGeneIdentityBank(num_nodes=5, embed_dim=3) + assert bank().shape == (5, 3) + assert bank().requires_grad + + with pytest.raises(ValueError, match='num_nodes must be positive'): + LearnableGeneIdentityBank(num_nodes=0, embed_dim=3) + with pytest.raises(ValueError, match='embed_dim must be positive'): + LearnableGeneIdentityBank(num_nodes=5, embed_dim=0) + + +def test_concat_before_all_cell_encoder(): + num_nodes, feature_dim, embed_dim, hidden_dim = 5, 1, 3, 8 + bank = LearnableGeneIdentityBank(num_nodes, embed_dim) + base = AllCellFeatureEncoder( + in_channels=[feature_dim + embed_dim], + out_channels=hidden_dim, + ) + encoder = GeneIdentityFeatureEncoder(base, bank, combine='concat') + + data = Data(x=torch.randn(num_nodes, feature_dim)) + data.batch_0 = torch.zeros(num_nodes, dtype=torch.long) + output = encoder(data) + + assert data.x.shape == (num_nodes, feature_dim + embed_dim) + assert output.x_0.shape == (num_nodes, hidden_dim) + + +def test_concat_repeats_identity_for_batched_graphs(): + num_nodes, feature_dim, embed_dim, batch_size = 4, 2, 3, 3 + bank = LearnableGeneIdentityBank(num_nodes, embed_dim) + base = AllCellFeatureEncoder( + in_channels=[feature_dim + embed_dim], + out_channels=6, + ) + encoder = GeneIdentityFeatureEncoder(base, bank, combine='concat') + + data = Data(x=torch.randn(batch_size * num_nodes, feature_dim)) + data.batch_0 = torch.arange(batch_size).repeat_interleave(num_nodes) + output = encoder(data) + + expected = bank().detach() + for graph_index in range(batch_size): + start = graph_index * num_nodes + assert torch.allclose( + data.x[start : start + num_nodes, feature_dim:], + expected, + ) + assert output.x_0.shape == (batch_size * num_nodes, 6) + + +def test_rejects_non_shared_node_count(): + bank = LearnableGeneIdentityBank(num_nodes=4, embed_dim=2) + base = AllCellFeatureEncoder(in_channels=[3], out_channels=5) + encoder = GeneIdentityFeatureEncoder(base, bank, combine='concat') + data = Data(x=torch.randn(5, 1)) + data.batch_0 = torch.zeros(5, dtype=torch.long) + + with pytest.raises(ValueError, match='fixed node order'): + encoder(data) + + +def test_flat_encoder_with_identity(): + num_nodes, feature_dim, embed_dim, batch_size = 3, 1, 2, 2 + bank = LearnableGeneIdentityBank(num_nodes, embed_dim) + base = FlatEncoder( + in_channels=feature_dim, + out_channels=num_nodes * (feature_dim + embed_dim), + ) + encoder = GeneIdentityFeatureEncoder(base, bank, combine='concat') + + data = Data(x=torch.randn(batch_size * num_nodes, feature_dim)) + data.y = torch.zeros(batch_size, dtype=torch.long) + data.batch_size = batch_size + output = encoder(data) + + assert output.x_0.shape == (batch_size, num_nodes * (feature_dim + embed_dim)) + + +def test_add_combine_preserves_feature_shape(): + num_nodes, feature_dim, embed_dim = 4, 2, 5 + bank = LearnableGeneIdentityBank(num_nodes, embed_dim) + base = AllCellFeatureEncoder(in_channels=[feature_dim], out_channels=7) + encoder = GeneIdentityFeatureEncoder( + base, + bank, + combine='add', + add_in_channels=feature_dim, + ) + data = Data(x=torch.randn(num_nodes, feature_dim)) + data.batch_0 = torch.zeros(num_nodes, dtype=torch.long) + + output = encoder(data) + + assert data.x.shape == (num_nodes, feature_dim) + assert output.x_0.shape == (num_nodes, 7) + + +def _config(encoder_name: str, in_channels, out_channels: int): + return OmegaConf.create( + { + 'gene_identity': { + 'mode': 'learnable', + 'embed_dim': 3, + 'combine': 'concat', + }, + 'dataset': {'parameters': {'num_nodes': 4, 'num_features': 2}}, + 'model': { + 'feature_encoder': { + 'encoder_name': encoder_name, + 'in_channels': in_channels, + 'out_channels': out_channels, + } + }, + } + ) + + +def test_setup_bumps_graph_encoder_input_channels(): + cfg = _config('AllCellFeatureEncoder', [2], 8) + + bank = setup_gene_identity(cfg, [Data(x=torch.randn(4, 2))]) + + assert bank is not None + assert cfg.model.feature_encoder.in_channels == [5] + assert cfg.gene_identity._runtime.num_nodes == 4 + + +def test_setup_recomputes_flat_encoder_output_channels(): + cfg = _config('FlatEncoder', 2, 8) + + setup_gene_identity(cfg, [Data(x=torch.randn(4, 2))]) + + assert cfg.model.feature_encoder.out_channels == 20 + + +def test_setup_rejects_configured_node_mismatch(): + cfg = _config('AllCellFeatureEncoder', [2], 8) + + with pytest.raises(ValueError, match='configured=4, actual=5'): + setup_gene_identity(cfg, [Data(x=torch.randn(5, 2))]) + + +def test_apply_wraps_model_feature_encoder(): + cfg = _config('AllCellFeatureEncoder', [2], 8) + bank = LearnableGeneIdentityBank(num_nodes=4, embed_dim=3) + model = SimpleNamespace(feature_encoder=AllCellFeatureEncoder(in_channels=[5], out_channels=8)) + + wrapped = apply_gene_identity_to_model(model, bank, cfg) + + assert isinstance(wrapped.feature_encoder, GeneIdentityFeatureEncoder) + assert wrapped.feature_encoder.bank is bank