Add src/data.py

src/data.py

@@ -0,0 +1,270 @@
+"""Data generation and loading for experiments."""
+import numpy as np
+from typing import List, Tuple, Dict, Optional
+from collections import defaultdict
+from src.graph_utils import (
+    generate_bounded_degree_graph, generate_erdos_renyi_graph,
+    generate_power_law_graph, build_adjacency
+)
+
+
+def generate_gamma_poisson_data(N, M, K, graph_type, avg_degree,
+                                 count_scale, a0, b0, c0, d0,
+                                 seed=0, keep_zeros=False):
+    """Generate synthetic Gamma-Poisson matrix factorization data."""
+    rng = np.random.RandomState(seed)
+    U_true = rng.gamma(a0, 1.0 / b0, size=(N, K))
+    V_true = rng.gamma(c0, 1.0 / d0, size=(M, K))
+    
+    if graph_type == 'bounded_degree':
+        graph_edges = generate_bounded_degree_graph(N, M, avg_degree, seed)
+    elif graph_type == 'erdos_renyi':
+        graph_edges = generate_erdos_renyi_graph(N, M, avg_degree, seed)
+    elif graph_type == 'power_law':
+        graph_edges = generate_power_law_graph(N, M, avg_degree, seed)
+    else:
+        raise ValueError(f"Unknown graph type: {graph_type}")
+    
+    edges = []
+    for i, j in graph_edges:
+        rate = count_scale * np.dot(U_true[i], V_true[j])
+        x = rng.poisson(max(rate, 1e-10))
+        if x > 0 or keep_zeros:
+            edges.append((i, j, int(x)))
+    
+    return edges, U_true, V_true, graph_edges
+
+
+def generate_gaussian_gaussian_data(N, M, K, graph_type, avg_degree,
+                                     sigma_U, sigma_V, sigma_x, seed=0):
+    """Generate synthetic Gaussian-Gaussian MF data."""
+    rng = np.random.RandomState(seed)
+    U_true = rng.normal(0, sigma_U, size=(N, K))
+    V_true = rng.normal(0, sigma_V, size=(M, K))
+    
+    if graph_type == 'bounded_degree':
+        graph_edges = generate_bounded_degree_graph(N, M, avg_degree, seed)
+    elif graph_type == 'erdos_renyi':
+        graph_edges = generate_erdos_renyi_graph(N, M, avg_degree, seed)
+    elif graph_type == 'power_law':
+        graph_edges = generate_power_law_graph(N, M, avg_degree, seed)
+    else:
+        raise ValueError(f"Unknown graph type: {graph_type}")
+    
+    edges = []
+    for i, j in graph_edges:
+        mean = np.dot(U_true[i], V_true[j])
+        x = rng.normal(mean, sigma_x)
+        edges.append((i, j, float(x)))
+    
+    return edges, U_true, V_true, graph_edges
+
+
+def generate_gaussian_gamma_data(N, M, K, graph_type, avg_degree,
+                                  a0, b0, c0, d0, sigma_x, seed=0):
+    """Generate synthetic Gaussian likelihood + Gamma prior data."""
+    rng = np.random.RandomState(seed)
+    U_true = rng.gamma(a0, 1.0 / b0, size=(N, K))
+    V_true = rng.gamma(c0, 1.0 / d0, size=(M, K))
+    
+    if graph_type == 'bounded_degree':
+        graph_edges = generate_bounded_degree_graph(N, M, avg_degree, seed)
+    elif graph_type == 'erdos_renyi':
+        graph_edges = generate_erdos_renyi_graph(N, M, avg_degree, seed)
+    elif graph_type == 'power_law':
+        graph_edges = generate_power_law_graph(N, M, avg_degree, seed)
+    else:
+        raise ValueError(f"Unknown graph type: {graph_type}")
+    
+    edges = []
+    for i, j in graph_edges:
+        mean = np.dot(U_true[i], V_true[j])
+        x = rng.normal(mean, sigma_x)
+        edges.append((i, j, float(x)))
+    
+    return edges, U_true, V_true, graph_edges
+
+
+def load_lastfm_data(max_users=2000, max_items=2000, max_edges=100000,
+                     min_user_degree=5, min_item_degree=5, max_count=100, seed=42):
+    """Load Last.fm user-artist counts from HF dataset."""
+    from datasets import load_dataset
+    
+    print("Loading Last.fm dataset...")
+    ds = load_dataset("matthewfranglen/lastfm-1k", split="train")
+    
+    user_artist_counts = defaultdict(lambda: defaultdict(int))
+    for row in ds:
+        uid = row['user_index']
+        aid = row['artist_index']
+        user_artist_counts[uid][aid] += 1
+    
+    user_degrees = {u: len(v) for u, v in user_artist_counts.items()}
+    valid_users = [u for u, d in user_degrees.items() if d >= min_user_degree]
+    
+    item_degree = defaultdict(int)
+    for u in valid_users:
+        for a in user_artist_counts[u]:
+            item_degree[a] += 1
+    valid_items = set(a for a, d in item_degree.items() if d >= min_item_degree)
+    
+    rng = np.random.RandomState(seed)
+    valid_users = sorted(valid_users)
+    if len(valid_users) > max_users:
+        valid_users = list(rng.choice(valid_users, max_users, replace=False))
+    valid_users_set = set(valid_users)
+    
+    all_items = set()
+    for u in valid_users:
+        for a in user_artist_counts[u]:
+            if a in valid_items:
+                all_items.add(a)
+    all_items = sorted(all_items)
+    if len(all_items) > max_items:
+        all_items = list(rng.choice(all_items, max_items, replace=False))
+    valid_items_set = set(all_items)
+    
+    user_map = {u: idx for idx, u in enumerate(sorted(valid_users_set))}
+    item_map = {a: idx for idx, a in enumerate(sorted(valid_items_set))}
+    
+    edges = []
+    for u in valid_users_set:
+        for a, count in user_artist_counts[u].items():
+            if a in valid_items_set:
+                c = min(count, max_count)
+                if c > 0:
+                    edges.append((user_map[u], item_map[a], int(c)))
+    
+    if len(edges) > max_edges:
+        indices = rng.choice(len(edges), max_edges, replace=False)
+        edges = [edges[i] for i in indices]
+    
+    N = len(user_map)
+    M = len(item_map)
+    preprocessing = {
+        'dataset': 'matthewfranglen/lastfm-1k', 'N': N, 'M': M,
+        'n_edges': len(edges), 'max_count': max_count, 'seed': seed,
+    }
+    
+    print(f"Last.fm loaded: N={N}, M={M}, edges={len(edges)}")
+    return edges, N, M, preprocessing
+
+
+def load_movielens_data(mode='rating_count', max_users=2000, max_items=2000,
+                        max_edges=100000, min_user_degree=5, min_item_degree=5, seed=42):
+    """Load MovieLens ratings from HF dataset."""
+    from datasets import load_dataset
+    
+    print("Loading MovieLens dataset...")
+    ds = load_dataset("ashraq/movielens_ratings", split="train")
+    
+    rng = np.random.RandomState(seed)
+    user_item_ratings = defaultdict(dict)
+    for row in ds:
+        uid = row['user_id']
+        mid = row['movie_id']
+        rating = row['rating']
+        user_item_ratings[uid][mid] = rating
+    
+    user_degrees = {u: len(v) for u, v in user_item_ratings.items()}
+    valid_users = [u for u, d in user_degrees.items() if d >= min_user_degree]
+    
+    item_degree = defaultdict(int)
+    for u in valid_users:
+        for m in user_item_ratings[u]:
+            item_degree[m] += 1
+    valid_items = set(m for m, d in item_degree.items() if d >= min_item_degree)
+    
+    if len(valid_users) > max_users:
+        valid_users = list(rng.choice(valid_users, max_users, replace=False))
+    valid_users_set = set(valid_users)
+    
+    all_items = set()
+    for u in valid_users_set:
+        for m in user_item_ratings[u]:
+            if m in valid_items:
+                all_items.add(m)
+    all_items = sorted(all_items)
+    if len(all_items) > max_items:
+        all_items = list(rng.choice(all_items, max_items, replace=False))
+    valid_items_set = set(all_items)
+    
+    user_map = {u: idx for idx, u in enumerate(sorted(valid_users_set))}
+    item_map = {m: idx for idx, m in enumerate(sorted(valid_items_set))}
+    
+    edges = []
+    for u in valid_users_set:
+        for m, rating in user_item_ratings[u].items():
+            if m in valid_items_set:
+                if mode == 'rating_count':
+                    x = int(np.ceil(rating))
+                elif mode == 'binary':
+                    x = 1
+                else:
+                    raise ValueError(f"Unknown mode: {mode}")
+                if x > 0:
+                    edges.append((user_map[u], item_map[m], x))
+    
+    if len(edges) > max_edges:
+        indices = rng.choice(len(edges), max_edges, replace=False)
+        edges = [edges[i] for i in indices]
+    
+    N = len(user_map)
+    M = len(item_map)
+    preprocessing = {
+        'dataset': 'ashraq/movielens_ratings', 'mode': mode,
+        'N': N, 'M': M, 'n_edges': len(edges), 'seed': seed,
+    }
+    
+    print(f"MovieLens ({mode}) loaded: N={N}, M={M}, edges={len(edges)}")
+    return edges, N, M, preprocessing
+
+
+def sample_deletions(edges, user_to_items, item_to_users, num_deletions, seed=0):
+    """Sample deletions with 25% each: random, high-count, hub-adjacent, low-degree."""
+    rng = np.random.RandomState(seed)
+    n_per_type = num_deletions // 4
+    remainder = num_deletions - 4 * n_per_type
+    
+    counts = np.array([e[2] for e in edges], dtype=float)
+    
+    user_degrees = defaultdict(int)
+    item_degrees = defaultdict(int)
+    for i, j, x in edges:
+        user_degrees[i] += 1
+        item_degrees[j] += 1
+    
+    hub_scores = np.array([max(user_degrees[e[0]], item_degrees[e[1]]) for e in edges], dtype=float)
+    low_deg_scores = np.array([min(user_degrees[e[0]], item_degrees[e[1]]) for e in edges], dtype=float)
+    
+    sampled = []
+    used = set()
+    
+    def _sample(scores, n, dtype, high=True):
+        avail = [i for i in range(len(edges)) if i not in used]
+        if not avail or n <= 0:
+            return
+        sc = scores[avail]
+        if high:
+            ranked = np.argsort(-sc)
+        else:
+            ranked = np.argsort(sc)
+        pool = ranked[:min(len(avail), max(n * 3, 20))]
+        chosen = rng.choice(pool, size=min(n, len(pool)), replace=False)
+        for idx in chosen:
+            eidx = avail[idx]
+            used.add(eidx)
+            sampled.append((edges[eidx], dtype))
+    
+    # Random
+    avail = list(range(len(edges)))
+    rng.shuffle(avail)
+    for idx in avail[:n_per_type + remainder]:
+        used.add(idx)
+        sampled.append((edges[idx], 'random'))
+    
+    _sample(counts, n_per_type, 'high_count', high=True)
+    _sample(hub_scores, n_per_type, 'hub_adjacent', high=True)
+    _sample(low_deg_scores, n_per_type, 'low_degree', high=False)
+    
+    return sampled