Add scripts/run_real.py

scripts/run_real.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""Run real-data experiments."""
+import os
+import sys
+import json
+import time
+import argparse
+import yaml
+import numpy as np
+from datetime import datetime
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.data import load_lastfm_data, load_movielens_data, sample_deletions
+from src.model import PoissonGammaVI
+from src.graph_utils import build_adjacency, compute_graph_stats
+from src.metrics import (compute_all_metrics, compute_deletion_influence_by_distance,
+                         fit_exponential_decay, compute_local_error, compute_chi_poisson_gamma,
+                         compute_gradient_interference)
+from src.unlearning import one_step_downdate_poisson_gamma
+from src.utils import generate_run_id, generate_config_id, save_jsonl, ensure_dir
+
+
+def run_real_dataset(dataset_name, edges, N, M, preprocessing, config):
+    """Run deletion experiments on a real dataset."""
+    K_values = config.get('K_values', [5, 10])
+    num_deletions = config.get('num_deletions', 50)
+    radii = config.get('radii', [1, 2, 3, 4])
+    prior = config.get('prior', {})
+    a0 = prior.get('a0', 0.3)
+    b0 = prior.get('b0', 1.0)
+    c0 = prior.get('c0', 0.3)
+    d0 = prior.get('d0', 1.0)
+    max_iter = config.get('max_iter', 300)
+    tol = config.get('tol', 1e-4)
+    seed = config.get('seed', 42)
+    
+    all_records = []
+    
+    for K in K_values:
+        print(f"\n  K={K}")
+        run_id = generate_run_id()
+        config_id = generate_config_id({**config, 'K': K, 'dataset': dataset_name})
+        
+        model = PoissonGammaVI(N, M, K, a0, b0, c0, d0, max_iter=max_iter, tol=tol, seed=seed)
+        
+        print(f"    Fitting full model...")
+        t0 = time.time()
+        full_result = model.fit_full(edges)
+        t_full = time.time() - t0
+        full_params = full_result.params
+        print(f"    Full fit: {full_result.n_iterations} iters, {t_full:.1f}s")
+        
+        user_to_items, item_to_users, edge_dict = build_adjacency(edges, N, M)
+        deletion_samples = sample_deletions(edges, user_to_items, item_to_users, num_deletions, seed=seed)
+        
+        print(f"    Running {len(deletion_samples)} deletions...")
+        
+        for del_idx, (edge_to_del, del_type) in enumerate(deletion_samples):
+            if del_idx % 10 == 0:
+                print(f"      Deletion {del_idx+1}/{len(deletion_samples)}")
+            
+            i_del, j_del, x_del = edge_to_del
+            
+            # Exact
+            exact_result = model.fit_without_edge(edges, edge_to_del, init_params=full_params)
+            exact_params = exact_result.params
+            
+            # Local
+            local_results = {}
+            local_params = {}
+            for R in radii:
+                lr = model.fit_local(edges, edge_to_del, R, init_params=full_params)
+                local_results[R] = lr
+                local_params[R] = lr.params
+            
+            # Warm-start
+            ws_result = model.fit_warm_start_global(edges, edge_to_del, init_params=full_params)
+            
+            # One-step
+            os_result = one_step_downdate_poisson_gamma(
+                edges, edge_to_del, full_params, N, M, K, a0, b0, c0, d0)
+            
+            # Metrics
+            model_kwargs = {'a0': a0, 'b0': b0, 'c0': c0, 'd0': d0}
+            metrics = compute_all_metrics(
+                full_params, exact_params, local_params,
+                ws_result.params, os_result.params,
+                edge_to_del, edges, N, M, K,
+                'poisson_gamma', model=model, radii=radii,
+                model_kwargs=model_kwargs)
+            
+            record = {
+                'run_id': run_id,
+                'config_id': config_id,
+                'dataset_type': 'real',
+                'dataset_name': dataset_name,
+                'model_family': 'poisson_gamma',
+                'inference_type': 'vi',
+                'likelihood': 'poisson',
+                'prior': 'gamma',
+                'N': N, 'M': M, 'K': K,
+                'n_edges': len(edges),
+                'deletion_edge': [int(i_del), int(j_del), float(x_del)],
+                'deletion_type': del_type,
+                'deletion_index': del_idx,
+                'runtime_full': t_full,
+                'runtime_exact': exact_result.runtime_sec,
+                'runtime_warm_start': ws_result.runtime_sec,
+                'runtime_one_step': os_result.runtime_sec,
+                'exact_converged': exact_result.converged,
+                'a0': a0, 'b0': b0, 'c0': c0, 'd0': d0,
+            }
+            
+            for R in radii:
+                record[f'runtime_local_R{R}'] = local_results[R].runtime_sec
+                record[f'local_R{R}_converged'] = local_results[R].converged
+            
+            record.update(metrics)
+            
+            if 'influence_by_distance' in record:
+                for d_str, val in record['influence_by_distance'].items():
+                    record[f'influence_d{d_str}'] = val
+            
+            all_records.append(record)
+    
+    return all_records
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, default='config/real_data.yaml')
+    parser.add_argument('--datasets', nargs='*', default=None)
+    args = parser.parse_args()
+    
+    with open(args.config) as f:
+        real_cfg = yaml.safe_load(f)
+    
+    output_dir = ensure_dir('results/raw')
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_file = os.path.join(output_dir, f'real_{timestamp}.jsonl')
+    
+    datasets_to_run = args.datasets or list(real_cfg.keys())
+    
+    for ds_name in datasets_to_run:
+        if ds_name not in real_cfg:
+            print(f"Unknown dataset config: {ds_name}")
+            continue
+        
+        ds_cfg = real_cfg[ds_name]
+        print(f"\n{'='*60}")
+        print(f"Dataset: {ds_name}")
+        print(f"{'='*60}")
+        
+        # Load data
+        if 'lastfm' in ds_name:
+            edges, N, M, preproc = load_lastfm_data(
+                max_users=ds_cfg.get('max_users', 1000),
+                max_items=ds_cfg.get('max_items', 1000),
+                max_edges=ds_cfg.get('max_edges', 50000),
+                min_user_degree=ds_cfg.get('min_user_degree', 5),
+                min_item_degree=ds_cfg.get('min_item_degree', 5),
+                max_count=ds_cfg.get('max_count', 100),
+                seed=ds_cfg.get('seed', 42))
+        elif 'movielens' in ds_name:
+            mode = ds_cfg.get('mode', 'rating_count')
+            edges, N, M, preproc = load_movielens_data(
+                mode=mode,
+                max_users=ds_cfg.get('max_users', 1000),
+                max_items=ds_cfg.get('max_items', 1000),
+                max_edges=ds_cfg.get('max_edges', 50000),
+                min_user_degree=ds_cfg.get('min_user_degree', 5),
+                min_item_degree=ds_cfg.get('min_item_degree', 5),
+                seed=ds_cfg.get('seed', 42))
+        else:
+            print(f"  Unsupported dataset: {ds_name}")
+            continue
+        
+        # Save preprocessing
+        preproc_dir = ensure_dir('results/reports')
+        with open(os.path.join(preproc_dir, f'dataset_card_{ds_name}.json'), 'w') as f:
+            json.dump(preproc, f, indent=2)
+        
+        graph_stats = compute_graph_stats([(e[0], e[1]) for e in edges], N, M)
+        print(f"  Graph stats: {json.dumps(graph_stats, indent=2)}")
+        
+        records = run_real_dataset(ds_name, edges, N, M, preproc, ds_cfg)
+        save_jsonl(records, output_file)
+        print(f"  Saved {len(records)} records for {ds_name}")
+    
+    print(f"\nOutput: {output_file}")
+
+
+if __name__ == '__main__':
+    main()