Add src/metrics.py

src/metrics.py

@@ -0,0 +1,357 @@
+"""Metrics computation for unlearning experiments."""
+import numpy as np
+from scipy import stats
+from scipy.special import digamma, polygamma
+from typing import Dict, List, Tuple, Optional
+from collections import defaultdict
+
+from src.graph_utils import build_adjacency, get_deletion_neighborhood, get_blocks_at_distance
+
+
+# ============================================================
+# Parameter distance / influence
+# ============================================================
+
+def compute_block_param_vector(params: dict, node_type: str, idx: int, model_family: str) -> np.ndarray:
+    """Extract flat parameter vector for a single block."""
+    if model_family == 'poisson_gamma':
+        if node_type == 'user':
+            return np.concatenate([params['a'][idx], params['b'][idx]])
+        else:
+            return np.concatenate([params['c'][idx], params['d'][idx]])
+    elif model_family == 'gaussian_gaussian':
+        if node_type == 'user':
+            return np.concatenate([params['m_U'][idx], params['s_U'][idx]])
+        else:
+            return np.concatenate([params['m_V'][idx], params['s_V'][idx]])
+    elif model_family == 'gaussian_gamma_map':
+        if node_type == 'user':
+            return params['alpha'][idx].copy()
+        else:
+            return params['beta'][idx].copy()
+    else:
+        raise ValueError(f"Unknown model family: {model_family}")
+
+
+def compute_all_param_vector(params: dict, model_family: str) -> np.ndarray:
+    """Flatten all parameters into a single vector."""
+    if model_family == 'poisson_gamma':
+        return np.concatenate([params['a'].ravel(), params['b'].ravel(), 
+                               params['c'].ravel(), params['d'].ravel()])
+    elif model_family == 'gaussian_gaussian':
+        return np.concatenate([params['m_U'].ravel(), params['s_U'].ravel(),
+                               params['m_V'].ravel(), params['s_V'].ravel()])
+    elif model_family == 'gaussian_gamma_map':
+        return np.concatenate([params['alpha'].ravel(), params['beta'].ravel()])
+    else:
+        raise ValueError(f"Unknown model family: {model_family}")
+
+
+def compute_deletion_influence_by_distance(full_params, exact_params, edge_to_remove,
+                                            edges, N, M, model_family, max_radius=6):
+    """Compute deletion influence Delta_u(z) = ||lambda_u^* - lambda_u^{\\z}||
+    grouped by graph distance from seed set."""
+    user_to_items, item_to_users, _ = build_adjacency(edges, N, M)
+    distances = get_deletion_neighborhood(edge_to_remove, user_to_items, item_to_users, 
+                                           N, M, max_radius)
+    blocks_by_dist = get_blocks_at_distance(distances, N)
+    
+    influence_by_dist = {}
+    for dist, blocks in blocks_by_dist.items():
+        influences = []
+        for node_type, idx in blocks:
+            v_full = compute_block_param_vector(full_params, node_type, idx, model_family)
+            v_exact = compute_block_param_vector(exact_params, node_type, idx, model_family)
+            delta = np.linalg.norm(v_full - v_exact)
+            influences.append(delta)
+        influence_by_dist[dist] = {
+            'mean': float(np.mean(influences)),
+            'std': float(np.std(influences)),
+            'median': float(np.median(influences)),
+            'max': float(np.max(influences)),
+            'n_blocks': len(influences),
+            'values': [float(v) for v in influences],
+        }
+    
+    return influence_by_dist
+
+
+def fit_exponential_decay(influence_by_dist: dict, min_shells: int = 3, eps: float = 1e-12):
+    """Fit log(Delta(r)) = alpha - mu * r."""
+    distances = sorted(influence_by_dist.keys())
+    r_vals = []
+    log_vals = []
+    
+    for r in distances:
+        mean_inf = influence_by_dist[r]['mean']
+        if mean_inf > eps and influence_by_dist[r]['n_blocks'] >= 2:
+            r_vals.append(r)
+            log_vals.append(np.log(mean_inf + eps))
+    
+    if len(r_vals) < min_shells:
+        return {
+            'mu_emp': None,
+            'intercept': None,
+            'r_squared': None,
+            'n_shells': len(r_vals),
+            'valid': False,
+        }
+    
+    r_arr = np.array(r_vals, dtype=float)
+    log_arr = np.array(log_vals)
+    
+    slope, intercept, r_value, p_value, std_err = stats.linregress(r_arr, log_arr)
+    
+    return {
+        'mu_emp': float(-slope),
+        'intercept': float(intercept),
+        'r_squared': float(r_value ** 2),
+        'n_shells': len(r_vals),
+        'p_value': float(p_value),
+        'std_err': float(std_err),
+        'valid': True,
+    }
+
+
+# ============================================================
+# Local approximation error
+# ============================================================
+
+def compute_local_error(local_params, exact_params, model_family):
+    """Err_R(z) = ||lambda^(R)_local - lambda^{\\z}||"""
+    v_local = compute_all_param_vector(local_params, model_family)
+    v_exact = compute_all_param_vector(exact_params, model_family)
+    
+    err = np.linalg.norm(v_local - v_exact)
+    rel_err = err / (1.0 + np.linalg.norm(v_exact))
+    
+    return {
+        'param_error': float(err),
+        'relative_error': float(rel_err),
+    }
+
+
+# ============================================================
+# ELBO / objective gap
+# ============================================================
+
+def compute_objective_gap(model, edges_without, exact_params, approx_params):
+    """Gap_R = L_{\\z}(lambda^{\\z}) - L_{\\z}(lambda^(R))."""
+    try:
+        obj_exact = model.compute_elbo(edges_without, exact_params) if hasattr(model, 'compute_elbo') \
+            else model.compute_objective(edges_without, exact_params)
+        obj_approx = model.compute_elbo(edges_without, approx_params) if hasattr(model, 'compute_elbo') \
+            else model.compute_objective(edges_without, approx_params)
+        return float(obj_exact - obj_approx)
+    except Exception as e:
+        return None
+
+
+# ============================================================
+# Weighted interaction statistics (chi)
+# ============================================================
+
+def compute_chi_poisson_gamma(edge_to_remove, edges, params, N, M, K,
+                               a0=0.3, b0=1.0, c0=0.3, d0=1.0):
+    """Compute chi statistics for Poisson-Gamma model."""
+    a, b, c, d = params['a'], params['b'], params['c'], params['d']
+    
+    # Compute constants
+    a_min = np.min(a[a > 0]) if np.any(a > 0) else a0
+    b_min = np.min(b[b > 0]) if np.any(b > 0) else b0
+    c_min = np.min(c[c > 0]) if np.any(c > 0) else c0
+    d_min = np.min(d[d > 0]) if np.any(d > 0) else d0
+    a_max = np.max(a)
+    c_max = np.max(c)
+    
+    C_x = 0.5 * polygamma(1, max(c_min, 1e-3)) + 0.5 / max(d_min, 1e-6)
+    C_0 = 1.0 / max(d_min, 1e-6) + c_max / max(d_min**2, 1e-12)
+    
+    C_tilde_x = 0.5 * polygamma(1, max(a_min, 1e-3)) + 0.5 / max(b_min, 1e-6)
+    C_tilde_0 = 1.0 / max(b_min, 1e-6) + a_max / max(b_min**2, 1e-12)
+    
+    # Build adjacency
+    user_to_items = defaultdict(list)
+    item_to_users = defaultdict(list)
+    edge_dict = {}
+    for i, j, x in edges:
+        user_to_items[i].append(j)
+        item_to_users[j].append(i)
+        edge_dict[(i, j)] = x
+    
+    i_del, j_del, x_del = edge_to_remove
+    
+    # chi_i
+    chi_i = sum(C_x * edge_dict.get((i_del, j), 0) + C_0 for j in user_to_items.get(i_del, []))
+    
+    # chi_tilde_j
+    chi_tilde_j = sum(C_tilde_x * edge_dict.get((i, j_del), 0) + C_tilde_0 
+                      for i in item_to_users.get(j_del, []))
+    
+    chi_max = max(chi_i, chi_tilde_j)
+    chi_sum = chi_i + chi_tilde_j
+    
+    # Empirical alternatives
+    seed_degree = len(user_to_items.get(i_del, [])) + len(item_to_users.get(j_del, []))
+    seed_count_sum = sum(edge_dict.get((i_del, j), 0) for j in user_to_items.get(i_del, [])) + \
+                     sum(edge_dict.get((i, j_del), 0) for i in item_to_users.get(j_del, []))
+    
+    return {
+        'chi_i': float(chi_i),
+        'chi_tilde_j': float(chi_tilde_j),
+        'chi_max': float(chi_max),
+        'chi_sum': float(chi_sum),
+        'seed_degree': int(seed_degree),
+        'seed_count_sum': float(seed_count_sum),
+        'C_x': float(C_x),
+        'C_0': float(C_0),
+        'C_tilde_x': float(C_tilde_x),
+        'C_tilde_0': float(C_tilde_0),
+    }
+
+
+def compute_chi_gaussian(edge_to_remove, edges, params, N, M, K, sigma_x,
+                          model_family='gaussian_gaussian'):
+    """Compute interaction proxy for Gaussian models."""
+    user_to_items = defaultdict(list)
+    item_to_users = defaultdict(list)
+    edge_dict = {}
+    for i, j, x in edges:
+        user_to_items[i].append(j)
+        item_to_users[j].append(i)
+        edge_dict[(i, j)] = x
+    
+    i_del, j_del, x_del = edge_to_remove
+    prec_x = 1.0 / (sigma_x ** 2)
+    
+    if model_family == 'gaussian_gaussian':
+        m_U, s_U = params['m_U'], params['s_U']
+        m_V, s_V = params['m_V'], params['s_V']
+        
+        chi_i = sum(prec_x * np.sum(m_V[j]**2 + s_V[j]) for j in user_to_items.get(i_del, []))
+        chi_tilde_j = sum(prec_x * np.sum(m_U[i]**2 + s_U[i]) for i in item_to_users.get(j_del, []))
+    elif model_family == 'gaussian_gamma_map':
+        from src.model import GaussianGammaMAP
+        sp = lambda x: np.log1p(np.exp(np.clip(x, -20, 20)))
+        U = sp(params['alpha'])
+        V = sp(params['beta'])
+        
+        chi_i = sum(prec_x * np.sum(V[j]**2) for j in user_to_items.get(i_del, []))
+        chi_tilde_j = sum(prec_x * np.sum(U[i]**2) for i in item_to_users.get(j_del, []))
+    
+    chi_max = max(chi_i, chi_tilde_j)
+    chi_sum = chi_i + chi_tilde_j
+    
+    seed_degree = len(user_to_items.get(i_del, [])) + len(item_to_users.get(j_del, []))
+    
+    return {
+        'chi_i': float(chi_i),
+        'chi_tilde_j': float(chi_tilde_j),
+        'chi_max': float(chi_max),
+        'chi_sum': float(chi_sum),
+        'seed_degree': int(seed_degree),
+    }
+
+
+# ============================================================
+# Gradient interference proxy
+# ============================================================
+
+def compute_gradient_interference(full_params, exact_params, local_params, model_family):
+    """Compute gradient interference proxy.
+    
+    g_del = lambda^* - lambda^{\\z}
+    g_ret = lambda^{\\z} - lambda^(R)_local
+    I(z) = |sum_u <g_del_u, g_ret_u>|
+    """
+    v_full = compute_all_param_vector(full_params, model_family)
+    v_exact = compute_all_param_vector(exact_params, model_family)
+    v_local = compute_all_param_vector(local_params, model_family)
+    
+    g_del = v_full - v_exact
+    g_ret = v_exact - v_local
+    
+    raw_interference = float(np.abs(np.dot(g_del, g_ret)))
+    
+    norm_del = np.linalg.norm(g_del)
+    norm_ret = np.linalg.norm(g_ret)
+    eps = 1e-12
+    cosine_interference = raw_interference / (norm_del * norm_ret + eps)
+    
+    return {
+        'interference_raw': float(raw_interference),
+        'interference_cosine': float(cosine_interference),
+        'g_del_norm': float(norm_del),
+        'g_ret_norm': float(norm_ret),
+    }
+
+
+# ============================================================
+# Full metrics for one deletion
+# ============================================================
+
+def compute_all_metrics(full_params, exact_params, local_params_by_radius,
+                        warm_start_params, one_step_params,
+                        edge_to_remove, edges, N, M, K,
+                        model_family, model=None, radii=[1, 2, 3, 4],
+                        model_kwargs=None):
+    """Compute all metrics for one deletion."""
+    results = {}
+    
+    # Influence by distance
+    influence = compute_deletion_influence_by_distance(
+        full_params, exact_params, edge_to_remove, edges, N, M, model_family, 
+        max_radius=max(radii) + 2)
+    results['influence_by_distance'] = {str(k): v['mean'] for k, v in influence.items()}
+    results['influence_by_distance_full'] = influence
+    
+    # Decay fit
+    decay = fit_exponential_decay(influence)
+    results['empirical_decay_mu'] = decay['mu_emp']
+    results['empirical_decay_r2'] = decay['r_squared']
+    results['decay_valid'] = decay['valid']
+    
+    # Chi statistics
+    if model_kwargs is None:
+        model_kwargs = {}
+    
+    if model_family == 'poisson_gamma':
+        chi = compute_chi_poisson_gamma(
+            edge_to_remove, edges, full_params, N, M, K,
+            a0=model_kwargs.get('a0', 0.3), b0=model_kwargs.get('b0', 1.0),
+            c0=model_kwargs.get('c0', 0.3), d0=model_kwargs.get('d0', 1.0))
+    else:
+        chi = compute_chi_gaussian(
+            edge_to_remove, edges, full_params, N, M, K,
+            sigma_x=model_kwargs.get('sigma_x', 1.0), model_family=model_family)
+    
+    results['chi_seed_max'] = chi['chi_max']
+    results['chi_seed_sum'] = chi['chi_sum']
+    results['seed_degree'] = chi['seed_degree']
+    
+    # Local errors by radius
+    for R in radii:
+        if R in local_params_by_radius:
+            err = compute_local_error(local_params_by_radius[R], exact_params, model_family)
+            results[f'error_R{R}'] = err['param_error']
+            results[f'rel_error_R{R}'] = err['relative_error']
+            
+            # Interference
+            interf = compute_gradient_interference(
+                full_params, exact_params, local_params_by_radius[R], model_family)
+            results[f'interference_raw_R{R}'] = interf['interference_raw']
+            results[f'interference_cosine_R{R}'] = interf['interference_cosine']
+    
+    # Warm start error
+    if warm_start_params is not None:
+        ws_err = compute_local_error(warm_start_params, exact_params, model_family)
+        results['error_warm_start'] = ws_err['param_error']
+        results['rel_error_warm_start'] = ws_err['relative_error']
+    
+    # One-step error
+    if one_step_params is not None:
+        os_err = compute_local_error(one_step_params, exact_params, model_family)
+        results['error_one_step'] = os_err['param_error']
+        results['rel_error_one_step'] = os_err['relative_error']
+    
+    return results