@pikaa-ai/pikaa 0.2.4 → 0.3.0

package/skills/torch-geometric/references/message_passing.md

@@ -0,0 +1,121 @@
+# Custom Message Passing Layers
+
+Full reference for implementing custom GNN layers via the `MessagePassing` base class.
+
+## MessagePassing API
+
+```python
+MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)
+```
+
+- `aggr`: Aggregation scheme — `"add"`, `"mean"`, or `"max"`
+- `flow`: Message direction — `"source_to_target"` (default) or `"target_to_source"`
+- `node_dim`: Axis along which to propagate
+
+### Methods to override
+
+- `message(...)`: Constructs messages for each edge. Access source/target node features via `_j`/`_i` suffixes.
+- `aggregate(inputs, index)`: Aggregates messages (usually handled by `aggr` parameter).
+- `update(aggr_out, ...)`: Post-aggregation transform on each node.
+- `propagate(edge_index, size=None, **kwargs)`: Orchestrates the full pipeline. Call this from `forward()`.
+
+Any tensor passed to `propagate()` can be auto-indexed in `message()` by appending `_i` (target) or `_j` (source). E.g., passing `x=features` lets you use `x_i` and `x_j` in the message function.
+
+For bipartite graphs, pass `size=(N, M)` to `propagate()` and provide features as tuples: `x=(x_src, x_dst)`.
+
+## Example: GCN Layer from Scratch
+
+```python
+import torch
+from torch.nn import Linear, Parameter
+from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import add_self_loops, degree
+
+class GCNConv(MessagePassing):
+    def __init__(self, in_channels, out_channels):
+        super().__init__(aggr='add')
+        self.lin = Linear(in_channels, out_channels, bias=False)
+        self.bias = Parameter(torch.empty(out_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.lin.reset_parameters()
+        self.bias.data.zero_()
+
+    def forward(self, x, edge_index):
+        # 1. Add self-loops
+        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
+        # 2. Linear transform
+        x = self.lin(x)
+        # 3. Compute normalization coefficients
+        row, col = edge_index
+        deg = degree(col, x.size(0), dtype=x.dtype)
+        deg_inv_sqrt = deg.pow(-0.5)
+        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
+        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
+        # 4-5. Message passing
+        out = self.propagate(edge_index, x=x, norm=norm)
+        # 6. Add bias
+        return out + self.bias
+
+    def message(self, x_j, norm):
+        # x_j: source node features for each edge [num_edges, out_channels]
+        # norm: normalization coefficients [num_edges]
+        return norm.view(-1, 1) * x_j
+```
+
+## Example: EdgeConv Layer
+
+```python
+import torch
+from torch.nn import Sequential as Seq, Linear, ReLU
+from torch_geometric.nn import MessagePassing
+
+class EdgeConv(MessagePassing):
+    def __init__(self, in_channels, out_channels):
+        super().__init__(aggr='max')
+        self.mlp = Seq(
+            Linear(2 * in_channels, out_channels),
+            ReLU(),
+            Linear(out_channels, out_channels),
+        )
+
+    def forward(self, x, edge_index):
+        return self.propagate(edge_index, x=x)
+
+    def message(self, x_i, x_j):
+        # x_i: target node features [num_edges, in_channels]
+        # x_j: source node features [num_edges, in_channels]
+        return self.mlp(torch.cat([x_i, x_j - x_i], dim=1))
+```
+
+## Example: Dynamic EdgeConv (recomputes graph each layer)
+
+```python
+from torch_geometric.nn import knn_graph
+
+class DynamicEdgeConv(EdgeConv):
+    def __init__(self, in_channels, out_channels, k=6):
+        super().__init__(in_channels, out_channels)
+        self.k = k
+
+    def forward(self, x, batch=None):
+        edge_index = knn_graph(x, self.k, batch, loop=False, flow=self.flow)
+        return super().forward(x, edge_index)
+```
+
+## Utility Functions
+
+```python
+from torch_geometric.utils import (
+    add_self_loops,      # Add self-loop edges
+    remove_self_loops,   # Remove self-loop edges
+    degree,              # Compute node degrees
+    softmax,             # Sparse softmax over neighborhoods
+    to_dense_adj,        # Convert edge_index to dense adjacency matrix
+    to_undirected,       # Make edge_index undirected
+    contains_self_loops, # Check for self-loops
+    is_undirected,       # Check if graph is undirected
+    scatter,             # Scatter operations (sum, mean, max)
+)
+```

package/skills/torch-geometric/references/scaling.md

@@ -0,0 +1,269 @@
+# Scaling GNNs — Full Reference
+
+Techniques for training GNNs on large graphs that don't fit in GPU memory, multi-GPU training, and performance optimization.
+
+## Table of Contents
+1. Neighbor Sampling (NeighborLoader)
+2. Other Sampling Strategies
+3. Multi-GPU / Distributed Training
+4. torch.compile Support
+5. Performance Tips
+
+---
+
+## 1. Neighbor Sampling (NeighborLoader)
+
+The primary approach for large single-graph training. Recursively samples a fixed number of neighbors per hop, bounding the computation graph.
+
+```python
+from torch_geometric.loader import NeighborLoader
+
+loader = NeighborLoader(
+    data,
+    num_neighbors=[15, 10],       # Max neighbors per hop (hop 1: 15, hop 2: 10)
+    batch_size=1024,               # Number of seed nodes per batch
+    input_nodes=data.train_mask,   # Which nodes to sample from
+    shuffle=True,
+    num_workers=4,                 # Parallel data loading
+    replace=False,                 # Sample without replacement
+)
+```
+
+**Key parameters:**
+- `num_neighbors`: List of max neighbors per hop. Length should match GNN depth. Use `-1` to sample all neighbors for a hop.
+- `input_nodes`: Seed nodes — can be a mask, tensor of indices, or `('node_type', mask)` tuple for hetero graphs.
+- `subgraph_type`: `"directional"` (default), `"bidirectional"` (add reverse edges), or `"induced"` (full induced subgraph).
+- `disjoint`: If `True`, don't fuse neighborhoods across seed nodes (uses more memory but can be needed).
+
+**Training pattern:**
+```python
+model = GraphSAGE(in_channels, hidden_channels, out_channels, num_layers=2)
+
+for batch in loader:
+    batch = batch.to(device)
+    out = model(batch.x, batch.edge_index)
+    # CRITICAL: only first batch_size nodes are seed nodes
+    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
+```
+
+**Important details:**
+- Nodes are sorted: first `batch.batch_size` nodes are the seed nodes
+- `batch.n_id` maps local indices back to original node IDs
+- Sampling >2-3 hops is generally infeasible (exponential neighborhood growth)
+- Keep `len(num_neighbors) == num_gnn_layers` for efficiency
+- PyG 2.7 adds `BidirectionalSampler` and forward+reverse edge sampling on `NeighborSampler` for undirected graphs
+
+### LinkNeighborLoader (for link prediction)
+
+Samples subgraphs around supervision edges:
+
+```python
+from torch_geometric.loader import LinkNeighborLoader
+
+loader = LinkNeighborLoader(
+    data,
+    num_neighbors=[20, 10],
+    edge_label_index=train_data.edge_label_index,
+    edge_label=train_data.edge_label,
+    batch_size=256,
+    neg_sampling_ratio=1.0,
+    shuffle=True,
+)
+```
+
+### HGTLoader (type-aware heterogeneous sampling)
+
+Samples a fixed number of nodes per type per hop, following HGT paper:
+
+```python
+from torch_geometric.loader import HGTLoader
+
+loader = HGTLoader(
+    data,
+    num_samples=[512] * 2,        # Nodes per type per hop
+    batch_size=128,
+    input_nodes=('paper', data['paper'].train_mask),
+)
+```
+
+## 2. Other Sampling Strategies
+
+### ClusterLoader (ClusterGCN)
+
+Partitions the graph into clusters, trains on full subgraphs. Better for deeper GNNs since messages flow freely within clusters:
+
+```python
+from torch_geometric.loader import ClusterData, ClusterLoader
+
+cluster_data = ClusterData(data, num_parts=1500)
+loader = ClusterLoader(cluster_data, batch_size=20, shuffle=True)
+
+for batch in loader:
+    # batch is a full subgraph — no slicing needed
+    out = model(batch.x, batch.edge_index)
+    loss = F.cross_entropy(out[batch.train_mask], batch.y[batch.train_mask])
+```
+
+### GraphSAINTSampler
+
+Samples subgraphs via random walks, nodes, or edges with importance-based normalization:
+
+```python
+from torch_geometric.loader import GraphSAINTRandomWalkSampler
+
+loader = GraphSAINTRandomWalkSampler(
+    data, batch_size=6000, walk_length=2, num_steps=5,
+)
+```
+
+### ShaDowKHopSampler
+
+Extracts K-hop induced subgraphs around seed nodes — decouples depth from scope:
+
+```python
+from torch_geometric.loader import ShaDowKHopSampler
+
+loader = ShaDowKHopSampler(
+    data, depth=2, num_neighbors=5, batch_size=64,
+    input_nodes=data.train_mask,
+)
+```
+
+## 3. Multi-GPU / Distributed Training
+
+`torch_geometric.distributed` is deprecated as of PyG 2.7 — use standard PyTorch DDP below.
+
+### DistributedDataParallel (DDP)
+
+Standard PyTorch DDP works with PyG. Each GPU gets a partition of the seed nodes:
+
+```python
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+import torch.multiprocessing as mp
+
+def run(rank, world_size, dataset):
+    # Initialize process group
+    os.environ['MASTER_ADDR'] = 'localhost'
+    os.environ['MASTER_PORT'] = '12345'
+    dist.init_process_group('nccl', rank=rank, world_size=world_size)
+
+    data = dataset[0]
+
+    # Split training nodes across GPUs
+    train_idx = data.train_mask.nonzero().view(-1)
+    train_idx = train_idx.split(train_idx.size(0) // world_size)[rank]
+
+    loader = NeighborLoader(
+        data,
+        input_nodes=train_idx,
+        num_neighbors=[25, 10],
+        batch_size=1024,
+        num_workers=4,
+        shuffle=True,
+    )
+
+    # Wrap model in DDP
+    model = GraphSAGE(...).to(rank)
+    model = DistributedDataParallel(model, device_ids=[rank])
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    for epoch in range(10):
+        model.train()
+        for batch in loader:
+            batch = batch.to(rank)
+            optimizer.zero_grad()
+            out = model(batch.x, batch.edge_index)[:batch.batch_size]
+            loss = F.cross_entropy(out, batch.y[:batch.batch_size])
+            loss.backward()
+            optimizer.step()
+
+        # Synchronize before evaluation
+        dist.barrier()
+
+        if rank == 0:
+            # Evaluate on rank 0 only
+            ...
+
+    dist.destroy_process_group()
+
+# Launch
+if __name__ == '__main__':
+    dataset = Reddit('./data/Reddit')
+    world_size = torch.cuda.device_count()
+    mp.spawn(run, args=(world_size, dataset), nprocs=world_size, join=True)
+```
+
+**Key points:**
+- Initialize dataset before `mp.spawn()` — data auto-moves to shared memory
+- Each rank creates its own NeighborLoader with a subset of seed nodes
+- Call `dist.barrier()` to synchronize before evaluation
+- Evaluate on rank 0 only for simplicity
+- Clean up with `dist.destroy_process_group()`
+
+### PyTorch Lightning Integration
+
+PyG provides Lightning wrappers for minimal boilerplate. PyG 2.7 supports the `lightning` package (not only legacy `pytorch-lightning`):
+
+```python
+import lightning as L
+from torch_geometric.data import LightningNodeData
+
+datamodule = LightningNodeData(
+    data,
+    input_train_nodes=data.train_mask,
+    input_val_nodes=data.val_mask,
+    input_test_nodes=data.test_mask,
+    loader='neighbor',
+    num_neighbors=[25, 10],
+    batch_size=1024,
+)
+
+# Use with any Lightning Trainer
+trainer = L.Trainer(devices=4, accelerator='gpu', strategy='ddp')
+trainer.fit(model, datamodule)
+```
+
+Also available: `LightningLinkData` for link prediction, `LightningDataset` for graph-level tasks.
+
+## 4. torch.compile Support
+
+PyG supports `torch.compile` for faster execution:
+
+```python
+model = GCN(...)
+model = torch.compile(model)
+
+# Works with standard training loops
+out = model(data.x, data.edge_index)
+```
+
+**What works:**
+- Most GNN layers (GCNConv, SAGEConv, GATConv, etc.)
+- Standard training/inference pipelines
+- Both CPU and CUDA backends
+
+**Limitations:**
+- Dynamic shapes (varying graph sizes per batch) may trigger recompilation
+- Some specialized layers or custom MessagePassing subclasses may not compile
+- Use `torch.compile(model, dynamic=True)` if batch graph sizes vary significantly
+
+## 5. Performance Tips
+
+- **num_workers**: Set `num_workers=4` (or more) in data loaders for CPU-side parallelism
+- **pin_memory**: Use `pin_memory=True` in loaders for faster CPU-to-GPU transfer
+- **Sparse tensors**: Use `SparseTensor` from `torch_sparse` instead of `edge_index` for faster message passing on some layers
+- **Profiling**: Use `torch_geometric.profile` to measure time and memory of individual layers
+- **Mixed precision**: Standard PyTorch AMP works with PyG:
+  ```python
+  from torch.amp import autocast, GradScaler
+  scaler = GradScaler()
+  with autocast('cuda'):
+      out = model(batch.x, batch.edge_index)
+      loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
+  scaler.scale(loss).backward()
+  scaler.step(optimizer)
+  scaler.update()
+  ```
+- **Reduce sampling**: Fewer neighbors per hop = faster but noisier. Start with `[15, 10]` for 2-layer GNNs.
+- **Avoid unnecessary computation**: With NeighborLoader, only the first `batch_size` outputs matter — don't compute metrics on sampled-only nodes.

package/skills/torchdrug/SKILL.md

@@ -0,0 +1,241 @@
+---
+name: torchdrug
+description: Build and troubleshoot TorchDrug 0.2.1 workflows for molecular graphs, property prediction, self-supervised pretraining, molecule generation, retrosynthesis, protein representation learning, and knowledge graph reasoning. Use when code imports torchdrug or needs its datasets, models, tasks, or Engine.
+license: Apache-2.0 license
+compatibility: TorchDrug 0.2.1 requires Python 3.7-3.10 and supports PyTorch 1.8-2.0. Apple Silicon is CPU-only; MPS is unsupported.
+allowed-tools: Read Write Edit Bash
+metadata:
+  version: "1.1"
+  skill-author: K-Dense Inc.
+---
+
+# TorchDrug
+
+Use TorchDrug as a modular PyTorch graph-learning stack:
+
+1. load a `datasets.*` dataset,
+2. choose a `models.*` representation model,
+3. wrap it in a `tasks.*` objective,
+4. train and evaluate it with `core.Engine`.
+
+The current official documentation and latest release are both **0.2.1**. Treat
+newer Python or PyTorch combinations as unverified rather than silently assuming
+compatibility.
+
+## Start with the version guard
+
+Before generating or debugging code, inspect the environment:
+
+```bash
+python --version
+python -c "import torch; print(torch.__version__)"
+python -c "import torchdrug; print(torchdrug.__version__)"
+```
+
+The supported matrix for TorchDrug 0.2.1 is:
+
+- Python 3.7 through 3.10
+- PyTorch 1.8 through 2.0
+- Linux, Windows, or macOS
+- Apple Silicon: PyTorch 1.13 or later, CPU only; no MPS support
+
+If the project uses Python 3.11+ or PyTorch 2.1+, create a compatible environment
+or explicitly test a source build. Do not present such combinations as supported.
+
+## Installation
+
+Prefer a dedicated Python 3.10 environment and pin the TorchDrug release:
+
+```bash
+uv venv --python 3.10
+source .venv/bin/activate
+uv pip install "torch==2.0.0"
+```
+
+Install `torch-scatter` and `torch-cluster` wheels matched to the exact PyTorch
+and CUDA pair, following the
+[official installation page](https://torchdrug.ai/docs/installation.html). For a
+CPU-only PyTorch 2.0 environment, one reproducible wheel combination is:
+
+```bash
+uv pip install "torch-scatter==2.1.1" "torch-cluster==1.6.1" \
+  --find-links "https://data.pyg.org/whl/torch-2.0.0+cpu.html"
+uv pip install "torchdrug==0.2.1"
+```
+
+Do not copy a CUDA wheel URL between environments. Match the PyTorch version,
+CUDA build, Python ABI, and platform. On Apple Silicon, the official docs require
+building `torch-scatter` and `torch-cluster` from source; pin reviewed source
+revisions and expect CPU execution.
+
+## Canonical property-prediction workflow
+
+Use the documented ClinTox → GIN → `PropertyPrediction` → `Engine` pattern:
+
+```python
+import torch
+from torchdrug import core, datasets, models, tasks
+
+dataset = datasets.ClinTox("~/molecule-datasets/")
+lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
+lengths.append(len(dataset) - sum(lengths))
+train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)
+
+model = models.GIN(
+    input_dim=dataset.node_feature_dim,
+    hidden_dims=[256, 256, 256, 256],
+    short_cut=True,
+    batch_norm=True,
+    concat_hidden=True,
+)
+task = tasks.PropertyPrediction(
+    model,
+    task=dataset.tasks,
+    criterion="bce",
+    metric=("auprc", "auroc"),
+)
+
+optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
+solver = core.Engine(
+    task,
+    train_set,
+    valid_set,
+    test_set,
+    optimizer,
+    batch_size=1024,
+)
+solver.train(num_epoch=100)
+solver.evaluate("valid")
+```
+
+Add `gpus=[0]` only when a supported CUDA device is available. Omit `gpus` for
+CPU execution.
+
+For binary classification, `task.predict(batch)` returns logits; apply
+`torch.sigmoid` when probabilities are needed. In 0.2.1, normalized regression
+predictions are returned on the original target scale, which is a breaking change
+from older releases.
+
+## Choose the official workflow
+
+### Molecular property prediction
+
+- Dataset: `datasets.ClinTox`, `BBBP`, `Tox21`, `QM9`, or another documented
+  molecule dataset.
+- Model: start with `models.GIN`; use `edge_input_dim` when the selected feature
+  configuration supplies edge features.
+- Task: `tasks.PropertyPrediction`.
+- Read [molecular property prediction](references/molecular_property_prediction.md).
+
+### Self-supervised molecular pretraining
+
+- InfoGraph: `models.InfoGraph(gin_model, separate_model=False)` wrapped by
+  `tasks.Unsupervised`.
+- Attribute masking: `tasks.AttributeMasking(model, mask_rate=0.15)`.
+- Recreate the same encoder for fine-tuning, then load the checkpoint with
+  `strict=False` before training `tasks.PropertyPrediction`.
+- Read [molecular property prediction](references/molecular_property_prediction.md).
+
+### Molecule generation
+
+- Dataset: `datasets.ZINC250k(..., kekulize=True, atom_feature="symbol")`.
+- GCPN: an `models.RGCN` encoder wrapped by `tasks.GCPNGeneration`.
+- GraphAF: node and edge `models.GraphAF` flows wrapped by
+  `tasks.AutoregressiveGeneration`.
+- Supported optimization tasks in the tutorial are `"qed"` and `"plogp"`;
+  criteria are `"nll"` and/or `"ppo"`.
+- Read [molecular generation](references/molecular_generation.md).
+
+### Retrosynthesis
+
+- Create two synchronized `datasets.USPTO50k` views: reaction mode for center
+  identification and `as_synthon=True` for synthon completion.
+- Train `tasks.CenterIdentification` and `tasks.SynthonCompletion` separately.
+- Combine the trained tasks with `tasks.Retrosynthesis`; do not pass raw models
+  directly to the end-to-end task.
+- Read [retrosynthesis](references/retrosynthesis.md).
+
+### Knowledge graph reasoning
+
+- Embedding workflow: `datasets.FB15k237` → `models.RotatE` →
+  `tasks.KnowledgeGraphCompletion`.
+- Neural reasoning workflow: `models.NeuralLP` with `fact_ratio=0.75`.
+- Read [knowledge graph reasoning](references/knowledge_graphs.md).
+
+### Protein modeling
+
+- Build proteins with `data.Protein.from_sequence`, `from_pdb`, or
+  `from_molecule`.
+- Sequence encoders include `models.ESM`, `ProteinCNN`, `ProteinResNet`,
+  `ProteinLSTM`, and `ProteinBERT`; structure encoders include `models.GearNet`.
+- Use documented graph-construction layers rather than a nonexistent
+  `protein.residue_graph()` convenience method.
+- Read [protein modeling](references/protein_modeling.md).
+
+## Rules for reliable TorchDrug code
+
+1. **Follow the 0.2.1 API.** The official docs are not a rolling latest-version
+   site.
+2. **Prefer documented feature names.** Use `atom_feature`, `bond_feature`,
+   `residue_feature`, and `mol_feature`; `node_feature`, `edge_feature`, and
+   `graph_feature` are deprecated aliases in relevant dataset constructors.
+3. **Let `Engine` preprocess tasks.** If composing pre-trained tasks without
+   constructing their solvers, call each task's `preprocess()` manually.
+4. **Keep paired splits synchronized.** For retrosynthesis, reset the same random
+   seed before splitting reaction and synthon datasets.
+5. **Use TorchDrug collation.** Use `data.graph_collate` or `core.Engine`;
+   generic PyTorch collation does not know how to pack TorchDrug graphs.
+6. **Separate model, task, and engine arguments.** A common source of invented
+   code is passing task options to a model or passing raw models where a composed
+   task is required.
+7. **Validate generated chemistry.** Treat model outputs as candidates, not as
+   experimentally valid or synthesizable compounds.
+
+## Troubleshooting
+
+### Installation or import failure
+
+Check Python, PyTorch, `torch-scatter`, and `torch-cluster` as one compatibility
+set. Most failures are binary-wheel mismatches, unsupported Python versions, or
+attempts to use MPS.
+
+### Feature dimension mismatch
+
+Build model dimensions from the loaded dataset:
+
+- `dataset.node_feature_dim`
+- `dataset.edge_feature_dim`
+- `dataset.num_bond_type`
+- `dataset.num_entity` and `dataset.num_relation` for knowledge graphs
+
+Do not hard-code dimensions copied from a different feature configuration.
+
+### Device mismatch
+
+Pass `gpus=[0]` to `core.Engine` for supported CUDA execution. For manual
+prediction, collate first and move the entire nested batch with `utils.cuda`.
+
+### Checkpoint mismatch
+
+Recreate the same model and feature configuration. For pretraining-to-fine-tuning
+transfer, load the checkpoint's `"model"` state with `strict=False`; for a complete
+solver, use `solver.save()` and `solver.load()`.
+
+## Reference index
+
+- [Core concepts and data structures](references/core_concepts.md)
+- [Datasets](references/datasets.md)
+- [Models and architectures](references/models_architectures.md)
+- [Molecular property prediction and pretraining](references/molecular_property_prediction.md)
+- [Protein modeling](references/protein_modeling.md)
+- [Molecular generation](references/molecular_generation.md)
+- [Retrosynthesis](references/retrosynthesis.md)
+- [Knowledge graph reasoning](references/knowledge_graphs.md)
+
+## Upstream sources
+
+- [TorchDrug 0.2.1 documentation](https://torchdrug.ai/docs/)
+- [Tutorial index](https://torchdrug.ai/docs/tutorials/)
+- [Installation](https://torchdrug.ai/docs/installation.html)
+- [Package reference](https://torchdrug.ai/docs/api/)
+- [TorchDrug 0.2.1 release notes](https://github.com/DeepGraphLearning/torchdrug/releases/tag/v0.2.1)