v2: model.py

src/model.py

@@ -191,7 +191,14 @@ class PoissonGammaVI:
                                                self.N, self.M, radius)
         users_in_R, items_in_R = get_user_item_sets_in_radius(distances, self.N, radius)
         
-        I, J, X, n_edges = self._prepare_edges(filtered)
+        # KEY OPTIMIZATION: filter edges to only those touching neighborhood
+        # For user i update: need all edges (i, j, x) where i in users_in_R
+        # For item j update: need all edges (i, j, x) where j in items_in_R
+        # Union: edges where i in users_in_R OR j in items_in_R
+        local_edges = [(i, j, x) for i, j, x in filtered 
+                       if i in users_in_R or j in items_in_R]
+        
+        I, J, X, n_edges = self._prepare_edges(local_edges)
         
         converged = False
         for it in range(self.max_iter):
@@ -405,7 +412,10 @@ class GaussianGaussianVI:
                                                self.N, self.M, radius)
         users_in_R, items_in_R = get_user_item_sets_in_radius(distances, self.N, radius)
         
-        I, J, X, n_edges = self._prepare_edges(filtered)
+        # Filter edges to neighborhood
+        local_edges = [(i, j, x) for i, j, x in filtered 
+                       if i in users_in_R or j in items_in_R]
+        I, J, X, n_edges = self._prepare_edges(local_edges)
         converged = False
         for it in range(self.max_iter):
             old_params = {k: v.copy() for k, v in params.items()}
@@ -436,10 +446,14 @@ class GaussianGaussianVI:
 # ============================================================
 
 class GaussianGammaMAP:
-    """Gaussian likelihood + Gamma prior, MAP via softplus parameterization. Vectorized."""
+    """Gaussian likelihood + Gamma prior, MAP via softplus parameterization.
+    
+    Uses Adam optimizer with gradient clipping for stable convergence.
+    """
     
     def __init__(self, N, M, K, a0=0.3, b0=1.0, c0=0.3, d0=1.0,
-                 sigma_x=1.0, lr=0.01, max_iter=200, tol=1e-5, seed=0):
+                 sigma_x=1.0, lr=0.01, max_iter=500, tol=1e-5, seed=0,
+                 grad_clip=5.0, adam_beta1=0.9, adam_beta2=0.999):
         self.N = N
         self.M = M
         self.K = K
@@ -452,6 +466,9 @@ class GaussianGammaMAP:
         self.max_iter = max_iter
         self.tol = tol
         self.seed = seed
+        self.grad_clip = grad_clip
+        self.adam_beta1 = adam_beta1
+        self.adam_beta2 = adam_beta2
     
     def _softplus(self, x):
         return np.log1p(np.exp(np.clip(x, -20, 20)))
@@ -459,12 +476,28 @@ class GaussianGammaMAP:
     def _softplus_grad(self, x):
         return 1.0 / (1.0 + np.exp(-np.clip(x, -20, 20)))
     
-    def _init_params(self, rng=None):
+    def _inv_softplus(self, y):
+        """Inverse of softplus: log(exp(y) - 1)."""
+        return np.log(np.exp(np.clip(y, 1e-8, 20)) - 1 + 1e-30)
+    
+    def _init_params(self, rng=None, edges=None):
         if rng is None:
             rng = np.random.RandomState(self.seed)
+        # Data-informed initialization: use NMF-style init from mean values
+        if edges is not None:
+            I = np.array([e[0] for e in edges], dtype=np.int32)
+            J = np.array([e[1] for e in edges], dtype=np.int32)
+            X = np.array([e[2] for e in edges], dtype=np.float64)
+            # Compute user/item means
+            x_mean = np.abs(X).mean()
+            init_scale = np.sqrt(np.abs(x_mean) / self.K + 0.1)
+        else:
+            init_scale = 0.5
+        U_init = np.abs(rng.randn(self.N, self.K)) * init_scale + 0.1
+        V_init = np.abs(rng.randn(self.M, self.K)) * init_scale + 0.1
         return {
-            'alpha': rng.randn(self.N, self.K) * 0.5,
-            'beta': rng.randn(self.M, self.K) * 0.5,
+            'alpha': self._inv_softplus(U_init),
+            'beta': self._inv_softplus(V_init),
         }
     
     def _prepare_edges(self, edges):
@@ -487,71 +520,98 @@ class GaussianGammaMAP:
         obj += np.sum((self.c0 - 1) * np.log(V + 1e-30) - self.d0 * V)
         return float(obj)
     
-    def _gradient_step(self, I, J, X, params, update_users=None, update_items=None):
-        """One gradient step."""
+    def _compute_gradients(self, I, J, X, params, update_users=None, update_items=None):
+        """Compute gradients with clipping."""
         U = self._softplus(params['alpha'])
         V = self._softplus(params['beta'])
         prec_x = 1.0 / (self.sigma_x ** 2)
         
-        pred = np.sum(U[I] * V[J], axis=1)  # (n_edges,)
-        residual = X - pred  # (n_edges,)
+        pred = np.sum(U[I] * V[J], axis=1)
+        residual = X - pred
         
         sp_grad_alpha = self._softplus_grad(params['alpha'])
         sp_grad_beta = self._softplus_grad(params['beta'])
         
-        # Gradient for alpha (user params)
-        # dL/dU[i,k] = prec_x * sum_j residual[e] * V[j,k] + (a0-1)/U[i,k] - b0
         grad_U = np.zeros_like(U)
         for k in range(self.K):
             contrib = prec_x * residual * V[J, k]
             np.add.at(grad_U[:, k], I, contrib)
         
-        # Prior gradient
-        prior_grad_U = (self.a0 - 1) / (U + 1e-30) - self.b0
+        prior_grad_U = (self.a0 - 1) / (U + 1e-6) - self.b0
         grad_U += prior_grad_U
-        
-        # Chain rule through softplus
         grad_alpha = grad_U * sp_grad_alpha
         
-        # Gradient for beta (item params)
         grad_V = np.zeros_like(V)
         for k in range(self.K):
             contrib = prec_x * residual * U[I, k]
             np.add.at(grad_V[:, k], J, contrib)
         
-        prior_grad_V = (self.c0 - 1) / (V + 1e-30) - self.d0
+        prior_grad_V = (self.c0 - 1) / (V + 1e-6) - self.d0
         grad_V += prior_grad_V
-        
         grad_beta = grad_V * sp_grad_beta
         
-        alpha_new = params['alpha'].copy()
-        beta_new = params['beta'].copy()
-        
-        if update_users is not None:
-            ul = list(update_users)
-            alpha_new[ul] += self.lr * grad_alpha[ul]
-        else:
-            alpha_new += self.lr * grad_alpha
+        # Gradient clipping
+        if self.grad_clip > 0:
+            gnorm_a = np.linalg.norm(grad_alpha)
+            if gnorm_a > self.grad_clip:
+                grad_alpha *= self.grad_clip / gnorm_a
+            gnorm_b = np.linalg.norm(grad_beta)
+            if gnorm_b > self.grad_clip:
+                grad_beta *= self.grad_clip / gnorm_b
         
-        if update_items is not None:
-            il = list(update_items)
-            beta_new[il] += self.lr * grad_beta[il]
-        else:
-            beta_new += self.lr * grad_beta
-        
-        return {'alpha': alpha_new, 'beta': beta_new}
+        return grad_alpha, grad_beta
     
-    def fit_full(self, edges, config=None, init_params=None):
+    def _fit_internal(self, edges, params, max_iter=None, 
+                      update_users=None, update_items=None):
+        """Internal fit with Adam optimizer."""
         t0 = time.time()
+        if max_iter is None:
+            max_iter = self.max_iter
         I, J, X, n_edges = self._prepare_edges(edges)
         
-        params = {k: v.copy() for k, v in (init_params or self._init_params()).items()}
+        # Adam state
+        m_alpha = np.zeros_like(params['alpha'])
+        v_alpha = np.zeros_like(params['alpha'])
+        m_beta = np.zeros_like(params['beta'])
+        v_beta = np.zeros_like(params['beta'])
+        eps_adam = 1e-8
+        
         obj_trace = []
         converged = False
         
-        for it in range(self.max_iter):
+        for it in range(max_iter):
             old_params = {k: v.copy() for k, v in params.items()}
-            params = self._gradient_step(I, J, X, params)
+            
+            grad_alpha, grad_beta = self._compute_gradients(
+                I, J, X, params, update_users, update_items)
+            
+            # Adam updates
+            t_adam = it + 1
+            m_alpha = self.adam_beta1 * m_alpha + (1 - self.adam_beta1) * grad_alpha
+            v_alpha = self.adam_beta2 * v_alpha + (1 - self.adam_beta2) * grad_alpha**2
+            m_hat_a = m_alpha / (1 - self.adam_beta1**t_adam)
+            v_hat_a = v_alpha / (1 - self.adam_beta2**t_adam)
+            
+            m_beta = self.adam_beta1 * m_beta + (1 - self.adam_beta1) * grad_beta
+            v_beta = self.adam_beta2 * v_beta + (1 - self.adam_beta2) * grad_beta**2
+            m_hat_b = m_beta / (1 - self.adam_beta1**t_adam)
+            v_hat_b = v_beta / (1 - self.adam_beta2**t_adam)
+            
+            step_alpha = self.lr * m_hat_a / (np.sqrt(v_hat_a) + eps_adam)
+            step_beta = self.lr * m_hat_b / (np.sqrt(v_hat_b) + eps_adam)
+            
+            if update_users is not None:
+                ul = list(update_users)
+                params['alpha'][ul] += step_alpha[ul]
+            else:
+                params['alpha'] = params['alpha'] + step_alpha
+            
+            if update_items is not None:
+                il = list(update_items)
+                params['beta'][il] += step_beta[il]
+            else:
+                params['beta'] = params['beta'] + step_beta
+            
             change = relative_param_change(old_params, params)
             if it % 50 == 0:
                 obj_trace.append(self.compute_objective(edges, params))
@@ -559,10 +619,20 @@ class GaussianGammaMAP:
                 converged = True
                 break
         
+        return params, obj_trace, it + 1, converged, time.time() - t0
+    
+    def fit_full(self, edges, config=None, init_params=None):
+        if init_params is not None:
+            params = {k: v.copy() for k, v in init_params.items()}
+        else:
+            params = self._init_params(edges=edges)
+        
+        params, obj_trace, n_iter, converged, runtime = self._fit_internal(edges, params)
+        
         return FitResult(
             params=params, objective_trace=obj_trace,
-            n_iterations=it + 1, converged=converged,
-            runtime_sec=time.time() - t0,
+            n_iterations=n_iter, converged=converged,
+            runtime_sec=runtime,
             model_family='gaussian_gamma_map', inference_type='map',
             likelihood='gaussian', prior='gamma',
         )
@@ -573,7 +643,6 @@ class GaussianGammaMAP:
         return self.fit_full(filtered, config, init_params)
     
     def fit_local(self, edges, edge_to_remove, radius, config=None, init_params=None):
-        t0 = time.time()
         i_del, j_del = int(edge_to_remove[0]), int(edge_to_remove[1])
         filtered = [(i, j, x) for i, j, x in edges if not (i == i_del and j == j_del)]
         
@@ -586,24 +655,21 @@ class GaussianGammaMAP:
                                                self.N, self.M, radius)
         users_in_R, items_in_R = get_user_item_sets_in_radius(distances, self.N, radius)
         
-        I, J, X, n_edges = self._prepare_edges(filtered)
-        converged = False
-        for it in range(self.max_iter):
-            old_params = {k: v.copy() for k, v in params.items()}
-            params = self._gradient_step(I, J, X, params,
-                                          update_users=users_in_R, update_items=items_in_R)
-            change = relative_param_change(old_params, params)
-            if change < self.tol:
-                converged = True
-                break
+        # Filter edges to neighborhood
+        local_edges = [(i, j, x) for i, j, x in filtered 
+                       if i in users_in_R or j in items_in_R]
+        
+        params, obj_trace, n_iter, converged, runtime = self._fit_internal(
+            local_edges, params, update_users=users_in_R, update_items=items_in_R)
         
         return FitResult(
-            params=params, objective_trace=[],
-            n_iterations=it + 1, converged=converged,
-            runtime_sec=time.time() - t0,
+            params=params, objective_trace=obj_trace,
+            n_iterations=n_iter, converged=converged,
+            runtime_sec=runtime,
             model_family='gaussian_gamma_map', inference_type='map',
             likelihood='gaussian', prior='gamma',
-            diagnostics={'radius': radius}
+            diagnostics={'radius': radius, 'n_users_updated': len(users_in_R),
+                         'n_items_updated': len(items_in_R)}
         )
     
     def fit_warm_start_global(self, edges, edge_to_remove, config=None, init_params=None):
@@ -621,7 +687,7 @@ def get_model(model_family, N, M, K, **kwargs):
         valid = {'sigma_U', 'sigma_V', 'sigma_x', 'max_iter', 'tol', 'damping', 'seed'}
         return GaussianGaussianVI(N, M, K, **{k: v for k, v in kwargs.items() if k in valid})
     elif model_family == 'gaussian_gamma_map':
-        valid = {'a0', 'b0', 'c0', 'd0', 'sigma_x', 'lr', 'max_iter', 'tol', 'seed'}
+        valid = {'a0', 'b0', 'c0', 'd0', 'sigma_x', 'lr', 'max_iter', 'tol', 'seed', 'grad_clip'}
         return GaussianGammaMAP(N, M, K, **{k: v for k, v in kwargs.items() if k in valid})
     else:
         raise ValueError(f"Unknown model family: {model_family}")