scripts/run_real.py

#!/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()