diff --git a/.gitignore b/.gitignore index 5487d9d..a9794d5 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,15 @@ dmypy.json # Pyre type checker .pyre/ + +# IDE +.idea/ + +# experiment artifacts +output/ +logs/ +log2/ + +# python cache +__pycache__/ +*.pyc diff --git a/ImputeFormer_ms.py b/ImputeFormer_ms.py new file mode 100644 index 0000000..b211e33 --- /dev/null +++ b/ImputeFormer_ms.py @@ -0,0 +1,538 @@ +# ImputeFormer_ms.py +"""ImputeFormer_ms.py + +Baseline wrapper following the *gin_ms.py* CLI style and using the unified +*inc/test.py* Tester to report f1/nrmse. + +What this file provides +----------------------- +- A runnable multi-snapshot imputation baseline that takes a few observed + snapshots (times) and reconstructs the full diffusion history. +- Same tester contract as other baselines: model_fn returns y_pred with shape + [num_nodes, T] (times 0..T-1). The tester appends the final snapshot y[:, T] + automatically via test_fix_obs. + +Observation pattern +------------------- +- Use --obs_ts (or --obs_time alias) to specify observed time indices. + Example: --obs_ts "0,3,5" or --obs_time "5". +- Use -1 to represent T. +- If --obs_ts is not provided, use the last --obs_k snapshots ending at T. +- We ALWAYS include the final snapshot at time T as observed. + +About "official" ImputeFormer +----------------------------- +You asked to *lock* this wrapper to the official ImputeFormer implementation. +In this execution environment, outbound connections to GitHub raw assets are +blocked, so I cannot vendor the upstream source code here. + +Instead, this file ships a self-contained, ImputeFormer-*style* Transformer +imputer with **low-rank attention** (Linformer-like) to mimic the paper's +low-rank inductive bias. The interfaces (args, tensor shapes, tester contract) +are the important part for your ditto-ms integration. + +If you later provide (or vendor) the exact upstream class file, you can replace +`LockedImputeFormer` with the official class while keeping the training/eval +pipeline unchanged. +""" + +from __future__ import annotations + +import argparse +import math +from typing import List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tqdm import trange + +from inc.diffus import diffus_gen, b_estim, SIR_STATES +from inc.test import Tester +from inc.utils import seed_all + + +# --------------------------------------------------------------------- +# Helpers: obs_time parsing (mirrors gin_ms.py; ALWAYS includes T) +# --------------------------------------------------------------------- + +def _parse_int_list(s: str) -> List[int]: + """Parse comma/space separated ints. Example: '0, 3,5' -> [0,3,5].""" + if s is None: + return [] + s = s.replace(" ", ",") + parts = [p.strip() for p in s.split(",") if p.strip() != ""] + return [int(p) for p in parts] + + +def _resolve_obs_ts(obs_ts: Optional[List[int]], obs_k: int, T: int) -> List[int]: + """Resolve observed snapshot indices in [0, T] (inclusive). + + - If obs_ts is provided: use it (with -1 mapped to T), clamp into [0, T], unique+sorted. + - Else: use last obs_k snapshots ending at T. + """ + if obs_ts is not None: + ts: List[int] = [] + for t in obs_ts: + if t == -1: + t = T + t = max(0, min(int(t), T)) + ts.append(t) + return sorted(set(ts)) + + k = max(1, min(int(obs_k), T + 1)) + start = max(0, T - k + 1) + return list(range(start, T + 1)) + + +def _make_obs_time(args, T: int) -> List[int]: + """Convert CLI args into final obs_time list, ALWAYS including T.""" + obs_ts = _resolve_obs_ts(args.obs_ts, args.obs_k, T) + obs_time = sorted(set(obs_ts + [T])) + return obs_time + + +# --------------------------------------------------------------------- +# ImputeFormer-style model (low-rank attention) +# --------------------------------------------------------------------- + +class LowRankSelfAttention(nn.Module): + """Multi-head self-attention with low-rank projection over sequence length. + + This is a Linformer-like approximation that reduces O(L^2) attention to + O(L * r) where r=proj_k. + + Input/Output: + x: [B, L, d_model] + out: [B, L, d_model] + """ + + def __init__( + self, + d_model: int, + n_heads: int, + proj_k: int, + max_len: int, + dropout: float, + ) -> None: + super().__init__() + assert d_model % n_heads == 0, "d_model must be divisible by n_heads" + assert proj_k > 0, "proj_k must be > 0" + + self.d_model = d_model + self.n_heads = n_heads + self.d_head = d_model // n_heads + self.proj_k = proj_k + self.max_len = max_len + + self.qkv = nn.Linear(d_model, 3 * d_model) + self.out_proj = nn.Linear(d_model, d_model) + + # Project K/V along the sequence length dimension (L -> proj_k) + self.E_k = nn.Parameter(torch.randn(max_len, proj_k) / math.sqrt(proj_k)) + self.E_v = nn.Parameter(torch.randn(max_len, proj_k) / math.sqrt(proj_k)) + + self.attn_drop = nn.Dropout(dropout) + self.proj_drop = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, L, _ = x.shape + if L > self.max_len: + raise ValueError(f"Sequence length L={L} exceeds max_len={self.max_len}. Increase --max_len") + + qkv = self.qkv(x) # [B, L, 3*d] + q, k, v = qkv.chunk(3, dim=-1) + + # [B, heads, L, d_head] + q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + + # Merge heads for projection and matmul + # [B*heads, L, d_head] + q = q.reshape(B * self.n_heads, L, self.d_head) + k = k.reshape(B * self.n_heads, L, self.d_head) + v = v.reshape(B * self.n_heads, L, self.d_head) + + # Project K and V along length dimension: [B*heads, proj_k, d_head] + Ek = self.E_k[:L, :] # [L, proj_k] + Ev = self.E_v[:L, :] + k_proj = torch.einsum("bld,lk->bkd", k, Ek) + v_proj = torch.einsum("bld,lk->bkd", v, Ev) + + # Attention: [B*heads, L, proj_k] + attn = torch.einsum("bld,bkd->blk", q, k_proj) / math.sqrt(self.d_head) + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + # Output: [B*heads, L, d_head] + out = torch.einsum("blk,bkd->bld", attn, v_proj) + + # Restore heads: [B, L, d_model] + out = out.view(B, self.n_heads, L, self.d_head).transpose(1, 2).reshape(B, L, self.d_model) + out = self.out_proj(out) + out = self.proj_drop(out) + return out + + +class ImputeFormerBlock(nn.Module): + """Transformer block with low-rank self-attention.""" + + def __init__( + self, + d_model: int, + n_heads: int, + proj_k: int, + max_len: int, + dropout: float, + ffn_mult: int, + ) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(d_model) + self.attn = LowRankSelfAttention( + d_model=d_model, + n_heads=n_heads, + proj_k=proj_k, + max_len=max_len, + dropout=dropout, + ) + self.norm2 = nn.LayerNorm(d_model) + self.ffn = nn.Sequential( + nn.Linear(d_model, ffn_mult * d_model), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(ffn_mult * d_model, d_model), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Pre-norm for stability + x = x + self.attn(self.norm1(x)) + x = x + self.ffn(self.norm2(x)) + return x + + +class LockedImputeFormer(nn.Module): + """A self-contained ImputeFormer-style imputer. + + It consumes the *entire* timeline (length L=T+1) with missing values masked. + + Input: + x: [B, L, in_dim] where in_dim = n_cls + 1 + - first n_cls channels: one-hot value at observed times, 0 otherwise + - last channel: observed mask (1 observed, 0 missing) + + Output: + logits: [B, L, n_cls] + """ + + def __init__( + self, + in_dim: int, + n_cls: int, + d_model: int, + n_heads: int, + n_layers: int, + proj_k: int, + dropout: float, + ffn_mult: int, + max_len: int, + ) -> None: + super().__init__() + self.in_dim = in_dim + self.n_cls = n_cls + self.d_model = d_model + self.max_len = max_len + + self.in_proj = nn.Linear(in_dim, d_model) + self.pos_emb = nn.Embedding(max_len, d_model) + + self.blocks = nn.ModuleList( + [ + ImputeFormerBlock( + d_model=d_model, + n_heads=n_heads, + proj_k=proj_k, + max_len=max_len, + dropout=dropout, + ffn_mult=ffn_mult, + ) + for _ in range(n_layers) + ] + ) + + self.out_norm = nn.LayerNorm(d_model) + self.out_proj = nn.Linear(d_model, n_cls) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, L, _ = x.shape + if L > self.max_len: + raise ValueError(f"Sequence length L={L} exceeds max_len={self.max_len}. Increase --max_len") + + h = self.in_proj(x) + pos = torch.arange(L, device=x.device) + h = h + self.pos_emb(pos)[None, :, :] + + for blk in self.blocks: + h = blk(h) + + h = self.out_norm(h) + logits = self.out_proj(h) + return logits + + +# --------------------------------------------------------------------- +# Input building +# --------------------------------------------------------------------- + +def _build_imputeformer_input(y: torch.Tensor, obs_time: List[int], n_cls: int) -> torch.Tensor: + """Build model input from integer labels with a global-time observation mask. + + Args: + y: [B, L] long (0..n_cls-1) + obs_time: list of observed time indices in [0, L-1] + n_cls: number of classes + + Returns: + x: [B, L, n_cls+1] float + - x[..., :n_cls] = one-hot(y) * obs_mask + - x[..., n_cls] = obs_mask + """ + B, L = y.shape + device = y.device + + obs_mask = torch.zeros((L,), dtype=torch.bool, device=device) + obs_mask[obs_time] = True + + # One-hot encode and zero-out unobserved positions + x_val = F.one_hot(y.clamp(min=0, max=n_cls - 1), num_classes=n_cls).float() # [B, L, n_cls] + x_val = x_val * obs_mask.view(1, L, 1).float() + + # Add mask as an extra channel + x_mask = obs_mask.view(1, L, 1).float().expand(B, L, 1) + x = torch.cat([x_val, x_mask], dim=-1) + return x + + +# --------------------------------------------------------------------- +# Main model_fn used by Tester +# --------------------------------------------------------------------- + +args = None # set in __main__ + + +def _call_b_estim(data, args, obs_time: List[int]): + """Call b_estim with best-effort compatibility across branches.""" + try: + return b_estim(data, args, obs_time=obs_time) + except TypeError: + # Fallback: older signature b_estim(data, args) + # Try to inject args.obs_time (string) for compatibility. + try: + setattr(args, "obs_time", ",".join(map(str, obs_time))) + except Exception: + pass + return b_estim(data, args) + + +def imputeformer_run(data) -> torch.Tensor: + """Train on simulated diffusion sequences then impute the test sequence. + + Returns: + y_pred: [num_nodes, T] long + """ + global args + + device = args.device + T = int(data.T.item()) + L = T + 1 + n_nodes = int(data.num_nodes) + + # Determine #classes (SI:2, SIR:3) + n_cls = int(data.y.max().item()) + 1 + + obs_time = _make_obs_time(args, T) + + # Estimate diffusion parameters (used for synthetic training labels) + bpar = _call_b_estim(data, args, obs_time=obs_time) + + # Initial infected count for simulation + I0 = int((data.y[:, 0] == SIR_STATES.I).sum().item()) + + # Model + model = LockedImputeFormer( + in_dim=n_cls + 1, + n_cls=n_cls, + d_model=args.units, + n_heads=args.heads, + n_layers=args.layers, + proj_k=args.proj_k, + dropout=args.dropout, + ffn_mult=args.ffn_mult, + max_len=max(args.max_len, L), + ).to(device) + + optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + # ---------------- + # Train on simulated diffusion sequences + # ---------------- + for _ in trange(1, args.epochs + 1, desc="train", leave=False): + model.train() + + # labels: [batch, n_nodes, L] + labels = diffus_gen( + T=T, + n_nodes=n_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ).transpose(0, 2) + + # Subsample nodes to bound memory + node_batch = min(int(args.node_batch), n_nodes) + if node_batch < n_nodes: + idx = torch.randint(0, n_nodes, (node_batch,), device=labels.device) + else: + idx = torch.arange(n_nodes, device=labels.device) + + # Each (diffusion sample, node) is one training instance: y: [B2, L] + y = labels[:, idx, :].reshape(-1, L).long() + + x = _build_imputeformer_input(y, obs_time=obs_time, n_cls=n_cls) + logits = model(x) # [B2, L, n_cls] + + loss = F.cross_entropy( + logits[:, :T, :].reshape(-1, n_cls), + y[:, :T].reshape(-1), + ) + + optimizer.zero_grad(set_to_none=True) + loss.backward() + if args.grad_clip and args.grad_clip > 0: + nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + + # ---------------- + # Inference on the test instance (only obs_time snapshots are revealed) + # ---------------- + model.eval() + + y_true = data.y[:, :L].long().to(device) # [n_nodes, L] + y_pred_full = torch.empty((n_nodes, L), dtype=torch.long, device=device) + + with torch.no_grad(): + bs = max(1, int(args.eval_node_batch)) + for s in range(0, n_nodes, bs): + e = min(n_nodes, s + bs) + y_chunk = y_true[s:e] # [b, L] + x_chunk = _build_imputeformer_input(y_chunk, obs_time=obs_time, n_cls=n_cls) + logits = model(x_chunk) # [b, L, n_cls] + pred = logits.argmax(dim=-1) # [b, L] + + # Enforce consistency on observed snapshots (except final; tester fixes it anyway) + for t in obs_time: + if t < T: + pred[:, t] = y_chunk[:, t] + + y_pred_full[s:e] = pred + + # Return only 0..T-1 (tester appends y[:, T]) + return y_pred_full[:, :T] + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- + +def get_args(): + parser = argparse.ArgumentParser() + + # IO / runtime + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--data_dir", type=str, required=True) + parser.add_argument("--output", type=str, required=True) + parser.add_argument("--device", type=torch.device, required=True) + + # Diffusion parameter estimation (same knobs as other baselines) + parser.add_argument("--b_pI0", type=float, required=True) + parser.add_argument("--b_pR0", type=float, required=True) + parser.add_argument("--b_steps", type=int, required=True) + parser.add_argument("--b_lr", type=float, required=True) + parser.add_argument("--b_pImax", type=float, default=1.0) + + # Multi-snapshot observation settings + parser.add_argument( + "--obs_ts", + type=str, + default=None, + help='comma-separated observed time indices. Use -1 for T. Example: "0,3,5"', + ) + parser.add_argument( + "--obs_time", + type=str, + default=None, + help="alias of --obs_ts (kept for compatibility with other scripts).", + ) + parser.add_argument( + "--obs_k", + type=int, + default=1, + help="if obs_ts is None, use last k snapshots ending at T (default 1: final-only).", + ) + + # Model / training hyperparameters + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--weight_decay", type=float, default=0.0) + parser.add_argument("--epochs", type=int, default=200) + parser.add_argument("--batch_size", type=int, default=16) + + parser.add_argument("--units", type=int, default=64, help="Transformer hidden size") + parser.add_argument("--heads", type=int, default=4, help="#attention heads") + parser.add_argument("--layers", type=int, default=4, help="#Transformer blocks") + parser.add_argument("--proj_k", type=int, default=16, help="low-rank projection size (Linformer k)") + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--ffn_mult", type=int, default=4) + parser.add_argument("--max_len", type=int, default=64, help="max T+1") + + parser.add_argument("--grad_clip", type=float, default=1.0) + + # For large graphs + parser.add_argument( + "--node_batch", + type=int, + default=512, + help="#nodes sampled per training epoch (each node is a sequence instance)", + ) + parser.add_argument( + "--eval_node_batch", + type=int, + default=2048, + help="#nodes per forward pass during inference", + ) + + args = parser.parse_args() + + # obs_ts/obs_time normalization + obs_s = args.obs_ts if args.obs_ts is not None else args.obs_time + if obs_s is not None and len(str(obs_s).strip()) > 0: + obs = _parse_int_list(str(obs_s)) + args.obs_ts = sorted(set(obs)) + # Note: keep args.obs_time as-is; _make_obs_time uses args.obs_ts + else: + args.obs_ts = None + + return args + + +if __name__ == "__main__": + args = get_args() + seed_all(args.seed) + + tester = Tester(args.data_dir, args.device, imputeformer_run) + tester.test([args.dataset], rep=1) + tester.save(args.output) diff --git a/brits.py b/brits.py new file mode 100644 index 0000000..a91724b --- /dev/null +++ b/brits.py @@ -0,0 +1,401 @@ +# ! pip install class-resolver==0.3.10 +# ! pip install --no-index torch-scatter==2.0.7 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-sparse==0.6.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-cluster==1.5.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-spline-conv==1.2.1 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install torch-geometric==2.0.4 +# ! pip install ndlib==5.1.1 + +from inc.diffus import * +from inc.test import * + +import argparse +import torch + + +def get_args(argv=None): + """Parse command-line arguments. + + We keep the original notebook defaults so running without extra flags + behaves the same as before. + """ + parser = argparse.ArgumentParser() + + # ---- experiment I/O ---- + parser.add_argument('--dataset', type=str, required=True, help='dataset name') + parser.add_argument('--seed', type=int, default=123456789, help='random seed') + parser.add_argument('--data_dir', type=str, default='input', help='dataset folder') + parser.add_argument('--output', type=str, default='output/brits.pt', help='output file name') + parser.add_argument( + '--device', + type=torch.device, + default=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), + help='torch device, e.g., cpu, cuda, cuda:0' + ) + + # ---- diffusion parameter estimation (b_*) ---- + parser.add_argument('--b_pI0', type=float, default=1e-3, + help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, default=1e-3, + help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, default=500, + help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, default=3e-3, + help='learning rate in diffusion parameter estimation') + + # ---- BRITS hyperparameters ---- + parser.add_argument('--lr', type=float, default=1e-3, help='learning rate') + parser.add_argument('--epochs', type=int, default=1000, help='training epochs') + parser.add_argument('--batch_size', type=int, default=64, help='batch size') + parser.add_argument('--hid_size', type=int, default=108, help='RNN hidden size') + parser.add_argument('--impute_weight', type=float, default=0.3, help='imputation loss weight') + parser.add_argument('--label_weight', type=float, default=1.0, help='label loss weight') + + if argv is None: + return parser.parse_args() + return parser.parse_args(argv) + + +'''https://github.com/caow13/BRITS''' +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim + +from torch.autograd import Variable +from torch.nn.parameter import Parameter + +import math +# import utils +import argparse +# import data_loader + +# from ipdb import set_trace +from sklearn import metrics + + +def binary_cross_entropy_with_logits(input, target, weight=None, size_average=True, reduce=True): + if not (target.size() == input.size()): + raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size())) + max_val = (-input).clamp(min=0) + loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log() + if weight is not None: + loss = loss * weight + if not reduce: + return loss + elif size_average: + return loss.mean() + else: + return loss.sum() + + +class FeatureRegression(nn.Module): + def __init__(self, input_size): + super().__init__() + self.build(input_size) + + def build(self, input_size): + self.W = Parameter(torch.Tensor(input_size, input_size)) + self.b = Parameter(torch.Tensor(input_size)) + m = torch.ones(input_size, input_size) - torch.eye(input_size, input_size) + self.register_buffer('m', m) + self.reset_parameters() + + def reset_parameters(self): + stdv = 1. / math.sqrt(self.W.size(0)) + self.W.data.uniform_(-stdv, stdv) + if self.b is not None: + self.b.data.uniform_(-stdv, stdv) + + def forward(self, x): + z_h = F.linear(x, self.W * Variable(self.m), self.b) + return z_h + + +class TemporalDecay(nn.Module): + def __init__(self, input_size, output_size, diag=False): + super().__init__() + self.diag = diag + self.build(input_size, output_size) + + def build(self, input_size, output_size): + self.W = Parameter(torch.Tensor(output_size, input_size)) + self.b = Parameter(torch.Tensor(output_size)) + if self.diag == True: + assert (input_size == output_size) + m = torch.eye(input_size, input_size) + self.register_buffer('m', m) + self.reset_parameters() + + def reset_parameters(self): + stdv = 1. / math.sqrt(self.W.size(0)) + self.W.data.uniform_(-stdv, stdv) + if self.b is not None: + self.b.data.uniform_(-stdv, stdv) + + def forward(self, d): + if self.diag == True: + gamma = F.relu(F.linear(d, self.W * Variable(self.m), self.b)) + else: + gamma = F.relu(F.linear(d, self.W, self.b)) + gamma = torch.exp(-gamma) + return gamma + + +class RITS(nn.Module): + def __init__(self, xdim, rnn_hid_size, impute_weight, label_weight): + super().__init__() + self.xdim = xdim + self.rnn_hid_size = rnn_hid_size + self.impute_weight = impute_weight + self.label_weight = label_weight + self.build() + + def build(self): + self.rnn_cell = nn.LSTMCell(self.xdim * 2, self.rnn_hid_size) + self.temp_decay_h = TemporalDecay(input_size=self.xdim, output_size=self.rnn_hid_size, diag=False) + self.temp_decay_x = TemporalDecay(input_size=self.xdim, output_size=self.xdim, diag=True) + self.hist_reg = nn.Linear(self.rnn_hid_size, self.xdim) + self.feat_reg = FeatureRegression(self.xdim) + self.weight_combine = nn.Linear(self.xdim * 2, self.xdim) + self.dropout = nn.Dropout(p=0.25) + self.out = nn.Linear(self.rnn_hid_size, 1) + + def forward(self, data, direct): + values = data[direct]['values'] + masks = data[direct]['masks'] + deltas = data[direct]['deltas'] + evals = data[direct]['evals'] + eval_masks = data[direct]['eval_masks'] + labels = data['labels'].reshape((-1, 1)) + is_train = data['is_train'].reshape((-1, 1)) + h = Variable(torch.zeros((values.size(0), self.rnn_hid_size))) + c = Variable(torch.zeros((values.size(0), self.rnn_hid_size))) + if torch.cuda.is_available(): + h, c = h.cuda(), c.cuda() + x_loss = 0.0 + y_loss = 0.0 + imputations = [] + for t in range(min(values.size(1), masks.size(1), deltas.size(1))): + x = values[:, t, :] + m = masks[:, t, :] + d = deltas[:, t, :] + gamma_h = self.temp_decay_h(d) + gamma_x = self.temp_decay_x(d) + h = h * gamma_h + x_h = self.hist_reg(h) + x_loss += torch.sum(torch.abs(x - x_h) * m) / (torch.sum(m) + 1e-5) + x_c = m * x + (1 - m) * x_h + z_h = self.feat_reg(x_c) + x_loss += torch.sum(torch.abs(x - z_h) * m) / (torch.sum(m) + 1e-5) + alpha = self.weight_combine(torch.cat([gamma_x, m], dim=1)) + c_h = alpha * z_h + (1 - alpha) * x_h + x_loss += torch.sum(torch.abs(x - c_h) * m) / (torch.sum(m) + 1e-5) + c_c = m * x + (1 - m) * c_h + inputs = torch.cat([c_c, m], dim=1) + h, c = self.rnn_cell(inputs, (h, c)) + imputations.append(c_c.unsqueeze(dim=1)) + imputations = torch.cat(imputations, dim=1) + y_h = self.out(h) + y_loss = binary_cross_entropy_with_logits(y_h, labels, reduce=False) + y_loss = torch.sum(y_loss * is_train) / (torch.sum(is_train) + 1e-5) + y_h = torch.sigmoid(y_h) + return {'loss': x_loss * self.impute_weight + y_loss * self.label_weight, 'predictions': y_h, \ + 'imputations': imputations, 'labels': labels, 'is_train': is_train, \ + 'evals': evals, 'eval_masks': eval_masks} + + def run_on_batch(self, data, optimizer, epoch=None): + ret = self(data, direct='forward') + if optimizer is not None: + optimizer.zero_grad() + ret['loss'].backward() + optimizer.step() + return ret + + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim + +from torch.autograd import Variable +from torch.nn.parameter import Parameter + +import math +# import utils +import argparse +# import data_loader + +# import rits +from sklearn import metrics + + +# from ipdb import set_trace + +class BRITS(nn.Module): + def __init__(self, xdim, rnn_hid_size, impute_weight, label_weight): + super().__init__() + self.xdim = xdim + self.rnn_hid_size = rnn_hid_size + self.impute_weight = impute_weight + self.label_weight = label_weight + self.build() + + def build(self): + self.rits_f = RITS(self.xdim, self.rnn_hid_size, self.impute_weight, self.label_weight) + self.rits_b = RITS(self.xdim, self.rnn_hid_size, self.impute_weight, self.label_weight) + + def forward(self, data): + ret_f = self.rits_f(data, 'forward') + ret_b = self.reverse(self.rits_b(data, 'backward')) + ret = self.merge_ret(ret_f, ret_b) + return ret + + def merge_ret(self, ret_f, ret_b): + loss_f = ret_f['loss'] + loss_b = ret_b['loss'] + loss_c = self.get_consistency_loss(ret_f['imputations'], ret_b['imputations']) + loss = loss_f + loss_b + loss_c + predictions = (ret_f['predictions'] + ret_b['predictions']) / 2 + imputations = (ret_f['imputations'] + ret_b['imputations']) / 2 + ret_f['loss'] = loss + ret_f['predictions'] = predictions + ret_f['imputations'] = imputations + return ret_f + + def get_consistency_loss(self, pred_f, pred_b): + loss = torch.abs(pred_f - pred_b).mean() * 1e-1 + return loss + + def reverse(self, ret): + def reverse_tensor(tensor_): + if tensor_.dim() <= 1: + return tensor_ + indices = range(tensor_.size()[1])[::-1] + indices = Variable(torch.LongTensor(indices), requires_grad=False) + if torch.cuda.is_available(): + indices = indices.cuda() + return tensor_.index_select(1, indices) + + for key in ret: + ret[key] = reverse_tensor(ret[key]) + return ret + + def run_on_batch(self, data, optimizer, epoch=None): + ret = self(data) + if optimizer is not None: + optimizer.zero_grad() + ret['loss'].backward() + optimizer.step() + return ret + + +import copy +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.optim.lr_scheduler import StepLR + +import numpy as np + +import time +# import utils +# import models +import argparse +# import data_loader +import pandas as pd +import ujson as json + +from sklearn import metrics + + +# from ipdb import set_trace + +def to_var(var): + if torch.is_tensor(var): + device = var.device + var = torch.autograd.Variable(var) + var = var.to(device) + return var + if isinstance(var, int) or isinstance(var, float) or isinstance(var, str): + return var + if isinstance(var, dict): + return {key: to_var(val) for key, val in var.items()} + if isinstance(var, list): + return [to_var(x) for x in var] + + +@torch.no_grad() +def brits_prep_rec(y, back): # (samples, T + 1, nodes) + evals = y.float().clone() + values = evals.clone() + if back: + values[:, 1:] = 0 + else: + values[:, : -1] = 0 + masks = torch.zeros_like(y) + if back: + masks[:, 0] = True + else: + masks[:, -1] = True + eval_masks = masks.clone() + deltas = torch.cat([torch.arange(y.size(dim=1) - 1, dtype=torch.float, device=y.device), + torch.zeros(1, dtype=torch.float, device=y.device)], dim=0).unsqueeze(dim=-1).expand(*y.size()) + return dict(values=values.contiguous(), masks=masks.contiguous(), evals=evals.contiguous(), + eval_masks=eval_masks.contiguous(), deltas=deltas.contiguous()) + + +@torch.no_grad() +def brits_prep(y, is_train): # y: (samples, nodes, T + 1) + n_samples = y.size(dim=0) + y = y.transpose(1, 2) # (samples, T + 1, nodes) + return to_var(dict( + forward=brits_prep_rec(y, back=False), + backward=brits_prep_rec(y.flip(dims=[1]), back=True), + labels=torch.zeros(n_samples, 1, dtype=torch.long, device=y.device), + is_train=torch.tensor([is_train] * n_samples, dtype=torch.float, device=y.device), + )) + + +def brits_run(data, args): + """Run BRITS on a single dataset instance (loaded by Tester).""" + bpar = b_estim(data, args) + + model = BRITS(data.num_nodes, args.hid_size, args.impute_weight, args.label_weight) + model = model.to(args.device) + + # train + I0 = (data.y[:, 0] == SIR_STATES.I).long().sum().item() + optimizer = optim.Adam(model.parameters(), lr=args.lr) + pbar = trange(args.epochs) + for epoch in pbar: + model.train() + batch = brits_prep(diffus_gen(T=data.T.item(), n_nodes=data.num_nodes, edge_index=data.edge_index, I0=I0, + n_samples=args.batch_size, pI=bpar.pI, pR=bpar.pR).transpose(0, 2).clone(), + is_train=1) + ret = model.run_on_batch(batch, optimizer, epoch) + pbar.set_description(f'epoch={epoch + 1} loss={ret["loss"].item():.4f}') + + # infer + with torch.no_grad(): + model.eval() + rec = brits_prep(data.y.unsqueeze(dim=0).clone(), is_train=0) + ret = model.run_on_batch(rec, None) + y_pred = ret['imputations'].long().clamp(0, data.y.max()).squeeze(dim=0).T.clone() + return y_pred.clone() + + +def main(argv=None): + args = get_args(argv) + + # Match the other runners (e.g., hermes.py): run one dataset specified by CLI. + seed_all(args.seed) + tester = Tester(args.data_dir, args.device, lambda data: brits_run(data, args)) + tester.test([args.dataset], seed=args.seed, rep=1) + tester.save(args.output) + + +if __name__ == '__main__': + main() + diff --git a/brits_ms.py b/brits_ms.py new file mode 100644 index 0000000..1b48825 --- /dev/null +++ b/brits_ms.py @@ -0,0 +1,532 @@ + + +from __future__ import annotations + +import argparse +import math +from typing import List, Optional, Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.autograd import Variable +from torch.nn.parameter import Parameter + +from inc.test import Tester + +from tqdm import trange + +# Project utilities (diffusion simulator, parameter estimator, seeding, etc.) +from inc.diffus import * + + +# ----------------------------------------------------------------------------- +# CLI +# ----------------------------------------------------------------------------- + +def get_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description='BRITS baseline (multi-snapshot)') + + # ---- experiment I/O ---- + parser.add_argument('--dataset', type=str, required=True, help='dataset name') + parser.add_argument('--seed', type=int, default=123456789, help='random seed') + parser.add_argument('--data_dir', type=str, default='input', help='dataset folder') + parser.add_argument('--output', type=str, default='output/brits.pt', help='output file name') + parser.add_argument( + '--device', + type=torch.device, + default=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), + help='torch device, e.g., cpu, cuda, cuda:0' + ) + + # ---- multi-snapshot observed times ---- + parser.add_argument( + '--obs_time', '--snapshot', + dest='obs_time', + type=str, + default='', + help='extra observed snapshot times, comma-separated, e.g., 5,7,9. ' + 'Final time T is always observed.' + ) + + # ---- diffusion parameter estimation (b_*) ---- + parser.add_argument('--b_pI0', type=float, default=1e-3, + help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, default=1e-3, + help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, default=500, + help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, default=3e-3, + help='learning rate in diffusion parameter estimation') + + # ---- BRITS hyperparameters ---- + parser.add_argument('--lr', type=float, default=1e-3, help='learning rate') + parser.add_argument('--epochs', type=int, default=1000, help='training epochs') + parser.add_argument('--batch_size', type=int, default=64, help='batch size') + parser.add_argument('--hid_size', type=int, default=108, help='RNN hidden size') + parser.add_argument('--impute_weight', type=float, default=0.3, help='imputation loss weight') + parser.add_argument('--label_weight', type=float, default=1.0, help='label loss weight') + + # ---- evaluation ---- + parser.add_argument('--rep', type=int, default=1, help='number of test repetitions') + + if argv is None: + return parser.parse_args() + return parser.parse_args(argv) + + +def _parse_obs_time_str(s: str) -> List[int]: + """Parse comma-separated time indices. + + Empty string -> []. Whitespace is ignored. + """ + s = '' if s is None else str(s) + out: List[int] = [] + for tok in s.split(','): + tok = tok.strip() + if not tok: + continue + out.append(int(tok)) + return out + + +def get_obs_time(data, args: argparse.Namespace) -> List[int]: + """Return sorted unique observed times (always includes final time T).""" + T = int(data.T.item()) + obs_time = _parse_obs_time_str(getattr(args, 'obs_time', '')) + obs_time.append(T) + obs_time = sorted({t for t in obs_time if 0 <= int(t) <= T}) + if len(obs_time) == 0 or obs_time[-1] != T: + obs_time.append(T) + return obs_time + + +@torch.no_grad() +def make_obs_mask(T: int, obs_time: Sequence[int], device: torch.device) -> torch.Tensor: + """Make boolean mask of shape (T+1,) indicating observed time steps.""" + mask = torch.zeros(T + 1, dtype=torch.bool, device=device) + if len(obs_time) == 0: + mask[T] = True + return mask + ts = torch.as_tensor(list(obs_time), dtype=torch.long, device=device).clamp(0, T) + ts = ts.unique() + mask[ts] = True + if not bool(mask[T].item()): + mask[T] = True + return mask + + +# ----------------------------------------------------------------------------- +# BRITS core (largely copied from https://github.com/caow13/BRITS) +# ----------------------------------------------------------------------------- + +def binary_cross_entropy_with_logits( + input: torch.Tensor, + target: torch.Tensor, + weight: Optional[torch.Tensor] = None, + size_average: bool = True, + reduce: bool = True, +) -> torch.Tensor: + """A numerically-stable BCE-with-logits implementation. + + Kept to match the original notebook/code. + """ + if target.size() != input.size(): + raise ValueError(f'Target size ({target.size()}) must be the same as input size ({input.size()})') + max_val = (-input).clamp(min=0) + loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log() + if weight is not None: + loss = loss * weight + if not reduce: + return loss + if size_average: + return loss.mean() + return loss.sum() + + +class FeatureRegression(nn.Module): + def __init__(self, input_size: int): + super().__init__() + self.W = Parameter(torch.Tensor(input_size, input_size)) + self.b = Parameter(torch.Tensor(input_size)) + m = torch.ones(input_size, input_size) - torch.eye(input_size, input_size) + self.register_buffer('m', m) + self.reset_parameters() + + def reset_parameters(self) -> None: + stdv = 1.0 / math.sqrt(self.W.size(0)) + self.W.data.uniform_(-stdv, stdv) + if self.b is not None: + self.b.data.uniform_(-stdv, stdv) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Mask diagonal to avoid trivial self-regression + z_h = F.linear(x, self.W * Variable(self.m), self.b) + return z_h + + +class TemporalDecay(nn.Module): + def __init__(self, input_size: int, output_size: int, diag: bool = False): + super().__init__() + self.diag = diag + self.W = Parameter(torch.Tensor(output_size, input_size)) + self.b = Parameter(torch.Tensor(output_size)) + if self.diag: + assert input_size == output_size + m = torch.eye(input_size, input_size) + self.register_buffer('m', m) + self.reset_parameters() + + def reset_parameters(self) -> None: + stdv = 1.0 / math.sqrt(self.W.size(0)) + self.W.data.uniform_(-stdv, stdv) + if self.b is not None: + self.b.data.uniform_(-stdv, stdv) + + def forward(self, d: torch.Tensor) -> torch.Tensor: + if self.diag: + gamma = F.relu(F.linear(d, self.W * Variable(self.m), self.b)) + else: + gamma = F.relu(F.linear(d, self.W, self.b)) + gamma = torch.exp(-gamma) + return gamma + + +class RITS(nn.Module): + def __init__(self, xdim: int, rnn_hid_size: int, impute_weight: float, label_weight: float): + super().__init__() + self.xdim = int(xdim) + self.rnn_hid_size = int(rnn_hid_size) + self.impute_weight = float(impute_weight) + self.label_weight = float(label_weight) + + self.rnn_cell = nn.LSTMCell(self.xdim * 2, self.rnn_hid_size) + self.temp_decay_h = TemporalDecay(input_size=self.xdim, output_size=self.rnn_hid_size, diag=False) + self.temp_decay_x = TemporalDecay(input_size=self.xdim, output_size=self.xdim, diag=True) + self.hist_reg = nn.Linear(self.rnn_hid_size, self.xdim) + self.feat_reg = FeatureRegression(self.xdim) + self.weight_combine = nn.Linear(self.xdim * 2, self.xdim) + self.dropout = nn.Dropout(p=0.25) + self.out = nn.Linear(self.rnn_hid_size, 1) + + def forward(self, data: dict, direct: str): + values = data[direct]['values'] + masks = data[direct]['masks'] + deltas = data[direct]['deltas'] + evals = data[direct]['evals'] + eval_masks = data[direct]['eval_masks'] + + labels = data['labels'].reshape((-1, 1)) + is_train = data['is_train'].reshape((-1, 1)) + + device = values.device + h = Variable(torch.zeros((values.size(0), self.rnn_hid_size), device=device)) + c = Variable(torch.zeros((values.size(0), self.rnn_hid_size), device=device)) + + x_loss = 0.0 + imputations = [] + + T_seq = min(values.size(1), masks.size(1), deltas.size(1)) + for t in range(T_seq): + x = values[:, t, :] + m = masks[:, t, :] + d = deltas[:, t, :] + + gamma_h = self.temp_decay_h(d) + gamma_x = self.temp_decay_x(d) + + h = h * gamma_h + x_h = self.hist_reg(h) + + x_loss += torch.sum(torch.abs(x - x_h) * m) / (torch.sum(m) + 1e-5) + + x_c = m * x + (1 - m) * x_h + z_h = self.feat_reg(x_c) + x_loss += torch.sum(torch.abs(x - z_h) * m) / (torch.sum(m) + 1e-5) + + alpha = self.weight_combine(torch.cat([gamma_x, m], dim=1)) + c_h = alpha * z_h + (1 - alpha) * x_h + x_loss += torch.sum(torch.abs(x - c_h) * m) / (torch.sum(m) + 1e-5) + + c_c = m * x + (1 - m) * c_h + inputs = torch.cat([c_c, m], dim=1) + h, c = self.rnn_cell(inputs, (h, c)) + + imputations.append(c_c.unsqueeze(dim=1)) + + imputations = torch.cat(imputations, dim=1) + + # (unused in our setting, but keep original structure) + y_h = self.out(h) + y_loss = binary_cross_entropy_with_logits(y_h, labels, reduce=False) + y_loss = torch.sum(y_loss * is_train) / (torch.sum(is_train) + 1e-5) + y_h = torch.sigmoid(y_h) + + loss = x_loss * self.impute_weight + y_loss * self.label_weight + return { + 'loss': loss, + 'predictions': y_h, + 'imputations': imputations, + 'labels': labels, + 'is_train': is_train, + 'evals': evals, + 'eval_masks': eval_masks, + } + + def run_on_batch(self, data: dict, optimizer: Optional[optim.Optimizer], epoch: Optional[int] = None): + ret = self(data, direct='forward') + if optimizer is not None: + optimizer.zero_grad() + ret['loss'].backward() + optimizer.step() + return ret + + +class BRITSModel(nn.Module): + def __init__(self, input_size: int, hidden_size: int, impute_weight: float, label_weight: float): + super().__init__() + self.rits_f = RITS(input_size, hidden_size, impute_weight, label_weight) + self.rits_b = RITS(input_size, hidden_size, impute_weight, label_weight) + + def forward(self, data: dict): + ret_f = self.rits_f(data, direct='forward') + ret_b = self.rits_b(data, direct='backward') + ret_b = self.reverse(ret_b) + + loss_f = ret_f['loss'] + loss_b = ret_b['loss'] + loss_c = self.get_consistency_loss(ret_f['imputations'], ret_b['imputations']) + loss = loss_f + loss_b + loss_c + + predictions = (ret_f['predictions'] + ret_b['predictions']) / 2 + imputations = (ret_f['imputations'] + ret_b['imputations']) / 2 + + ret_f['loss'] = loss + ret_f['predictions'] = predictions + ret_f['imputations'] = imputations + return ret_f + + @staticmethod + def get_consistency_loss(pred_f: torch.Tensor, pred_b: torch.Tensor) -> torch.Tensor: + return torch.abs(pred_f - pred_b).mean() * 1e-1 + + @staticmethod + def reverse(ret: dict) -> dict: + def reverse_tensor(tensor_: torch.Tensor) -> torch.Tensor: + if tensor_.dim() <= 1: + return tensor_ + # reverse along time dimension (dim=1) + idx = torch.arange(tensor_.size(1) - 1, -1, -1, device=tensor_.device) + return tensor_.index_select(1, idx) + + return {k: reverse_tensor(v) for k, v in ret.items()} + + def run_on_batch(self, data: dict, optimizer: Optional[optim.Optimizer], epoch: Optional[int] = None): + ret = self(data) + if optimizer is not None: + optimizer.zero_grad() + ret['loss'].backward() + optimizer.step() + return ret + + +class BRITS(nn.Module): + def __init__(self, input_size: int, hidden_size: int, impute_weight: float = 1.0, label_weight: float = 1.0): + super().__init__() + self.model = BRITSModel(input_size, hidden_size, impute_weight, label_weight) + + def forward(self, data: dict): + return self.model(data) + + def run_on_batch(self, data: dict, optimizer: Optional[optim.Optimizer], epoch: Optional[int] = None): + return self.model.run_on_batch(data, optimizer, epoch) + + +# ----------------------------------------------------------------------------- +# Data preparation (multi-snapshot masking) +# ----------------------------------------------------------------------------- + +def to_var(var): + """Recursively move nested structures into torch Variables on the same device.""" + if torch.is_tensor(var): + dev = var.device + var = torch.autograd.Variable(var) + var = var.to(dev) + return var + if isinstance(var, (int, float, str)): + return var + if isinstance(var, dict): + return {key: to_var(val) for key, val in var.items()} + if isinstance(var, list): + return [to_var(x) for x in var] + return var + + +@torch.no_grad() +def _brits_make_deltas(obs_mask_1d: torch.Tensor) -> torch.Tensor: + """Make a (T+1,) float tensor of time gaps since the last observation. + + In forward direction: + delta[t] = 0 if t is observed else delta[t-1] + 1. + In backward direction we apply the same logic on the reversed mask. + """ + if obs_mask_1d.dim() != 1: + raise ValueError(f'obs_mask_1d must be 1D, got shape={tuple(obs_mask_1d.shape)}') + + L = obs_mask_1d.numel() + d = torch.zeros(L, dtype=torch.float, device=obs_mask_1d.device) + for t in range(1, L): + d[t] = 0.0 if bool(obs_mask_1d[t].item()) else (d[t - 1] + 1.0) + return d + + +@torch.no_grad() +def brits_prep_rec(y: torch.Tensor, obs_mask_1d: torch.Tensor) -> dict: + """Prepare one direction (forward OR backward) input for BRITS. + + Args: + y: (samples, T+1, nodes) full ground-truth sequence (will be masked). + obs_mask_1d: (T+1,) bool mask for which time steps are observed. + + Returns: + dict with keys values/masks/evals/eval_masks/deltas, all shaped + (samples, T+1, nodes). + """ + if y.dim() != 3: + raise ValueError(f'y must be 3D (samples,T+1,nodes), got {tuple(y.shape)}') + + # Eval targets (full sequence, used only for reporting; NOT fed as observed). + evals = y.float().clone() + + # Observation mask broadcast to (samples, T+1, nodes) + obs_mask = obs_mask_1d.to(device=y.device, dtype=torch.bool) + masks = obs_mask.view(1, -1, 1).expand(y.size(0), -1, y.size(2)).float() + + # Only keep observed frames in values; unobserved are set to 0. + values = evals * masks + + # delta features + deltas_1d = _brits_make_deltas(obs_mask) + deltas = deltas_1d.view(1, -1, 1).expand_as(values).contiguous() + + eval_masks = masks.clone() + + return { + 'values': values.contiguous(), + 'masks': masks.contiguous(), + 'evals': evals.contiguous(), + 'eval_masks': eval_masks.contiguous(), + 'deltas': deltas.contiguous(), + } + + +@torch.no_grad() +def brits_prep(y: torch.Tensor, is_train: int, obs_mask_1d: torch.Tensor) -> dict: + """Prepare BRITS batch dict. + + Args: + y: (samples, nodes, T+1) + is_train: 1 or 0 + obs_mask_1d: (T+1,) bool in *forward* time. + + Returns: + dict consumable by BRITSModel. + """ + if y.dim() != 3: + raise ValueError(f'y must be 3D (samples,nodes,T+1), got {tuple(y.shape)}') + + n_samples = int(y.size(0)) + y = y.transpose(1, 2) # (samples, T+1, nodes) + + obs_mask_1d = obs_mask_1d.to(device=y.device, dtype=torch.bool) + obs_mask_bwd = obs_mask_1d.flip(dims=[0]) + + return to_var({ + 'forward': brits_prep_rec(y, obs_mask_1d), + 'backward': brits_prep_rec(y.flip(dims=[1]), obs_mask_bwd), + # BRITS classification head is unused; keep placeholders + 'labels': torch.zeros(n_samples, 1, dtype=torch.long, device=y.device), + 'is_train': torch.full((n_samples, 1), float(is_train), dtype=torch.float, device=y.device), + }) + + +# ----------------------------------------------------------------------------- +# Experiment runner +# ----------------------------------------------------------------------------- + +def brits_run(data, args: argparse.Namespace) -> torch.Tensor: + """Train BRITS on synthetic histories and infer a history for the given data. + + Returns: + y_pred: (nodes, T+1) long tensor of reconstructed states. + """ + device = args.device + data = data.to(device) + + T = int(data.T.item()) + obs_time = get_obs_time(data, args) + + # Attach observation times for the tester (so metrics exclude observed frames) + # `inc.test_ms` uses `data.obs_ts` or `data.obs_mask` if present. + data.obs_ts = obs_time + + obs_mask = make_obs_mask(T, obs_time, device=device) + + # Estimate diffusion parameters (supports multi-snapshot via obs_time) + bpar = b_estim(data, args, obs_time=obs_time) + + model = BRITS(data.num_nodes, args.hid_size, args.impute_weight, args.label_weight).to(device) + + # -------------------- train -------------------- + I0 = int((data.y[:, 0] == SIR_STATES.I).long().sum().item()) + optimizer = optim.Adam(model.parameters(), lr=args.lr) + + pbar = trange(args.epochs) + for epoch in pbar: + model.train() + # Generate synthetic training batch (full history), then mask it. + Y = diffus_gen( + T=T, + n_nodes=data.num_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ) # (T+1, nodes, samples) + + batch_y = Y.transpose(0, 2).clone() # -> (samples, nodes, T+1) + batch = brits_prep(batch_y, is_train=1, obs_mask_1d=obs_mask) + ret = model.run_on_batch(batch, optimizer, epoch) + pbar.set_description(f'epoch={epoch + 1} loss={ret["loss"].item():.4f}') + + # -------------------- infer -------------------- + with torch.no_grad(): + model.eval() + rec = brits_prep(data.y.unsqueeze(dim=0).clone(), is_train=0, obs_mask_1d=obs_mask) + ret = model.run_on_batch(rec, None) + + # ret['imputations']: (1, T+1, nodes) float + y_pred = ret['imputations'] + y_pred = y_pred.long().clamp(0, int(data.y.max().item())) + y_pred = y_pred.squeeze(dim=0).T.contiguous() # (nodes, T+1) + return y_pred.detach().clone() + + +def main(argv: Optional[Sequence[str]] = None) -> None: + args = get_args(argv) + + # Reproducibility + seed_all(args.seed) + + def _model_fn(data): + return brits_run(data, args) + + tester = Tester(args.data_dir, args.device, _model_fn) + tester.test([args.dataset], seed=args.seed, rep=args.rep) + tester.save(args.output) + + +if __name__ == '__main__': + main() diff --git a/cri_ms.py b/cri_ms.py new file mode 100644 index 0000000..9091023 --- /dev/null +++ b/cri_ms.py @@ -0,0 +1,292 @@ +from inc.diffus import * +from inc.test import * + +import argparse +import numpy as np +import networkx as nx + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type=str, help='dataset name') + parser.add_argument('--seed', type=int, help='random seed') + parser.add_argument('--data_dir', type=str, help='dataset folder') + parser.add_argument('--output', type=str, help='output file name') + parser.add_argument('--device', type=torch.device, help='torch device') + + # Multi-snapshot options (aligned with inc/ditto_ms.py). + parser.add_argument( + '--obs_ts', type=str, default=None, + help='comma-separated observed time indices, e.g. "0,3,5"; ' + 'None means single-snapshot (use only the final snapshot)' + ) + parser.add_argument( + '--obs_k', type=int, default=None, + help='number of observed snapshots; if None, set to len(obs_ts) when obs_ts ' + 'is given, otherwise 1 (single-snapshot)' + ) + + args = parser.parse_args() + # normalize obs_ts / obs_k + if args.obs_ts is not None and len(args.obs_ts.strip()) > 0: + obs = [int(x) for x in args.obs_ts.split(',') if x.strip() != ''] + obs = sorted(set(obs)) + args.obs_ts = obs + if args.obs_k is None: + args.obs_k = len(obs) + else: + args.obs_ts = None + if args.obs_k is None: + args.obs_k = 1 + return args + + +def _bfs_dist(G: nx.Graph, s: int, cutoff: int = None) -> dict: + """ + Return {node: hop_distance} for nodes reachable from s. + + cutoff: optional BFS depth limit. When we only care whether dist <= T_thr, + setting cutoff=T_thr can be much faster on large graphs. + """ + if cutoff is None: + return dict(nx.single_source_shortest_path_length(G, s)) + return dict(nx.single_source_shortest_path_length(G, s, cutoff=int(cutoff))) + + +def cri_cluster(G: nx.Graph, obs_mask: np.ndarray, T_thr: int): + """ + CRI clustering (greedy k-center) — optimized for speed. + + Original bottleneck: + The legacy version repeatedly called BFS from *every infected node* inside + the greedy loop, causing huge runtimes on large VI. + + Fix (same concept, much faster): + Run BFS only from the current centers (with cutoff=T_thr), and maintain + each infected node's distance to its nearest center incrementally. + + Output: + clusters: list of infected-node lists (one per center) + VI: list of all infected nodes + """ + n = int(obs_mask.shape[0]) + VI = [u for u in range(n) if obs_mask[u] == 1] + if len(VI) == 0: + return [], VI + if len(VI) == 1: + return [VI], VI + + # If T_thr <= 0, every infected node must be its own center (no BFS needed). + if T_thr <= 0: + return [[u] for u in VI], VI + + VI_arr = np.asarray(VI, dtype=np.int64) + idx_of = {int(u): i for i, u in enumerate(VI_arr)} # node -> index in VI + + # --- 2-BFS "double sweep" heuristic to pick an initial far-apart pair --- + # We do NOT use cutoff here; only 2 BFS calls, so it's cheap and gives a better pair. + u0 = int(VI_arr[0]) + dist0 = _bfs_dist(G, u0, cutoff=None) + + u1 = next((int(u) for u in VI_arr if int(u) not in dist0), None) + if u1 is None: + u1 = max((int(u) for u in VI_arr), key=lambda u: dist0.get(u, -1)) + + dist1 = _bfs_dist(G, u1, cutoff=None) + u2 = next((int(u) for u in VI_arr if int(u) not in dist1), None) + if u2 is None: + u2 = max((int(u) for u in VI_arr), key=lambda u: dist1.get(u, -1)) + + centers = [u1] + if u2 != u1: + centers.append(u2) + + # --- Greedy k-center with incremental nearest-center distances --- + is_center = np.zeros(len(VI_arr), dtype=bool) + d_near = np.full(len(VI_arr), np.inf, dtype=np.float32) + + # For each center s, store only distances to infected nodes (shape: [|VI|]). + dist_to_VI = {} # center -> np.ndarray(|VI|,) + + def add_center(s: int): + nonlocal d_near + if s in dist_to_VI: + return + + # Key speedup: cutoff=T_thr (we only need to know if dist <= T_thr) + dist_s = _bfs_dist(G, s, cutoff=T_thr) + ds = np.fromiter((dist_s.get(int(u), np.inf) for u in VI_arr), + dtype=np.float32, count=len(VI_arr)) + dist_to_VI[s] = ds + + if s in idx_of: + is_center[idx_of[s]] = True + + d_near = np.minimum(d_near, ds) + d_near[is_center] = 0.0 + + for s in centers: + add_center(int(s)) + + with tqdm(desc='cluster', leave=False) as pbar: + while True: + # farthest infected node from current centers (excluding centers) + if (~is_center).any(): + tmp = d_near.copy() + tmp[is_center] = -1.0 + far_i = int(np.argmax(tmp)) + far_d = float(tmp[far_i]) + else: + far_i, far_d = -1, -1.0 + + if far_d <= float(T_thr): + break + + far_node = int(VI_arr[far_i]) + centers.append(far_node) + add_center(far_node) + pbar.update(1) + + # --- Assign each infected node to its nearest center (vectorized) --- + center_list = list(dict.fromkeys(centers)) # unique, keep order + dist_mat = np.vstack([dist_to_VI[s] for s in center_list]) # (k, |VI|) + assign = dist_mat.argmin(axis=0) # (|VI|,) + + clusters_map = {s: [] for s in center_list} + for i, u in enumerate(VI_arr): + s = center_list[int(assign[i])] + clusters_map[s].append(int(u)) + + clusters = [lst for lst in clusters_map.values() if len(lst) > 0] + return clusters, VI + + +def cri_rev_infect(G: nx.Graph, VI_all, Vi, y): + """ + Reverse infection for a cluster Vi (same as legacy version, but faster): + + - Expand BFS "wavefronts" tagged by sources x in Vi (pairs (u, x)). + - Stop once any node receives all tags. + - Choose the best center s (min sum of tag distances). + - For each x in Vi: set predicted infection time tI[x] = dist(s, x), + and mark y[x, tI[x]:] = 1. + + Speed fixes: + - Avoid allocating n empty dicts: create per-node dicts on-demand. + - Avoid O(n) scan each layer (track max_seen incrementally). + - Avoid final O(n) scan for candidates (track candidates during BFS). + """ + n = G.number_of_nodes() + ni = len(Vi) + if ni == 0: + return + + # g[u] is a dict {x: dist(u, x)}; allocate lazily + g = [None] * n + + def has_label(u: int, x: int) -> bool: + du = g[u] + return (du is not None) and (x in du) + + frontier = set() + for x in Vi: + x = int(x) + frontier.add((x, x)) + for v in G[x]: + frontier.add((int(v), x)) + + t_layer = 0 + max_seen = 0 + candidates = set() + + with tqdm(desc='rev_infect.expand', leave=False) as pbar: + while frontier and max_seen < ni: + next_frontier = set() + for u, x in frontier: + if has_label(u, x): + continue + if g[u] is None: + g[u] = {} + g[u][x] = t_layer + + lu = len(g[u]) + if lu > max_seen: + max_seen = lu + if lu == ni: + candidates.add(u) + + for v in G[u]: + v = int(v) + if not has_label(v, x): + next_frontier.add((v, x)) + + frontier = next_frontier + t_layer += 1 + pbar.update(1) + + if not candidates: + # Fallback (should be rare if clustering worked): + x0 = int(Vi[0]) + y[x0, 0:] = 1 + return + + s = min(candidates, key=lambda u: sum(g[u].values())) + + # Set first-infection time for each tagged x in the cluster and mark trajectory + tI_map = g[s] + for x, t0 in tI_map.items(): + t0 = int(t0) + if 0 <= x < y.shape[0]: + t0 = max(0, min(t0, y.shape[1] - 1)) # clamp to [0, T] + y[x, t0:] = 1 + + +@torch.no_grad() +def cri_ms_run(data): + """ + Multi-snapshot CRI: + - For each observed time t_obs: + * extract infected mask (I=1, others=0), + * cluster with radius threshold = t_obs, + * reverse-infect per cluster to get first-infection times, + * mark y[:, tI: ] = 1 for those nodes, + - Merge across snapshots by OR (sum since we fill with 1's from tI onward). + """ + T = int(data.T.item()) + n_nodes = int(data.num_nodes) + n_cls = int(data.y.max().item() + 1) + + # Determine observed times to use + if args.obs_ts is not None and len(args.obs_ts) > 0: + obs_times = [t for t in args.obs_ts if 0 <= t <= T] + obs_times = sorted(set(obs_times)) + if len(obs_times) == 0: + obs_times = [T] + else: + obs_times = [T] + + # Build graph + G = pyg.utils.to_networkx(data, to_undirected=True, remove_self_loops=True) + + # Accumulator for predictions + y_pred = np.zeros((n_nodes, T + 1), dtype=np.int32) + + for t_obs in tqdm(obs_times, desc='obs_times', leave=False): + obs_vec = (data.y[:, t_obs].cpu().detach().numpy() & 1).astype(np.int32) + if obs_vec.sum() == 0: + continue + + clusters, VI_all = cri_cluster(G, obs_vec, T_thr=int(t_obs)) + for Vi in tqdm(clusters, desc=f'rev_infect@t={t_obs}', leave=False): + cri_rev_infect(G, VI_all, Vi, y_pred) + + return torch.tensor(np.minimum(y_pred, n_cls - 1), dtype=torch.long, device=data.y.device) + + +# ---- entry point ---- +args = get_args() +seed_all(args.seed) +tester = Tester(args.data_dir, args.device, cri_ms_run) +tester.test([args.dataset], rep=1) +tester.save(args.output) + diff --git a/dhrec_ms.py b/dhrec_ms.py new file mode 100644 index 0000000..dd73648 --- /dev/null +++ b/dhrec_ms.py @@ -0,0 +1,224 @@ +from inc.diffus import * +from inc.test import * + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type=str, help='dataset name') + parser.add_argument('--seed', type=int, help='random seed') + parser.add_argument('--data_dir', type=str, help='dataset folder') + parser.add_argument('--output', type=str, help='output file name') + parser.add_argument('--device', type=torch.device, help='torch device') + + # diffusion parameter estimation (kept as in the baseline) + parser.add_argument('--b_pI0', type=float, help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, help='learning rate in diffusion parameter estimation') + + # multi-snapshot options (aligned with ditto_ms.py) + parser.add_argument( + '--obs_ts', type=str, default=None, + help='comma-separated observed time indices, e.g. "0,3,5"; ' + 'None means single-snapshot (use only the final snapshot)' + ) + parser.add_argument( + '--obs_k', type=int, default=None, + help='number of observed snapshots; if None, set to len(obs_ts) when obs_ts is given, ' + 'otherwise 1 (single-snapshot)' + ) + args = parser.parse_args() + + # parse obs_ts in the same spirit as inc/ditto_ms.py + if args.obs_ts is not None and len(args.obs_ts.strip()) > 0: + obs = [int(x) for x in args.obs_ts.split(',') if x.strip() != ''] + obs = sorted(set(obs)) + args.obs_ts = obs + if args.obs_k is None: + args.obs_k = len(obs) + else: + args.obs_ts = None + if args.obs_k is None: + args.obs_k = 1 + return args + + +def pcdsvc_greedy(bpar, G, y): + """ + One-step backward inference (t -> t-1) under the PCDSVC-style greedy rule. + This is the original single-snapshot kernel kept unchanged; it maps a snapshot y(t) + to a previous snapshot x(t-1). + + Fix: + - Some SI datasets legitimately have pR = 0 (no recovery). In that case, lr = -log(pR) + should NOT be evaluated (it is unused anyway because there is no R state in y). + - Also guard against boundary numeric issues (pI -> 1, pR -> 0) to avoid log(0). + """ + n = G.number_of_nodes() + + # ---- numeric guards / SI-safe handling ---- + eps = 1e-12 + + # pI is used in all cases; only clip the upper bound to avoid log(0) at (1 - pI). + pI = float(getattr(bpar, 'pI', 0.0)) + if not np.isfinite(pI): + pI = 0.0 + pI = min(max(pI, 0.0), 1.0 - eps) + l1s = -np.log(1.0 - pI) + + # pR is only needed when there are recovered nodes in the current snapshot. + # For SI datasets, y never contains state=2, so skip log(pR) entirely. + if np.any(y == 2): + pR = float(getattr(bpar, 'pR', 0.0)) + if not np.isfinite(pR): + pR = eps + pR = min(max(pR, eps), 1.0) + lr = -np.log(pR) + else: + lr = 0.0 + # ------------------------------------------ + + # x: 2 means "unknown yet / start from R" in the original code, then moves to 1 or 0 + x = np.where(y == 2, 2, 0) + + we = l1s + ws = np.zeros(n, dtype=np.float32) + wi = np.zeros(n, dtype=np.float32) + + for u in range(n): + if y[u] == 0: # S + for v in G.neighbors(u): + wi[v] += l1s + elif y[u] == 1: # I + ws[u] -= 1. + else: # R + ws[u] += lr - 1. + wi[u] += lr + + # R --> I + pbar = tqdm(disable=True) + while True: + mvs = [] + for u in range(n): + if y[u] >= 1 and x[u] != 1: + cur = wi[u] + for v in G.neighbors(u): + if y[v] >= 1 and x[v] != 1: + cur -= we + mvs.append((cur, u)) + if len(mvs) == 0: + break + mv = min(mvs, key=lambda mv: mv[0]) + if mv[0] >= 0.: + dom = True + for u in range(n): + if y[u] == 1: + dm = (x[u] == 1) + for v in G.neighbors(u): + dm |= (x[v] == 1) + if dm: + break + dom &= dm + if dom: + break + x[mv[1]] = 1 + pbar.update(1) + + # I --> S + for u in range(n): + if x[u] == 2 and ws[u] < 0: + dm = False + for v in G.neighbors(u): + dm |= (x[v] == 1) + if dm: + break + if dm: + x[u] = 0 + pbar.update(1) + pbar.close() + return x + + +def _parse_obs_ts(args, T): + """ + Build a sorted & unique list of observed time indices within [0, T], + ensuring the final snapshot T is always included as the anchor. + """ + if args.obs_ts is None: + obs_ts = [] + else: + obs_ts = [t for t in args.obs_ts if 0 <= t <= T] + + if T not in obs_ts: + obs_ts.append(T) + obs_ts = sorted(set(obs_ts)) + return obs_ts + + +def _build_observation_map(data, obs_ts): + """ + Return {t: np.array states at time t} for all observed times t. + """ + y = data.y # (nodes, T+1) + obs = {} + for t in obs_ts: + obs[t] = y[:, t].cpu().detach().numpy() + return obs + + +def pcdsvc_run(data): + """ + Multi-snapshot backward reconstruction: + - Split the timeline by observed time points (including T as anchor). + - For each segment [t_prev, t_next], start from the observed y(t_next) + and repeatedly apply the single-step greedy kernel to obtain y(t_prev+1),...,y(t_prev), + clamping to ground-truth observation whenever we hit an observed time. + """ + bpar = b_estim(data, args) # estimate diffusion parameters once + + with torch.no_grad(): + T = int(data.T.item()) + n_nodes = data.num_nodes + n_cls = int(data.y.max().item() + 1) + + # build graph & observations + G = pyg.utils.to_networkx(data, to_undirected=True, remove_self_loops=True) + obs_ts = _parse_obs_ts(args, T) # ensure T is present + obs_map = _build_observation_map(data, obs_ts) + + # init buffer and seed observed frames + y_pred = np.zeros((n_nodes, T + 1), dtype=np.int32) + for t in obs_ts: + y_pred[:, t] = obs_map[t] + + # include 0 to cover the entire span; we will reconstruct down to t=0 + time_cuts = sorted(set(obs_ts + [0])) + + # process segments from later to earlier + for idx in trange(len(time_cuts) - 1, 0, -1, desc='ms-backtrack'): + t_prev = time_cuts[idx - 1] + t_next = time_cuts[idx] # t_next > t_prev + y_cur = y_pred[:, t_next] # this is observed or already clamped + + # step-by-step: t_next -> t_prev + for t in range(t_next, t_prev, -1): + y_prev = pcdsvc_greedy(bpar, G, y_cur) + + # if the new time (t-1) is observed, clamp to the observation + if (t - 1) in obs_map: + y_prev = obs_map[t - 1] + + y_pred[:, t - 1] = y_prev + y_cur = y_prev + + # clip to valid classes and return torch tensor on the original device + return torch.tensor(np.minimum(y_pred, n_cls - 1), + dtype=torch.long, device=data.y.device) + + +args = get_args() +seed_all(args.seed) +tester = Tester(args.data_dir, args.device, pcdsvc_run) +tester.test([args.dataset], rep=1) +tester.save(args.output) + diff --git a/ditto.py b/ditto.py index b4e482c..438d8fe 100644 --- a/ditto.py +++ b/ditto.py @@ -121,7 +121,16 @@ def lik(self, Y): # Y: (T+1, nodes, samples) rem.flatten()[vids] = torch.where(mski.repeat_interleave(repeats = degi), torch.where(trsi.repeat_interleave(repeats = degi), rems - 1, self.n_inf), rems) # (T * sum neighbs) mskI[uidi] &= opti # (T * samples) # likR + likI - lik = (torch.where(mskR, torch.where(trsR, lR1, lR0), self.zero).view(-1, n_samples) + torch.where(mskI.view(-1, n_samples), torch.where(trsI.view(-1, n_samples), lI1.view(-1, n_samples), lI0.view(-1, n_samples)), self.zero)).sum(dim = 0) # (samples,) + likR = torch.where(mskR, torch.where(trsR, lR1, lR0), self.zero).reshape(-1, n_samples) + + mskI2 = mskI.reshape(-1, n_samples) + trsI2 = trsI.reshape(-1, n_samples) + lI1_2 = lI1.reshape(-1, n_samples) + lI0_2 = lI0.reshape(-1, n_samples) + + likI = torch.where(mskI2, torch.where(trsI2, lI1_2, lI0_2), self.zero) + + lik = (likR + likI).sum(dim=0) return lik, zI0, zR0, zI, zR # (samples,) @torch.no_grad() def clamp_grad(self, z0, grad): @@ -217,26 +226,39 @@ def t_mcmc(data, bpar, q_net, args, keepdim = True): tR_avg = args.t_keep * tR_avg + (1. - args.t_keep) * tR # (nodes, 1) return tI_avg, tR_avg # (nodes, 1) -def main(data): +def run_ditto_on_data(data, args): # estimate diffusion parameters bpar = b_estim(data, args) print(f'[est] pI={bpar.pI:.4f}, pR={bpar.pR:.4f}', flush = True) + # train a proposal network q_net = q_train(data, bpar, args) + # estimate transition times - tI, tR = t_mcmc(data, bpar, q_net, args, keepdim = True) # (nodes, 1) - T = data.T.item() + tI, tR = t_mcmc(data, bpar, q_net, args, keepdim = True) # (nodes, 1) tI = tI.round().long() tR = tR.round().long() + # compose a history with torch.no_grad(): - y_pred = torch.zeros_like(data.y) # (nodes, T+1) + y_pred = torch.zeros_like(data.y) # (nodes, T+1) y_pred.scatter_(dim = 1, index = torch.minimum(tI, data.T), src = torch.full_like(tI, 1)) y_pred.scatter_(dim = 1, index = torch.minimum(tR, data.T), src = torch.full_like(tR, 2)) y_pred = y_pred[:, : data.T.item()].cummax(dim = 1).values return y_pred -args = get_args() -tester = Tester(args.data_dir, args.device, main) -tester.test([args.dataset], seed = args.seed, rep = 1) -tester.save(args.output) + +def main(data): + return run_ditto_on_data(data, args) + + +if __name__ == '__main__': + args = get_args() + if args.device is None: + args.device = torch_device() + + tester = Tester(args.data_dir, args.device, main) + tester.test([args.dataset], seed = args.seed, rep = 1) + + if args.output is not None: + tester.save(args.output) \ No newline at end of file diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..cc1e9af --- /dev/null +++ b/environment.yml @@ -0,0 +1,8 @@ +name: ditto-gpu +channels: + - defaults +dependencies: + - python=3.11.14 + - pip + - git +prefix: C:\Users\z1585\anaconda3\envs\ditto-gpu diff --git a/experiments/ditto_seg_twosnap.py b/experiments/ditto_seg_twosnap.py new file mode 100644 index 0000000..20b78dc --- /dev/null +++ b/experiments/ditto_seg_twosnap.py @@ -0,0 +1,142 @@ +from inc.diffus import * +from inc.test import * +from ditto import run_ditto_on_data + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type = str, default = None, help = 'single dataset name') + parser.add_argument('--datasets', type = str, default = None, help = 'comma-separated dataset names') + parser.add_argument('--seed', type = int, help = 'random seed') + parser.add_argument('--data_dir', type = str, help = 'dataset folder') + parser.add_argument('--output', type = str, help = 'output file name') + parser.add_argument('--device', type = torch.device, help = 'torch device') + + # segment-and-stitch specific + parser.add_argument('--split_time', type = int, default = None, + help = 'global split time; default=floor(T/2)') + + # same DITTO args as ditto.py + parser.add_argument('--b_pI0', type = float, help = 'initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type = float, help = 'initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type = int, help = 'optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type = float, help = 'learning rate in diffusion parameter estimation') + + parser.add_argument('--q_steps', type = int, help = 'training steps for the proposal model') + parser.add_argument('--q_lr', type = float, help = 'learning rate for the proposal model') + parser.add_argument('--q_hid', type = int, help = 'hidden size of the proposal model') + parser.add_argument('--q_gnn', type = int, help = 'number of layers of the GNN in the proposal model') + parser.add_argument('--q_mlp', type = int, help = 'number of layers of the MLP in the proposal model') + parser.add_argument('--q_samples', type = int, help = 'sample size to estimate the loss function of the proposal model') + parser.add_argument('--q_zlim', type = int, help = 'a hyperparameter to stabilize gradient') + + parser.add_argument('--p_coef', type = float, help = 'the coefficient gamma in the initial distribution P[y_0]') + parser.add_argument('--t_samples', type = int, help = 'MCMC sample size') + parser.add_argument('--t_steps', type = int, help = 'MCMC steps') + parser.add_argument('--t_keep', type = float, help = 'moving average in MCMC') + + args = parser.parse_args() + + if args.device is None: + args.device = torch_device() + + return args + + +def resolve_datasets(args): + if args.datasets is not None: + datasets = [x.strip() for x in args.datasets.split(',') if x.strip()] + assert len(datasets) > 0, '--datasets is empty' + return datasets + + assert args.dataset is not None, 'either --dataset or --datasets is required' + return [args.dataset] + + +@torch.no_grad() +def make_segment_data(data, t_start, t_end): + """ + Build a local segment subproblem from global times [t_start, t_end]. + + Local time 0 <-> global time t_start + Local time seg_T <-> global time t_end + """ + assert 0 <= t_start < t_end <= data.T.item(), 'invalid segment boundary' + + seg_y = data.y[:, t_start : t_end + 1].detach().clone() # (nodes, seg_T+1) + + seg = Dict( + edge_index = data.edge_index, + num_nodes = data.num_nodes, + y = seg_y, + T = torch.tensor(t_end - t_start, dtype = data.T.dtype, device = data.T.device), + name = f'{data.name}[{t_start},{t_end}]', + ) + return seg + + +@torch.no_grad() +def stitch_two_segments(y1_full, y2_full): + """ + y1_full: global [0, ..., t_split] + y2_full: global [t_split, ..., T] + + Direct stitching rule: + - keep segment 1's terminal snapshot at t_split + - append segment 2 from t_split+1 onward + """ + y_full = torch.cat([y1_full, y2_full[:, 1:]], dim = 1) + return y_full + + +def run_ditto_seg_on_data(data, args): + """ + Naive two-snapshot extension baseline: + 1) run single-snapshot DITTO on [0, t_split], using y_{t_split} as final snapshot + 2) run single-snapshot DITTO on [t_split, T], using y_T as final snapshot + 3) directly stitch the two histories + """ + T = data.T.item() + split_time = args.split_time if args.split_time is not None else (T // 2) + assert 1 <= split_time < T, f'split_time must be in [1, {T - 1}]' + + seg1 = make_segment_data(data, 0, split_time) + seg2 = make_segment_data(data, split_time, T) + + print(f'[split] dataset={data.name} T={T} split={split_time}', flush = True) + + # segment 1: [0, split_time] + print(f'[seg1] run DITTO on [0, {split_time}] with final snapshot y_{split_time}', flush = True) + y1_pred = run_ditto_on_data(seg1, args) # (nodes, split_time) + y1_full = torch.cat([y1_pred, seg1.y[:, -1 :]], dim = 1) # (nodes, split_time+1) + + # segment 2: [split_time, T] + print(f'[seg2] run DITTO on [{split_time}, {T}] with final snapshot y_{T}', flush = True) + y2_pred = run_ditto_on_data(seg2, args) # (nodes, T-split_time) + y2_full = torch.cat([y2_pred, seg2.y[:, -1 :]], dim = 1) # (nodes, T-split_time+1) + + # stitch + y_full = stitch_two_segments(y1_full, y2_full) # (nodes, T+1) + + assert y_full.size(1) == T + 1, 'stitched history has wrong length' + assert torch.equal(y_full[:, split_time], data.y[:, split_time]), 'split snapshot mismatch after stitching' + assert torch.equal(y_full[:, -1], data.y[:, -1]), 'final snapshot mismatch after stitching' + + # Tester expects shape (nodes, T), i.e. all times except the final snapshot + return y_full[:, : T] + + +if __name__ == '__main__': + args = get_args() + datasets = resolve_datasets(args) + + tester = Tester( + args.data_dir, + args.device, + lambda data: run_ditto_seg_on_data(data, args), + ) + + tester.test(datasets, seed = args.seed, rep = 1) + + if args.output is not None: + tester.save(args.output) \ No newline at end of file diff --git a/experiments/exp1_scalability.py b/experiments/exp1_scalability.py new file mode 100644 index 0000000..d39e655 --- /dev/null +++ b/experiments/exp1_scalability.py @@ -0,0 +1,511 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Experiment 1: Scalability (HERMES) + +Runtime profiling for: + (1) vs T (fixed n, vary T) + (2) vs n (fixed T, vary n) + +Dataset: ba-sir +Observation: only TWO observed frames for each run: + obs_time = [floor(T/2), T] + +Defaults for vsT: + T in {3,4,5,6,7,8,9,10} (skip T=1,2) + +Defaults for vsN: + T_fixed = 10 + scale n by tiling disjoint copies (factors) + +Run from repo root: + python experiments/exp1_scalability.py --method hermes --dataset ba-sir --data_dir input --device cuda + +Optional: + --save_datasets to export generated .pt datasets (for traceability) +""" + +from __future__ import annotations + +import os +import sys +import gc +import csv +import time +import argparse +import platform +from dataclasses import dataclass, asdict +from datetime import datetime +from typing import List, Dict, Any + +import torch +from torch_geometric.data import Data + +# --------------------------------------------------------------------- +# Make repo root importable (so `import inc.*` and `import hermes` works) +# --------------------------------------------------------------------- +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +# project imports (re-use existing code) +from inc.data import data_load, data_make_states # type: ignore +import hermes # hermes.py must be import-safe (guarded by __main__) + +# seeding util (fallback if inc.utils not present) +try: + from inc.utils import seed_all # type: ignore +except Exception: # pragma: no cover + import random + import numpy as np + + def seed_all(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +CSV_FIELDS = [ + "timestamp", + "exp", + "method", + "base_dataset", + "dataset_pt", + "seed", + "device", + "status", + "n_nodes", + "n_edges", + "T", + "obs_time_mid", + "obs_time", + "n_factor", + # HERMES HP + "b_pI0", + "b_pR0", + "b_steps", + "b_lr", + "q_steps", + "q_lr", + "q_hid", + "q_gnn", + "q_mlp", + "q_samples", + "q_zlim", + "p_coef", + "t_samples", + "t_steps", + "t_keep", + # timings + "b_estim_sec", + "q_train_sec", + "t_mcmc_sec", + "total_sec", + # meta + "host", +] + + +def _sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _parse_int_list(s: str) -> List[int]: + if s is None: + return [] + s = s.strip() + if not s: + return [] + parts = s.replace(",", " ").split() + return [int(x) for x in parts] + + +def _ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def _safe_save_pt(data: Data, path: str) -> None: + _ensure_dir(os.path.dirname(path)) + torch.save(data.cpu(), path) + + +def _variant_timespan(base: Data, T_new: int) -> Data: + """ + Truncate/clip diffusion timespan to T_new and recompute y from (clipped) tI/tR. + """ + assert T_new >= 1, "T_new should be >= 1" + device = base.tI.device + T_tensor = torch.tensor(T_new, dtype=base.T.dtype, device=device) + + # Clip hitting times beyond T_new to 'never' = T_new+1 + cap = torch.full_like(base.tI, T_new + 1) + tI_new = torch.minimum(base.tI, cap) + + tR_new = None + if hasattr(base, "tR") and getattr(base, "tR") is not None: + capR = torch.full_like(base.tR, T_new + 1) + tR_new = torch.minimum(base.tR, capR) + + y_new = data_make_states(T_new, tI_new, tR_new) + + out = Data(edge_index=base.edge_index.clone(), y=y_new, T=T_tensor, tI=tI_new.clone()) + if tR_new is not None: + out.tR = tR_new.clone() + + out.num_nodes = y_new.size(0) + return out + + +def _variant_scale_n(base: Data, factor: int) -> Data: + """ + Tile disjoint copies of a base graph/history to scale n. + This avoids rewriting the synthetic generator and is sufficient for runtime scaling. + """ + assert factor >= 1 + if factor == 1: + out = Data(edge_index=base.edge_index.clone(), y=base.y.clone(), T=base.T.clone(), tI=base.tI.clone()) + if hasattr(base, "tR") and getattr(base, "tR") is not None: + out.tR = base.tR.clone() + out.num_nodes = base.num_nodes + return out + + n0 = int(base.num_nodes) + eidx_list = [base.edge_index + k * n0 for k in range(factor)] + edge_index = torch.cat(eidx_list, dim=1) + + y = base.y.repeat(factor, 1) + tI = base.tI.repeat(factor) + out = Data(edge_index=edge_index, y=y, T=base.T.clone(), tI=tI) + if hasattr(base, "tR") and getattr(base, "tR") is not None: + out.tR = base.tR.repeat(factor) + + out.num_nodes = y.size(0) + return out + + +@dataclass +class HermesHP: + """ + Defaults match your provided ba-sir command: + --b_pI0 0.001 --b_pR0 0.001 --b_steps 500 --b_lr 0.003 + --q_steps 250 --q_lr 0.003 --q_hid 16 --q_gnn 3 --q_mlp 2 --q_samples 10 --q_zlim 16 + --p_coef 1.0 --t_samples 100 --t_steps 100 --t_keep 0.5 + """ + b_pI0: float = 0.001 + b_pR0: float = 0.001 + b_steps: int = 500 + b_lr: float = 0.003 + q_steps: int = 250 + q_lr: float = 0.003 + q_hid: int = 16 + q_gnn: int = 3 + q_mlp: int = 2 + q_samples: int = 10 + q_zlim: int = 16 + p_coef: float = 1.0 + t_samples: int = 100 + t_steps: int = 100 + t_keep: float = 0.5 + + def to_namespace(self, *, device: torch.device, seed: int, obs_time_mid: int) -> argparse.Namespace: + ns = argparse.Namespace(**asdict(self)) + ns.device = device + ns.seed = seed + # for compatibility; we pass obs_time explicitly in calls anyway + ns.obs_time = str(obs_time_mid) + return ns + + +def _run_hermes_once_timed(data: Data, args: argparse.Namespace, obs_time: List[int]) -> Dict[str, float]: + """ + Run HERMES pipeline once and return per-stage runtimes (seconds): + b_estim, q_train, t_mcmc, total + """ + obs_time = sorted(set(int(t) for t in obs_time)) + + # move to device (exclude copy time from timing) + data = data.to(args.device) + + _sync(args.device) + t0 = time.perf_counter() + + # 1) diffusion parameter estimation + b0 = time.perf_counter() + bpar = hermes.b_estim(data, args, obs_time=obs_time) + _sync(args.device) + t_b = time.perf_counter() - b0 + + # 2) proposal training + q0 = time.perf_counter() + q_net = hermes.q_train(data, obs_time, bpar, args) + _sync(args.device) + t_q = time.perf_counter() - q0 + + # 3) MCMC inference + m0 = time.perf_counter() + _ = hermes.t_mcmc(data, bpar, q_net, args, obs_time=obs_time, keepdim=True) + _sync(args.device) + t_m = time.perf_counter() - m0 + + t_total = time.perf_counter() - t0 + + # cleanup (outside timing) + del q_net, bpar + gc.collect() + if args.device.type == "cuda": + torch.cuda.empty_cache() + + return {"b_estim_sec": t_b, "q_train_sec": t_q, "t_mcmc_sec": t_m, "total_sec": t_total} + + +def _append_csv(path: str, row: Dict[str, Any]) -> None: + _ensure_dir(os.path.dirname(path)) + file_exists = os.path.exists(path) + with open(path, "a", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=CSV_FIELDS, extrasaction="ignore") + if not file_exists: + writer.writeheader() + fixed_row = {k: row.get(k, "") for k in CSV_FIELDS} + writer.writerow(fixed_row) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, default="hermes", choices=["hermes"]) + parser.add_argument("--dataset", type=str, default="ba-sir") + parser.add_argument("--data_dir", type=str, default="input") + parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--seed", type=int, default=123456789) + + # Exp controls + parser.add_argument("--run_vsT", action="store_true", help="run scalability vs T") + parser.add_argument("--run_vsN", action="store_true", help="run scalability vs n") + + # ba-sir: T=3..10 (skip 1,2) + parser.add_argument("--T_list", type=str, default="3,4,5,6,7,8,9,10", + help="comma-separated T list for vsT") + parser.add_argument("--T_fixed", type=int, default=10, help="fixed T for vsN") + + # scale-n factors (n = factor * n0). You can override to larger factors, script will mark OOM if happens. + parser.add_argument("--n_factors", type=str, default="1,2,4,8,16", + help="tile factors for vsN (n = factor * n0)") + parser.add_argument("--n_factor_for_vsT", type=int, default=1, + help="optionally scale n before varying T") + + # IO + parser.add_argument("--gen_dir", type=str, default="", + help="where to save generated .pt datasets (default: /exp1_scalability)") + parser.add_argument("--save_datasets", action="store_true", help="save generated datasets as .pt") + parser.add_argument("--out_csv", type=str, default="output/exp1_scalability_runtime_ba_sir.csv") + parser.add_argument("--overwrite_csv", action="store_true") + + # HERMES hyperparameters (defaults match your provided ba-sir command) + parser.add_argument("--b_pI0", type=float, default=0.001) + parser.add_argument("--b_pR0", type=float, default=0.001) + parser.add_argument("--b_steps", type=int, default=500) + parser.add_argument("--b_lr", type=float, default=0.003) + parser.add_argument("--q_steps", type=int, default=250) + parser.add_argument("--q_lr", type=float, default=0.003) + parser.add_argument("--q_hid", type=int, default=16) + parser.add_argument("--q_gnn", type=int, default=3) + parser.add_argument("--q_mlp", type=int, default=2) + parser.add_argument("--q_samples", type=int, default=10) + parser.add_argument("--q_zlim", type=int, default=16) + parser.add_argument("--p_coef", type=float, default=1.0) + parser.add_argument("--t_samples", type=int, default=100) + parser.add_argument("--t_steps", type=int, default=100) + parser.add_argument("--t_keep", type=float, default=0.5) + + args_cli = parser.parse_args() + + if not args_cli.run_vsT and not args_cli.run_vsN: + # default: run both + args_cli.run_vsT = True + args_cli.run_vsN = True + + # device resolve + if args_cli.device.startswith("cuda") and not torch.cuda.is_available(): + print("[warn] CUDA not available, fallback to CPU.") + device = torch.device("cpu") + else: + device = torch.device(args_cli.device) + + # output csv + if args_cli.overwrite_csv and os.path.exists(args_cli.out_csv): + os.remove(args_cli.out_csv) + + # gen dir + gen_dir = args_cli.gen_dir.strip() or os.path.join(args_cli.data_dir, "exp1_scalability") + _ensure_dir(gen_dir) + + # load base dataset on CPU (generation happens on CPU, then move to device for timing) + base = data_load(args_cli.dataset, args_cli.data_dir, torch.device("cpu")) + n0 = int(base.num_nodes) + m0 = int(base.edge_index.size(1)) + T0 = int(base.T.item()) + print(f"[load] {args_cli.dataset}: n={n0}, m={m0}, T={T0}") + + # build HP template + hp = HermesHP( + b_pI0=args_cli.b_pI0, + b_pR0=args_cli.b_pR0, + b_steps=args_cli.b_steps, + b_lr=args_cli.b_lr, + q_steps=args_cli.q_steps, + q_lr=args_cli.q_lr, + q_hid=args_cli.q_hid, + q_gnn=args_cli.q_gnn, + q_mlp=args_cli.q_mlp, + q_samples=args_cli.q_samples, + q_zlim=args_cli.q_zlim, + p_coef=args_cli.p_coef, + t_samples=args_cli.t_samples, + t_steps=args_cli.t_steps, + t_keep=args_cli.t_keep, + ) + + host = platform.node() + + # ------------------------- + # run vsT + # ------------------------- + if args_cli.run_vsT: + T_list = _parse_int_list(args_cli.T_list) + assert len(T_list) > 0, "T_list is empty" + print(f"[exp] vsT: T_list={T_list}, n_factor_for_vsT={args_cli.n_factor_for_vsT}") + + base_scaled = _variant_scale_n(base, args_cli.n_factor_for_vsT) + + for T in T_list: + if T > T0: + print(f"[skip] T={T} > base T0={T0}.") + continue + if T <= 2: + print(f"[skip] T={T} (skip T<=2 for this setting).") + continue + + data_T = _variant_timespan(base_scaled, T) + obs_mid = T // 2 + obs_time = [obs_mid, T] # exactly TWO frames + + pt_path = "" + if args_cli.save_datasets: + pt_path = os.path.join(gen_dir, "vsT", f"{args_cli.dataset}_n{data_T.num_nodes}_T{T}.pt") + _safe_save_pt(data_T, pt_path) + + seed_all(args_cli.seed) + run_args = hp.to_namespace(device=device, seed=args_cli.seed, obs_time_mid=obs_mid) + + status = "ok" + times = {"b_estim_sec": float("nan"), "q_train_sec": float("nan"), + "t_mcmc_sec": float("nan"), "total_sec": float("nan")} + try: + times = _run_hermes_once_timed(data_T, run_args, obs_time) + except RuntimeError as e: + if "out of memory" in str(e).lower(): + status = "oom" + if device.type == "cuda": + torch.cuda.empty_cache() + else: + raise + + row = { + "timestamp": datetime.now().isoformat(timespec="seconds"), + "exp": "vsT", + "method": args_cli.method, + "base_dataset": args_cli.dataset, + "dataset_pt": pt_path, + "seed": args_cli.seed, + "device": str(device), + "status": status, + "n_nodes": int(data_T.num_nodes), + "n_edges": int(data_T.edge_index.size(1)), + "T": int(T), + "obs_time_mid": int(obs_mid), + "obs_time": ",".join(map(str, obs_time)), + "n_factor": int(args_cli.n_factor_for_vsT), + **asdict(hp), + **times, + "host": host, + } + _append_csv(args_cli.out_csv, row) + print(f"[done] vsT T={T} n={row['n_nodes']} total={row['total_sec']:.3f}s status={status}") + + del data_T + gc.collect() + + # ------------------------- + # run vsN + # ------------------------- + if args_cli.run_vsN: + factors = _parse_int_list(args_cli.n_factors) + assert len(factors) > 0, "n_factors is empty" + T_fixed = int(args_cli.T_fixed) + if T_fixed > T0: + print(f"[warn] T_fixed={T_fixed} > base T0={T0}, truncate to T0={T0}.") + T_fixed = T0 + + print(f"[exp] vsN: factors={factors}, T_fixed={T_fixed}") + + base_T = _variant_timespan(base, T_fixed) + obs_mid = T_fixed // 2 + obs_time = [obs_mid, T_fixed] # exactly TWO frames + + for fac in factors: + data_N = _variant_scale_n(base_T, fac) + + pt_path = "" + if args_cli.save_datasets: + pt_path = os.path.join(gen_dir, "vsN", f"{args_cli.dataset}_n{data_N.num_nodes}_T{T_fixed}.pt") + _safe_save_pt(data_N, pt_path) + + seed_all(args_cli.seed) + run_args = hp.to_namespace(device=device, seed=args_cli.seed, obs_time_mid=obs_mid) + + status = "ok" + times = {"b_estim_sec": float("nan"), "q_train_sec": float("nan"), + "t_mcmc_sec": float("nan"), "total_sec": float("nan")} + try: + times = _run_hermes_once_timed(data_N, run_args, obs_time) + except RuntimeError as e: + if "out of memory" in str(e).lower(): + status = "oom" + if device.type == "cuda": + torch.cuda.empty_cache() + else: + raise + + row = { + "timestamp": datetime.now().isoformat(timespec="seconds"), + "exp": "vsN", + "method": args_cli.method, + "base_dataset": args_cli.dataset, + "dataset_pt": pt_path, + "seed": args_cli.seed, + "device": str(device), + "status": status, + "n_nodes": int(data_N.num_nodes), + "n_edges": int(data_N.edge_index.size(1)), + "T": int(T_fixed), + "obs_time_mid": int(obs_mid), + "obs_time": ",".join(map(str, obs_time)), + "n_factor": int(fac), + **asdict(hp), + **times, + "host": host, + } + _append_csv(args_cli.out_csv, row) + print(f"[done] vsN fac={fac} n={row['n_nodes']} total={row['total_sec']:.3f}s status={status}") + + del data_N + gc.collect() + + print(f"[ok] wrote CSV -> {args_cli.out_csv}") + + +if __name__ == "__main__": + main() diff --git a/experiments/exp2_timespan.py b/experiments/exp2_timespan.py new file mode 100644 index 0000000..5c866b2 --- /dev/null +++ b/experiments/exp2_timespan.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Experiment 2: Effect of Timespan (ba-sir) + +- Observations: only TWO frames per run: + obs_time = { floor(T/2), T } + +- Sweep T from 3 to 10 (skip 1,2). +- For each T, regenerate a dataset file: + /exp2_timespan/T{T}/synthetic/ba-sir.pt + +- Compare: + 1) HERMES (hermes.py) + 2) CRI-MS (cri_ms.py) + 3) DHREC-MS (dhrec_ms.py) + +- Evaluation: + Use the existing inc/tester inside each script (they already do), + then this orchestrator reads the saved .pt result and writes a CSV + with f1 and nrmse (no plotting). + +Run from repo root: + python experiments/exp2_timespan.py --data_dir input --device_hermes cuda --device_dhrec cuda --device_cri cpu +""" + +from __future__ import annotations + +import os +import sys +import csv +import time +import argparse +import subprocess +from datetime import datetime +from typing import Dict, Any, List + +import torch +from torch_geometric.data import Data + +# --------------------------------------------------------------------- +# Make repo root importable (so `from inc.data import ...` works) +# --------------------------------------------------------------------- +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +from inc.data import data_load, data_make_states # type: ignore + + +CSV_FIELDS = [ + "timestamp", + "dataset", + "T", + "obs_time_mid", + "obs_time", + "method", + "device", + "seed", + "status", + "f1", + "nrmse", + "result_pt", + "data_dir_used", + "data_pt_used", + "cmd", + "elapsed_sec", +] + + +def _ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def _append_csv(path: str, row: Dict[str, Any]) -> None: + _ensure_dir(os.path.dirname(path)) + file_exists = os.path.exists(path) + with open(path, "a", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=CSV_FIELDS, extrasaction="ignore") + if not file_exists: + w.writeheader() + fixed_row = {k: row.get(k, "") for k in CSV_FIELDS} + w.writerow(fixed_row) + + +def _variant_timespan(base: Data, T_new: int) -> Data: + """ + Create a new Data with timespan clipped to T_new: + - tI, tR clipped to <= T_new+1 + - y rebuilt by data_make_states(T_new, tI_new, tR_new) + """ + assert T_new >= 1 + device = base.tI.device # typically CPU here + T_tensor = torch.tensor(T_new, dtype=base.T.dtype, device=device) + + cap = torch.full_like(base.tI, T_new + 1) + tI_new = torch.minimum(base.tI, cap) + + tR_new = None + if hasattr(base, "tR") and getattr(base, "tR") is not None: + capR = torch.full_like(base.tR, T_new + 1) + tR_new = torch.minimum(base.tR, capR) + + y_new = data_make_states(T_new, tI_new, tR_new) + + out = Data(edge_index=base.edge_index.clone(), y=y_new, T=T_tensor, tI=tI_new.clone()) + if tR_new is not None: + out.tR = tR_new.clone() + out.num_nodes = y_new.size(0) + return out + + +def _save_ba_sir_variant(data_dir_T: str, data_T: Data) -> str: + """ + Save variant dataset to: + /synthetic/ba-sir.pt + so that data_load('ba-sir', data_dir_T, device) will load this file. + """ + pt_path = os.path.join(data_dir_T, "synthetic", "ba-sir.pt") + _ensure_dir(os.path.dirname(pt_path)) + torch.save(data_T.cpu(), pt_path) + return pt_path + + +def _load_metrics_from_tester_pt(result_pt: str, dataset: str) -> Dict[str, float]: + """ + Each method script saves tester.res via torch.save(res, output). + res format: res[dataset]['f1'] = [..], res[dataset]['nrmse'] = [..] + """ + res = torch.load(result_pt, map_location="cpu", weights_only=False) + f1 = float(res[dataset]["f1"][-1]) + nrmse = float(res[dataset]["nrmse"][-1]) + return {"f1": f1, "nrmse": nrmse} + + +def _run_cmd(cmd: List[str]) -> float: + """ + Run a command and return wall-clock seconds. + """ + t0 = time.perf_counter() + subprocess.run(cmd, check=True) + return time.perf_counter() - t0 + + +def main() -> None: + p = argparse.ArgumentParser() + + # sweep setting + p.add_argument("--dataset", type=str, default="ba-sir") + p.add_argument("--seed", type=int, default=123456789) + p.add_argument("--data_dir", type=str, default="input") + p.add_argument("--T_min", type=int, default=3) + p.add_argument("--T_max", type=int, default=10) + + # where to store per-T datasets and per-run result .pt + p.add_argument("--gen_root", type=str, default="", help="default: /exp2_timespan") + p.add_argument("--out_root", type=str, default="output/exp2_timespan") + p.add_argument("--out_csv", type=str, default="output/exp2_timespan_metrics_ba_sir.csv") + p.add_argument("--overwrite_csv", action="store_true") + + # devices + p.add_argument("--device_hermes", type=str, default="cuda") + p.add_argument("--device_dhrec", type=str, default="cuda") + p.add_argument("--device_cri", type=str, default="cpu") + + # HERMES hyperparameters (keep fixed) + p.add_argument("--b_pI0", type=float, default=0.001) + p.add_argument("--b_pR0", type=float, default=0.001) + p.add_argument("--b_steps", type=int, default=500) + p.add_argument("--b_lr", type=float, default=0.003) + p.add_argument("--q_steps", type=int, default=250) + p.add_argument("--q_lr", type=float, default=0.003) + p.add_argument("--q_hid", type=int, default=16) + p.add_argument("--q_gnn", type=int, default=3) + p.add_argument("--q_mlp", type=int, default=2) + p.add_argument("--q_samples", type=int, default=10) + p.add_argument("--q_zlim", type=int, default=16) + p.add_argument("--p_coef", type=float, default=1.0) + p.add_argument("--t_samples", type=int, default=100) + p.add_argument("--t_steps", type=int, default=100) + p.add_argument("--t_keep", type=float, default=0.5) + + args = p.parse_args() + + if args.overwrite_csv and os.path.exists(args.out_csv): + os.remove(args.out_csv) + + gen_root = args.gen_root.strip() or os.path.join(args.data_dir, "exp2_timespan") + _ensure_dir(gen_root) + _ensure_dir(args.out_root) + + # Load base ba-sir once (CPU). For synthetic ba-sir, this is typically T=10 cached at /synthetic/ba-sir.pt. + base = data_load(args.dataset, args.data_dir, torch.device("cpu")) + base_T = int(base.T.item()) + + print(f"[load base] dataset={args.dataset} base_T={base_T} (expect >= {args.T_max})") + + py = sys.executable + + # sweep + for T in range(args.T_min, args.T_max + 1): + if T <= 2: + continue + if T > base_T: + print(f"[skip] T={T} > base_T={base_T}") + continue + + obs_mid = T // 2 + obs_time = [obs_mid, T] # exactly two frames + obs_time_str = ",".join(map(str, obs_time)) + + # 1) regenerate data for this T into a dedicated data_dir + data_dir_T = os.path.join(gen_root, f"T{T}") + data_T = _variant_timespan(base, T) + data_pt = _save_ba_sir_variant(data_dir_T, data_T) + print(f"[data] T={T} saved -> {data_pt} (obs={obs_time_str})") + + # 2) run methods + # 2.1 HERMES: obs_time arg is "extra observed times"; final T always included inside hermes.py + hermes_out = os.path.join(args.out_root, f"hermes_T{T}.pt") + cmd_hermes = [ + py, "hermes.py", + "--dataset", args.dataset, + "--seed", str(args.seed), + "--data_dir", data_dir_T, + "--output", hermes_out, + "--device", args.device_hermes, + "--obs_time", str(obs_mid), + "--b_pI0", str(args.b_pI0), "--b_pR0", str(args.b_pR0), "--b_steps", str(args.b_steps), "--b_lr", str(args.b_lr), + "--q_steps", str(args.q_steps), "--q_lr", str(args.q_lr), "--q_hid", str(args.q_hid), "--q_gnn", str(args.q_gnn), + "--q_mlp", str(args.q_mlp), "--q_samples", str(args.q_samples), "--q_zlim", str(args.q_zlim), + "--p_coef", str(args.p_coef), + "--t_samples", str(args.t_samples), "--t_steps", str(args.t_steps), "--t_keep", str(args.t_keep), + ] + + # 2.2 CRI-MS: pass BOTH mid and T explicitly (cri_ms.py doesn't auto-append T) + cri_out = os.path.join(args.out_root, f"cri_T{T}.pt") + cmd_cri = [ + py, "cri_ms.py", + "--dataset", args.dataset, + "--seed", str(args.seed), + "--data_dir", data_dir_T, + "--output", cri_out, + "--device", args.device_cri, + "--obs_ts", obs_time_str, + ] + + # 2.3 DHREC-MS: pass BOTH mid and T (dhrec_ms.py will ensure T included anyway) + dhrec_out = os.path.join(args.out_root, f"dhrec_T{T}.pt") + cmd_dhrec = [ + py, "dhrec_ms.py", + "--dataset", args.dataset, + "--seed", str(args.seed), + "--data_dir", data_dir_T, + "--output", dhrec_out, + "--device", args.device_dhrec, + "--b_pI0", str(args.b_pI0), "--b_pR0", str(args.b_pR0), "--b_steps", str(args.b_steps), "--b_lr", str(args.b_lr), + "--obs_ts", obs_time_str, + ] + + runs = [ + ("hermes", args.device_hermes, cmd_hermes, hermes_out), + ("cri", args.device_cri, cmd_cri, cri_out), + ("dhrec", args.device_dhrec, cmd_dhrec, dhrec_out), + ] + + for method, device, cmd, out_pt in runs: + row = { + "timestamp": datetime.now().isoformat(timespec="seconds"), + "dataset": args.dataset, + "T": T, + "obs_time_mid": obs_mid, + "obs_time": obs_time_str, + "method": method, + "device": device, + "seed": args.seed, + "data_dir_used": data_dir_T, + "data_pt_used": data_pt, + "result_pt": out_pt, + "cmd": " ".join(cmd), + } + + status = "ok" + elapsed = float("nan") + f1 = float("nan") + nrmse = float("nan") + + try: + elapsed = _run_cmd(cmd) + mets = _load_metrics_from_tester_pt(out_pt, args.dataset) + f1, nrmse = mets["f1"], mets["nrmse"] + except subprocess.CalledProcessError: + status = "error" + except FileNotFoundError: + status = "missing_output" + except Exception: + status = "error" + + row.update({ + "status": status, + "elapsed_sec": elapsed, + "f1": f1, + "nrmse": nrmse, + }) + _append_csv(args.out_csv, row) + + print(f"[done] T={T} method={method:6s} status={status} " + f"f1={f1 if f1==f1 else float('nan'):.4f} " + f"nrmse={nrmse if nrmse==nrmse else float('nan'):.4f} " + f"sec={elapsed if elapsed==elapsed else float('nan'):.2f}") + + print(f"[ok] wrote CSV -> {args.out_csv}") + + +if __name__ == "__main__": + main() diff --git a/experiments/exp3_beta_scatter.py b/experiments/exp3_beta_scatter.py new file mode 100644 index 0000000..4948668 --- /dev/null +++ b/experiments/exp3_beta_scatter.py @@ -0,0 +1,212 @@ +import os +import sys +import csv +import argparse +from typing import Dict, Any, List, Tuple + +import torch + +HERE = os.path.abspath(os.path.dirname(__file__)) +CANDIDATE_ROOTS = [HERE, os.path.abspath(os.path.join(HERE, ".."))] +for cand in CANDIDATE_ROOTS: + if os.path.isdir(os.path.join(cand, "inc")): + if cand not in sys.path: + sys.path.insert(0, cand) + break + +from inc.data import data_load # type: ignore +from inc.diffus import b_estim # type: ignore + + +def _ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def parse_datasets_arg(s: str) -> List[str]: + s = (s or "").strip() + if s.lower() in {"main", "default"}: + # 8 datasets used in main experiments where true parameters are known by construction. + return [ + "ba-si", "ba-sir", + "er-si", "er-sir", + "oregon2-si", "oregon2-sir", + "prost-si", "prost-sir", + ] + return [x.strip() for x in s.split(",") if x.strip()] + + +def parse_obs_time_arg(s: str, T: int) -> List[int]: + """ + Parse --obs_time. If empty, use main protocol {floor(T/2), T}. + Input format examples: + --obs_time "" -> [T//2, T] + --obs_time "5" -> [5, T] + --obs_time "5,10" -> [5, 10] (and ensures T included) + """ + s = (s or "").strip() + if not s: + return [T // 2, T] + ts = [int(x) for x in s.split(",") if x.strip() != ""] + ts.append(T) + ts = sorted(set(t for t in ts if 0 <= t <= T)) + if ts[-1] != T: + ts.append(T) + return ts + + +# ------------------------------- +# Justified true parameters for MAIN synthetic datasets +# ------------------------------- +def true_params_for_dataset(dataset: str) -> Tuple[float, float]: + d = dataset.lower() + is_sir = d.endswith("-sir") + is_si = d.endswith("-si") + if not (is_sir or is_si): + raise ValueError(f"dataset must end with -si or -sir, got: {dataset}") + + # D1: synthetic graphs (BA/ER) + if d.startswith("ba-") or d.startswith("er-"): + pI = 0.1 + pR = 0.1 if is_sir else 0.0 + return pI, pR + + # D2: real graphs (Oregon2/Prost) with synthetic diffusion + if d.startswith("oregon2-") or d.startswith("prost-"): + pI = 0.1 + pR = 0.05 if is_sir else 0.0 + return pI, pR + + raise ValueError( + f"Unsupported dataset for exp3 load-only scatter (needs known true beta/gamma): {dataset}" + ) + + +def save_points_csv(path: str, rows: List[Dict[str, Any]]) -> None: + _ensure_dir(os.path.dirname(path)) + fieldnames = ["dataset", "param", "beta_true", "beta_hat", "T", "obs_time"] + with open(path, "w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=fieldnames) + w.writeheader() + for r in rows: + w.writerow({k: r.get(k, "") for k in fieldnames}) + + +def main() -> None: + p = argparse.ArgumentParser() + + p.add_argument( + "--datasets", + type=str, + default="main", + help="Comma-separated datasets, or 'main' for BA/ER/Oregon2/Prost with SI+SIR.", + ) + p.add_argument("--data_dir", type=str, default="input") + p.add_argument("--device", type=str, default="cuda") + + # Observation times: if empty, use main protocol {floor(T/2), T} + p.add_argument( + "--obs_time", + type=str, + default="", + help='Observation frames used for estimation. "" means {floor(T/2), T}. Example: "5" or "5,10".', + ) + + # b_estim hyperparams + p.add_argument("--b_pI0", type=float, default=0.05) + p.add_argument("--b_pR0", type=float, default=0.05) + p.add_argument("--b_steps", type=int, default=300) + p.add_argument("--b_lr", type=float, default=0.001) + + # output + p.add_argument("--out_dir", type=str, default="output/exp3_beta_scatter_main") + p.add_argument("--csv_name", type=str, default="points.csv") + + # whether to also output pR (gamma) for SIR datasets + p.add_argument( + "--only_betaI", + action="store_true", + help="If set, only output infection beta (pI).", + ) + + args = p.parse_args() + + # device resolve + if args.device.startswith("cuda") and (not torch.cuda.is_available()): + print("[warn] CUDA not available, fallback to CPU.") + device = torch.device("cpu") + else: + device = torch.device(args.device) + + datasets = parse_datasets_arg(args.datasets) + if not datasets: + raise ValueError("Empty --datasets") + + _ensure_dir(args.out_dir) + + # b_estim expects these fields on args; we pass obs_time explicitly anyway. + b_args = argparse.Namespace( + b_pI0=float(args.b_pI0), + b_pR0=float(args.b_pR0), + b_steps=int(args.b_steps), + b_lr=float(args.b_lr), + obs_time="", # not used when obs_time is passed explicitly + device=device, + ) + + rows: List[Dict[str, Any]] = [] + + for dataset in datasets: + data = data_load(dataset, args.data_dir, device) + T = int(data.T.item()) + obs_time = parse_obs_time_arg(args.obs_time, T) + obs_time_str = ",".join(str(t) for t in obs_time) + + pI_true, pR_true = true_params_for_dataset(dataset) + + est = b_estim(data, b_args, obs_time=obs_time) # dict with keys pI, pR + + # always output infection beta (pI) + rows.append( + dict( + dataset=dataset, + param="pI", + beta_true=float(pI_true), + beta_hat=float(est.get("pI", float("nan"))), + T=T, + obs_time=obs_time_str, + ) + ) + + # optionally output recovery beta (pR) for SIR datasets + if (not args.only_betaI) and dataset.lower().endswith("-sir"): + rows.append( + dict( + dataset=dataset, + param="pR", + beta_true=float(pR_true), + beta_hat=float(est.get("pR", float("nan"))), + T=T, + obs_time=obs_time_str, + ) + ) + + print( + "[ok] %s: T=%d, obs={%s}, true(pI=%.3f, pR=%.3f) -> hat(pI=%.4f, pR=%.4f)" + % ( + dataset, + T, + obs_time_str, + pI_true, + pR_true, + float(est.get("pI", 0.0)), + float(est.get("pR", 0.0)), + ) + ) + + out_csv = os.path.join(args.out_dir, args.csv_name) + save_points_csv(out_csv, rows) + print("[ok] saved csv -> %s (rows=%d)" % (out_csv, len(rows))) + + +if __name__ == "__main__": + main() diff --git a/experiments/exp_mcmc_diagnostics.py b/experiments/exp_mcmc_diagnostics.py new file mode 100644 index 0000000..e64f43b --- /dev/null +++ b/experiments/exp_mcmc_diagnostics.py @@ -0,0 +1,212 @@ +import gc +import os +import sys +import os.path as osp +from copy import deepcopy + +import pandas as pd +import matplotlib.pyplot as plt + +ROOT = osp.dirname(osp.dirname(osp.abspath(__file__))) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +from inc.data import * +from inc.test import TEST_METRICS +from hermes import b_estim, q_train, t_mcmc + +DEFAULT_T_STEPS = [25, 50, 100, 200] +DEFAULT_REPS = 4 + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type = str, required = True, help = 'dataset name') + parser.add_argument('--seed', type = int, default = 12345, help = 'base seed used to fit bpar/q_net and derive MCMC seeds') + parser.add_argument('--data_dir', type = str, required = True, help = 'dataset folder') + parser.add_argument('--output_prefix', type = str, required = True, help = 'output prefix for csv/png files') + parser.add_argument('--device', type = torch.device, default = torch_device(), help = 'torch device') + + # same HERMES hyperparameters as hermes.py + parser.add_argument('--b_pI0', type = float, help = 'initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type = float, help = 'initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type = int, help = 'optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type = float, help = 'learning rate in diffusion parameter estimation') + parser.add_argument('--q_steps', type = int, help = 'training steps for the proposal model') + parser.add_argument('--q_lr', type = float, help = 'learning rate for the proposal model') + parser.add_argument('--q_hid', type = int, help = 'hidden size of the proposal model') + parser.add_argument('--q_gnn', type = int, help = 'number of layers of the GNN in the proposal model') + parser.add_argument('--q_mlp', type = int, help = 'number of layers of the MLP in the proposal model') + parser.add_argument('--q_samples', type = int, help = 'sample size to estimate the loss function of the proposal model') + parser.add_argument('--q_zlim', type = int, help = 'a hyperparameter to stablize gradient') + parser.add_argument('--p_coef', type = float, help = 'the coefficient gamma in the initial distribution P[y_0]') + parser.add_argument('--t_samples', type = int, help = 'MCMC sample size') + parser.add_argument('--t_steps', type = int, default = 100, help = 'unused default; sweep values are controlled internally') + parser.add_argument('--t_keep', type = float, help = 'moving average in MCMC') + parser.add_argument('--obs_time', type = str, default = '', help = 'extra observed snapshot times, comma-separated, e.g., 5,7,9') + return parser.parse_args() + + +def parse_obs_time(data, obs_time): + out = [int(t) for t in str(obs_time).split(',') if t] + out.append(data.T.item()) + return sorted(set(out)) + + +@torch.no_grad() +def compose_history(data, tI, tR): + tI = tI.round().long() + tR = tR.round().long() + y_pred = torch.zeros_like(data.y) + y_pred.scatter_(dim = 1, index = torch.minimum(tI, data.T), src = torch.full_like(tI, SIR_STATES.I)) + y_pred.scatter_(dim = 1, index = torch.minimum(tR, data.T), src = torch.full_like(tR, SIR_STATES.R)) + return y_pred[:, : data.T.item()].cummax(dim = 1).values + + +def save_trace_plot(df, y_col, ylabel, title, fpath): + plt.figure() + for t_steps, group in df.groupby('t_steps'): + mean_trace = group.groupby('step')[y_col].mean().reset_index() + plt.plot(mean_trace['step'], mean_trace[y_col], label = f'S={t_steps}') + plt.xlabel('MCMC step') + plt.ylabel(ylabel) + plt.title(title) + plt.legend() + plt.tight_layout() + plt.savefig(fpath, dpi = 200) + plt.close() + + +def save_metric_plot(df, metric, ylabel, title, fpath): + stats = df.groupby('t_steps')[metric].agg(['mean', 'std']).reset_index() + plt.figure() + plt.errorbar(stats['t_steps'], stats['mean'], yerr = stats['std'], marker = 'o') + plt.xlabel('t_steps (S)') + plt.ylabel(ylabel) + plt.title(title) + plt.tight_layout() + plt.savefig(fpath, dpi = 200) + plt.close() + + +def main(): + args = get_args() + out_dir = osp.dirname(args.output_prefix) + if out_dir: + os.makedirs(out_dir, exist_ok = True) + + data = data_load(args.dataset, args.data_dir, args.device) + obs_time = parse_obs_time(data, args.obs_time) + + # Fix parameter estimation + proposal training. + seed_all(args.seed) + bpar = b_estim(data, args, obs_time = obs_time) + print(f'[est] pI={bpar.pI:.4f}, pR={bpar.pR:.4f}', flush = True) + q_net = q_train(data, obs_time, bpar, args) + + trace_rows = [] + metric_rows = [] + mcmc_seeds = [args.seed ^ (rep + 1) for rep in range(DEFAULT_REPS)] + + for t_steps in DEFAULT_T_STEPS: + args_run = deepcopy(args) + args_run.t_steps = t_steps + for rep, mcmc_seed in enumerate(mcmc_seeds): + print(f'[dataset={args.dataset}] [S={t_steps}] [rep={rep}] [seed={mcmc_seed}]', flush = True) + seed_all(mcmc_seed) + tI, tR, diag_rows = t_mcmc( + data, + bpar, + q_net, + args_run, + obs_time = obs_time, + keepdim = True, + diagnostics = True, + ) + + y_pred = compose_history(data, tI, tR) + f1 = TEST_METRICS['f1'](data, y_pred) + nrmse = TEST_METRICS['nrmse'](data, y_pred) + + metric_rows.append(dict( + dataset = args.dataset, + t_steps = t_steps, + rep = rep, + seed = mcmc_seed, + estimated_pI = bpar.pI, + estimated_pR = bpar.pR, + F1 = f1, + NRMSE = nrmse, + )) + + for row in diag_rows: + trace_rows.append(dict( + dataset = args.dataset, + t_steps = t_steps, + rep = rep, + seed = mcmc_seed, + **row, + )) + + gc.collect() + if getattr(args.device, 'type', None) == 'cuda': + torch.cuda.empty_cache() + + traces = pd.DataFrame(trace_rows, columns = [ + 'dataset', 't_steps', 'rep', 'seed', 'step', 'accept_rate', + 'mean_tI', 'mean_tR', 'mean_tI_avg', 'mean_tR_avg', 'mean_lp' + ]) + metrics = pd.DataFrame(metric_rows, columns = [ + 'dataset', 't_steps', 'rep', 'seed', 'estimated_pI', 'estimated_pR', 'F1', 'NRMSE' + ]) + + traces.to_csv(f'{args.output_prefix}_traces.csv', index = False) + metrics.to_csv(f'{args.output_prefix}_metrics.csv', index = False) + + save_trace_plot( + traces, + y_col = 'accept_rate', + ylabel = 'acceptance rate', + title = f'MCMC acceptance trajectory ({args.dataset})', + fpath = f'{args.output_prefix}_accept.png', + ) + save_trace_plot( + traces, + y_col = 'mean_tI_avg', + ylabel = 'mean infection hitting time', + title = f'MCMC infection-time trace ({args.dataset})', + fpath = f'{args.output_prefix}_mean_tI.png', + ) + save_trace_plot( + traces, + y_col = 'mean_tR_avg', + ylabel = 'mean recovery hitting time', + title = f'MCMC recovery-time trace ({args.dataset})', + fpath = f'{args.output_prefix}_mean_tR.png', + ) + save_metric_plot( + metrics, + metric = 'F1', + ylabel = 'F1', + title = f'Final F1 vs MCMC steps ({args.dataset})', + fpath = f'{args.output_prefix}_F1.png', + ) + save_metric_plot( + metrics, + metric = 'NRMSE', + ylabel = 'NRMSE', + title = f'Final NRMSE vs MCMC steps ({args.dataset})', + fpath = f'{args.output_prefix}_NRMSE.png', + ) + + print(f'[saved] {args.output_prefix}_traces.csv', flush = True) + print(f'[saved] {args.output_prefix}_metrics.csv', flush = True) + print(f'[saved] {args.output_prefix}_accept.png', flush = True) + print(f'[saved] {args.output_prefix}_mean_tI.png', flush = True) + print(f'[saved] {args.output_prefix}_mean_tR.png', flush = True) + print(f'[saved] {args.output_prefix}_F1.png', flush = True) + print(f'[saved] {args.output_prefix}_NRMSE.png', flush = True) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/experiments/exp_nI0_sensitivity.py b/experiments/exp_nI0_sensitivity.py new file mode 100644 index 0000000..c98ade4 --- /dev/null +++ b/experiments/exp_nI0_sensitivity.py @@ -0,0 +1,107 @@ +import os +import sys +import os.path as osp + +ROOT = osp.dirname(osp.dirname(osp.abspath(__file__))) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +from inc.data import * +from inc.test import TEST_METRICS +from hermes import run_hermes + +DEFAULT_DATASETS = ['ba-si', 'ba-sir'] +DEFAULT_RATIOS = [0.5, 0.75, 1.0, 1.25, 1.5] + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--datasets', type = str, default = ','.join(DEFAULT_DATASETS), help = 'comma-separated dataset names') + parser.add_argument('--ratios', type = str, default = ','.join(map(str, DEFAULT_RATIOS)), help = 'comma-separated I0 ratios') + parser.add_argument('--seed', type = int, default = 12345, help = 'random seed reused across all runs') + parser.add_argument('--data_dir', type = str, required = True, help = 'dataset folder') + parser.add_argument('--output', type = str, required = True, help = 'output csv file name') + parser.add_argument('--device', type = torch.device, default = torch_device(), help = 'torch device') + + # same HERMES hyperparameters as hermes.py + parser.add_argument('--b_pI0', type = float, help = 'initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type = float, help = 'initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type = int, help = 'optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type = float, help = 'learning rate in diffusion parameter estimation') + parser.add_argument('--q_steps', type = int, help = 'training steps for the proposal model') + parser.add_argument('--q_lr', type = float, help = 'learning rate for the proposal model') + parser.add_argument('--q_hid', type = int, help = 'hidden size of the proposal model') + parser.add_argument('--q_gnn', type = int, help = 'number of layers of the GNN in the proposal model') + parser.add_argument('--q_mlp', type = int, help = 'number of layers of the MLP in the proposal model') + parser.add_argument('--q_samples', type = int, help = 'sample size to estimate the loss function of the proposal model') + parser.add_argument('--q_zlim', type = int, help = 'a hyperparameter to stablize gradient') + parser.add_argument('--p_coef', type = float, help = 'the coefficient gamma in the initial distribution P[y_0]') + parser.add_argument('--t_samples', type = int, help = 'MCMC sample size') + parser.add_argument('--t_steps', type = int, help = 'MCMC steps') + parser.add_argument('--t_keep', type = float, help = 'moving average in MCMC') + parser.add_argument('--obs_time', type = str, default = '', help = 'extra observed snapshot times, comma-separated, e.g., 5,7,9') + return parser.parse_args() + + +def parse_list(raw, cast_fn): + return [cast_fn(x.strip()) for x in raw.split(',') if x.strip()] + + +def ratio_to_I0(true_I0, ratio, n_nodes): + assumed_I0 = int(np.floor(true_I0 * ratio + 0.5)) + return max(0, min(int(n_nodes), assumed_I0)) + + +def run_one(data, args, ratio): + true_I0 = resolve_assumed_I0(data, None) + assumed_I0 = ratio_to_I0(true_I0, ratio, data.num_nodes) + if args.seed is not None: + seed_all(args.seed) + + y_pred, extra = run_hermes(data, args, assumed_I0 = assumed_I0, return_extra = True) + + return Dict( + dataset = data.name, + assumed_I0 = assumed_I0, + ratio = ratio, + estimated_pI = extra.pI, + estimated_pR = extra.pR, + F1 = TEST_METRICS['f1'](data, y_pred), + NRMSE = TEST_METRICS['nrmse'](data, y_pred), + ) + + +def main(): + args = get_args() + datasets = parse_list(args.datasets, str) + ratios = parse_list(args.ratios, float) + + rows = [] + for dataset in datasets: + data = data_load(dataset, args.data_dir, args.device) + for ratio in ratios: + print(f'[dataset={dataset}] [ratio={ratio}]', flush = True) + row = run_one(data, args, ratio) + rows.append(dict(row)) + print( + f"[dataset={dataset}] [ratio={ratio}] assumed_I0={row.assumed_I0} " + f"pI={row.estimated_pI:.4f} pR={row.estimated_pR:.4f} " + f"F1={row.F1:.4f} NRMSE={row.NRMSE:.4f}", + flush = True, + ) + gc.collect() + if getattr(args.device, 'type', None) == 'cuda': + torch.cuda.empty_cache() + + cols = ['dataset', 'assumed_I0', 'ratio', 'estimated_pI', 'estimated_pR', 'F1', 'NRMSE'] + df = pd.DataFrame(rows, columns = cols) + + out_dir = osp.dirname(args.output) + if out_dir: + os.makedirs(out_dir, exist_ok = True) + df.to_csv(args.output, index = False) + print(f'[saved] {args.output}', flush = True) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/experiments/runtime_breakdown_ba.py b/experiments/runtime_breakdown_ba.py new file mode 100644 index 0000000..f4ee0ca --- /dev/null +++ b/experiments/runtime_breakdown_ba.py @@ -0,0 +1,484 @@ +# experiments/runtime_breakdown_ba.py +# -*- coding: utf-8 -*- + +import os +import sys +import gc +import time +import argparse +import traceback +from collections import OrderedDict + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import networkx as nx +import pandas as pd +import torch + +# Make the project root importable when this script is placed under experiments/ +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "")) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +from inc.data import data_simulate +from inc.diffus import b_estim +from inc.utils import seed_all +from hermes import q_train, t_mcmc + + +# ---------------------------- +# helpers +# ---------------------------- +def parse_int_list(text): + if text is None: + return [] + text = str(text).strip() + if text == "": + return [] + return [int(x.strip()) for x in text.split(",") if x.strip()] + + +def resolve_device(device_str): + device_str = (device_str or "auto").strip().lower() + if device_str == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + if device_str.startswith("cuda") and not torch.cuda.is_available(): + print("[warn] CUDA requested but not available; falling back to CPU.", flush=True) + return torch.device("cpu") + return torch.device(device_str) + + +def cleanup(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def sync_device(device): + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def timed_call(fn, device): + sync_device(device) + t0 = time.perf_counter() + out = fn() + sync_device(device) + return out, time.perf_counter() - t0 + + +def format_size_label(n): + n = int(n) + if n % 1000 == 0: + return f"{n // 1000}k" + return f"{n:,}" + + +def parse_obs_time(obs_time_str, T): + obs = parse_int_list(obs_time_str) + obs = [int(t) for t in obs if 0 <= int(t) <= int(T)] + if int(T) not in obs: + obs.append(int(T)) + obs = sorted(set(obs)) + return obs + + +# ---------------------------- +# data + compose +# ---------------------------- +def make_ba_sir_data(num_nodes, args): + """ + Create BA-SIR locally in this experiment script, without touching inc/data.py. + Uses the same BA-SIR setting as the current 1k synthetic BA experiment, + except graph size is user-controlled. + """ + Gnx = nx.barabasi_albert_graph( + n=int(num_nodes), + m=int(args.ba_m), + seed=int(args.seed), + ) + + # BA with m>=1 is connected in practice, but keep this for safety. + if not nx.is_connected(Gnx): + Gnx = Gnx.subgraph(max(nx.connected_components(Gnx), key=len)).copy() + + meta = { + "graph_size_actual": int(Gnx.number_of_nodes()), + "num_edges": int(Gnx.number_of_edges()), + } + + params = dict( + fraction_infected=float(args.fraction_infected), + beta=float(args.sim_pI), + gamma=float(args.sim_pR), + ) + + data = data_simulate( + Gnx=Gnx, + seed=int(args.seed), + T=int(args.T), + diffus="sir", + params=params, + ).to(args.device) + + data.name = f"ba-sir-n{meta['graph_size_actual']}" + return data, meta + + +@torch.no_grad() +def compose_history(data, tI, tR): + """ + Same compose logic as run_hermes(), timed separately. + Output shape follows the current project convention: (nodes, T) + and the final observed snapshot is not duplicated here. + """ + tI = tI.round().long() + tR = tR.round().long() + + y_pred = torch.zeros_like(data.y) # (nodes, T+1) + y_pred.scatter_( + dim=1, + index=torch.minimum(tI, data.T), + src=torch.full_like(tI, 1), + ) + y_pred.scatter_( + dim=1, + index=torch.minimum(tR, data.T), + src=torch.full_like(tR, 2), + ) + y_pred = y_pred[:, : data.T.item()].cummax(dim=1).values + return y_pred + + +# ---------------------------- +# one run +# ---------------------------- +def init_row(requested_graph_size, run_order, obs_time): + row = OrderedDict() + row["run_order"] = int(run_order) + row["graph_size_requested"] = int(requested_graph_size) + row["graph_size_actual"] = None + row["num_edges"] = None + row["T"] = None + row["obs_time"] = ",".join(str(t) for t in obs_time) + row["status"] = "pending" + row["fail_stage"] = "" + row["error"] = "" + + row["graph_build_sec"] = None + row["beta_est_sec"] = None + row["proposal_train_sec"] = None + row["mcmc_sec"] = None + row["compose_sec"] = None + row["total_algo_sec"] = None + row["total_wall_sec"] = None + + row["est_pI"] = None + row["est_pR"] = None + return row + + +def run_one_size(requested_graph_size, run_order, args): + obs_time = parse_obs_time(args.obs_time, args.T) + row = init_row( + requested_graph_size=requested_graph_size, + run_order=run_order, + obs_time=obs_time, + ) + + data = None + meta = None + bpar = None + q_net = None + tI = None + tR = None + y_pred = None + stage = "start" + + try: + seed_all(int(args.seed)) + cleanup() + + stage = "graph_build" + (data, meta), graph_build_sec = timed_call( + lambda: make_ba_sir_data(requested_graph_size, args), + args.device, + ) + row["graph_size_actual"] = int(meta["graph_size_actual"]) + row["num_edges"] = int(meta["num_edges"]) + row["T"] = int(data.T.item()) + row["graph_build_sec"] = float(graph_build_sec) + + stage = "beta_est" + bpar, beta_est_sec = timed_call( + lambda: b_estim( + data=data, + args=args, + obs_time=obs_time, + assumed_I0=args.assumed_I0, + ), + args.device, + ) + row["beta_est_sec"] = float(beta_est_sec) + row["est_pI"] = float(bpar.pI) + row["est_pR"] = float(bpar.pR) + + stage = "proposal_train" + q_net, proposal_train_sec = timed_call( + lambda: q_train( + data=data, + obs_time=obs_time, + bpar=bpar, + args=args, + assumed_I0=args.assumed_I0, + ), + args.device, + ) + row["proposal_train_sec"] = float(proposal_train_sec) + + stage = "mcmc" + (tI, tR), mcmc_sec = timed_call( + lambda: t_mcmc( + data=data, + bpar=bpar, + q_net=q_net, + args=args, + obs_time=obs_time, + keepdim=True, + assumed_I0=args.assumed_I0, + diagnostics=False, + ), + args.device, + ) + row["mcmc_sec"] = float(mcmc_sec) + + stage = "compose" + y_pred, compose_sec = timed_call( + lambda: compose_history(data, tI, tR), + args.device, + ) + row["compose_sec"] = float(compose_sec) + + row["total_algo_sec"] = float( + row["beta_est_sec"] + + row["proposal_train_sec"] + + row["mcmc_sec"] + + row["compose_sec"] + ) + row["total_wall_sec"] = float( + row["graph_build_sec"] + row["total_algo_sec"] + ) + row["status"] = "ok" + + except Exception as exc: + row["status"] = "failed" + row["fail_stage"] = stage + row["error"] = f"{type(exc).__name__}: {exc}" + print(f"[failed] n={requested_graph_size}, stage={stage}, error={row['error']}", flush=True) + traceback.print_exc() + + finally: + del data, meta, bpar, q_net, tI, tR, y_pred + cleanup() + + return row + + +# ---------------------------- +# save / plot +# ---------------------------- +def save_rows(rows, csv_path): + df = pd.DataFrame(rows) + df.to_csv(csv_path, index=False) + + +def make_plot(rows, fig_path, q_steps): + df = pd.DataFrame(rows) + df = df[df["status"] == "ok"].copy() + if df.empty: + print(f"[warn] no successful runs; skip figure: {fig_path}", flush=True) + return + + df = df.sort_values("run_order") + + stage_cols = [ + ("beta_est_sec", "beta estimation"), + ("proposal_train_sec", "proposal training"), + ("mcmc_sec", "MCMC"), + ("compose_sec", "compose"), + ] + + x = list(range(len(df))) + xticklabels = [format_size_label(n) for n in df["graph_size_actual"].tolist()] + bottoms = [0.0] * len(df) + + fig, ax = plt.subplots(figsize=(8, 5)) + + for col, label in stage_cols: + vals = df[col].fillna(0.0).astype(float).tolist() + ax.bar(x, vals, bottom=bottoms, label=label) + bottoms = [b + v for b, v in zip(bottoms, vals)] + + totals = df["total_algo_sec"].fillna(0.0).astype(float).tolist() + for i, total in enumerate(totals): + ax.text( + i, + total, + f"{total:.1f}s", + ha="center", + va="bottom", + fontsize=9, + ) + + ax.set_xticks(x) + ax.set_xticklabels(xticklabels) + ax.set_xlabel("BA graph size n") + ax.set_ylabel("running time (sec)") + ax.set_title(f"HERMES runtime breakdown on BA-SIR (q_steps={q_steps})") + ax.legend(frameon=False) + plt.tight_layout() + plt.savefig(fig_path, dpi=200, bbox_inches="tight") + plt.close(fig) + + +# ---------------------------- +# CLI +# ---------------------------- +def get_args(): + parser = argparse.ArgumentParser( + description="Runtime breakdown on BA-SIR for 50k -> 30k -> 1k in one run.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + # run order: do NOT sort, preserve the exact user input order + parser.add_argument( + "--graph_sizes", + type=str, + default="50000,30000,1000", + help="Comma-separated BA sizes; order is preserved exactly.", + ) + parser.add_argument("--seed", type=int, default=123456789) + parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--output_dir", type=str, default="output/runtime_breakdown_ba") + + # BA-SIR generation + parser.add_argument("--T", type=int, default=10) + parser.add_argument( + "--obs_time", + type=str, + default="5", + help="Extra observed times excluding T; T is appended automatically.", + ) + parser.add_argument("--ba_m", type=int, default=4) + parser.add_argument("--fraction_infected", type=float, default=0.05) + parser.add_argument("--sim_pI", type=float, default=0.1, help="Ground-truth infection rate for simulation.") + parser.add_argument("--sim_pR", type=float, default=0.1, help="Ground-truth recovery rate for simulation.") + + # diffusion parameter estimation + parser.add_argument("--b_pI0", type=float, default=0.001) + parser.add_argument("--b_pR0", type=float, default=0.001) + parser.add_argument("--b_steps", type=int, default=500) + parser.add_argument("--b_lr", type=float, default=0.003) + + # proposal training + parser.add_argument("--q_steps", type=int, default=200) # user-requested change + parser.add_argument("--q_lr", type=float, default=0.003) + parser.add_argument("--q_hid", type=int, default=16) + parser.add_argument("--q_gnn", type=int, default=3) + parser.add_argument("--q_mlp", type=int, default=2) + parser.add_argument("--q_samples", type=int, default=10) + parser.add_argument("--q_zlim", type=int, default=16) + + # MCMC + parser.add_argument("--p_coef", type=float, default=1.0) + parser.add_argument("--t_samples", type=int, default=100) + parser.add_argument("--t_steps", type=int, default=100) + parser.add_argument("--t_keep", type=float, default=0.5) + + # optional initial infected prior + parser.add_argument( + "--assumed_I0", + type=int, + default=None, + help="If None, use the current code behavior (read I0 from data.y[:,0]).", + ) + + args = parser.parse_args() + args.device = resolve_device(args.device) + return args + + +# ---------------------------- +# main +# ---------------------------- +def main(): + args = get_args() + graph_sizes = parse_int_list(args.graph_sizes) + if len(graph_sizes) == 0: + raise ValueError("--graph_sizes is empty.") + + os.makedirs(args.output_dir, exist_ok=True) + + csv_path = os.path.join(args.output_dir, "runtime_breakdown_ba.csv") + fig_path = os.path.join(args.output_dir, "runtime_breakdown_ba.png") + + print("=" * 80, flush=True) + print("BA-SIR runtime breakdown experiment", flush=True) + print(f"device : {args.device}", flush=True) + print(f"graph_sizes : {graph_sizes}", flush=True) + print(f"T : {args.T}", flush=True) + print(f"obs_time : {parse_obs_time(args.obs_time, args.T)}", flush=True) + print(f"q_steps : {args.q_steps}", flush=True) + print("=" * 80, flush=True) + + rows = [] + total_runs = len(graph_sizes) + + for run_order, requested_graph_size in enumerate(graph_sizes, start=1): + print( + f"\n=== [{run_order}/{total_runs}] running BA-SIR with n={requested_graph_size} ===", + flush=True, + ) + row = run_one_size( + requested_graph_size=requested_graph_size, + run_order=run_order, + args=args, + ) + rows.append(row) + save_rows(rows, csv_path) + + if row["status"] == "ok": + print( + "[ok] " + f"n={row['graph_size_actual']:,}, " + f"m={row['num_edges']:,}, " + f"build={row['graph_build_sec']:.2f}s, " + f"beta={row['beta_est_sec']:.2f}s, " + f"q_train={row['proposal_train_sec']:.2f}s, " + f"mcmc={row['mcmc_sec']:.2f}s, " + f"compose={row['compose_sec']:.2f}s, " + f"total_algo={row['total_algo_sec']:.2f}s, " + f"total_wall={row['total_wall_sec']:.2f}s, " + f"est_pI={row['est_pI']:.4f}, " + f"est_pR={row['est_pR']:.4f}", + flush=True, + ) + else: + print( + "[failed] " + f"n={requested_graph_size:,}, " + f"stage={row['fail_stage']}, " + f"error={row['error']}", + flush=True, + ) + + make_plot(rows, fig_path, q_steps=args.q_steps) + + print("\nSaved:") + print(f" CSV : {csv_path}", flush=True) + print(f" FIG : {fig_path}", flush=True) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/runtime_breakdown_ba_sparse.py b/experiments/runtime_breakdown_ba_sparse.py new file mode 100644 index 0000000..cf1286d --- /dev/null +++ b/experiments/runtime_breakdown_ba_sparse.py @@ -0,0 +1,160 @@ +# experiments/runtime_breakdown_ba_sparse.py +# -*- coding: utf-8 -*- + +import os +import sys +import torch + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.abspath(os.path.join(THIS_DIR, "..")) + +if THIS_DIR not in sys.path: + sys.path.insert(0, THIS_DIR) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +import runtime_breakdown_ba as base +from hermes import QNet +from inc.diffus import SIR_STATES + + +def apply_sparse_rhs_patch(): + @torch.no_grad() + def _lik_step_sparse_rhs(self, y0, y1, lI1, lI0, lR1, lR0, yL=None, reach=None): + n_nodes, n_samples = y0.shape + lik = self.zero + + # R -> I part (unchanged) + msk = (y1 == SIR_STATES.R) + if yL is not None: + msk = msk & (yL != SIR_STATES.R) & reach + lik = lik + torch.where( + msk, + torch.where(y0 != SIR_STATES.R, lR1, lR0), + self.zero, + ) + + # I -> S part + uid = lI1.argsort(dim=0, descending=True) # (nodes, samples) + msk = (y1 == SIR_STATES.I) | (msk & (y0 != SIR_STATES.R)) + + rem = torch.where( + msk, + (reach.long() if reach is not None else 1) + + torch.sparse.mm(self.adj, msk.float()).long(), + self.n_inf, + ) # (nodes, samples) + + cols = torch.arange(n_samples, dtype=torch.long, device=rem.device) + + for i, u in enumerate(uid): + rem_v = torch.full( + (self.n_nodes, n_samples), + self.n_inf, + dtype=rem.dtype, + device=rem.device, + ) + rem_v = rem_v.scatter_reduce( + dim=0, + index=self.eidx[0, :, None].expand(-1, n_samples), + src=rem[self.eidx[1]], + reduce="amin", + include_self=True, + ) + rem_v = rem_v.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + rem_u = rem.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + + opt = (rem_u > 1) & (rem_v > 1) + if yL is not None: + opt = opt & ( + yL.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) != SIR_STATES.I + ) + + msk_u = msk.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + msk_opt = msk_u & opt + + lik = lik + torch.where( + msk_opt, + torch.where(y0 == SIR_STATES.S, lI1, lI0), + self.zero, + ) + + trs = ( + y0.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) + == SIR_STATES.S + ) # (samples,) + + # ========================================================== + # OLD: + # rem = rem - (torch.sparse.mm( + # self.adj, + # torch.zeros(rem.size(), ...).scatter(...) + # ) > 0).to(rem.dtype) + # + # NEW: + # build RHS directly as sparse COO and use sparse indices. + # ========================================================== + active_cols = torch.nonzero(msk_u & trs, as_tuple=False).flatten() + if active_cols.numel() > 0: + rhs_row = u.index_select(0, active_cols) # row index varies by sample + rhs_col = active_cols + rhs_idx = torch.stack([rhs_row, rhs_col], dim=0) + rhs_val = torch.ones( + active_cols.numel(), + dtype=self.adj.dtype, + device=rem.device, + ) + + rhs = torch.sparse_coo_tensor( + indices=rhs_idx, + values=rhs_val, + size=rem.size(), # (nodes, samples) + dtype=self.adj.dtype, + device=rem.device, + ).coalesce() + + nbr_hit = torch.sparse.mm(self.adj, rhs) + if nbr_hit.layout != torch.sparse_coo: + nbr_hit = nbr_hit.to_sparse_coo() + nbr_hit = nbr_hit.coalesce() + + # IMPORTANT: + # sparse tensor usually should not continue with `> 0` here. + # Use sparse indices directly. + if nbr_hit._nnz() > 0: + hit_idx = nbr_hit.indices() + rem.index_put_( + (hit_idx[0], hit_idx[1]), + -torch.ones( + hit_idx.size(1), + dtype=rem.dtype, + device=rem.device, + ), + accumulate=True, + ) + + # self-row update (unchanged logic; in-place for less allocation) + rem.index_put_( + (u, cols), + torch.where( + msk_u, + torch.where(trs, rem_u - 1, self.n_inf), + rem.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0), + ), + accumulate=False, + ) + msk.index_put_((u, cols), msk_opt, accumulate=False) + + return lik + + QNet._lik_step = _lik_step_sparse_rhs + print("[patch] QNet._lik_step -> sparse RHS version", flush=True) + + +def main(): + apply_sparse_rhs_patch() + base.main() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/gcn_ms.py b/gcn_ms.py new file mode 100644 index 0000000..b1fa075 --- /dev/null +++ b/gcn_ms.py @@ -0,0 +1,178 @@ +from inc.diffus import * +from inc.test import * +import argparse + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type=str, help='dataset name') + parser.add_argument('--seed', type=int, help='random seed') + parser.add_argument('--data_dir', type=str, help='dataset folder') + parser.add_argument('--output', type=str, help='output file name') + parser.add_argument('--device', type=torch.device, help='torch device') + + # Diffusion parameter estimation (same as baseline) + parser.add_argument('--b_pI0', type=float, help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, help='learning rate in diffusion parameter estimation') + + # GCN hyperparameters (same as baseline) + parser.add_argument('--lr', type=float, help='learning rate for GCN') + parser.add_argument('--epochs', type=int, help='training epochs for GCN') + parser.add_argument('--batch_size', type=int, help='batch size when training GCN') + parser.add_argument('--units', type=int, help='hidden size of GCN') + parser.add_argument('--layers', type=int, help='number of layers in GCN') + parser.add_argument('--dropout', type=float, help='dropout rate in GCN') + + # === New: multi-snapshot controls (align with inc/ditto_ms.py). === + parser.add_argument( + '--obs_ts', type=str, default=None, + help='comma-separated observed time indices, e.g. "0,3,5"; ' + 'None means single-snapshot (use only the final snapshot)' + ) + parser.add_argument( + '--obs_k', type=int, default=None, + help='number of observed snapshots; if None, set to len(obs_ts) when obs_ts is given, ' + 'otherwise 1 (single-snapshot)' + ) + + args = parser.parse_args() + + # Normalize obs_ts / obs_k exactly like inc/ditto_ms.py does. :contentReference[oaicite:8]{index=8} + if args.obs_ts is not None and len(args.obs_ts.strip()) > 0: + obs = [int(x) for x in args.obs_ts.split(',') if x.strip() != ''] + obs = sorted(set(obs)) + args.obs_ts = obs + if args.obs_k is None: + args.obs_k = len(obs) + else: + args.obs_ts = None + if args.obs_k is None: + args.obs_k = 1 + + return args + + +def _select_obs_times(T: int, obs_ts: list | None) -> list: + """ + Choose which snapshots to treat as 'observed' inputs. + - If obs_ts is provided, we use it (clamped to [0, T]). + - Otherwise, we default to the final snapshot only (t = T). + """ + if obs_ts is None or len(obs_ts) == 0: + return [T] + times = [] + for t in obs_ts: + if t < 0: + t = 0 + if t > T: + t = T + times.append(t) + times = sorted(set(times)) + if len(times) == 0: + times = [T] + return times + + +def _build_input_from_labels(labels: torch.Tensor, obs_times: list) -> torch.Tensor: + """ + labels: (batch, nodes, T+1), dtype=long + returns x: (batch*nodes, K) where K = len(obs_times) + """ + # Stack observed snapshots as feature channels. + xs = [labels[:, :, t] for t in obs_times] # each: (batch, nodes) + x = torch.stack(xs, dim=2) # (batch, nodes, K) + x = x.reshape(-1, x.size(2)) # (batch*nodes, K) + return x + + +def _build_input_from_data_y(data, obs_times: list) -> torch.Tensor: + """ + data.y: (nodes, T+1), dtype=long + returns x: (nodes, K) where K = len(obs_times) + """ + xs = [data.y[:, t] for t in obs_times] # each: (nodes,) + x = torch.stack(xs, dim=1) # (nodes, K) + return x + + +def gcn_run(data): + # Estimate diffusion parameters as in the baseline GCN. :contentReference[oaicite:9]{index=9} + bpar = b_estim(data, args) + + # Shapes / counts + T = data.T.item() + n_nodes = data.num_nodes + n_cls = data.y.max().item() + 1 + + # Observed time indices used as inputs (multi-snapshot). :contentReference[oaicite:10]{index=10} + obs_times = _select_obs_times(T, args.obs_ts) + k_in = len(obs_times) + + # Model: input dim = #observed snapshots, output dim = T * n_cls (class per time). :contentReference[oaicite:11]{index=11} + model = gnn.GCN(k_in, args.units, args.layers, T * n_cls, args.dropout) + model = model.to(args.device) + + # ---------------------------- + # Train (self-supervised on synthetic data from bpar). + # ---------------------------- + model.train() + I0 = (data.y[:, 0] == SIR_STATES.I).long().sum().item() + opt = optim.Adam(model.parameters(), lr=args.lr) + pbar = trange(1, args.epochs + 1) + for epoch in pbar: + opt.zero_grad() + + # Simulate histories with estimated parameters; arrange as (batch, nodes, T+1). + labels = diffus_gen( + T=data.T.item(), n_nodes=data.num_nodes, edge_index=data.edge_index, + I0=I0, n_samples=args.batch_size, pI=bpar.pI, pR=bpar.pR + ).transpose(0, 2) # (batch, nodes, T + 1) + + # Build per-node, multi-snapshot inputs from the chosen observed times. + x = _build_input_from_labels(labels, obs_times) # (batch*nodes, K) + x = x.float() + + # Replicate edges across batch. (2, E*batch) + edge_index = ( + data.edge_index.unsqueeze(dim=2) + + n_nodes * torch.arange(args.batch_size, dtype=torch.long, device=x.device) + ).flatten(start_dim=1) + + # Predict all previous T states for each node (0..T-1). + logits = F.log_softmax(model(x, edge_index).view(-1, n_cls), dim=-1) + target = labels[:, :, :T].flatten() # (batch*nodes*T,) + loss = F.nll_loss(logits, target) + + pbar.set_description(f'epoch={epoch} loss={loss.item():.4f}') + + # Standard backward/update. + loss.backward() + opt.step() + + # ---------------------------- + # Inference (condition on the provided observed snapshots in data.y). + # ---------------------------- + with torch.no_grad(): + model.eval() + + # Prepare inference inputs from observed times. + x_inf = _build_input_from_data_y(data, obs_times).float() # (nodes, K) + y_logits = model(x_inf, data.edge_index).view(n_nodes, T, n_cls) # (nodes, T, n_cls) + y_pred = y_logits.argmax(dim=2).contiguous() # (nodes, T), dtype=long + + # Optional: enforce hard consistency on any observed snapshots that lie inside [0, T-1]. + # (If an observed snapshot includes the final T, it is *input* only; we do not predict y_T.) + for t in obs_times: + if 0 <= t < T: + y_pred[:, t] = data.y[:, t] + + return y_pred.clone() + + +args = get_args() +seed_all(args.seed) +tester = Tester(args.data_dir, args.device, gcn_run) +tester.test([args.dataset], rep=1) +tester.save(args.output) diff --git a/gin_ms.py b/gin_ms.py new file mode 100644 index 0000000..9ab56d4 --- /dev/null +++ b/gin_ms.py @@ -0,0 +1,156 @@ +from inc.diffus import * # diffus_gen, SIR_STATES, b_estim, etc. +from inc.test import * # Tester +import argparse +import torch +import torch.nn.functional as F + + +def get_args(): + parser = argparse.ArgumentParser() + # dataset & runtime + parser.add_argument('--dataset', type=str, help='dataset name') + parser.add_argument('--seed', type=int, help='random seed') + parser.add_argument('--data_dir', type=str, help='dataset folder') + parser.add_argument('--output', type=str, help='output file name') + parser.add_argument('--device', type=torch.device, help='torch device') + + # diffusion parameter estimation (used to synthesize training labels) + parser.add_argument('--b_pI0', type=float, help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, help='learning rate in diffusion parameter estimation') + parser.add_argument('--b_pImax', type=float, default=1.0, help='upper bound to clamp pI during estimation (safety)') + + # GIN model & training + parser.add_argument('--lr', type=float, help='learning rate for GIN') + parser.add_argument('--epochs', type=int, help='training epochs for GIN') + parser.add_argument('--batch_size', type=int, help='batch size when training GIN') + parser.add_argument('--units', type=int, help='hidden size of GIN') + parser.add_argument('--layers', type=int, help='number of layers in GIN') + parser.add_argument('--dropout', type=float, help='dropout rate in GIN') + + # multi-snapshot settings (align with ditto_ms.py) :contentReference[oaicite:4]{index=4} + parser.add_argument( + '--obs_ts', type=str, default=None, + help='comma-separated observed time indices, e.g. "0,3,5"; ' + 'None means use obs_k snapshots ending at T (default: final-only)' + ) + parser.add_argument( + '--obs_k', type=int, default=None, + help='number of observed snapshots; if None, set to len(obs_ts) when obs_ts is given, ' + 'otherwise 1 (single-snapshot)' + ) + args = parser.parse_args() + + # normalize obs_ts / obs_k following ditto_ms.py conventions :contentReference[oaicite:5]{index=5} + if args.obs_ts is not None and len(args.obs_ts.strip()) > 0: + obs = [int(x) for x in args.obs_ts.split(',') if x.strip() != ''] + obs = sorted(set(obs)) + args.obs_ts = obs + if args.obs_k is None: + args.obs_k = len(obs) + else: + args.obs_ts = None + if args.obs_k is None: + args.obs_k = 1 + return args + + +def _resolve_obs_ts(args, T: int): + """ + Resolve the list of observed snapshot indices for a given horizon T. + Rules: + - If args.obs_ts is provided: use it directly (values in [0, T] allowed; -1 means T). + - Else: use the last args.obs_k snapshots ending at T (inclusive), + i.e., max(0, T - obs_k + 1) ... T. + """ + if args.obs_ts is not None: + ts = [(T if (t == -1) else int(t)) for t in args.obs_ts] + ts = [min(max(0, t), T) for t in ts] # clamp into [0, T] + ts = sorted(set(ts)) + return ts + # default: use last k snapshots ending at T + k = int(max(1, min(args.obs_k, T + 1))) + start = max(0, T - k + 1) + return list(range(start, T + 1)) + + +def gin_run(data): + # Estimate diffusion parameters for synthetic training labels + bpar = b_estim(data, args) + + # Problem sizes + T = int(data.T.item()) + n_nodes = int(data.num_nodes) + n_cls = int(data.y.max().item() + 1) + + # Resolve observed snapshot indices for training/inference + obs_ts = _resolve_obs_ts(args, T) # indices in [0, T], inclusive + k_in = int(len(obs_ts)) # number of observed snapshots (channels) + + # Build model: GIN maps k_in-channel node features to T * n_cls logits per node + model = gnn.GIN(k_in, args.units, args.layers, T * n_cls, args.dropout) # :contentReference[oaicite:6]{index=6} + model = model.to(args.device) + + # ------------------------- + # Train on simulated labels + # ------------------------- + model.train() + I0 = int((data.y[:, 0] == SIR_STATES.I).long().sum().item()) + opt = optim.Adam(model.parameters(), lr=args.lr) + pbar = trange(1, args.epochs + 1) + + for epoch in pbar: + opt.zero_grad() + + # Simulate training batch: (T+1, nodes, batch) -> transpose -> (batch, nodes, T+1) + labels = diffus_gen( + T=T, n_nodes=n_nodes, edge_index=data.edge_index, + I0=I0, n_samples=args.batch_size, pI=bpar.pI, pR=bpar.pR + ).transpose(0, 2) # (batch, nodes, T+1) + + # Collect multi-snapshot inputs at obs_ts -> x: (batch * nodes, k_in) + # Note: obs_ts indices are in [0, T] inclusive; labels' last dim matches that. + x = labels[:, :, obs_ts] # (batch, nodes, k_in) + x = x.reshape((-1, k_in)) # (batch*nodes, k_in) + + # Batch-edge indexing: replicate the graph 'batch' times with node-ID offsets + edge_index = ( + data.edge_index.unsqueeze(dim=2) + + n_nodes * torch.arange(args.batch_size, dtype=torch.long, device=x.device) + ).flatten(start_dim=1) # (2, batch * n_edges) + + # Forward & loss: predict states for times 0..T-1 (exclude the final observed time T) + logits = model(x.float(), edge_index).view(-1, n_cls) # (batch*nodes*T, n_cls) after view below + logits = F.log_softmax(logits, dim=-1) + + target = labels[:, :, :T].flatten() # (batch*nodes*T,) + loss = F.nll_loss(logits, target) + + loss.backward() + opt.step() + + pbar.set_description(f'epoch={epoch} loss={loss.item():.4f}') + + # ------------- + # Inference + # ------------- + with torch.no_grad(): + model.eval() + + # Build input features from the observed snapshots of the test instance + # data.y: (nodes, T+1) + obs = data.y[:, obs_ts] # (nodes, k_in) + y_pred = model(obs.float(), data.edge_index) \ + .view(n_nodes, T, n_cls) \ + .argmax(dim=2) # (nodes, T), states for times [0..T-1] + + return y_pred.clone() + + +if __name__ == '__main__': + args = get_args() + seed_all(args.seed) + tester = Tester(args.data_dir, args.device, gin_run) + tester.test([args.dataset], rep=1) + tester.save(args.output) diff --git a/grin.py b/grin.py new file mode 100644 index 0000000..ed90ba4 --- /dev/null +++ b/grin.py @@ -0,0 +1,219 @@ +# -*- coding: utf-8 -*- +# NOTE: This file was extracted from `grin.ipynb`. +# - Spatiotemporal 0.1.1: https://github.com/TorchSpatiotemporal/tsl/tree/1ae3289e00b28d0e84dfd54799561162df1917cd +# - SPIN: https://github.com/Graph-Machine-Learning-Group/spin + +from __future__ import annotations + +import argparse + +import torch +from tsl.nn.models.stgn import GRINModel + +from inc.diffus import * +from inc.test import * + + +def get_args() -> argparse.Namespace: + """Parse command line arguments.""" + + parser = argparse.ArgumentParser(description='GRIN baseline (single-snapshot)') + + # ---- standard experiment args (align with other runners, e.g., hermes.py) ---- + parser.add_argument('--dataset', type=str, required=True, help='dataset name') + parser.add_argument('--seed', type=int, default=123456789, help='random seed') + parser.add_argument('--data_dir', type=str, default='input', help='dataset folder') + parser.add_argument('--output', type=str, default='output/grin.pt', help='output file name') + parser.add_argument( + '--device', + type=str, + default='cuda' if torch.cuda.is_available() else 'cpu', + help='torch device, e.g., cpu | cuda | cuda:0', + ) + + # ---- diffusion parameter estimation (b_*) ---- + parser.add_argument( + '--b_pI0', + type=float, + default=1e-3, + help='initial infection rate in diffusion parameter estimation', + ) + parser.add_argument( + '--b_pR0', + type=float, + default=1e-3, + help='initial recovery rate in diffusion parameter estimation', + ) + parser.add_argument( + '--b_steps', + type=int, + default=500, + help='optimization steps in diffusion parameter estimation', + ) + parser.add_argument( + '--b_lr', + type=float, + default=3e-3, + help='learning rate in diffusion parameter estimation', + ) + + # ---- GRINModel hyper-parameters ---- + # Ref: https://github.com/Graph-Machine-Learning-Group/spin/blob/main/config/imputation/grin.yaml + parser.add_argument('--hidden_size', type=int, default=64) + parser.add_argument('--ff_size', type=int, default=64) + parser.add_argument('--embedding_size', type=int, default=8) + parser.add_argument('--n_layers', type=int, default=1) + parser.add_argument('--kernel_size', type=int, default=2) + parser.add_argument('--decoder_order', type=int, default=1) + parser.add_argument( + '--layer_norm', + action='store_true', + help='enable layer norm inside the GRIN model', + ) + parser.add_argument('--dropout', type=float, default=0.0) + parser.add_argument('--ff_dropout', type=float, default=0.0) + parser.add_argument('--merge_mode', type=str, default='mlp') + + # ---- optimizer / training ---- + parser.add_argument('--lr', type=float, default=1e-3, help='Adam learning rate') + parser.add_argument('--l2_reg', type=float, default=0.0, help='Adam weight decay') + parser.add_argument('--epochs', type=int, default=300) + parser.add_argument('--batch_size', type=int, default=1) + + # ---- evaluation ---- + parser.add_argument('--rep', type=int, default=1, help='repetitions per dataset') + + args = parser.parse_args() + + # Normalize device to torch.device (keeps compatibility with other modules). + args.device = torch.device(args.device) + + return args + + +def grin_prep(y: torch.Tensor, edge_index: torch.Tensor, device: torch.device): + """Prepare inputs for GRIN. + + Args: + y: (samples, nodes, T+1) integer states. + edge_index: (2, edges) + device: torch device for mask tensor + + Returns: + x: (samples, T+1, nodes, 1) + mask: (samples, T+1, nodes, 1) with only last snapshot observed + ei: edge_index (kept for API symmetry) + """ + + n_samples, n_nodes, T1 = y.size() + T = T1 - 1 + + # (samples, T+1, nodes, 1) + x = y.transpose(1, 2).unsqueeze(dim=3).float() + + # only last snapshot observed + mask = ( + torch.tensor([[0]] * T + [[1]], device=device) + .expand(n_samples, n_nodes, -1, 1) + .transpose(1, 2) + ) + + ei = edge_index + return x, mask, ei + + +def grin_run(data, args: argparse.Namespace): + """Train GRIN on synthetic histories generated from estimated diffusion params, + then impute the unobserved history for `data`. + + This is a minimal refactor of the original notebook logic. + """ + + bpar = b_estim(data, args) # Dict(pI=..., pR=...) + + T = data.T.item() + n_nodes = data.num_nodes + n_out = data.y[:, -1].max().item() + 1 + + # train + model = GRINModel( + input_size=1, + hidden_size=args.hidden_size, + ff_size=args.ff_size, + embedding_size=args.embedding_size, + n_layers=args.n_layers, + n_nodes=n_nodes, + kernel_size=args.kernel_size, + decoder_order=args.decoder_order, + layer_norm=args.layer_norm, + dropout=args.dropout, + ff_dropout=args.ff_dropout, + merge_mode=args.merge_mode, + ).to(args.device) + + I0 = (data.y[:, 0] == SIR_STATES.I).long().sum().item() + + opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2_reg) + model.train() + + pbar = trange(1, args.epochs + 1) + for epoch in pbar: + opt.zero_grad() + + # generate synthetic histories for training + Y_true = diffus_gen( + T=data.T.item(), + n_nodes=data.num_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ) # (T+1, nodes, samples) + + # GRIN expects (samples, T+1, nodes, features) + x, mask, _ = grin_prep(Y_true.transpose(0, 2), data.edge_index, device=args.device) + + z = model(x=x, mask=mask, edge_index=data.edge_index)[0] # (samples, T+1, nodes, 1) + + # L1 loss on the unobserved part (t < T) + loss = ( + z[:, :-1].flatten() + - Y_true[:-1].transpose(1, 2).transpose(0, 1).flatten() + ).abs().mean() + + pbar.set_description(f'[epoch={epoch}] loss={loss.item():.4f}') + loss.backward() + opt.step() + + # infer + with torch.no_grad(): + model.eval() + x, mask, ei = grin_prep(data.y.unsqueeze(dim=0).clone(), data.edge_index, device=args.device) + z = model(x=x, mask=mask, edge_index=ei)[0] # (1, T+1, nodes, 1) + + y_pred = data.y.clone() + y_pred[:, :-1] = ( + z[0, :-1, :, 0] + .clamp(0, n_out - 1) + .T + .round() + .long() + ) # (nodes, T) + + return y_pred.clone() + + +def main() -> None: + args = get_args() + + # Tester expects a callable: model_fn(data) -> y_pred + model_fn = lambda data: grin_run(data, args) + + tester = Tester(args.data_dir, args.device, model_fn) + tester.test([args.dataset], seed=args.seed, rep=args.rep) + tester.save(args.output) + + +if __name__ == '__main__': + main() diff --git a/grin_ms.py b/grin_ms.py new file mode 100644 index 0000000..afc7415 --- /dev/null +++ b/grin_ms.py @@ -0,0 +1,270 @@ + +from __future__ import annotations + +import argparse +import os +from typing import List + +import torch +from tqdm import trange +from tsl.nn.models.stgn import GRINModel + +# Project utilities (same style as other runners) +from inc.diffus import SIR_STATES, b_estim, diffus_gen, seed_all + +from inc.test import Tester + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="GRIN baseline (multi-snapshot)") + + # ---- standard experiment args (align with hermes.py style) ---- + parser.add_argument("--dataset", type=str, required=True, help="dataset name") + parser.add_argument("--seed", type=int, default=123456789, help="random seed") + parser.add_argument("--data_dir", type=str, default="input", help="dataset folder") + parser.add_argument("--output", type=str, default="output/grin.pt", help="output file name") + parser.add_argument( + "--device", + type=str, + default="cuda" if torch.cuda.is_available() else "cpu", + help="torch device, e.g., cpu | cuda | cuda:0", + ) + + # ---- multi-snapshot control ---- + # Keep hermes-style naming (obs_time) and also provide snapshot as an alias. + parser.add_argument( + "--obs_time", + "--snapshot", + dest="obs_time", + type=str, + default="", + help="extra observed snapshot times, comma-separated (e.g., 5 or 3,7,9). " + "Final time T is always observed automatically.", + ) + + # ---- diffusion parameter estimation (b_*) ---- + parser.add_argument("--b_pI0", type=float, default=1e-3, + help="initial infection rate in diffusion parameter estimation") + parser.add_argument("--b_pR0", type=float, default=1e-3, + help="initial recovery rate in diffusion parameter estimation") + parser.add_argument("--b_steps", type=int, default=500, + help="optimization steps in diffusion parameter estimation") + parser.add_argument("--b_lr", type=float, default=3e-3, + help="learning rate in diffusion parameter estimation") + + # ---- GRINModel hyper-parameters ---- + # Ref: SPIN config/imputation/grin.yaml + parser.add_argument("--hidden_size", type=int, default=64) + parser.add_argument("--ff_size", type=int, default=64) + parser.add_argument("--embedding_size", type=int, default=8) + parser.add_argument("--n_layers", type=int, default=1) + parser.add_argument("--kernel_size", type=int, default=2) + parser.add_argument("--decoder_order", type=int, default=1) + parser.add_argument("--layer_norm", action="store_true", + help="enable layer norm inside GRIN") + parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument("--ff_dropout", type=float, default=0.0) + parser.add_argument("--merge_mode", type=str, default="mlp") + + # ---- optimizer / training ---- + parser.add_argument("--lr", type=float, default=1e-3, help="Adam learning rate") + parser.add_argument("--l2_reg", type=float, default=0.0, help="Adam weight decay") + parser.add_argument("--epochs", type=int, default=300) + parser.add_argument("--batch_size", type=int, default=1) + + # ---- evaluation ---- + parser.add_argument("--rep", type=int, default=1, help="repetitions per dataset") + + return parser + + +def parse_obs_time(obs_time_str: str, T: int) -> List[int]: + """ + Parse comma-separated times from CLI, clamp to [0, T], and always include T. + + hermes.py behavior: + obs_time = [int(t) for t in args.obs_time.split(',') if t] + obs_time.append(T) + obs_time = sorted(set(obs_time)) + """ + times: List[int] = [] + if obs_time_str: + for s in str(obs_time_str).split(","): + s = s.strip() + if not s: + continue + try: + t = int(s) + except ValueError: + continue + if 0 <= t <= T: + times.append(t) + + times.append(T) # final snapshot always observed + times = sorted(set(times)) + return times + + +def build_time_mask(obs_time: List[int], T: int, device: torch.device) -> torch.Tensor: + """ + Return a 1D mask over time: (T+1,) with 1 at observed times, else 0. + """ + m = torch.zeros(T + 1, dtype=torch.long, device=device) + if len(obs_time) > 0: + idx = torch.tensor(obs_time, dtype=torch.long, device=device).clamp(0, T) + m[idx.unique()] = 1 + return m + + +def grin_prep(y: torch.Tensor, obs_time: List[int], device: torch.device): + """ + Prepare inputs for GRIN. + + Args: + y: (samples, nodes, T+1) integer states. + obs_time: list of observed snapshot times (must include T). + device: torch device for mask. + + Returns: + x: (samples, T+1, nodes, 1) float, with missing frames zeroed + mask: (samples, T+1, nodes, 1) long {0,1}, 1 means observed + """ + n_samples, n_nodes, T1 = y.size() + T = T1 - 1 + + # time mask: (T+1,) + tmask = build_time_mask(obs_time, T=T, device=device) # long (T+1,) + + # GRIN input layout: (samples, T+1, nodes, 1) + x = y.transpose(1, 2).unsqueeze(dim=3).float() + + # mask layout: (samples, T+1, nodes, 1) + mask = tmask.view(1, T1, 1, 1).expand(n_samples, T1, n_nodes, 1) + + # IMPORTANT: avoid leaking ground-truth values at missing times + x = x * mask.float() + + return x, mask + + +def grin_run_ms(data, args: argparse.Namespace) -> torch.Tensor: + """ + Train GRIN on synthetic histories generated from estimated diffusion params, + then impute the unobserved history for `data` under the multi-snapshot mask. + """ + # ---- parse observed snapshots ---- + T = int(data.T.item()) + obs_time = parse_obs_time(args.obs_time, T) + + # attach obs info for the ms tester (so metrics only evaluate unobserved frames) + # test_ms.py will look for data.obs_ts / obs_mask / obs_masks + data.obs_ts = obs_time + + # ---- estimate diffusion params using observed snapshots ---- + bpar = b_estim(data, args, obs_time=obs_time) # Dict(pI=..., pR=...) + + n_nodes = int(data.num_nodes) + n_out = int(data.y[:, -1].max().item() + 1) + + # ---- build model ---- + model = GRINModel( + input_size=1, + hidden_size=args.hidden_size, + ff_size=args.ff_size, + embedding_size=args.embedding_size, + n_layers=args.n_layers, + n_nodes=n_nodes, + kernel_size=args.kernel_size, + decoder_order=args.decoder_order, + layer_norm=args.layer_norm, + dropout=args.dropout, + ff_dropout=args.ff_dropout, + merge_mode=args.merge_mode, + ).to(args.device) + + # prior info used in original baseline (kept) + I0 = int((data.y[:, 0] == SIR_STATES.I).long().sum().item()) + + opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2_reg) + model.train() + + # ---- training ---- + pbar = trange(1, args.epochs + 1) + for epoch in pbar: + opt.zero_grad() + + # synthetic training histories: (T+1, nodes, samples) + Y_true = diffus_gen( + T=T, + n_nodes=n_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ) + + # GRIN expects (samples, T+1, nodes, 1) + # y for prep: (samples, nodes, T+1) + y_snt = Y_true.transpose(0, 2) # (samples, nodes, T+1) + x, mask = grin_prep(y_snt, obs_time=obs_time, device=args.device) + + # model output: (samples, T+1, nodes, 1) + z = model(x=x, mask=mask, edge_index=data.edge_index)[0] + + # ground-truth in same layout: (samples, T+1, nodes, 1) + y_true_seq = Y_true.permute(2, 0, 1).unsqueeze(-1).float() + + # loss only on missing frames (mask == 0) + unobs = (mask == 0) + if bool(unobs.any()): + loss = (z - y_true_seq).abs()[unobs].mean() + else: + # degenerate case: everything observed (rare), fallback to full loss + loss = (z - y_true_seq).abs().mean() + + pbar.set_description(f"[epoch={epoch}] loss={loss.item():.4f}") + loss.backward() + opt.step() + + # ---- inference ---- + with torch.no_grad(): + model.eval() + + # build masked input from observed snapshots + tmask_1d = build_time_mask(obs_time, T=T, device=args.device).bool() # (T+1,) + y_in = data.y.clone() + y_in[:, ~tmask_1d] = 0 # hide unobserved frames + + x, mask = grin_prep(y_in.unsqueeze(0), obs_time=obs_time, device=args.device) + + z = model(x=x, mask=mask, edge_index=data.edge_index)[0] # (1, T+1, nodes, 1) + + pred = z[0, :, :, 0].T # (nodes, T+1) + pred = pred.clamp(0, n_out - 1).round().long() + + y_pred = data.y.clone() + y_pred[:, ~tmask_1d] = pred[:, ~tmask_1d] # only fill missing times + return y_pred + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + # normalize device + args.device = torch.device(args.device) + + # make sure output dir exists + out_dir = os.path.dirname(args.output) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + + # run + tester = Tester(args.data_dir, args.device, lambda data: grin_run_ms(data, args)) + tester.test([args.dataset], seed=args.seed, rep=args.rep) + tester.save(args.output) + + +if __name__ == "__main__": + main() diff --git a/hermes.py b/hermes.py new file mode 100644 index 0000000..2f02f83 --- /dev/null +++ b/hermes.py @@ -0,0 +1,383 @@ +import sys + +from inc.diffus import * +from inc.nn import * +from inc.test import * + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--dataset', type = str, help = 'dataset name') + parser.add_argument('--seed', type = int, help = 'random seed') + parser.add_argument('--data_dir', type = str, help = 'dataset folder') + parser.add_argument('--output', type = str, help = 'output file name') + parser.add_argument('--device', type = torch.device, help = 'torch device') + parser.add_argument('--b_pI0', type = float, help = 'initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type = float, help = 'initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type = int, help = 'optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type = float, help = 'learning rate in diffusion parameter estimation') + parser.add_argument('--q_steps', type = int, help = 'training steps for the proposal model') + parser.add_argument('--q_lr', type = float, help = 'learning rate for the proposal model') + parser.add_argument('--q_hid', type = int, help = 'hidden size of the proposal model') + parser.add_argument('--q_gnn', type = int, help = 'number of layers of the GNN in the proposal model') + parser.add_argument('--q_mlp', type = int, help = 'number of layers of the MLP in the proposal model') + parser.add_argument('--q_samples', type = int, help = 'sample size to estimate the loss function of the proposal model') + parser.add_argument('--q_zlim', type = int, help = 'a hyperparameter to stablize gradient') + parser.add_argument('--p_coef', type = float, help = 'the coefficient gamma in the initial distribution P[y_0]') + parser.add_argument('--t_samples', type = int, help = 'MCMC sample size') + parser.add_argument('--t_steps', type = int, help = 'MCMC steps') + parser.add_argument('--t_keep', type = float, help = 'moving average in MCMC') + parser.add_argument('--obs_time', type = str, default = '', help = 'extra observed snapshot times, comma-separated, e.g., 5,7,9') + parser.add_argument('--assumed_I0', type=int, default=None, + help='assumed initial infected count; default uses the current behavior (true I0 from data)') + args = parser.parse_args() + return args + +class QNet(nn.Module): + @classmethod + def make(cls, data, obs_time, args): + return cls( + eidx = data.edge_index, + T = data.T.item(), + n_obs = len(obs_time), + hid = args.q_hid, + gnn = args.q_gnn, + mlp = args.q_mlp, + n_nodes = data.num_nodes, + zlim = args.q_zlim, + ).to(args.device) + def __init__(self, eidx, T, n_obs, hid, gnn, mlp, n_nodes, zlim): + super().__init__() + self.eidx = eidx + self.device = self.eidx.device + self.n_nodes = n_nodes + self.n_inf = self.n_nodes + 2 + self.n_edges = self.eidx.size(dim = 1) + self.zlim = zlim + # Maximum rejection rounds per backward step in multi-snapshot segment sampling. + # This avoids infinite loops when a segment has empty/tiny support under hard constraints. + self.T = T + self.hid = int(hid) + self.gnn_dep = int(gnn) + self.mlp_dep = int(mlp) + self.w = nn.Parameter(data = torch.randn((self.n_edges, self.hid), dtype = torch.float32, device = self.device), requires_grad = True) + self.gnn = GNN(v_in = n_obs, e_in = self.hid, hid = self.hid, dep = self.gnn_dep) + self.mlp = MLP([self.hid] * self.mlp_dep + [2 * self.T]) + self.rem = (pyg.utils.degree(self.eidx[1], num_nodes = self.n_nodes).long().unsqueeze(dim = 1) + 1).detach().clone() # (nodes, 1) + self.neighbs = [[] for u in range(self.n_nodes)] + for i in range(self.n_edges): + self.neighbs[self.eidx[0, i].item()].append(self.eidx[1, i].item()) + for u in range(self.n_nodes): + self.neighbs[u] = torch.tensor(self.neighbs[u], dtype = torch.long, device = self.device) + self.adj = torch.sparse_coo_tensor( + indices = torch.stack([self.eidx[1], self.eidx[0]], dim = 0), + values = torch.ones(self.n_edges, dtype = torch.float, device = self.device), + size = (self.n_nodes, self.n_nodes), + ).coalesce() + self.zero = torch.tensor(0., dtype = torch.float, device = self.device) + def clamp_z(self, z): + return z.clamp(-self.zlim, self.zlim) + def forward(self, y, orig = False): # y: (nodes, samples, obs) + n_nodes, n_samples, n_obs = y.size() + y = y.permute(1, 0, 2).flatten(end_dim = 1) # (samples*nodes, obs) + eidx = (self.eidx.unsqueeze(dim = 1) + n_nodes * torch.arange(n_samples, dtype = torch.long, device = y.device).unsqueeze(dim = -1)).reshape((2, -1)) # (2, samples*edges) + w = self.w.repeat(n_samples, 1) # (samples, hid) + z, e = self.gnn(y.float(), eidx, w) + z = self.mlp(z) # (samples*nodes, 2*T) + z = z.T.reshape((2 * self.T, n_samples, -1)) # (2*T, samples, nodes) + zI, zR = z[: self.T], z[self.T :] # (T, samples, nodes) + zI, zR = zI.transpose(1, 2), zR.transpose(1, 2) # (T, nodes, samples) + if orig: + return zI, zR, self.clamp_z(zI), self.clamp_z(zR) + else: + return self.clamp_z(zI), self.clamp_z(zR) + def _lik_step(self, y0, y1, lI1, lI0, lR1, lR0, yL=None, reach=None): # y*, l*, reach: (nodes, sampls) + n_nodes, n_samples = y0.shape + lik = self.zero + # unreachable has log1=0 + # R->I + msk = (y1 == SIR_STATES.R) # (nodes, samples) + if yL is not None: + msk = msk & (yL != SIR_STATES.R) & reach + lik = lik + torch.where(msk, torch.where(y0 != SIR_STATES.R, lR1, lR0), self.zero) + # I->S + uid = lI1.argsort(dim = 0, descending = True) # (nodes, samples) + msk = (y1 == SIR_STATES.I) | (msk & (y0 != SIR_STATES.R)) # (nodes, samples) + rem = torch.where(msk, (reach.long() if reach is not None else 1) + torch.sparse.mm(self.adj, msk.float()).long(), self.n_inf) # (nodes, samples) + cols = torch.arange(n_samples, dtype = torch.long, device = rem.device) + for i, u in enumerate(uid): + rem_v = torch.full((self.n_nodes, n_samples), self.n_inf, dtype=rem.dtype, device=rem.device) + rem_v = rem_v.scatter_reduce(dim=0, index=self.eidx[0, :, None].expand(-1, n_samples), + src=rem[self.eidx[1]], reduce="amin", include_self=True) + rem_v = rem_v.gather(dim=0, index=u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + rem_u = rem.gather(dim = 0, index = u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + opt = (rem_u > 1) & (rem_v > 1) # (samples,) + if yL is not None: + opt = opt & (yL.gather(dim = 0, index = u.unsqueeze(dim=0)).squeeze(dim=0) != SIR_STATES.I) + msk_u = msk.gather(dim = 0, index = u.unsqueeze(dim=0)).squeeze(dim=0) # (samples,) + msk_opt = msk_u & opt # (samples,) + lik = lik + torch.where(msk_opt, torch.where(y0 == SIR_STATES.S, lI1, lI0), self.zero) + trs = (y0.gather(dim = 0, index = u.unsqueeze(dim=0)).squeeze(dim=0) == SIR_STATES.S) # (samples,) + rem = rem - (torch.sparse.mm(self.adj, torch.zeros(rem.size(), dtype=self.adj.dtype, device=rem.device).scatter(dim = 0, index = u.unsqueeze(dim=0), src = (msk_u & trs).unsqueeze(dim=0).to(self.adj.dtype))) > 0).to(rem.dtype) # (nodes, samples) + # rem = rem - torch.zeros_like(rem).index_put( + # (self.eidx[1, :, None].expand(-1, n_samples).flatten(), cols[None].expand(self.n_edges, -1).flatten()), + # ((msk_u & trs) & (self.eidx[0, :, None] == u)).flatten().long(), # (edges * samples,) + # accumulate = False) + rem = rem.index_put((u, cols), torch.where(msk_u, torch.where(trs, rem_u - 1, self.n_inf), rem.gather(dim = 0, index = u.unsqueeze(dim=0)).squeeze(dim=0))) + msk = msk.index_put((u, cols), msk_opt) + return lik + def lik_ms(self, Y, obs_time): #(T+1, nodes, samples) + assert Y.size(dim=0) == self.T + 1, "lik_ms expects Y with shape (T+1, nodes, samples)" + n_samples = Y.size(dim=2) + K = len(obs_time) + y_cond = torch.stack([Y[t] for t in obs_time], dim=2) # (nodes, samples, obs) + zI0, zR0, zI, zR = self.forward(y_cond, orig=True) # (T, nodes, samples) + zI = zI.clone().detach().requires_grad_(True); zI.retain_grad() + zR = zR.clone().detach().requires_grad_(True); zR.retain_grad() + lik = torch.zeros(n_samples, dtype=zI.dtype, device=self.device) + lI1, lI0 = F.logsigmoid(zI), F.logsigmoid(-zI) + lR1, lR0 = F.logsigmoid(zR), F.logsigmoid(-zR) + for i in range(K): + TL, TR = obs_time[i - 1] if i > 0 else None, obs_time[i] + if TL is None: + for t in range(TR - 1, -1, -1): + lik = lik + self._lik_step(Y[t], Y[t + 1], lI1[t], lI0[t], lR1[t], lR0[t]) + else: + yL = Y[TL] # (nodes, samples) + L = TR - TL + reach = torch.zeros(L, self.n_nodes, n_samples, dtype=torch.bool, device=self.device) + reach[0] = (yL == SIR_STATES.I) # (nodes, samples) + yL_not_R = (yL != SIR_STATES.R) # (nodes, samples) + for d in range(1, L): + reach[d] = reach[d - 1] | ((torch.sparse.mm(self.adj, reach[d - 1].float()) > 0) & yL_not_R) + for t in range(TR - 1, TL, -1): + lik = lik + self._lik_step(Y[t], Y[t + 1], lI1[t], lI0[t], lR1[t], lR0[t], yL = yL, reach = reach[t - TL]) + return lik, zI0, zR0, zI, zR + @torch.no_grad() + def clamp_grad(self, z0, grad): + return torch.where(z0 < self.zlim, torch.where(z0 > -self.zlim, grad, F.relu(grad)), -F.relu(-grad)) + def backward(self, loss, zI0, zR0, zI, zR): + loss.backward() + z0 = torch.stack([zI0, zR0], dim = 0) + z0.backward(torch.stack([self.clamp_grad(zI0, zI.grad), self.clamp_grad(zR0, zR.grad)], dim = 0)) + @torch.no_grad() + def _samp_step(self, y, zI, zR, compute_lik=False, yL=None, reach=None): + if reach is not None:# y: (nodes, samples); yL: (nodes,); z: (nodes, 1); reach: (nodes,) + reach = reach.unsqueeze(dim=1) # (nodes, 1) + n_samples = y.size(dim=1) + if compute_lik: + lik = self.zero + # unreachable + if reach is not None: + y = torch.where(reach, y, yL.unsqueeze(dim=1)) + # R->I + qR = torch.sigmoid(zR) # (nodes, 1) + xR = SIR_STATES.R - qR.expand(-1, n_samples).bernoulli().long() # (nodes, samples) + msk = (y == SIR_STATES.R) # (nodes, samples) + if yL is not None: + msk = msk & (yL != SIR_STATES.R).unsqueeze(dim=1) & reach + y = torch.where(msk, xR, y) + if compute_lik: + lik = lik + torch.where(msk, torch_log(torch.where(xR != SIR_STATES.R, qR, 1. - qR)), self.zero).sum(dim=0) # (samples,) + # I->S + zI, uid = zI.sort(dim = 0, descending = True) + uid = uid.squeeze(dim = -1) # (nodes,) + qI = torch.sigmoid(zI) # (nodes, 1) (already sorted by zI) + xI = SIR_STATES.I - qI.expand(-1, n_samples).bernoulli().long() # (nodes, samples) + if compute_lik: + lI = torch_log(torch.where(xI != SIR_STATES.I, qI, 1. - qI)) # (nodes, samples) + msk = (y == SIR_STATES.I) # (nodes, samples) + rem = torch.where(msk, (reach.long() if reach is not None else 1) + torch.sparse.mm(self.adj, msk.float()).long(), self.n_inf) # (nodes, samples) + for i, u in enumerate(uid): + if msk[u].max(): + vid = self.neighbs[u.item()] # (neighbs,) + rem_u = rem[u] # (samples,) + rem_v = rem[vid] # (neighbs, samples) + opt = (rem_u > 1) & (rem_v.min(dim=0).values > 1) # (samples,) + if yL is not None: + opt = opt & (yL[u] != SIR_STATES.I) + msk_opt = msk[u] & opt + y[u] = torch.where(msk_opt, xI[i], y[u]) # (samples,) + trs = (y[u] != SIR_STATES.I) # (samples,) + rem[vid] = torch.where(msk[u].unsqueeze(dim=0), torch.where(trs.unsqueeze(dim=0), rem_v - 1, rem_v), rem_v) + rem[u] = torch.where(msk[u], torch.where(trs, rem_u - 1, self.n_inf), rem[u]) + msk[u] = msk_opt + if compute_lik: + lik = lik + torch.where(msk[uid], lI, self.zero).sum(dim=0) # (samples,) + return y, lik + else: + return y, 0. + @torch.no_grad() + def _samp_seg(self, yR, zI, zR, n_samples, TL, TR, yL=None, compute_lik=False): # yR: (nodes,); zI: (T, nodes, 1); zR: (T, nodes, 1); return: Y: (TR - TL, nodes, samples), lik: (samples,) + if compute_lik: # segment log-likelihood under the proposal Q_theta + lik = torch.zeros(n_samples, dtype=torch.float, device=self.device) + y = yR.unsqueeze(dim=1).expand(-1, n_samples) # (nodes, samples) + if yL is None: # no left constraint, sample purely by original backward local-support steps + Y = torch.empty(TR, self.n_nodes, n_samples, dtype=torch.long, device=self.device) + for t in range(TR - 1, -1, -1): + y, lik_t = self._samp_step(y, zI[t], zR[t], compute_lik=compute_lik) + if compute_lik: + lik = lik + lik_t + Y[t] = y + else: + L = TR - TL + Y = torch.empty(L, self.n_nodes, n_samples, dtype=torch.long, device=self.device) + Y[0] = yL.unsqueeze(dim = -1) + reach = torch.zeros(L, self.n_nodes, dtype=torch.bool, device=yL.device) + reach[0] = (yL == SIR_STATES.I) + + yL_not_R = (yL != SIR_STATES.R) # (nodes,) FIX: keep as 1D boolean + for d in range(1, L): + nbr = (torch.sparse.mm(self.adj, reach[d - 1].float().unsqueeze(1)).squeeze(1) > 0) # (nodes,) + reach[d] = reach[d - 1] | (nbr & yL_not_R) + for t in range(TR - 1, TL, -1): + y, lik_t = self._samp_step(y, zI[t], zR[t], compute_lik=compute_lik, yL = yL, reach = reach[t - TL]) + if compute_lik: + lik = lik + lik_t + Y[t - TL] = y + if compute_lik: + return Y.detach().clone(), lik.detach().clone() + else: + return Y.detach().clone(), 0. + @torch.no_grad() + def samp_ms(self, y, zI, zR, n_samples, obs_time, compute_lik=False): # z: (T, nodes, 1) + # y: (nodes, T+1) + Y = torch.empty(self.T, self.n_nodes, n_samples, dtype=torch.long, device=self.device) + for t in obs_time: + if t < self.T: + Y[t] = y[:, t].unsqueeze(dim=1).expand(-1, n_samples) + if compute_lik: + lik = 0. # will become (samples,) after first addition + for i in range(len(obs_time) - 1, -1, -1): # Sample segments in reverse order (right endpoint always known). + TL, TR = obs_time[i - 1] if i > 0 else None, obs_time[i] + yL = None if TL is None else y[:, TL] + yR = y[:, TR] + Y[TL : TR], lik_seg = self._samp_seg(yR, zI, zR, n_samples, TL, TR, yL = yL, compute_lik = compute_lik) + if compute_lik: + lik = lik + lik_seg + if compute_lik: + return Y.detach().clone(), lik.detach().clone() + else: + return Y.detach().clone() + + +def q_loss(q_net, data, I0, bpar, n_samples, obs_time): + T = data.T.item() + n_nodes = data.num_nodes + Y = diffus_gen( + T=T, + n_nodes=n_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=n_samples, + pI=bpar.pI, + pR=bpar.pR, + ) # (T+1, nodes, samples) + + q_liks, zI0, zR0, zI, zR = q_net.lik_ms(Y=Y, obs_time=obs_time) + return -q_liks.mean(), zI0, zR0, zI, zR + +def q_train(data, obs_time, bpar, args, assumed_I0 = None): + I0 = resolve_assumed_I0(data, getattr(args, 'assumed_I0', None) if assumed_I0 is None else assumed_I0) + q_net = QNet.make(data, obs_time, args) + q_net.train() + opt = optim.AdamW(q_net.parameters(), lr=args.q_lr) + pbar = trange(1, args.q_steps + 1) + for step in pbar: + opt.zero_grad() + loss, zI0, zR0, zI, zR = q_loss(q_net, data, I0, bpar, args.q_samples, obs_time=obs_time) + pbar.set_description(f'[step={step}] loss={loss.item():.4f}') + q_net.backward(loss, zI0, zR0, zI, zR) + opt.step() + q_net.eval() + return q_net + +@torch.no_grad() +def t_mcmc(data, bpar, q_net, args, obs_time, keepdim=True, assumed_I0=None, diagnostics=False): + + I0 = resolve_assumed_I0( + data, + getattr(args, 'assumed_I0', None) if assumed_I0 is None else assumed_I0 + ) + + obs_time = sorted(list(obs_time)) + y_obs = torch.stack([data.y[:, t : t + 1] for t in obs_time], dim=2) # (nodes, 1, obs) + zI, zR = q_net(y_obs) # (T, nodes, 1) + + X, lqX = q_net.samp_ms(data.y, zI, zR, args.t_samples, obs_time=obs_time, compute_lik=True) + lpX = diffus_liks(Y=X, edge_index=data.edge_index, I0=I0, coef=args.p_coef, pI=bpar.pI, pR=bpar.pR) + + tI_avg = data_make_t(X, SIR_STATES.I, dim=0).float().mean(dim=1, keepdim=keepdim) + tR_avg = data_make_t(X, SIR_STATES.R, dim=0).float().mean(dim=1, keepdim=keepdim) + + diag_rows = [] if diagnostics else None + + pbar = trange(1, args.t_steps + 1) + for step in pbar: + Y, lqY = q_net.samp_ms(data.y, zI, zR, args.t_samples, obs_time=obs_time, compute_lik=True) + lpY = diffus_liks(Y=Y, edge_index=data.edge_index, I0=I0, coef=args.p_coef, pI=bpar.pI, pR=bpar.pR) + + # Hastings acceptance + a = torch.rand(args.t_samples, device=args.device) <= torch.exp(lpY + lqX - lpX - lqY) + acc = a.float().mean().item() + pbar.set_description(f"[step={step}] acc={acc:.3f}") + + X = torch.where(a, Y, X) + lqX = torch.where(a, lqY, lqX) + lpX = torch.where(a, lpY, lpX) + + tI = data_make_t(X, SIR_STATES.I, dim=0).float().mean(dim=1, keepdim=keepdim) + tR = data_make_t(X, SIR_STATES.R, dim=0).float().mean(dim=1, keepdim=keepdim) + tI_avg = args.t_keep * tI_avg + (1.0 - args.t_keep) * tI + tR_avg = args.t_keep * tR_avg + (1.0 - args.t_keep) * tR + + if diagnostics: + diag_rows.append(dict( + step=int(step), + accept_rate=float(acc), + mean_tI=float(tI.mean().item()), + mean_tR=float(tR.mean().item()), + mean_tI_avg=float(tI_avg.mean().item()), + mean_tR_avg=float(tR_avg.mean().item()), + mean_lp=float(lpX.mean().item()), + )) + + if diagnostics: + return tI_avg, tR_avg, diag_rows + return tI_avg, tR_avg +def run_hermes(data, args, assumed_I0 = None, return_extra = False): + # parse obs times + obs_time = [int(t) for t in args.obs_time.split(',') if t] + obs_time.append(data.T.item()) + obs_time = sorted(set(obs_time)) + assumed_I0 = resolve_assumed_I0(data, getattr(args, 'assumed_I0', None) if assumed_I0 is None else assumed_I0) + + # estimate diffusion parameters + bpar = b_estim(data, args, obs_time = obs_time, assumed_I0 = assumed_I0) + print(f'[est] pI={bpar.pI:.4f}, pR={bpar.pR:.4f}', flush = True) + + # train a proposal network + q_net = q_train(data, obs_time, bpar, args, assumed_I0 = assumed_I0) + + # estimate transition times + tI, tR = t_mcmc(data, bpar, q_net, args, obs_time = obs_time, keepdim = True, assumed_I0 = assumed_I0) # (nodes, 1) + + tI = tI.round().long() + tR = tR.round().long() + + # compose a history + with torch.no_grad(): + y_pred = torch.zeros_like(data.y) # (nodes, T+1) + y_pred.scatter_(dim = 1, index = torch.minimum(tI, data.T), src = torch.full_like(tI, 1)) + y_pred.scatter_(dim = 1, index = torch.minimum(tR, data.T), src = torch.full_like(tR, 2)) + y_pred = y_pred[:, : data.T.item()].cummax(dim = 1).values + if return_extra: + return y_pred, Dict(assumed_I0 = assumed_I0, pI = bpar.pI, pR = bpar.pR) + return y_pred + +def main(data): + return run_hermes(data, args) + +if __name__ == '__main__': + args = get_args() + tester = Tester(args.data_dir, args.device, main) + tester.test([args.dataset], seed = args.seed, rep = 1) \ No newline at end of file diff --git a/inc/data.py b/inc/data.py index c09bc56..41a2a8e 100644 --- a/inc/data.py +++ b/inc/data.py @@ -2,7 +2,6 @@ def data_make_t(y, x, dim = -1): return torch.where(*(y == x).max(dim), y.size(dim)) - def data_simulate(Gnx, seed, T, diffus, params): sir = (diffus == 'sir') cfg = ndmc.Configuration() @@ -67,7 +66,7 @@ def data_synthetic(graph, diffus, data_dir, device): data_dir = osp.join(data_dir, 'synthetic') f_data = osp.join(data_dir, f'{graph}-{diffus}.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: seed = 123456789 T = 10 @@ -89,7 +88,7 @@ def data_prost(diffus, data_dir, device): data_dir = osp.join(data_dir, 'prost') f_data = osp.join(data_dir, f'prost-{diffus}.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: seed = 123456789 T = 15 @@ -111,7 +110,7 @@ def data_oregon2(diffus, data_dir, device): data_dir = osp.join(data_dir, 'oregon2') f_data = osp.join(data_dir, f'oregon2-{diffus}.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: seed = 123456789 T = 15 @@ -131,7 +130,7 @@ def data_farmers_si(data_dir, device): data_dir = osp.join(data_dir, 'farmers') f_data = osp.join(data_dir, 'farmers-si.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: f_raw = file_require(None, data_dir, 'brfarmers.rdata') df = pyreadr.read_r('farmers/brfarmers.rdata')['brfarmers'] @@ -161,7 +160,7 @@ def data_pol_si(data_dir, device): data_dir = osp.join(data_dir, 'pol') f_data = osp.join(data_dir, 'pol-si.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: f_edge = file_require('https://nrvis.com/download/data/rt/rt-pol.zip', data_dir, 'rt-pol.txt', z = 'zip') df_fr, df_to, df_time = [], [], [] @@ -190,7 +189,7 @@ def data_covid_sir(data_dir, device): data_dir = osp.join(data_dir, 'covid') f_data = osp.join(data_dir, 'covid-sir.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: COVID_KNN = 10 f_s2a = file_require(None, data_dir, 'state2abbr.pyon') @@ -251,7 +250,7 @@ def data_heb_sir(data_dir, device): data_dir = osp.join(data_dir, 'heb') f_data = osp.join(data_dir, 'heb-sir.pt') if osp.exists(f_data): - return torch.load(f_data, map_location = device) + return torch.load(f_data, map_location=device, weights_only=False) else: f_edge = file_require(url = None, fdir = data_dir, fname = 'DS1_NON_VIRAL_Gtw.tsv') df = pd.read_csv(f_edge, sep = '\t', header = None, names = ['time', 'to', 'fr'], dtype = dict(time = str, to = int, fr = int), parse_dates = ['time']) diff --git a/inc/diffus.py b/inc/diffus.py index 40dc77f..a452857 100644 --- a/inc/diffus.py +++ b/inc/diffus.py @@ -24,12 +24,19 @@ def diffus_sim(edge_index, y0, WI, WR = None): # y0: (nodes, samples); WI: (T, e @torch.no_grad() def diffus_gen(T, n_nodes, edge_index, I0, n_samples, pI, pR): n_edges = edge_index.size(dim = 1) - idx = torch.ones(n_samples, n_nodes, device = edge_index.device).multinomial(I0, replacement = False).T # (I0, samples) y0 = torch.full((n_nodes, n_samples), SIR_STATES.S, dtype = torch.long, device = edge_index.device) - y0.scatter_(index = idx, dim = 0, src = torch.full_like(idx, SIR_STATES.I)) + if I0 > 0: + idx = torch.ones(n_samples, n_nodes, device = edge_index.device).multinomial(I0, replacement = False).T # (I0, samples) + y0.scatter_(index = idx, dim = 0, src = torch.full_like(idx, SIR_STATES.I)) WI = (torch.rand(T, n_edges, n_samples, device = y0.device) < pI).long() WR = (torch.rand(T, n_nodes, n_samples, device = y0.device) < pR).long() if pR > 0 else None return diffus_sim(edge_index, y0, WI, WR) # (T+1, nodes, samples) +def resolve_assumed_I0(data, assumed_I0 = None): + if assumed_I0 is None: + I0 = (data.y[:, 0] == SIR_STATES.I).long().sum().item() + else: + I0 = int(assumed_I0) + return max(0, min(int(data.num_nodes), I0)) def diffus_liks(Y, edge_index, I0, coef, pI, pR): # Y: (T+1, nodes, samples) # assuming Y feasible log1pI = torch_log(1. - pI) if isinstance(pI, torch.Tensor) else math_log(1. - pI) @@ -62,50 +69,96 @@ def __repr__(self, digits = 4): def dict(self): return Dict(pI = self.pI.item(), pR = self.pR.item()) -def b_lik(bpar, data): # yT: (nodes,) +def _mf_init_from_prior(data, n_nodes, device, assumed_I0 = None): + # prior: only use I0 count (same as current code) + I0 = resolve_assumed_I0(data, assumed_I0) + lI = torch.full((n_nodes,), I0 / n_nodes, dtype=torch.float, device=device) + lS = torch.full((n_nodes,), 1. - I0 / n_nodes, dtype=torch.float, device=device) + lR = torch.zeros(n_nodes, dtype=torch.float, device=device) + return lS, lI, lR + +def _mf_init_from_snapshot(y, device): + # hard clamp to observed snapshot y (nodes,) + lS = (y == SIR_STATES.S).float().to(device) + lI = (y == SIR_STATES.I).float().to(device) + lR = (y == SIR_STATES.R).float().to(device) + return lS, lI, lR + +def b_lik(bpar, data, obs_time=None, assumed_I0 = None): device = data.y.device n_nodes = data.num_nodes T = data.T.item() ei = data.edge_index - I0 = (data.y[:, 0] == 1).sum() pI, pR = bpar.pI, bpar.pR - lSs, lIs, lRs = [], [], [] - lIs.append(torch.full((n_nodes,), I0 / n_nodes, dtype = torch.float, device = device)) - lSs.append(torch.full((n_nodes,), 1. - I0 / n_nodes, dtype = torch.float, device = device)) - lRs.append(torch.zeros(n_nodes, dtype = torch.float, device = device)) - for t in range(T): - aI = pysc.scatter_mul(src = (1. - lIs[-1] * pI)[ei[0]], dim = 0, index = ei[1], dim_size = n_nodes) - lS = lSs[-1] * aI - kI = lIs[-1] + lSs[-1] * (1. - aI) - lI = kI * (1. - pR) - lR = lRs[-1] + kI * pR - lSs.append(lS) - lIs.append(lI) - lRs.append(lR) - lik = torch.stack([lSs[-1], lIs[-1], lRs[-1]], dim = 0) # (states, nodes) - yT = data.y[:, -1].unsqueeze(dim = 0) # (1, nodes) - lik = lik.gather(dim = 0, index = yT) # (1, nodes) - lik = torch_log(lik).mean() - return lik - -def b_estim(data, args): + + # -------- obs times -------- + if obs_time is None: + obs_time = [T] + obs_time = sorted(set(int(t) for t in obs_time if 0 <= int(t) <= T)) + if len(obs_time) == 0 or obs_time[-1] != T: + obs_time.append(T) + + # -------- segmented mean-field -------- + lS, lI, lR = _mf_init_from_prior(data, n_nodes, device, assumed_I0 = assumed_I0) + t_prev = 0 + lik_total = 0.0 + + for t_obs in obs_time: + # forward from t_prev -> t_obs + for _ in range(t_obs - t_prev): + aI = pysc.scatter_mul( + src=(1. - lI * pI)[ei[0]], + dim=0, + index=ei[1], + dim_size=n_nodes + ) + lS_new = lS * aI + kI = lI + lS * (1. - aI) + lI_new = kI * (1. - pR) + lR_new = lR + kI * pR + lS, lI, lR = lS_new, lI_new, lR_new + + # score snapshot at t_obs + lik = torch.stack([lS, lI, lR], dim=0) # (3, nodes) + y = data.y[:, t_obs].unsqueeze(dim=0) # (1, nodes) + prob = lik.gather(dim=0, index=y).squeeze(dim=0) # (nodes,) + lik_total = lik_total + torch_log(prob).mean() + + # clamp for next segment (if any) + lS, lI, lR = _mf_init_from_snapshot(data.y[:, t_obs], device) + t_prev = t_obs + + # optional: normalize by number of observed frames, keeps loss scale stable + lik_total = lik_total / len(obs_time) + return lik_total + +def b_estim(data, args, obs_time=None, assumed_I0 = None): T = data.T.item() - n_nodes = data.num_nodes - n_edges = data.edge_index.size(dim = 1) n_cls = data.y[:, T].max().item() + 1 pI, pR = args.b_pI0, (args.b_pR0 if n_cls == 3 else 0.) - #print(f'[ini] pI={pI:.4f}, pR={pR:.4f}', flush = True) device = data.y.device - bpar = BPar(pI = pI, pR = pR, device = device) + + # if caller doesn't pass obs_time, parse from args (compatible with current main) + if obs_time is None and hasattr(args, "obs_time"): + tmp = [int(t) for t in str(args.obs_time).split(',') if t] + tmp.append(T) + obs_time = tmp + if assumed_I0 is None and hasattr(args, 'assumed_I0'): + assumed_I0 = args.assumed_I0 + + bpar = BPar(pI=pI, pR=pR, device=device) bpar.train() - opt = optim.AdamW(bpar.parameters(), lr = args.b_lr, betas = (0.5, 0.5)) + opt = optim.AdamW(bpar.parameters(), lr=args.b_lr, betas=(0.5, 0.5)) pbar = trange(1, args.b_steps + 1) + for step in pbar: opt.zero_grad() - loss = -b_lik(bpar, data) + loss = -b_lik(bpar, data, obs_time=obs_time, assumed_I0 = assumed_I0) loss.backward() opt.step() bpar.clamp_() pbar.set_description(f'[step={step}] {bpar}') + bpar.eval() return bpar.dict() + diff --git a/input/covid/dist-covid.pt b/input/covid/dist-covid.pt new file mode 100644 index 0000000..57fae13 Binary files /dev/null and b/input/covid/dist-covid.pt differ diff --git a/input/exp1_scalability/ba-si-n1000-T1-g0-d1.pt b/input/exp1_scalability/ba-si-n1000-T1-g0-d1.pt new file mode 100644 index 0000000..f3d6852 Binary files /dev/null and b/input/exp1_scalability/ba-si-n1000-T1-g0-d1.pt differ diff --git a/input/exp1_scalability/vsN/farmers-si_n82_T16.pt b/input/exp1_scalability/vsN/farmers-si_n82_T16.pt new file mode 100644 index 0000000..e8793b3 Binary files /dev/null and b/input/exp1_scalability/vsN/farmers-si_n82_T16.pt differ diff --git a/input/exp1_scalability/vsT/farmers-si_n82_T4.pt b/input/exp1_scalability/vsT/farmers-si_n82_T4.pt new file mode 100644 index 0000000..5151d6b Binary files /dev/null and b/input/exp1_scalability/vsT/farmers-si_n82_T4.pt differ diff --git a/input/exp1_scalability/vsT/farmers-si_n82_T6.pt b/input/exp1_scalability/vsT/farmers-si_n82_T6.pt new file mode 100644 index 0000000..82cbbec Binary files /dev/null and b/input/exp1_scalability/vsT/farmers-si_n82_T6.pt differ diff --git a/input/exp1_scalability/vsT/farmers-si_n82_T8.pt b/input/exp1_scalability/vsT/farmers-si_n82_T8.pt new file mode 100644 index 0000000..85cda61 Binary files /dev/null and b/input/exp1_scalability/vsT/farmers-si_n82_T8.pt differ diff --git a/input/exp2_timespan/T3/synthetic/ba-sir.pt b/input/exp2_timespan/T3/synthetic/ba-sir.pt new file mode 100644 index 0000000..8811e00 Binary files /dev/null and b/input/exp2_timespan/T3/synthetic/ba-sir.pt differ diff --git a/mkdir b/mkdir new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8a56907 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,65 @@ +aiohappyeyeballs==2.6.1 +aiohttp==3.13.3 +aiosignal==1.4.0 +attrs==25.4.0 +bokeh==3.8.2 +certifi==2026.1.4 +charset-normalizer==3.4.4 +colorama==0.4.6 +contourpy==1.3.3 +cycler==0.12.1 +decorator==5.2.1 +dynetx==0.3.2 +filelock==3.20.0 +fonttools==4.61.1 +frozenlist==1.8.0 +fsspec==2025.12.0 +future==1.0.0 +idna==3.11 +igraph==1.0.0 +Jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +MarkupSafe==2.1.5 +matplotlib==3.10.8 +mpmath==1.3.0 +multidict==6.7.0 +narwhals==2.15.0 +ndlib==5.1.1 +netdispatch==0.1.0 +networkx==3.6.1 +numpy==2.3.5 +packaging==26.0 +pandas==3.0.0 +pillow==12.0.0 +propcache==0.4.1 +psutil==7.2.1 +pyg-lib==0.5.0+pt28cu128 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +python-igraph==1.0.0 +PyYAML==6.0.3 +requests==2.32.5 +scikit-learn==1.8.0 +scipy==1.17.0 +seaborn==0.13.2 +six==1.17.0 +sympy==1.14.0 +texttable==1.7.0 +threadpoolctl==3.6.0 +torch==2.8.0+cu128 +torch-geometric==2.7.0 +torch_cluster==1.6.3+pt28cu128 +torch_scatter==2.1.2+pt28cu128 +torch_sparse==0.6.18+pt28cu128 +torch_spline_conv==1.2.2+pt28cu128 +torchaudio==2.8.0+cu128 +torchvision==0.23.0+cu128 +tornado==6.5.4 +tqdm==4.67.1 +typing_extensions==4.15.0 +tzdata==2025.3 +urllib3==2.6.3 +xxhash==3.6.0 +xyzservices==2025.11.0 +yarl==1.22.0 diff --git a/scripts/exp_sir_profile.py b/scripts/exp_sir_profile.py new file mode 100644 index 0000000..9bdea03 --- /dev/null +++ b/scripts/exp_sir_profile.py @@ -0,0 +1,112 @@ + +import os +import sys +from typing import List, Optional + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +import argparse +import torch + +from inc.test_ms import Tester +try: + from inc.header import SIR_STATES as _SIR_STATES + SIR_STATES = _SIR_STATES +except Exception: + class _SIR: + S, I, R = 0, 1, 2 + SIR_STATES = _SIR() + +def parse_obs_ts(s: Optional[str]) -> Optional[List[int]]: + if s is None: + return None + s = s.strip() + if not s: + return None + ts = [int(x) for x in s.split(",") if x.strip() != ""] + ts = sorted(set(ts)) + return ts if len(ts) > 0 else None + + +def main_factory(obs_ts: Optional[List[int]]): + def main(data): + # data.y: (nodes, T+1) + y0 = data.y + y = y0.detach().to("cpu", dtype=torch.long) + + T = int(data.T.item()) if torch.is_tensor(data.T) else int(data.T) + n = int(data.num_nodes) + + assert y.dim() == 2 and y.size(0) == n and y.size(1) == T + 1, \ + f"Expect y shape (nodes, T+1)=({n},{T+1}), got {tuple(y.shape)}" + + obs_set = set(obs_ts or []) + obs_set.add(T) + + def ratio_at(t: int): + yt = y[:, t] + cS = int((yt == SIR_STATES.S).sum().item()) + cI = int((yt == SIR_STATES.I).sum().item()) + cR = int((yt == SIR_STATES.R).sum().item()) + tot = cS + cI + cR + if tot == 0: + return (0, 0, 0, 0.0, 0.0, 0.0) + return (cS, cI, cR, cS / tot, cI / tot, cR / tot) + + print("=" * 80, flush=True) + print(f"[SIR PROFILE] nodes={n}, T={T}", flush=True) + print(f"[SIR PROFILE] observed ts = {sorted(obs_set)}", flush=True) + print("-" * 80, flush=True) + print(f"{'t':>3} {'tag':>6} {'S%':>8} {'I%':>8} {'R%':>8} {'(S,I,R counts)':>20}", flush=True) + + # per-time + for t in range(T + 1): + cS, cI, cR, rS, rI, rR = ratio_at(t) + tag = "OBS" if t in obs_set else "UNOBS" + print(f"{t:>3} {tag:>6} {rS:>8.4f} {rI:>8.4f} {rR:>8.4f} ({cS},{cI},{cR})", flush=True) + + # aggregate on unobserved + unobs_ts = [t for t in range(T + 1) if t not in obs_set] + if len(unobs_ts) == 0: + print("-" * 80, flush=True) + print("[SIR PROFILE] No unobserved time steps under current obs_ts.", flush=True) + print("=" * 80, flush=True) + return y0 + + cS_all = cI_all = cR_all = 0 + for t in unobs_ts: + cS, cI, cR, *_ = ratio_at(t) + cS_all += cS + cI_all += cI + cR_all += cR + tot_all = cS_all + cI_all + cR_all + rS_all = cS_all / tot_all + rI_all = cI_all / tot_all + rR_all = cR_all / tot_all + + print("-" * 80, flush=True) + print(f"[SIR PROFILE] UNOBS ts = {unobs_ts}", flush=True) + print(f"[SIR PROFILE] UNOBS weighted ratio: S={rS_all:.4f}, I={rI_all:.4f}, R={rR_all:.4f}", flush=True) + print("=" * 80, flush=True) + + return y0 + + return main + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--dataset", required=True) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--data_dir", type=str, default="input") + ap.add_argument("--device", type=str, default="cpu") + ap.add_argument("--obs_ts", type=str, default=None, help='e.g. "0,3,5,9"') + args = ap.parse_args() + + obs_ts = parse_obs_ts(args.obs_ts) + device = torch.device(args.device) + + tester = Tester(args.data_dir, device, main_factory(obs_ts)) + tester.test(datasets=[args.dataset], seed=args.seed, rep=1) diff --git a/spin.py b/spin.py new file mode 100644 index 0000000..0fd251d --- /dev/null +++ b/spin.py @@ -0,0 +1,672 @@ +# ! pip install --no-index torch-scatter==2.0.7 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-sparse==0.6.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-cluster==1.5.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +#! pip install --no-index torch-spline-conv==1.2.1 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install torch-geometric==2.0.4 +# ! pip install ndlib==5.1.1 + + +# ! pip install einops==0.6.0 +# ! pip install test_tube==0.7.5 + + +try: + from tsl.nn.base import StaticGraphEmbedding +except Exception: + import torch + from torch import nn + + class StaticGraphEmbedding(nn.Module): + def __init__(self, n_nodes, out_channels): + super().__init__() + self.emb = nn.Embedding(n_nodes, out_channels) + + def forward(self, token_index=None): + if token_index is None: + token_index = torch.arange(self.emb.num_embeddings, device=self.emb.weight.device) + return self.emb(token_index) +from tsl.nn.layers import PositionalEncoding +from tsl.nn.layers.norm import LayerNorm +from tsl.nn.blocks.encoders import MLP +from tsl.nn.functional import sparse_softmax +from tsl.engines import Imputer, Predictor +from tsl.ops.connectivity import weighted_degree +#from tsl.data import Batch, SpatioTemporalDataModule, ImputationDataset + +#SPINModel +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/postional_encoding.py''' +from typing import Optional + +from torch import nn + +class PositionalEncoder(nn.Module): + + def __init__(self, in_channels, out_channels, + n_layers: int = 1, + n_nodes: Optional[int] = None): + super(PositionalEncoder, self).__init__() + self.lin = nn.Linear(in_channels, out_channels) + self.activation = nn.LeakyReLU() + self.mlp = MLP(out_channels, out_channels, out_channels, + n_layers=n_layers, activation='relu') + self.positional = PositionalEncoding(out_channels) + if n_nodes is not None: + self.node_emb = StaticGraphEmbedding(n_nodes, out_channels) + else: + self.register_parameter('node_emb', None) + + def forward(self, x, node_emb=None, node_index=None): + if node_emb is None: + node_emb = self.node_emb(token_index=node_index) + # x: [b s c], node_emb: [n c] -> [b s n c] + x = self.lin(x) + x = self.activation(x.unsqueeze(-2) + node_emb) + #print('u:', tuple(x.shape), 'node_emb:', tuple(node_emb.shape))##### + out = self.mlp(x) + out = self.positional(out) + return out + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/additive_attention.py''' +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch import nn +from torch.nn import LayerNorm, functional as F +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.dense.linear import Linear +from torch_geometric.typing import Adj, OptTensor, PairTensor +from torch_scatter import scatter +from torch_scatter.utils import broadcast + + +class AdditiveAttention(MessagePassing): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + norm: bool = True, + dropout: float = 0.0, + dim: int = -2, + **kwargs): + kwargs.setdefault('aggr', 'add') + super().__init__(node_dim=dim, **kwargs) + + self.output_size = output_size + if isinstance(input_size, int): + self.src_size = self.tgt_size = input_size + else: + self.src_size, self.tgt_size = input_size + + self.msg_size = msg_size or self.output_size + self.msg_layers = msg_layers + + assert reweight in ['softmax', 'l1', None] + self.reweight = reweight + + self.root_weight = root_weight + self.dropout = dropout + + # key bias is discarded in softmax + self.lin_src = Linear(self.src_size, self.output_size, + weight_initializer='glorot', + bias_initializer='zeros') + self.lin_tgt = Linear(self.tgt_size, self.output_size, + weight_initializer='glorot', bias=False) + + if self.root_weight: + self.lin_skip = Linear(self.tgt_size, self.output_size, + bias=False) + else: + self.register_parameter('lin_skip', None) + + self.msg_nn = nn.Sequential( + nn.PReLU(init=0.2), + MLP(self.output_size, self.msg_size, self.output_size, + n_layers=self.msg_layers, dropout=self.dropout, + activation='prelu') + ) + + if self.reweight == 'softmax': + self.msg_gate = nn.Linear(self.output_size, 1, bias=False) + else: + self.msg_gate = nn.Sequential(nn.Linear(self.output_size, 1), + nn.Sigmoid()) + + if norm: + self.norm = LayerNorm(self.output_size) + else: + self.register_parameter('norm', None) + + self.reset_parameters() + + def reset_parameters(self): + self.lin_src.reset_parameters() + self.lin_tgt.reset_parameters() + if self.lin_skip is not None: + self.lin_skip.reset_parameters() + + def forward(self, x: PairTensor, edge_index: Adj, mask: OptTensor = None): + # if query/key not provided, defaults to x (e.g., for self-attention) + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + N_src, N_tgt = x_src.size(self.node_dim), x_tgt.size(self.node_dim) + + msg_src = self.lin_src(x_src) + msg_tgt = self.lin_tgt(x_tgt) + + msg = (msg_src, msg_tgt) + + # propagate_type: (msg: PairTensor, mask: OptTensor) + out = self.propagate(edge_index, msg=msg, mask=mask, + size=(N_src, N_tgt)) + + # skip connection + if self.root_weight: + out = out + self.lin_skip(x_tgt) + + if self.norm is not None: + out = self.norm(out) + + return out + + def normalize_weights(self, weights, index, num_nodes, mask=None): + # mask weights + if mask is not None: + fill_value = float("-inf") if self.reweight == 'softmax' else 0. + weights = weights.masked_fill(torch.logical_not(mask), fill_value) + # eventually reweight + if self.reweight == 'l1': + expanded_index = broadcast(index, weights, self.node_dim) + weights_sum = scatter(weights, expanded_index, self.node_dim, + dim_size=num_nodes, reduce='sum') + weights_sum = weights_sum.index_select(self.node_dim, index) + weights = weights / (weights_sum + 1e-5) + elif self.reweight == 'softmax': + weights = sparse_softmax(weights, index, num_nodes=num_nodes, + dim=self.node_dim) + return weights + + def message(self, msg_j: Tensor, msg_i: Tensor, index, size_i, + mask_j: OptTensor = None) -> Tensor: + msg = self.msg_nn(msg_j + msg_i) + gate = self.msg_gate(msg) + alpha = self.normalize_weights(gate, index, size_i, mask_j) + alpha = F.dropout(alpha, p=self.dropout, training=self.training) + out = alpha * msg + return out + + def __repr__(self) -> str: + return (f'{self.__class__.__name__}({self.output_size}, ' + f'dim={self.node_dim}, ' + f'root_weight={self.root_weight})') + + +class TemporalAdditiveAttention(AdditiveAttention): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + norm: bool = True, + dropout: float = 0.0, + **kwargs): + kwargs.setdefault('dim', 1) + super().__init__(input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + root_weight=root_weight, + reweight=reweight, + dropout=dropout, + norm=norm, + **kwargs) + + def forward(self, x: PairTensor, mask: OptTensor = None, + temporal_mask: OptTensor = None, + causal_lag: Optional[int] = None): + # x: [b s * c] query: [b l * c] key: [b s * c] + # mask: [b s * c] temporal_mask: [l s] + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + l, s = x_tgt.size(self.node_dim), x_src.size(self.node_dim) + i = torch.arange(l, dtype=torch.long, device=x_src.device) + j = torch.arange(s, dtype=torch.long, device=x_src.device) + + # compute temporal index, from j to i + if temporal_mask is None and isinstance(causal_lag, int): + temporal_mask = tuple(torch.tril_indices(l, l, offset=-causal_lag, + device=x_src.device)) + if temporal_mask is not None: + assert temporal_mask.size() == (l, s) + i, j = torch.meshgrid(i, j) + edge_index = torch.stack((j[temporal_mask], i[temporal_mask])) + else: + edge_index = torch.cartesian_prod(j, i).T + + return super(TemporalAdditiveAttention, self).forward(x, edge_index, + mask=mask) + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/temporal_graph_additive_attention.py''' +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.dense.linear import Linear +from torch_geometric.typing import Adj, OptTensor, OptPairTensor + +class TemporalGraphAdditiveAttention(MessagePassing): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + temporal_self_attention: bool = True, + mask_temporal: bool = True, + mask_spatial: bool = True, + norm: bool = True, + dropout: float = 0., + **kwargs): + kwargs.setdefault('aggr', 'add') + super(TemporalGraphAdditiveAttention, self).__init__(node_dim=-2, + **kwargs) + + # store dimensions + if isinstance(input_size, int): + self.src_size = self.tgt_size = input_size + else: + self.src_size, self.tgt_size = input_size + self.output_size = output_size + self.msg_size = msg_size or self.output_size + + self.mask_temporal = mask_temporal + self.mask_spatial = mask_spatial + + self.root_weight = root_weight + self.dropout = dropout + + if temporal_self_attention: + self.self_attention = TemporalAdditiveAttention( + input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + reweight=reweight, + dropout=dropout, + root_weight=False, + norm=False + ) + else: + self.register_parameter('self_attention', None) + + self.cross_attention = TemporalAdditiveAttention(input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + reweight=reweight, + dropout=dropout, + root_weight=False, + norm=False) + + if self.root_weight: + self.lin_skip = Linear(self.tgt_size, self.output_size, + bias_initializer='zeros') + else: + self.register_parameter('lin_skip', None) + + if norm: + self.norm = LayerNorm(output_size) + else: + self.register_parameter('norm', None) + + self.reset_parameters() + + def reset_parameters(self): + self.cross_attention.reset_parameters() + if self.self_attention is not None: + self.self_attention.reset_parameters() + if self.lin_skip is not None: + self.lin_skip.reset_parameters() + if self.norm is not None: + self.norm.reset_parameters() + + def forward(self, x: OptPairTensor, + edge_index: Adj, edge_weight: OptTensor = None, + mask: OptTensor = None): + # inputs: [batch, steps, nodes, channels] + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + n_src, n_tgt = x_src.size(-2), x_tgt.size(-2) + + # propagate query, key and value + #print('src:', x_src.shape, 'tgt:', x_tgt.shape, 'ei:', edge_index.shape, 'mask:', mask.shape, f'mask_spatial={self.mask_spatial}') + out = self.propagate(x=(x_src, x_tgt), + edge_index=edge_index, edge_weight=edge_weight, + mask=mask if self.mask_spatial else None, + size=(n_src, n_tgt)) + + if self.self_attention is not None: + s, l = x_src.size(1), x_tgt.size(1) + if s == l: + attn_mask = ~torch.eye(l, l, dtype=torch.bool, + device=x_tgt.device) + else: + attn_mask = None + temp = self.self_attention(x=(x_src, x_tgt), + mask=mask if self.mask_temporal else None, + temporal_mask=attn_mask) + out = out + temp + + # skip connection + if self.root_weight: + out = out + self.lin_skip(x_tgt) + + if self.norm is not None: + out = self.norm(out) + + return out + + def message(self, x_i: Tensor, x_j: Tensor, + edge_weight: OptTensor, mask_j: OptTensor) -> Tensor: + # [batch, steps, edges, channels] + + out = self.cross_attention((x_j, x_i), mask=mask_j) + #print('out:', out.shape) + + if edge_weight is not None: + out = out * edge_weight.view(-1, 1) + return out + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/models/spin.py''' +from typing import Optional + +import torch +from torch import nn, Tensor +from torch.nn import LayerNorm +from torch_geometric.typing import OptTensor + +class SPINModel(nn.Module): + + def __init__(self, input_size: int, + hidden_size: int, + n_nodes: int, + u_size: Optional[int] = None, + output_size: Optional[int] = None, + temporal_self_attention: bool = True, + reweight: Optional[str] = 'softmax', + n_layers: int = 4, + eta: int = 3, + message_layers: int = 1): + super(SPINModel, self).__init__() + + u_size = u_size or input_size + output_size = output_size or input_size + self.n_nodes = n_nodes + self.n_layers = n_layers + self.eta = eta + self.temporal_self_attention = temporal_self_attention + + self.u_enc = PositionalEncoder(in_channels=u_size, + out_channels=hidden_size, + n_layers=2, + n_nodes=n_nodes) + + self.h_enc = MLP(input_size, hidden_size, n_layers=2) + self.h_norm = LayerNorm(hidden_size) + + self.valid_emb = StaticGraphEmbedding(n_nodes, hidden_size) + self.mask_emb = StaticGraphEmbedding(n_nodes, hidden_size) + + self.x_skip = nn.ModuleList() + self.encoder, self.readout = nn.ModuleList(), nn.ModuleList() + for l in range(n_layers): + x_skip = nn.Linear(input_size, hidden_size) + encoder = TemporalGraphAdditiveAttention( + input_size=hidden_size, + output_size=hidden_size, + msg_size=hidden_size, + msg_layers=message_layers, + temporal_self_attention=temporal_self_attention, + reweight=reweight, + mask_temporal=True, + mask_spatial=l < eta, + norm=True, + root_weight=True, + dropout=0.0 + ) + readout = MLP(hidden_size, hidden_size, output_size, + n_layers=2) + self.x_skip.append(x_skip) + self.encoder.append(encoder) + self.readout.append(readout) + + def forward(self, x: Tensor, u: Tensor, mask: Tensor, + edge_index: Tensor, edge_weight: OptTensor = None, + node_index: OptTensor = None, target_nodes: OptTensor = None): + if target_nodes is None: + target_nodes = slice(None) + + # Whiten missing values + x = x * mask + + # POSITIONAL ENCODING ################################################# + # Obtain spatio-temporal positional encoding for every node-step pair # + # in both observed and target sets. Encoding are obtained by jointly # + # processing node and time positional encoding. # + + # Build (node, timestamp) encoding + q = self.u_enc(u, node_index=node_index) + # Condition value on key + h = self.h_enc(x) + q + + # ENCODER ############################################################# + # Obtain representations h^i_t for every (i, t) node-step pair by # + # only taking into account valid data in representation set. # + + # Replace H in missing entries with queries Q + h = torch.where(mask.bool(), h, q) + # Normalize features + h = self.h_norm(h) + + imputations = [] + + for l in range(self.n_layers): + if l == self.eta: + # Condition H on two different embeddings to distinguish + # valid values from masked ones + valid = self.valid_emb(token_index=node_index) + masked = self.mask_emb(token_index=node_index) + h = torch.where(mask.bool(), h + valid, h + masked) + # Masked Temporal GAT for encoding representation + h = h + self.x_skip[l](x) * mask # skip connection for valid x + #print(f'l={l}', 'h:', tuple(h.shape), 'x:', tuple(x.shape), 'mask:', tuple(mask.shape), 'ei:', edge_index) + h = self.encoder[l](h, edge_index, mask=mask) + # Read from H to get imputations + target_readout = self.readout[l](h[..., target_nodes, :]) + imputations.append(target_readout) + + # Get final layer imputations + x_hat = imputations.pop(-1) + + return x_hat, imputations + + +from inc.diffus import * +from inc.test import * + +import argparse + + +def get_args(): + parser = argparse.ArgumentParser() + # align with existing CLI style (e.g., hermes.py) + parser.add_argument('--dataset', type=str, default=None, + help='dataset name (default: run the built-in list used in the original spin.py)') + parser.add_argument('--seed', type=int, default=123456789, + help='random seed') + parser.add_argument('--data_dir', type=str, default='input', + help='dataset folder') + parser.add_argument('--output', type=str, default='output/spin.pt', + help='output file name') + parser.add_argument('--device', type=torch.device, + default=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), + help='torch device') + + # diffusion parameter estimation (same defaults as the original Dict args) + parser.add_argument('--b_pI0', type=float, default=1e-3, + help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, default=1e-3, + help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, default=500, + help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, default=3e-3, + help='learning rate in diffusion parameter estimation') + + # SPINModel hyperparameters (same defaults as the original Dict args) + parser.add_argument('--u_size', type=int, default=1, + help='size of exogenous input features u') + parser.add_argument('--hidden_size', type=int, default=32, + help='hidden size of SPIN') + parser.add_argument('--temporal_self_attention', dest='temporal_self_attention', + action='store_true', + help='enable temporal self-attention (default: enabled)') + parser.add_argument('--no_temporal_self_attention', dest='temporal_self_attention', + action='store_false', + help='disable temporal self-attention') + parser.set_defaults(temporal_self_attention=True) + parser.add_argument('--reweight', type=str, default='softmax', choices=['softmax', 'l1', 'none'], + help="edge reweighting in attention: 'softmax', 'l1', or 'none'") + parser.add_argument('--n_layers', type=int, default=4, + help='number of layers in SPIN') + parser.add_argument('--eta', type=int, default=3, + help='temporal window size eta in SPIN') + parser.add_argument('--message_layers', type=int, default=1, + help='number of message passing layers in SPIN') + + # Adam + parser.add_argument('--lr', type=float, default=8e-4, + help='learning rate') + parser.add_argument('--l2_reg', type=float, default=0.0, + help='weight decay') + + # training + parser.add_argument('--epochs', type=int, default=300, + help='training epochs') + parser.add_argument('--batch_size', type=int, default=1, + help='number of simulated histories per epoch') + + # evaluation repeats (kept for consistency with other scripts) + parser.add_argument('--rep', type=int, default=1, + help='number of repetitions in evaluation') + + args = parser.parse_args() + + # keep compatibility with the original SPIN code that expects reweight in {'softmax','l1',None} + if getattr(args, 'reweight', None) == 'none': + args.reweight = None + + return args + + +def spin_prep(y, edge_index, args): + n_samples, n_nodes, T = y.size() + T -= 1 + x = y.float().transpose(1, 2).unsqueeze(dim=3) # (samples, T + 1, nodes, 1) + u = torch.ones(n_samples, T + 1, args.u_size, dtype=torch.float32, device=args.device) # (samples, T + 1, u_size) + mask = ( + torch.tensor([[0]] * T + [[1]], device=args.device) + .expand(n_samples, n_nodes, -1, -1) + .transpose(1, 2) + ) # (samples, T + 1, nodes, 1) + ei = edge_index # keep the original behavior + return x, u, mask, ei + + +def spin_run(data, args): + bpar = b_estim(data, args) # diffusion parameter estimation + T = data.T.item() + n_nodes = data.num_nodes + n_out = data.y[:, -1].max().item() + 1 + + # train + model = SPINModel( + input_size=1, + u_size=args.u_size, + n_nodes=n_nodes, + hidden_size=args.hidden_size, + output_size=1, + temporal_self_attention=args.temporal_self_attention, + reweight=args.reweight, + n_layers=args.n_layers, + eta=args.eta, + message_layers=args.message_layers, + ).to(args.device) + + I0 = (data.y[:, 0] == SIR_STATES.I).long().sum().item() + opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2_reg) + model.train() + + pbar = trange(1, args.epochs + 1) + for epoch in pbar: + opt.zero_grad() + Y_true = diffus_gen( + T=data.T.item(), + n_nodes=data.num_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ) # (T + 1, nodes, samples) + x, u, mask, ei = spin_prep(Y_true.transpose(0, 2), data.edge_index, args) # (samples, T + 1, nodes, *) + z = model(x=x, u=u, mask=mask, edge_index=data.edge_index)[0] # (samples, T + 1, nodes, 1) + loss = (z[:, :-1].flatten() - Y_true[:-1].transpose(1, 2).transpose(0, 1).flatten()).abs().mean() + pbar.set_description(f'[epoch={epoch}] loss={loss.item():.4f}') + loss.backward() + opt.step() + + # infer + with torch.no_grad(): + model.eval() + x, u, mask, ei = spin_prep(data.y.unsqueeze(dim=0).clone(), data.edge_index, args) + z = model(x=x, u=u, mask=mask, edge_index=ei)[0] # (1, T + 1, nodes, 1) + y_pred = data.y.clone() + y_pred[:, :-1] = z[0, :-1, :, 0].clamp(0, n_out - 1).T.round().long() # (nodes, T) + return y_pred.clone() + + +def main(): + args = get_args() + + # Preserve the original default behavior (run a built-in list) when --dataset is not provided. + datasets = ( + ['heb-sir', 'ba-si', 'er-si', 'farmers-si', 'ba-sir', 'er-sir', 'covid-sir'] + if args.dataset is None + else [args.dataset] + ) + + model_fn = lambda data: spin_run(data, args) + tester = Tester(args.data_dir, args.device, model_fn) + tester.test(datasets, seed=args.seed, rep=args.rep) + tester.save(args.output) + + +if __name__ == '__main__': + main() diff --git a/spin_ms.py b/spin_ms.py new file mode 100644 index 0000000..4df7c37 --- /dev/null +++ b/spin_ms.py @@ -0,0 +1,706 @@ +# ! pip install --no-index torch-scatter==2.0.7 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-sparse==0.6.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install --no-index torch-cluster==1.5.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +#! pip install --no-index torch-spline-conv==1.2.1 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html +# ! pip install torch-geometric==2.0.4 +# ! pip install ndlib==5.1.1 + + +# ! pip install einops==0.6.0 +# ! pip install test_tube==0.7.5 + + +try: + from tsl.nn.base import StaticGraphEmbedding +except Exception: + import torch + from torch import nn + + class StaticGraphEmbedding(nn.Module): + def __init__(self, n_nodes, out_channels): + super().__init__() + self.emb = nn.Embedding(n_nodes, out_channels) + + def forward(self, token_index=None): + if token_index is None: + token_index = torch.arange(self.emb.num_embeddings, device=self.emb.weight.device) + return self.emb(token_index) +from tsl.nn.layers import PositionalEncoding +from tsl.nn.layers.norm import LayerNorm +from tsl.nn.blocks.encoders import MLP +from tsl.nn.functional import sparse_softmax +from tsl.engines import Imputer, Predictor +from tsl.ops.connectivity import weighted_degree +#from tsl.data import Batch, SpatioTemporalDataModule, ImputationDataset + +#SPINModel +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/postional_encoding.py''' +from typing import Optional + +from torch import nn + +class PositionalEncoder(nn.Module): + + def __init__(self, in_channels, out_channels, + n_layers: int = 1, + n_nodes: Optional[int] = None): + super(PositionalEncoder, self).__init__() + self.lin = nn.Linear(in_channels, out_channels) + self.activation = nn.LeakyReLU() + self.mlp = MLP(out_channels, out_channels, out_channels, + n_layers=n_layers, activation='relu') + self.positional = PositionalEncoding(out_channels) + if n_nodes is not None: + self.node_emb = StaticGraphEmbedding(n_nodes, out_channels) + else: + self.register_parameter('node_emb', None) + + def forward(self, x, node_emb=None, node_index=None): + if node_emb is None: + node_emb = self.node_emb(token_index=node_index) + # x: [b s c], node_emb: [n c] -> [b s n c] + x = self.lin(x) + x = self.activation(x.unsqueeze(-2) + node_emb) + #print('u:', tuple(x.shape), 'node_emb:', tuple(node_emb.shape))##### + out = self.mlp(x) + out = self.positional(out) + return out + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/additive_attention.py''' +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch import nn +from torch.nn import LayerNorm, functional as F +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.dense.linear import Linear +from torch_geometric.typing import Adj, OptTensor, PairTensor +from torch_scatter import scatter +from torch_scatter.utils import broadcast + + +class AdditiveAttention(MessagePassing): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + norm: bool = True, + dropout: float = 0.0, + dim: int = -2, + **kwargs): + kwargs.setdefault('aggr', 'add') + super().__init__(node_dim=dim, **kwargs) + + self.output_size = output_size + if isinstance(input_size, int): + self.src_size = self.tgt_size = input_size + else: + self.src_size, self.tgt_size = input_size + + self.msg_size = msg_size or self.output_size + self.msg_layers = msg_layers + + assert reweight in ['softmax', 'l1', None] + self.reweight = reweight + + self.root_weight = root_weight + self.dropout = dropout + + # key bias is discarded in softmax + self.lin_src = Linear(self.src_size, self.output_size, + weight_initializer='glorot', + bias_initializer='zeros') + self.lin_tgt = Linear(self.tgt_size, self.output_size, + weight_initializer='glorot', bias=False) + + if self.root_weight: + self.lin_skip = Linear(self.tgt_size, self.output_size, + bias=False) + else: + self.register_parameter('lin_skip', None) + + self.msg_nn = nn.Sequential( + nn.PReLU(init=0.2), + MLP(self.output_size, self.msg_size, self.output_size, + n_layers=self.msg_layers, dropout=self.dropout, + activation='prelu') + ) + + if self.reweight == 'softmax': + self.msg_gate = nn.Linear(self.output_size, 1, bias=False) + else: + self.msg_gate = nn.Sequential(nn.Linear(self.output_size, 1), + nn.Sigmoid()) + + if norm: + self.norm = LayerNorm(self.output_size) + else: + self.register_parameter('norm', None) + + self.reset_parameters() + + def reset_parameters(self): + self.lin_src.reset_parameters() + self.lin_tgt.reset_parameters() + if self.lin_skip is not None: + self.lin_skip.reset_parameters() + + def forward(self, x: PairTensor, edge_index: Adj, mask: OptTensor = None): + # if query/key not provided, defaults to x (e.g., for self-attention) + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + N_src, N_tgt = x_src.size(self.node_dim), x_tgt.size(self.node_dim) + + msg_src = self.lin_src(x_src) + msg_tgt = self.lin_tgt(x_tgt) + + msg = (msg_src, msg_tgt) + + # propagate_type: (msg: PairTensor, mask: OptTensor) + out = self.propagate(edge_index, msg=msg, mask=mask, + size=(N_src, N_tgt)) + + # skip connection + if self.root_weight: + out = out + self.lin_skip(x_tgt) + + if self.norm is not None: + out = self.norm(out) + + return out + + def normalize_weights(self, weights, index, num_nodes, mask=None): + # mask weights + if mask is not None: + fill_value = float("-inf") if self.reweight == 'softmax' else 0. + weights = weights.masked_fill(torch.logical_not(mask), fill_value) + # eventually reweight + if self.reweight == 'l1': + expanded_index = broadcast(index, weights, self.node_dim) + weights_sum = scatter(weights, expanded_index, self.node_dim, + dim_size=num_nodes, reduce='sum') + weights_sum = weights_sum.index_select(self.node_dim, index) + weights = weights / (weights_sum + 1e-5) + elif self.reweight == 'softmax': + weights = sparse_softmax(weights, index, num_nodes=num_nodes, + dim=self.node_dim) + return weights + + def message(self, msg_j: Tensor, msg_i: Tensor, index, size_i, + mask_j: OptTensor = None) -> Tensor: + msg = self.msg_nn(msg_j + msg_i) + gate = self.msg_gate(msg) + alpha = self.normalize_weights(gate, index, size_i, mask_j) + alpha = F.dropout(alpha, p=self.dropout, training=self.training) + out = alpha * msg + return out + + def __repr__(self) -> str: + return (f'{self.__class__.__name__}({self.output_size}, ' + f'dim={self.node_dim}, ' + f'root_weight={self.root_weight})') + + +class TemporalAdditiveAttention(AdditiveAttention): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + norm: bool = True, + dropout: float = 0.0, + **kwargs): + kwargs.setdefault('dim', 1) + super().__init__(input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + root_weight=root_weight, + reweight=reweight, + dropout=dropout, + norm=norm, + **kwargs) + + def forward(self, x: PairTensor, mask: OptTensor = None, + temporal_mask: OptTensor = None, + causal_lag: Optional[int] = None): + # x: [b s * c] query: [b l * c] key: [b s * c] + # mask: [b s * c] temporal_mask: [l s] + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + l, s = x_tgt.size(self.node_dim), x_src.size(self.node_dim) + i = torch.arange(l, dtype=torch.long, device=x_src.device) + j = torch.arange(s, dtype=torch.long, device=x_src.device) + + # compute temporal index, from j to i + if temporal_mask is None and isinstance(causal_lag, int): + temporal_mask = tuple(torch.tril_indices(l, l, offset=-causal_lag, + device=x_src.device)) + if temporal_mask is not None: + assert temporal_mask.size() == (l, s) + i, j = torch.meshgrid(i, j) + edge_index = torch.stack((j[temporal_mask], i[temporal_mask])) + else: + edge_index = torch.cartesian_prod(j, i).T + + return super(TemporalAdditiveAttention, self).forward(x, edge_index, + mask=mask) + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/layers/temporal_graph_additive_attention.py''' +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.dense.linear import Linear +from torch_geometric.typing import Adj, OptTensor, OptPairTensor + +class TemporalGraphAdditiveAttention(MessagePassing): + def __init__(self, input_size: Union[int, Tuple[int, int]], + output_size: int, + msg_size: Optional[int] = None, + msg_layers: int = 1, + root_weight: bool = True, + reweight: Optional[str] = None, + temporal_self_attention: bool = True, + mask_temporal: bool = True, + mask_spatial: bool = True, + norm: bool = True, + dropout: float = 0., + **kwargs): + kwargs.setdefault('aggr', 'add') + super(TemporalGraphAdditiveAttention, self).__init__(node_dim=-2, + **kwargs) + + # store dimensions + if isinstance(input_size, int): + self.src_size = self.tgt_size = input_size + else: + self.src_size, self.tgt_size = input_size + self.output_size = output_size + self.msg_size = msg_size or self.output_size + + self.mask_temporal = mask_temporal + self.mask_spatial = mask_spatial + + self.root_weight = root_weight + self.dropout = dropout + + if temporal_self_attention: + self.self_attention = TemporalAdditiveAttention( + input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + reweight=reweight, + dropout=dropout, + root_weight=False, + norm=False + ) + else: + self.register_parameter('self_attention', None) + + self.cross_attention = TemporalAdditiveAttention(input_size=input_size, + output_size=output_size, + msg_size=msg_size, + msg_layers=msg_layers, + reweight=reweight, + dropout=dropout, + root_weight=False, + norm=False) + + if self.root_weight: + self.lin_skip = Linear(self.tgt_size, self.output_size, + bias_initializer='zeros') + else: + self.register_parameter('lin_skip', None) + + if norm: + self.norm = LayerNorm(output_size) + else: + self.register_parameter('norm', None) + + self.reset_parameters() + + def reset_parameters(self): + self.cross_attention.reset_parameters() + if self.self_attention is not None: + self.self_attention.reset_parameters() + if self.lin_skip is not None: + self.lin_skip.reset_parameters() + if self.norm is not None: + self.norm.reset_parameters() + + def forward(self, x: OptPairTensor, + edge_index: Adj, edge_weight: OptTensor = None, + mask: OptTensor = None): + # inputs: [batch, steps, nodes, channels] + if isinstance(x, Tensor): + x_src = x_tgt = x + else: + x_src, x_tgt = x + x_tgt = x_tgt if x_tgt is not None else x_src + + n_src, n_tgt = x_src.size(-2), x_tgt.size(-2) + + # propagate query, key and value + #print('src:', x_src.shape, 'tgt:', x_tgt.shape, 'ei:', edge_index.shape, 'mask:', mask.shape, f'mask_spatial={self.mask_spatial}') + out = self.propagate(x=(x_src, x_tgt), + edge_index=edge_index, edge_weight=edge_weight, + mask=mask if self.mask_spatial else None, + size=(n_src, n_tgt)) + + if self.self_attention is not None: + s, l = x_src.size(1), x_tgt.size(1) + if s == l: + attn_mask = ~torch.eye(l, l, dtype=torch.bool, + device=x_tgt.device) + else: + attn_mask = None + temp = self.self_attention(x=(x_src, x_tgt), + mask=mask if self.mask_temporal else None, + temporal_mask=attn_mask) + out = out + temp + + # skip connection + if self.root_weight: + out = out + self.lin_skip(x_tgt) + + if self.norm is not None: + out = self.norm(out) + + return out + + def message(self, x_i: Tensor, x_j: Tensor, + edge_weight: OptTensor, mask_j: OptTensor) -> Tensor: + # [batch, steps, edges, channels] + + out = self.cross_attention((x_j, x_i), mask=mask_j) + #print('out:', out.shape) + + if edge_weight is not None: + out = out * edge_weight.view(-1, 1) + return out + +'''https://github.com/Graph-Machine-Learning-Group/spin/blob/main/spin/models/spin.py''' +from typing import Optional + +import torch +from torch import nn, Tensor +from torch.nn import LayerNorm +from torch_geometric.typing import OptTensor + +class SPINModel(nn.Module): + + def __init__(self, input_size: int, + hidden_size: int, + n_nodes: int, + u_size: Optional[int] = None, + output_size: Optional[int] = None, + temporal_self_attention: bool = True, + reweight: Optional[str] = 'softmax', + n_layers: int = 4, + eta: int = 3, + message_layers: int = 1): + super(SPINModel, self).__init__() + + u_size = u_size or input_size + output_size = output_size or input_size + self.n_nodes = n_nodes + self.n_layers = n_layers + self.eta = eta + self.temporal_self_attention = temporal_self_attention + + self.u_enc = PositionalEncoder(in_channels=u_size, + out_channels=hidden_size, + n_layers=2, + n_nodes=n_nodes) + + self.h_enc = MLP(input_size, hidden_size, n_layers=2) + self.h_norm = LayerNorm(hidden_size) + + self.valid_emb = StaticGraphEmbedding(n_nodes, hidden_size) + self.mask_emb = StaticGraphEmbedding(n_nodes, hidden_size) + + self.x_skip = nn.ModuleList() + self.encoder, self.readout = nn.ModuleList(), nn.ModuleList() + for l in range(n_layers): + x_skip = nn.Linear(input_size, hidden_size) + encoder = TemporalGraphAdditiveAttention( + input_size=hidden_size, + output_size=hidden_size, + msg_size=hidden_size, + msg_layers=message_layers, + temporal_self_attention=temporal_self_attention, + reweight=reweight, + mask_temporal=True, + mask_spatial=l < eta, + norm=True, + root_weight=True, + dropout=0.0 + ) + readout = MLP(hidden_size, hidden_size, output_size, + n_layers=2) + self.x_skip.append(x_skip) + self.encoder.append(encoder) + self.readout.append(readout) + + def forward(self, x: Tensor, u: Tensor, mask: Tensor, + edge_index: Tensor, edge_weight: OptTensor = None, + node_index: OptTensor = None, target_nodes: OptTensor = None): + if target_nodes is None: + target_nodes = slice(None) + + # Whiten missing values + x = x * mask + + # POSITIONAL ENCODING ################################################# + # Obtain spatio-temporal positional encoding for every node-step pair # + # in both observed and target sets. Encoding are obtained by jointly # + # processing node and time positional encoding. # + + # Build (node, timestamp) encoding + q = self.u_enc(u, node_index=node_index) + # Condition value on key + h = self.h_enc(x) + q + + # ENCODER ############################################################# + # Obtain representations h^i_t for every (i, t) node-step pair by # + # only taking into account valid data in representation set. # + + # Replace H in missing entries with queries Q + h = torch.where(mask.bool(), h, q) + # Normalize features + h = self.h_norm(h) + + imputations = [] + + for l in range(self.n_layers): + if l == self.eta: + # Condition H on two different embeddings to distinguish + # valid values from masked ones + valid = self.valid_emb(token_index=node_index) + masked = self.mask_emb(token_index=node_index) + h = torch.where(mask.bool(), h + valid, h + masked) + # Masked Temporal GAT for encoding representation + h = h + self.x_skip[l](x) * mask # skip connection for valid x + #print(f'l={l}', 'h:', tuple(h.shape), 'x:', tuple(x.shape), 'mask:', tuple(mask.shape), 'ei:', edge_index) + h = self.encoder[l](h, edge_index, mask=mask) + # Read from H to get imputations + target_readout = self.readout[l](h[..., target_nodes, :]) + imputations.append(target_readout) + + # Get final layer imputations + x_hat = imputations.pop(-1) + + return x_hat, imputations + + +import argparse +import torch + +from inc.diffus import * + +from inc.test import Tester # type: ignore + + +def get_args(): + parser = argparse.ArgumentParser() + + # -------------------- common experiment args -------------------- + parser.add_argument('--dataset', type=str, required=True, help='dataset name') + parser.add_argument('--seed', type=int, default=123456789, help='random seed') + parser.add_argument('--data_dir', type=str, default='input', help='dataset folder') + parser.add_argument('--output', type=str, default='output/spin.pt', help='output file name') + parser.add_argument( + '--device', type=str, + default='cuda' if torch.cuda.is_available() else 'cpu', + help='torch device, e.g., "cuda", "cuda:0", or "cpu"' + ) + parser.add_argument( + '--obs_time', '--snapshot', dest='obs_time', + type=str, default='', + help='extra observed snapshot times, comma-separated, e.g., "5,7,9". ' + 'Final time T will always be added automatically.' + ) + + # -------------------- diffusion parameter estimation -------------------- + parser.add_argument('--b_pI0', type=float, default=1e-3, + help='initial infection rate in diffusion parameter estimation') + parser.add_argument('--b_pR0', type=float, default=1e-3, + help='initial recovery rate in diffusion parameter estimation') + parser.add_argument('--b_steps', type=int, default=500, + help='optimization steps in diffusion parameter estimation') + parser.add_argument('--b_lr', type=float, default=3e-3, + help='learning rate in diffusion parameter estimation') + + # -------------------- SPIN model hyperparams -------------------- + parser.add_argument('--u_size', type=int, default=1, help='u feature size (kept as 1)') + parser.add_argument('--hidden_size', type=int, default=32, help='hidden size') + parser.add_argument('--reweight', type=str, default='softmax', + choices=['softmax', 'l1', 'none'], + help='attention reweighting: softmax | l1 | none') + parser.add_argument('--n_layers', type=int, default=4, help='number of SPIN layers') + parser.add_argument('--eta', type=int, default=3, help='layers before enabling spatial mask-off') + parser.add_argument('--message_layers', type=int, default=1, help='message MLP depth inside attention') + + parser.add_argument( + '--no_temporal_self_attention', action='store_false', + dest='temporal_self_attention', + help='disable temporal self-attention (default: enabled)' + ) + parser.set_defaults(temporal_self_attention=True) + + # -------------------- training hyperparams -------------------- + parser.add_argument('--lr', type=float, default=8e-4, help='Adam lr') + parser.add_argument('--l2_reg', type=float, default=0.0, help='Adam weight_decay') + parser.add_argument('--epochs', type=int, default=300, help='training epochs') + parser.add_argument('--batch_size', type=int, default=1, help='synthetic batch size per epoch') + + args = parser.parse_args() + args.device = torch.device(args.device) + + if args.reweight == 'none': + args.reweight = None + + return args + + +def parse_obs_time(obs_time_str: str, T: int): + """Parse comma-separated observed snapshot times and always include final time T.""" + times = [] + if obs_time_str: + for part in str(obs_time_str).split(','): + part = part.strip() + if part == '': + continue + times.append(int(part)) + # keep within [0, T] + times = [t for t in times if 0 <= t <= T] + if T not in times: + times.append(T) + return sorted(set(times)) + + +def build_obs_mask(n_samples: int, n_nodes: int, T: int, obs_time, device): + """Return float mask with shape (samples, T+1, nodes, 1), 1=observed, 0=missing.""" + mask = torch.zeros((n_samples, T + 1, n_nodes, 1), dtype=torch.float32, device=device) + if len(obs_time) > 0: + mask[:, obs_time, :, :] = 1.0 + return mask + + +def spin_prep(y, edge_index, obs_time, args): + """ + y: (samples, nodes, T+1) + return: + x: (samples, T+1, nodes, 1) + u: (samples, T+1, u_size) + mask: (samples, T+1, nodes, 1) + ei: edge_index + """ + n_samples, n_nodes, Tp1 = y.size() + T = Tp1 - 1 + + x = y.float().transpose(1, 2).unsqueeze(dim=3) # (samples, T+1, nodes, 1) + u = torch.ones(n_samples, T + 1, args.u_size, dtype=torch.float32, device=args.device) + mask = build_obs_mask(n_samples, n_nodes, T, obs_time, device=args.device) + + # NOTE: SPIN implementation here uses a single static edge_index shared across all steps. + ei = edge_index + return x, u, mask, ei + + +def spin_run(data, args): + """Train SPIN on synthetic histories and infer missing diffusion history.""" + T = int(data.T.item()) + obs_time = parse_obs_time(args.obs_time, T) + data.obs_ts = obs_time + bpar = b_estim(data, args, obs_time=obs_time) + + n_nodes = int(data.num_nodes) + n_out = int(data.y[:, -1].max().item() + 1) + + # -------------------- train -------------------- + model = SPINModel( + input_size=1, + u_size=args.u_size, + n_nodes=n_nodes, + hidden_size=args.hidden_size, + output_size=1, + temporal_self_attention=args.temporal_self_attention, + reweight=args.reweight, + n_layers=args.n_layers, + eta=args.eta, + message_layers=args.message_layers, + ).to(args.device) + + I0 = int((data.y[:, 0] == SIR_STATES.I).long().sum().item()) + opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2_reg) + + model.train() + pbar = trange(1, args.epochs + 1) + for epoch in pbar: + opt.zero_grad() + + # synthetic history + Y_true = diffus_gen( + T=T, + n_nodes=n_nodes, + edge_index=data.edge_index, + I0=I0, + n_samples=args.batch_size, + pI=bpar.pI, + pR=bpar.pR, + ) # (T+1, nodes, samples) + + x, u, mask, ei = spin_prep(Y_true.transpose(0, 2), data.edge_index, obs_time, args) + z = model(x=x, u=u, mask=mask, edge_index=ei)[0] # (samples, T+1, nodes, 1) + y_true = Y_true.permute(2, 0, 1).unsqueeze(-1) # (samples, T+1, nodes, 1) + unobs = ~mask.bool() + loss = (z - y_true).abs() + loss = loss[unobs].mean() if unobs.any() else loss.mean() + + pbar.set_description(f'[epoch={epoch}] loss={loss.item():.4f}') + loss.backward() + opt.step() + + # -------------------- infer -------------------- + with torch.no_grad(): + model.eval() + + x, u, mask, ei = spin_prep(data.y.unsqueeze(0).clone(), data.edge_index, obs_time, args) + z = model(x=x, u=u, mask=mask, edge_index=ei)[0] # (1, T+1, nodes, 1) + + y_pred = z[0, :, :, 0].clamp(0, n_out - 1).transpose(0, 1).round().long() # (nodes, T+1) + for t in obs_time: + y_pred[:, t] = data.y[:, t] + + return y_pred.detach().clone() + + +def main(): + args = get_args() + + def _run(data): + return spin_run(data, args) + + tester = Tester(args.data_dir, args.device, _run) + tester.test([args.dataset], seed=args.seed, rep=1) + tester.save(args.output) + + +if __name__ == '__main__': + main()