insurance_gam.pin

 60class PINModel(nn.Module):
 61    """
 62    Single Tree-like Pairwise Interaction Network.
 63
 64    Prediction:
 65        f_PIN(x) = exp( sum_{j<=k} w_{jk} * h_{jk}(x) + b )
 66
 67    Args:
 68        features: Dict mapping feature name to spec. Spec is 'continuous' or
 69            an int (number of categories). Order matters — features are indexed
 70            by position.
 71        embedding_dim: Feature embedding dimension d (default 10).
 72        hidden_dim: Hidden width for continuous embedding FNNs d' (default 20).
 73        token_dim: Interaction token dimension d0 (default 10).
 74        shared_dims: (d1, d2) widths for shared interaction network (default [30, 20]).
 75        loss: Loss name — 'poisson', 'gamma', or 'tweedie'.
 76        tweedie_p: Tweedie power (only used when loss='tweedie').
 77        lr: Adam learning rate (default 0.001).
 78        batch_size: Mini-batch size (default 128).
 79        max_epochs: Maximum training epochs (default 500).
 80        patience: Early stopping patience in epochs (default 20).
 81        lr_patience: ReduceLROnPlateau patience (default 5).
 82        lr_factor: ReduceLROnPlateau reduction factor (default 0.9).
 83        val_fraction: Fraction of training data for validation if X_val not given
 84            (default 0.1).
 85        device: Torch device string, or None to auto-detect.
 86        random_seed: Seed for reproducibility.
 87
 88    Examples:
 89        >>> model = PINModel(
 90        ...     features={"age": "continuous", "area": 5},
 91        ...     loss="poisson",
 92        ... )
 93    """
 94
 95    def __init__(
 96        self,
 97        features: FeatureSpec,
 98        embedding_dim: int = 10,
 99        hidden_dim: int = 20,
100        token_dim: int = 10,
101        shared_dims: Tuple[int, int] = (30, 20),
102        loss: str = "poisson",
103        tweedie_p: float = 1.5,
104        lr: float = 0.001,
105        batch_size: int = 128,
106        max_epochs: int = 500,
107        patience: int = 20,
108        lr_patience: int = 5,
109        lr_factor: float = 0.9,
110        val_fraction: float = 0.1,
111        device: Optional[str] = None,
112        random_seed: int = 42,
113    ) -> None:
114        super().__init__()
115
116        self.features = dict(features)
117        self.feature_names: List[str] = list(features.keys())
118        self.q = len(self.feature_names)
119        self.embedding_dim = embedding_dim
120        self.hidden_dim = hidden_dim
121        self.token_dim = token_dim
122        self.shared_dims = tuple(shared_dims)
123        self.loss_name = loss
124        self.tweedie_p = tweedie_p
125        self.lr = lr
126        self.batch_size = batch_size
127        self.max_epochs = max_epochs
128        self.patience = patience
129        self.lr_patience = lr_patience
130        self.lr_factor = lr_factor
131        self.val_fraction = val_fraction
132        self.random_seed = random_seed
133
134        # Resolve device
135        if device is None:
136            self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
137        else:
138            self._device = torch.device(device)
139
140        # --- Sub-modules ---
141        self.feature_embeddings = FeatureEmbeddings(
142            features=features,
143            embedding_dim=embedding_dim,
144            hidden_dim=hidden_dim,
145        )
146        self.interaction_tokens = InteractionTokens(
147            n_features=self.q,
148            token_dim=token_dim,
149        )
150        self.shared_net = SharedInteractionNet(
151            embedding_dim=embedding_dim,
152            token_dim=token_dim,
153            layer1_dim=shared_dims[0],
154            layer2_dim=shared_dims[1],
155        )
156
157        n_pairs = self.q * (self.q + 1) // 2
158        # Output weights w_{jk} and bias b — linear combination of pair terms
159        self.output_weights = nn.Parameter(torch.randn(n_pairs) * 0.01)
160        self.output_bias = nn.Parameter(torch.zeros(1))
161
162        # Loss function
163        loss_kwargs = {"p": tweedie_p} if loss == "tweedie" else {}
164        self._loss_fn = get_loss(loss, **loss_kwargs)
165
166        # Weight initialisation: smaller scales for numerical stability
167        self._init_weights()
168
169        # Training state
170        self._is_fitted = False
171        self.train_history: Dict[str, List[float]] = {"train_loss": [], "val_loss": []}
172
173        # Centering offsets (set post-hoc after fitting)
174        # h_{jk}^centered(x) = h_{jk}(x) - mean_train[h_{jk}]
175        # We store mean_train[w_{jk} * h_{jk}] per pair for efficiency
176        self._pair_means: Optional[torch.Tensor] = None
177
178    # ------------------------------------------------------------------
179    # Weight initialisation
180    # ------------------------------------------------------------------
181
182    def _init_weights(self) -> None:
183        """
184        Initialise all linear layers to smaller scales for training stability.
185
186        Default PyTorch Kaiming init can produce large initial activations
187        when many layers are composed. With smaller scales, early training
188        is more stable even on very small datasets.
189        """
190        for module in self.modules():
191            if isinstance(module, nn.Linear):
192                nn.init.normal_(module.weight, mean=0.0, std=0.01)
193                if module.bias is not None:
194                    nn.init.zeros_(module.bias)
195            elif isinstance(module, nn.Embedding):
196                nn.init.normal_(module.weight, mean=0.0, std=0.01)
197
198    # ------------------------------------------------------------------
199    # Forward pass
200    # ------------------------------------------------------------------
201
202    def _compute_linear_predictor(
203        self,
204        x_dict: Dict[str, torch.Tensor],
205        apply_centering: bool = True,
206    ) -> torch.Tensor:
207        """
208        Compute sum_{j<=k} w_{jk} * h_{jk}(x) + b.
209
210        Args:
211            x_dict: Feature tensors on self._device.
212            apply_centering: Subtract pair means (set during fit) for identifiability.
213
214        Returns:
215            Shape (batch_size,).
216        """
217        # Embed all features once
218        embeddings: Dict[str, torch.Tensor] = self.feature_embeddings.embed_all(x_dict)
219
220        pairs = self.interaction_tokens.pair_indices()
221        terms = []
222
223        for pair_idx, (j, k) in enumerate(pairs):
224            fname_j = self.feature_names[j]
225            fname_k = self.feature_names[k]
226
227            phi_j = embeddings[fname_j]
228            phi_k = embeddings[fname_k]
229            token = self.interaction_tokens.get_token(j, k)
230
231            raw = self.shared_net(phi_j, phi_k, token)  # (batch, 1)
232            h = centered_hard_sigmoid(raw).squeeze(-1)   # (batch,)
233
234            w = self.output_weights[pair_idx]
235            term = w * h  # (batch,)
236            terms.append(term)
237
238        # Stack to (batch, n_pairs) then sum over pairs
239        all_terms = torch.stack(terms, dim=1)  # (batch, n_pairs)
240
241        if apply_centering and self._pair_means is not None:
242            all_terms = all_terms - self._pair_means.unsqueeze(0)
243
244        linear_pred = all_terms.sum(dim=1) + self.output_bias.squeeze()
245        return linear_pred
246
247    def forward(
248        self,
249        x_dict: Dict[str, torch.Tensor],
250        exposure: Optional[torch.Tensor] = None,
251    ) -> torch.Tensor:
252        """
253        Compute predictions f_PIN(x) = exp(linear_predictor) * exposure.
254
255        When exposure is provided, the raw model output is frequency (claims per
256        year) and multiplying by exposure gives expected claim count. For
257        frequency models, typically you'd pass exposure=None and let the caller
258        multiply; this method supports both modes.
259
260        Args:
261            x_dict: Feature tensors. Dict of feature_name -> (batch,) tensor.
262            exposure: Optional per-sample exposure, shape (batch,).
263
264        Returns:
265            Predicted frequency, shape (batch,).
266        """
267        eta = self._compute_linear_predictor(x_dict)
268        # Clamp to avoid exp overflow (GLM link stabilisation)
269        eta = torch.clamp(eta, min=-20.0, max=20.0)
270        mu = torch.exp(eta)
271        if exposure is not None:
272            mu = mu * exposure
273        return mu
274
275    # ------------------------------------------------------------------
276    # Data preparation
277    # ------------------------------------------------------------------
278
279    def _prepare_features(
280        self,
281        X: Union[Dict, "pl.DataFrame", "pd.DataFrame"],  # noqa: F821
282    ) -> Dict[str, torch.Tensor]:
283        """
284        Convert input data to a dict of tensors.
285
286        Accepts:
287        - Dict[str, np.ndarray]
288        - Dict[str, list]
289        - polars.DataFrame
290        - pandas.DataFrame
291        """
292        # Try polars first
293        try:
294            import polars as pl
295            if isinstance(X, pl.DataFrame):
296                return self._polars_to_dict(X)
297        except ImportError:
298            pass
299
300        # Try pandas
301        try:
302            import pandas as pd
303            if isinstance(X, pd.DataFrame):
304                return {col: X[col].to_numpy() for col in self.feature_names}
305        except ImportError:
306            pass
307
308        # Assume it's already a dict
309        if isinstance(X, dict):
310            return X
311
312        raise TypeError(
313            f"X must be a dict, polars.DataFrame, or pandas.DataFrame. Got {type(X)}."
314        )
315
316    def _polars_to_dict(self, df) -> Dict[str, np.ndarray]:
317        result = {}
318        for name in self.feature_names:
319            result[name] = df[name].to_numpy()
320        return result
321
322    def _to_device_dict(
323        self, x_dict: Dict[str, Union[np.ndarray, torch.Tensor]]
324    ) -> Dict[str, torch.Tensor]:
325        """Convert dict of arrays to tensors on device."""
326        result = {}
327        for name, arr in x_dict.items():
328            spec = self.features[name]
329            if spec == "continuous":
330                result[name] = _to_tensor(arr).to(self._device)
331            else:
332                result[name] = _to_long_tensor(arr).to(self._device)
333        return result
334
335    # ------------------------------------------------------------------
336    # Training
337    # ------------------------------------------------------------------
338
339    def fit(
340        self,
341        X_train,
342        y_train: np.ndarray,
343        exposure: Optional[np.ndarray] = None,
344        X_val=None,
345        y_val: Optional[np.ndarray] = None,
346        exposure_val: Optional[np.ndarray] = None,
347        verbose: bool = True,
348    ) -> "PINModel":
349        """
350        Fit the model.
351
352        Args:
353            X_train: Training features. Dict, polars.DataFrame, or pandas.DataFrame.
354            y_train: Observed frequency (claims / exposure), shape (n,).
355            exposure: Per-sample exposure (years at risk), shape (n,).
356            X_val: Validation features (optional). If None, 10% of training data
357                is reserved.
358            y_val: Validation targets.
359            exposure_val: Validation exposure.
360            verbose: Print training progress.
361
362        Returns:
363            self (for chaining).
364        """
365        torch.manual_seed(self.random_seed)
366        np.random.seed(self.random_seed)
367
368        self.to(self._device)
369
370        # Prepare training data
371        x_dict = self._prepare_features(X_train)
372        x_dict = self._to_device_dict(x_dict)
373        y_t = _to_tensor(y_train).to(self._device)
374        exp_t = (
375            _to_tensor(exposure).to(self._device)
376            if exposure is not None
377            else torch.ones_like(y_t)
378        )
379
380        n = y_t.shape[0]
381
382        # Build or reserve validation set (must happen before bias init so we
383        # use training-only data for the mean frequency estimate).
384        if X_val is not None:
385            x_val_dict = self._prepare_features(X_val)
386            x_val_dict = self._to_device_dict(x_val_dict)
387            y_val_t = _to_tensor(y_val).to(self._device)
388            exp_val_t = (
389                _to_tensor(exposure_val).to(self._device)
390                if exposure_val is not None
391                else torch.ones_like(y_val_t)
392            )
393        else:
394            # Reserve val_fraction from training
395            val_size = max(1, int(n * self.val_fraction))
396            perm = torch.randperm(n, device=self._device)
397            val_idx = perm[:val_size]
398            train_idx = perm[val_size:]
399
400            x_val_dict = {k: v[val_idx] for k, v in x_dict.items()}
401            y_val_t = y_t[val_idx]
402            exp_val_t = exp_t[val_idx]
403
404            x_dict = {k: v[train_idx] for k, v in x_dict.items()}
405            y_t = y_t[train_idx]
406            exp_t = exp_t[train_idx]
407
408        n_train = y_t.shape[0]
409
410        # Initialise bias to log(mean_frequency) of the TRAINING split only.
411        # Computing this from the full dataset (including validation) would
412        # be a mild form of data leakage — small in practice but wrong in
413        # principle. We compute it here, after the split.
414        with torch.no_grad():
415            mean_freq = y_t.mean().clamp(min=1e-8)
416            self.output_bias.fill_(torch.log(mean_freq).item())
417
418        optimizer = optim.Adam(self.parameters(), lr=self.lr)
419        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
420            optimizer,
421            mode="min",
422            factor=self.lr_factor,
423            patience=self.lr_patience,
424        )
425
426        best_val_loss = float("inf")
427        best_state = None
428        epochs_no_improve = 0
429
430        self.train_history = {"train_loss": [], "val_loss": []}
431
432        for epoch in range(self.max_epochs):
433            self.train()
434            # Shuffle training data
435            perm = torch.randperm(n_train, device=self._device)
436
437            epoch_loss = 0.0
438            n_batches = 0
439
440            for start in range(0, n_train, self.batch_size):
441                end = min(start + self.batch_size, n_train)
442                batch_idx = perm[start:end]
443
444                x_batch = {k: v[batch_idx] for k, v in x_dict.items()}
445                y_batch = y_t[batch_idx]
446                exp_batch = exp_t[batch_idx]
447
448                optimizer.zero_grad()
449                mu = self.forward(x_batch)
450                loss = self._loss_fn(mu, y_batch, exp_batch)
451                if torch.isnan(loss):
452                    continue  # skip NaN batches
453                loss.backward()
454                # Gradient clipping prevents exploding gradients in early training
455                torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)
456                # Zero NaN gradients before step
457                for p in self.parameters():
458                    if p.grad is not None and torch.isnan(p.grad).any():
459                        p.grad.zero_()
460                optimizer.step()
461                # Restore any NaN params to their pre-step values
462                for p in self.parameters():
463                    if torch.isnan(p).any():
464                        torch.nan_to_num_(p, nan=0.0)
465
466                epoch_loss += loss.item()
467                n_batches += 1
468
469            train_loss = epoch_loss / n_batches
470
471            # Validation
472            self.eval()
473            with torch.no_grad():
474                mu_val = self.forward(x_val_dict)
475                val_loss = self._loss_fn(mu_val, y_val_t, exp_val_t).item()
476
477            scheduler.step(val_loss)
478
479            self.train_history["train_loss"].append(train_loss)
480            self.train_history["val_loss"].append(val_loss)
481
482            if verbose and (epoch % 50 == 0 or epoch < 10):
483                lr_now = optimizer.param_groups[0]["lr"]
484                print(
485                    f"Epoch {epoch:4d} | train={train_loss:.6f} | val={val_loss:.6f} | lr={lr_now:.2e}"
486                )
487
488            # Early stopping
489            if val_loss < best_val_loss:
490                best_val_loss = val_loss
491                best_state = copy.deepcopy(self.state_dict())
492                epochs_no_improve = 0
493            else:
494                epochs_no_improve += 1
495                if epochs_no_improve >= self.patience:
496                    if verbose:
497                        print(f"Early stopping at epoch {epoch}. Best val={best_val_loss:.6f}")
498                    break
499
500        # Restore best weights
501        if best_state is not None:
502            self.load_state_dict(best_state)
503
504        # Post-hoc centering: compute pair means on training data
505        self._compute_pair_centering(x_dict)
506
507        self._is_fitted = True
508        return self
509
510    def _compute_pair_centering(
511        self, x_dict: Dict[str, torch.Tensor], batch_size: int = 1024
512    ) -> None:
513        """
514        Compute mean of w_{jk} * h_{jk}(x) over training data for each pair.
515
516        Centering ensures that the bias b absorbs the intercept and the
517        interaction terms have zero mean — making w_{jk} interpretable as
518        the amplitude of the interaction effect.
519
520        Stores result in self._pair_means, shape (n_pairs,).
521        """
522        self.eval()
523        n = x_dict[self.feature_names[0]].shape[0]
524        n_pairs = self.q * (self.q + 1) // 2
525        pairs = self.interaction_tokens.pair_indices()
526
527        pair_sums = torch.zeros(n_pairs, device=self._device)
528        n_seen = 0
529
530        with torch.no_grad():
531            for start in range(0, n, batch_size):
532                end = min(start + batch_size, n)
533                x_batch = {k: v[start:end] for k, v in x_dict.items()}
534                embeddings = self.feature_embeddings.embed_all(x_batch)
535                batch_n = end - start
536
537                for pair_idx, (j, k) in enumerate(pairs):
538                    fname_j = self.feature_names[j]
539                    fname_k = self.feature_names[k]
540                    phi_j = embeddings[fname_j]
541                    phi_k = embeddings[fname_k]
542                    token = self.interaction_tokens.get_token(j, k)
543                    raw = self.shared_net(phi_j, phi_k, token)
544                    h = centered_hard_sigmoid(raw).squeeze(-1)
545                    w = self.output_weights[pair_idx]
546                    pair_sums[pair_idx] += (w * h).sum()
547
548                n_seen += batch_n
549
550        self._pair_means = pair_sums / n_seen
551
552    # ------------------------------------------------------------------
553    # Prediction
554    # ------------------------------------------------------------------
555
556    def predict(
557        self,
558        X,
559        exposure: Optional[np.ndarray] = None,
560    ) -> np.ndarray:
561        """
562        Predict frequency (or expected claims if exposure given).
563
564        Args:
565            X: Features. Dict, polars.DataFrame, or pandas.DataFrame.
566            exposure: Per-sample exposure.
567
568        Returns:
569            Predictions as numpy array, shape (n,).
570        """
571        if not self._is_fitted:
572            raise RuntimeError("Call fit() before predict().")
573
574        self.eval()
575        x_dict = self._prepare_features(X)
576        x_dict = self._to_device_dict(x_dict)
577
578        exp_t = (
579            _to_tensor(exposure).to(self._device)
580            if exposure is not None
581            else None
582        )
583
584        with torch.no_grad():
585            mu = self.forward(x_dict, exposure=exp_t)
586
587        return mu.cpu().numpy()
588
589    # ------------------------------------------------------------------
590    # Interpretability surfaces
591    # ------------------------------------------------------------------
592
593    def pair_contributions(
594        self,
595        X,
596    ) -> Dict[Tuple[str, str], np.ndarray]:
597        """
598        Compute w_{jk} * h_{jk}(x) for each pair and each sample.
599
600        These are the additive components on the linear predictor scale.
601        Useful for understanding which pairs drive predictions.
602
603        Args:
604            X: Features.
605
606        Returns:
607            Dict mapping (fname_j, fname_k) -> array of shape (n,).
608        """
609        if not self._is_fitted:
610            raise RuntimeError("Call fit() before pair_contributions().")
611
612        self.eval()
613        x_dict = self._prepare_features(X)
614        x_dict = self._to_device_dict(x_dict)
615        pairs = self.interaction_tokens.pair_indices()
616        result = {}
617
618        with torch.no_grad():
619            embeddings = self.feature_embeddings.embed_all(x_dict)
620
621            for pair_idx, (j, k) in enumerate(pairs):
622                fname_j = self.feature_names[j]
623                fname_k = self.feature_names[k]
624                phi_j = embeddings[fname_j]
625                phi_k = embeddings[fname_k]
626                token = self.interaction_tokens.get_token(j, k)
627                raw = self.shared_net(phi_j, phi_k, token)
628                h = centered_hard_sigmoid(raw).squeeze(-1)
629                w = self.output_weights[pair_idx]
630                contrib = (w * h)
631                if self._pair_means is not None:
632                    contrib = contrib - self._pair_means[pair_idx]
633                result[(fname_j, fname_k)] = contrib.cpu().numpy()
634
635        return result
636
637    def main_effects(
638        self,
639        X_background,
640        n_grid: int = 100,
641    ) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
642        """
643        Compute main effect curves for each feature.
644
645        For continuous features: evaluates h_{jj} over an evenly-spaced grid
646        while fixing all other features to background means.
647
648        Returns:
649            Dict mapping feature_name -> (grid_values, effect_values).
650            For categoricals: grid_values are integer category codes.
651        """
652        if not self._is_fitted:
653            raise RuntimeError("Call fit() before main_effects().")
654
655        self.eval()
656        bg_dict = self._prepare_features(X_background)
657        bg_dict = self._to_device_dict(bg_dict)
658
659        result = {}
660
661        with torch.no_grad():
662            for i, fname in enumerate(self.feature_names):
663                spec = self.features[fname]
664
665                if spec == "continuous":
666                    vals = bg_dict[fname].float()
667                    lo, hi = vals.min().item(), vals.max().item()
668                    grid = torch.linspace(lo, hi, n_grid, device=self._device)
669                elif isinstance(spec, int):
670                    grid = torch.arange(spec, device=self._device)
671                else:
672                    continue
673
674                # Build batch: vary feature i, fix others to background mean
675                n_grid_pts = grid.shape[0]
676                x_eval = {}
677                for fname2 in self.feature_names:
678                    spec2 = self.features[fname2]
679                    if fname2 == fname:
680                        x_eval[fname2] = grid.long() if isinstance(spec, int) else grid
681                    else:
682                        if spec2 == "continuous":
683                            mean_val = bg_dict[fname2].float().mean()
684                            x_eval[fname2] = mean_val.expand(n_grid_pts)
685                        else:
686                            mode_val = bg_dict[fname2].long().mode().values
687                            x_eval[fname2] = mode_val.expand(n_grid_pts)
688
689                # Get diagonal pair contribution (j=i, k=i)
690                phi_j = self.feature_embeddings.embed_feature(fname, x_eval[fname])
691                token = self.interaction_tokens.get_token(i, i)
692                raw = self.shared_net(phi_j, phi_j, token)
693                h = centered_hard_sigmoid(raw).squeeze(-1)
694                pair_idx = self.interaction_tokens._pair_to_idx[(i, i)]
695                w = self.output_weights[pair_idx]
696                effect = (w * h).cpu().numpy()
697
698                result[fname] = (grid.cpu().numpy(), effect)
699
700        return result
701
702    def interaction_surfaces(
703        self,
704        X_background,
705        n_grid: int = 30,
706        pairs: Optional[List[Tuple[str, str]]] = None,
707    ) -> Dict[Tuple[str, str], Dict]:
708        """
709        Compute 2D interaction surfaces for feature pairs.
710
711        For pair (j, k): evaluates w_{jk} * h_{jk}(x_j, x_k) over a 2D grid.
712        Main effect contributions (j=j and k=k pairs) are NOT subtracted here —
713        this shows the raw interaction term. Use pair_contributions() for
714        full decomposition.
715
716        Args:
717            X_background: Background data for range estimation.
718            n_grid: Grid resolution per axis (n_grid x n_grid for continuous).
719            pairs: List of (feature_j, feature_k) pairs to compute. If None,
720                computes all off-diagonal interaction pairs.
721
722        Returns:
723            Dict mapping (fname_j, fname_k) -> {
724                'grid_j': array of shape (n_grid,),
725                'grid_k': array of shape (n_grid,) or (n_cats,),
726                'surface': array of shape (n_grid_j, n_grid_k),
727            }
728        """
729        if not self._is_fitted:
730            raise RuntimeError("Call fit() before interaction_surfaces().")
731
732        self.eval()
733        bg_dict = self._prepare_features(X_background)
734        bg_dict = self._to_device_dict(bg_dict)
735
736        # All off-diagonal pairs if not specified
737        if pairs is None:
738            all_pairs = self.interaction_tokens.pair_indices()
739            pairs = [
740                (self.feature_names[j], self.feature_names[k])
741                for j, k in all_pairs
742                if j != k
743            ]
744
745        result = {}
746
747        with torch.no_grad():
748            for (fname_j, fname_k) in pairs:
749                j = self.feature_names.index(fname_j)
750                k = self.feature_names.index(fname_k)
751
752                spec_j = self.features[fname_j]
753                spec_k = self.features[fname_k]
754
755                if spec_j == "continuous":
756                    vals_j = bg_dict[fname_j].float()
757                    grid_j = torch.linspace(vals_j.min(), vals_j.max(), n_grid, device=self._device)
758                else:
759                    grid_j = torch.arange(spec_j, device=self._device)
760
761                if spec_k == "continuous":
762                    vals_k = bg_dict[fname_k].float()
763                    grid_k = torch.linspace(vals_k.min(), vals_k.max(), n_grid, device=self._device)
764                else:
765                    grid_k = torch.arange(spec_k, device=self._device)
766
767                nj, nk = grid_j.shape[0], grid_k.shape[0]
768
769                # Meshgrid
770                gj = grid_j.repeat_interleave(nk)  # (nj * nk,)
771                gk = grid_k.repeat(nj)              # (nj * nk,)
772
773                phi_j = self.feature_embeddings.embed_feature(
774                    fname_j, gj.long() if isinstance(spec_j, int) else gj
775                )
776                phi_k = self.feature_embeddings.embed_feature(
777                    fname_k, gk.long() if isinstance(spec_k, int) else gk
778                )
779
780                token = self.interaction_tokens.get_token(j, k)
781                raw = self.shared_net(phi_j, phi_k, token)
782                h = centered_hard_sigmoid(raw).squeeze(-1)
783                pair_idx = self.interaction_tokens._pair_to_idx[(j, k)]
784                w = self.output_weights[pair_idx]
785                surface = (w * h).reshape(nj, nk).cpu().numpy()
786
787                result[(fname_j, fname_k)] = {
788                    "grid_j": grid_j.cpu().numpy(),
789                    "grid_k": grid_k.cpu().numpy(),
790                    "surface": surface,
791                }
792
793        return result
794
795    def shapley_values(
796        self,
797        X,
798        X_background,
799        n_background: int = 100,
800    ) -> Dict[str, np.ndarray]:
801        """
802        Compute exact Shapley values using the pairwise additive structure.
803
804        Cost: 2(q+1) forward passes per test sample per background sample.
805        For q=9, n_background=100: ~2000 forward passes per test sample.
806        This is exact — no sampling approximation.
807
808        Args:
809            X: Test data.
810            X_background: Background data for baseline distribution.
811            n_background: Number of background samples to use.
812
813        Returns:
814            Dict mapping feature_name -> (n_test,) array of SHAP values.
815            Values are on the linear predictor scale.
816        """
817        from insurance_gam.pin.shapley import exact_shapley_values
818
819        if not self._is_fitted:
820            raise RuntimeError("Call fit() before shapley_values().")
821
822        x_dict = self._prepare_features(X)
823        x_dict = self._to_device_dict(x_dict)
824        bg_dict = self._prepare_features(X_background)
825        bg_dict = self._to_device_dict(bg_dict)
826
827        return exact_shapley_values(self, x_dict, bg_dict, n_background)
828
829    def interaction_weights(self) -> Dict[Tuple[str, str], float]:
830        """
831        Return output weights w_{jk} for all pairs as a dict.
832
833        Large absolute weights indicate important interactions or main effects.
834        Note: weights are on the scale of the linear predictor; compare
835        w_{jk} * h_{jk} range rather than w_{jk} alone for fair comparison.
836
837        Returns:
838            Dict mapping (fname_j, fname_k) -> float.
839        """
840        pairs = self.interaction_tokens.pair_indices()
841        weights = self.output_weights.detach().cpu().numpy()
842        return {
843            (self.feature_names[j], self.feature_names[k]): float(weights[idx])
844            for idx, (j, k) in enumerate(pairs)
845        }
846
847    def count_parameters(self) -> int:
848        """Count trainable parameters."""
849        return sum(p.numel() for p in self.parameters() if p.requires_grad)