insurance_gam.anam

 43class ANAM:
 44    """sklearn-compatible Actuarial Neural Additive Model.
 45
 46    Parameters
 47    ----------
 48    feature_configs:
 49        List of FeatureConfig objects defining features. If None, all
 50        features are treated as continuous with no constraints.
 51    feature_names:
 52        Feature names (used when feature_configs is None). Must match
 53        columns of X passed to fit().
 54    categorical_features:
 55        List of feature names that are categorical. Only used when
 56        feature_configs is None.
 57    monotone_increasing:
 58        Feature names to constrain as monotone increasing.
 59    monotone_decreasing:
 60        Feature names to constrain as monotone decreasing.
 61    link:
 62        Link function: 'log' (Poisson/Tweedie), 'identity' (Gaussian),
 63        'logit' (binary).
 64    loss:
 65        Distributional loss: 'poisson', 'tweedie', 'gamma', 'mse'.
 66    tweedie_p:
 67        Tweedie power parameter (only for loss='tweedie').
 68    interaction_pairs:
 69        List of (feature_i, feature_j) tuples for interaction subnetworks.
 70    hidden_sizes:
 71        Default hidden layer sizes for subnetworks.
 72    n_epochs:
 73        Maximum training epochs.
 74    batch_size:
 75        Mini-batch size.
 76    learning_rate:
 77        Adam learning rate.
 78    lambda_smooth:
 79        Smoothness regularisation weight.
 80    lambda_l2:
 81        L2 ridge weight.
 82    lambda_l1:
 83        L1 sparsity weight.
 84    patience:
 85        Early stopping patience (epochs).
 86    normalize:
 87        If True, standardise continuous features before training. The
 88        scaler is stored and applied automatically during predict().
 89    verbose:
 90        Training verbosity. 0 = silent.
 91    device:
 92        'cpu', 'cuda', or None (auto).
 93    """
 94
 95    def __init__(
 96        self,
 97        feature_configs: Optional[List[FeatureConfig]] = None,
 98        feature_names: Optional[List[str]] = None,
 99        categorical_features: Optional[List[str]] = None,
100        monotone_increasing: Optional[List[str]] = None,
101        monotone_decreasing: Optional[List[str]] = None,
102        link: Literal["log", "identity", "logit"] = "log",
103        loss: Literal["poisson", "tweedie", "gamma", "mse"] = "poisson",
104        tweedie_p: float = 1.5,
105        interaction_pairs: Optional[List[Tuple[str, str]]] = None,
106        hidden_sizes: Optional[List[int]] = None,
107        n_epochs: int = 100,
108        batch_size: int = 512,
109        learning_rate: float = 1e-3,
110        lambda_smooth: float = 1e-4,
111        lambda_l2: float = 1e-4,
112        lambda_l1: float = 0.0,
113        patience: int = 15,
114        normalize: bool = True,
115        verbose: int = 0,
116        device: Optional[str] = None,
117    ) -> None:
118        self.feature_configs = feature_configs
119        self.feature_names = feature_names
120        # Store exactly as passed — no coercion to [] here.
121        # sklearn clone() checks that get_params() round-trips through __init__
122        # exactly. Coercing None->[] would break that invariant.
123        # Defaults are applied at point-of-use in fit() and _build_feature_configs().
124        self.categorical_features = categorical_features
125        self.monotone_increasing = monotone_increasing
126        self.monotone_decreasing = monotone_decreasing
127        self.link = link
128        self.loss = loss
129        self.tweedie_p = tweedie_p
130        self.interaction_pairs = interaction_pairs
131        self.hidden_sizes = hidden_sizes
132        self.n_epochs = n_epochs
133        self.batch_size = batch_size
134        self.learning_rate = learning_rate
135        self.lambda_smooth = lambda_smooth
136        self.lambda_l2 = lambda_l2
137        self.lambda_l1 = lambda_l1
138        self.patience = patience
139        self.normalize = normalize
140        self.verbose = verbose
141        self.device = device
142
143        # Set after fit()
144        self.model_: Optional[ANAMModel] = None
145        self.scaler_: Optional[StandardScaler] = None
146        self.history_: Optional[TrainingHistory] = None
147        self.feature_names_in_: Optional[List[str]] = None
148        self._continuous_col_indices: List[int] = []
149        # P0-1 fix: track both the cached shapes and the n_points used to
150        # build them so we can detect stale cache on different n_points.
151        self._shapes_cache: Optional[Dict[str, ShapeFunction]] = None
152        self._shapes_cache_n_points: Optional[int] = None
153        self._X_train_scaled: Optional[np.ndarray] = None
154        # P0-2 fix: per-feature category remapping tables set during fit().
155        # Maps feature name -> array where remap[original_code] = new_code.
156        self._cat_remap_: Dict[str, np.ndarray] = {}
157
158    def fit(
159        self,
160        X: Union[np.ndarray, pl.DataFrame],
161        y: Union[np.ndarray, pl.Series],
162        sample_weight: Optional[Union[np.ndarray, pl.Series]] = None,
163    ) -> "ANAM":
164        """Fit the ANAM model.
165
166        Parameters
167        ----------
168        X:
169            Feature matrix, shape (n, p). Polars DataFrame accepted.
170        y:
171            Target vector (claim counts, rates, or severities).
172        sample_weight:
173            Exposure weights. For frequency models this is policy duration
174            (e.g. years on risk). If None, uniform exposure assumed.
175
176        Returns
177        -------
178        self
179        """
180        X_arr, y_arr, w_arr, names = self._validate_input(X, y, sample_weight)
181        self.feature_names_in_ = names
182
183        # Build feature configs if not provided. Also populates self._cat_remap_
184        # for any categorical columns that are not zero-indexed consecutive ints.
185        feat_configs = self.feature_configs or self._build_feature_configs(names, X_arr)
186
187        # P0-2 fix: apply category remapping to training data so that the
188        # trainer and scaler see zero-indexed codes, matching the embedding
189        # table sizes that _build_feature_configs computed.
190        if self._cat_remap_:
191            X_arr = self._apply_cat_remap(X_arr)
192
193        # Identify continuous column indices for normalisation
194        self._continuous_col_indices = [
195            i for i, cfg in enumerate(feat_configs)
196            if cfg.feature_type == "continuous"
197        ]
198
199        # Normalise continuous features
200        if self.normalize and self._continuous_col_indices:
201            self.scaler_ = StandardScaler()
202            cont_cols = np.array(self._continuous_col_indices)
203            X_arr_norm = X_arr.copy()
204            X_arr_norm[:, cont_cols] = self.scaler_.fit(
205                X_arr[:, cont_cols], feature_names=[names[i] for i in cont_cols]
206            ).transform(X_arr[:, cont_cols])
207        else:
208            X_arr_norm = X_arr
209
210        self._X_train_scaled = X_arr_norm
211
212        # Build interaction configs
213        interaction_configs: List[InteractionConfig] = []
214        for fi, fj in (self.interaction_pairs or []):
215            interaction_configs.append(
216                InteractionConfig(feature_i=fi, feature_j=fj)
217            )
218
219        # Build model
220        self.model_ = ANAMModel(
221            feature_configs=feat_configs,
222            link=self.link,
223            interaction_configs=interaction_configs,
224            hidden_sizes=self.hidden_sizes or [64, 32],
225        )
226
227        # Build trainer
228        train_cfg = TrainingConfig(
229            loss=self.loss,
230            tweedie_p=self.tweedie_p,
231            n_epochs=self.n_epochs,
232            batch_size=self.batch_size,
233            learning_rate=self.learning_rate,
234            lambda_smooth=self.lambda_smooth,
235            lambda_l2=self.lambda_l2,
236            lambda_l1=self.lambda_l1,
237            patience=self.patience,
238            verbose=self.verbose,
239            device=self.device,
240        )
241        trainer = ANAMTrainer(self.model_, train_cfg)
242        self.history_ = trainer.fit(X_arr_norm, y_arr, exposure=w_arr)
243        self.model_ = trainer.model  # may have moved to device
244
245        return self
246
247    def predict(
248        self,
249        X: Union[np.ndarray, pl.DataFrame],
250        exposure: Optional[Union[np.ndarray, pl.Series]] = None,
251    ) -> np.ndarray:
252        """Predict expected mean (mu) for each observation.
253
254        Parameters
255        ----------
256        X:
257            Feature matrix.
258        exposure:
259            Exposure for each row. If None, exposure=1 assumed (predicts
260            per-policy-year rate for log-link models).
261
262        Returns
263        -------
264        np.ndarray
265            Predicted means, shape (n,).
266        """
267        self._check_fitted()
268
269        X_arr = self._to_array(X)
270        X_arr = self._apply_cat_remap(X_arr)
271        X_arr = self._apply_scaling(X_arr)
272
273        if exposure is not None:
274            if isinstance(exposure, pl.Series):
275                exposure = exposure.to_numpy()
276            exp_arr = np.asarray(exposure, dtype=np.float32)
277            log_exp = torch.tensor(np.log(exp_arr.clip(min=1e-8)), dtype=torch.float32)
278        else:
279            log_exp = None
280
281        X_t = torch.tensor(X_arr, dtype=torch.float32)
282
283        self.model_.eval()
284        with torch.no_grad():
285            mu = self.model_(X_t, log_exposure=log_exp)
286
287        return mu.cpu().numpy()
288
289    def score(
290        self,
291        X: Union[np.ndarray, pl.DataFrame],
292        y: Union[np.ndarray, pl.Series],
293        sample_weight: Optional[Union[np.ndarray, pl.Series]] = None,
294    ) -> float:
295        """Return negative mean deviance (higher = better).
296
297        Follows sklearn convention where score() returns a value where
298        higher is better. We negate deviance so that model.score() can
299        be maximised by hyperparameter search.
300        """
301        self._check_fitted()
302
303        y_arr = np.asarray(y if not isinstance(y, pl.Series) else y.to_numpy(), dtype=np.float32)
304        y_pred = self.predict(X, exposure=sample_weight)
305
306        w = None
307        if sample_weight is not None:
308            w = np.asarray(
309                sample_weight if not isinstance(sample_weight, pl.Series)
310                else sample_weight.to_numpy(),
311                dtype=np.float32,
312            )
313
314        dev = compute_deviance_stat(
315            y_arr, y_pred, exposure=w,
316            loss=self.loss, tweedie_p=self.tweedie_p
317        )
318        return -dev
319
320    def shape_functions(
321        self,
322        n_points: int = 200,
323        category_labels: Optional[Dict[str, Dict[int, str]]] = None,
324    ) -> Dict[str, ShapeFunction]:
325        """Extract and cache shape functions for all features.
326
327        Returns a dict mapping feature name to ShapeFunction. Shape
328        functions are evaluated over the observed training data range.
329
330        The cache is keyed on n_points. Calling with a different n_points
331        than the cached value discards the stale cache and recomputes.
332        """
333        self._check_fitted()
334        assert self._X_train_scaled is not None
335
336        # P0-1 fix: invalidate cache when n_points changes.
337        if self._shapes_cache is None or self._shapes_cache_n_points != n_points:
338            self._shapes_cache = extract_shape_functions(
339                self.model_,
340                self._X_train_scaled,
341                n_points=n_points,
342                category_labels=category_labels,
343            )
344            self._shapes_cache_n_points = n_points
345        return self._shapes_cache
346
347    def feature_importance(self) -> pl.DataFrame:
348        """Feature importance as subnetwork weight norms.
349
350        Returns a Polars DataFrame with columns [feature, importance],
351        sorted by importance descending.
352        """
353        self._check_fitted()
354        imp = self.model_.feature_importance()
355        return pl.DataFrame(
356            {"feature": list(imp.keys()), "importance": list(imp.values())}
357        ).sort("importance", descending=True)
358
359    def get_params(self, deep: bool = True) -> Dict[str, Any]:
360        """sklearn get_params for grid search compatibility.
361
362        Returns ALL constructor parameters — including structural ones
363        (feature_configs, interaction_pairs, etc.) — so that
364        sklearn.base.clone() can reconstruct a functionally equivalent
365        model.
366        """
367        # P0-4 fix: include all constructor parameters, not just numerical ones.
368        return {
369            "feature_configs": self.feature_configs,
370            "feature_names": self.feature_names,
371            "categorical_features": self.categorical_features,
372            "monotone_increasing": self.monotone_increasing,
373            "monotone_decreasing": self.monotone_decreasing,
374            "link": self.link,
375            "loss": self.loss,
376            "tweedie_p": self.tweedie_p,
377            "interaction_pairs": self.interaction_pairs,
378            "hidden_sizes": self.hidden_sizes,
379            "n_epochs": self.n_epochs,
380            "batch_size": self.batch_size,
381            "learning_rate": self.learning_rate,
382            "lambda_smooth": self.lambda_smooth,
383            "lambda_l2": self.lambda_l2,
384            "lambda_l1": self.lambda_l1,
385            "patience": self.patience,
386            "normalize": self.normalize,
387            "verbose": self.verbose,
388            "device": self.device,
389        }
390
391    def set_params(self, **params: Any) -> "ANAM":
392        """sklearn set_params for grid search compatibility."""
393        for key, val in params.items():
394            setattr(self, key, val)
395        return self
396
397    # ------------------------------------------------------------------
398    # Internal helpers
399    # ------------------------------------------------------------------
400
401    def _check_fitted(self) -> None:
402        if self.model_ is None:
403            raise RuntimeError("Call fit() before predict() or score().")
404
405    def _to_array(self, X: Union[np.ndarray, pl.DataFrame]) -> np.ndarray:
406        if isinstance(X, pl.DataFrame):
407            return X.to_numpy()
408        return np.asarray(X, dtype=np.float64)
409
410    def _apply_scaling(self, X_arr: np.ndarray) -> np.ndarray:
411        if self.scaler_ is None or not self._continuous_col_indices:
412            return X_arr.astype(np.float32)
413        X_norm = X_arr.copy().astype(np.float32)
414        cont_cols = np.array(self._continuous_col_indices)
415        X_norm[:, cont_cols] = self.scaler_.transform(X_arr[:, cont_cols].astype(np.float64)).astype(np.float32)
416        return X_norm
417
418    def _apply_cat_remap(self, X_arr: np.ndarray) -> np.ndarray:
419        """Apply stored category remappings.
420
421        P0-2 fix: if any categorical features were remapped during fit
422        (because they were not zero-indexed consecutive integers), apply
423        the same remapping here. Called both during fit() (on training
424        data) and during predict() (on new data) so that codes always
425        align with the trained embedding table.
426        """
427        if not self._cat_remap_:
428            return X_arr
429        if self.feature_names_in_ is None:
430            return X_arr
431        X_out = X_arr.copy()
432        for feat_name, remap in self._cat_remap_.items():
433            if feat_name in self.feature_names_in_:
434                col_idx = self.feature_names_in_.index(feat_name)
435                codes = X_out[:, col_idx].astype(int)
436                X_out[:, col_idx] = remap[codes].astype(X_out.dtype)
437        return X_out
438
439    def _validate_input(
440        self,
441        X: Union[np.ndarray, pl.DataFrame],
442        y: Union[np.ndarray, pl.Series],
443        w: Optional[Union[np.ndarray, pl.Series]],
444    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], List[str]]:
445        """Convert inputs to numpy and extract feature names."""
446        if isinstance(X, pl.DataFrame):
447            names = X.columns
448            X_arr = X.to_numpy().astype(np.float64)
449        else:
450            X_arr = np.asarray(X, dtype=np.float64)
451            names = self.feature_names or [f"f{i}" for i in range(X_arr.shape[1])]
452
453        if isinstance(y, pl.Series):
454            y_arr = y.to_numpy().astype(np.float32)
455        else:
456            y_arr = np.asarray(y, dtype=np.float32)
457
458        if w is not None:
459            if isinstance(w, pl.Series):
460                w_arr: Optional[np.ndarray] = w.to_numpy().astype(np.float32)
461            else:
462                w_arr = np.asarray(w, dtype=np.float32)
463        else:
464            w_arr = None
465
466        return X_arr, y_arr, w_arr, list(names)
467
468    def _build_feature_configs(
469        self, names: List[str], X_arr: np.ndarray
470    ) -> List[FeatureConfig]:
471        """Auto-construct FeatureConfig list from names and constraints.
472
473        P0-2 fix: for categorical features, validate that column values are
474        zero-indexed consecutive integers. If not, auto-remap and warn. The
475        remap is stored in self._cat_remap_ and applied to the training
476        array in fit() before the scaler and trainer receive it.
477        """
478        configs: List[FeatureConfig] = []
479        self._cat_remap_ = {}
480
481        for i, name in enumerate(names):
482            if name in (self.categorical_features or []):
483                col = X_arr[:, i].astype(int)
484                unique_vals = np.unique(col)
485                expected = np.arange(len(unique_vals))
486                if not np.array_equal(unique_vals, expected):
487                    # Not zero-indexed consecutive integers. Auto-remap.
488                    warnings.warn(
489                        f"Categorical feature '{name}' has values {unique_vals.tolist()} "
490                        f"which are not zero-indexed consecutive integers. "
491                        f"Auto-remapping to 0..{len(unique_vals) - 1}. "
492                        f"The same remapping will be applied at predict time. "
493                        f"If you rely on specific category indices in post-hoc analysis, "
494                        f"remap the column yourself before passing to fit().",
495                        UserWarning,
496                        stacklevel=3,
497                    )
498                    # Build a lookup array: remap[original_code] = new_code.
499                    max_orig = int(unique_vals.max())
500                    remap = np.full(max_orig + 1, -1, dtype=int)
501                    for new_code, orig_code in enumerate(unique_vals):
502                        remap[orig_code] = new_code
503                    self._cat_remap_[name] = remap
504                    col = remap[col]
505
506                n_cats = int(col.max()) + 1
507                configs.append(
508                    FeatureConfig(
509                        name=name,
510                        feature_type="categorical",
511                        n_categories=n_cats,
512                    )
513                )
514            else:
515                mono: Literal["increasing", "decreasing", "none"] = "none"
516                if name in (self.monotone_increasing or []):
517                    mono = "increasing"
518                elif name in (self.monotone_decreasing or []):
519                    mono = "decreasing"
520                configs.append(
521                    FeatureConfig(
522                        name=name,
523                        feature_type="continuous",
524                        monotonicity=mono,
525                    )
526                )
527        return configs

 85class ANAMModel(nn.Module):
 86    """Actuarial Neural Additive Model.
 87
 88    Orchestrates one subnetwork per feature plus optional pairwise
 89    interaction networks. The output is:
 90
 91        eta = bias + sum_i f_i(x_i) + sum_{(i,j)} g_{ij}(x_i, x_j)
 92        mu  = link_inverse(eta + log_offset)
 93
 94    Parameters
 95    ----------
 96    feature_configs:
 97        Ordered list of FeatureConfig objects. The column order in X arrays
 98        passed to forward() must match this list.
 99    link:
100        Link function. 'log' for Poisson/Tweedie/Gamma, 'identity' for
101        Gaussian, 'logit' for binary.
102    interaction_configs:
103        Optional list of InteractionConfig for pairwise terms.
104    hidden_sizes:
105        Default hidden sizes for all subnetworks (overridden per-feature
106        by FeatureConfig.hidden_sizes).
107    dropout:
108        Dropout rate applied within subnetworks.
109    """
110
111    def __init__(
112        self,
113        feature_configs: List[FeatureConfig],
114        link: LinkFunction = "log",
115        interaction_configs: Optional[List[InteractionConfig]] = None,
116        hidden_sizes: Optional[List[int]] = None,
117        dropout: float = 0.0,
118    ) -> None:
119        super().__init__()
120
121        self.feature_configs = feature_configs
122        self.link = link
123        self.interaction_configs = interaction_configs or []
124
125        # Map feature names to column indices
126        self.feature_name_to_idx: Dict[str, int] = {
127            cfg.name: i for i, cfg in enumerate(feature_configs)
128        }
129
130        default_hidden = hidden_sizes or [64, 32]
131
132        # Build per-feature subnetworks
133        feature_nets: Dict[str, nn.Module] = {}
134        for cfg in feature_configs:
135            net_hidden = cfg.hidden_sizes or default_hidden
136            if cfg.feature_type == "continuous":
137                feature_nets[cfg.name] = FeatureNetwork(
138                    hidden_sizes=net_hidden,
139                    monotonicity=cfg.monotonicity,
140                    dropout=dropout,
141                )
142            elif cfg.feature_type == "categorical":
143                if cfg.n_categories is None:
144                    raise ValueError(
145                        f"Feature '{cfg.name}' is categorical but n_categories is not set."
146                    )
147                feature_nets[cfg.name] = CategoricalFeatureNetwork(
148                    n_categories=cfg.n_categories,
149                    embedding_dim=cfg.embedding_dim,
150                    hidden_sizes=net_hidden,
151                    dropout=dropout,
152                )
153            else:
154                raise ValueError(f"Unknown feature_type: {cfg.feature_type!r}")
155
156        self.feature_nets = nn.ModuleDict(feature_nets)
157
158        # Build interaction subnetworks
159        interaction_nets: Dict[str, nn.Module] = {}
160        for icfg in self.interaction_configs:
161            key = f"{icfg.feature_i}_x_{icfg.feature_j}"
162            i_idx = self.feature_name_to_idx[icfg.feature_i]
163            j_idx = self.feature_name_to_idx[icfg.feature_j]
164            i_hidden = icfg.hidden_sizes or [32, 16]
165            interaction_nets[key] = InteractionNetwork(
166                feature_indices=(i_idx, j_idx),
167                hidden_sizes=i_hidden,
168                dropout=dropout,
169            )
170
171        self.interaction_nets = nn.ModuleDict(interaction_nets)
172
173        # Scalar bias (learnable)
174        self.bias = nn.Parameter(torch.zeros(1))
175
176    def forward(
177        self,
178        X: torch.Tensor,
179        log_exposure: Optional[torch.Tensor] = None,
180    ) -> torch.Tensor:
181        """Forward pass through all subnetworks.
182
183        Parameters
184        ----------
185        X:
186            Feature matrix, shape (batch, n_features). Continuous features
187            should be pre-normalised. Categorical features should be integer
188            indices (will be cast to long internally).
189        log_exposure:
190            Log of exposure (e.g. log policy duration in years), shape
191            (batch,). Added as offset to the linear predictor before the
192            link function. If None, no offset applied (equivalent to
193            exposure=1 for all observations).
194
195        Returns
196        -------
197        torch.Tensor
198            Predicted means mu, shape (batch,).
199        """
200        batch_size = X.shape[0]
201
202        # Accumulate linear predictor starting from bias
203        eta = self.bias.expand(batch_size)
204
205        # Feature subnetwork contributions
206        for i, cfg in enumerate(self.feature_configs):
207            x_i = X[:, i]
208            net = self.feature_nets[cfg.name]
209
210            if cfg.feature_type == "categorical":
211                contrib = net(x_i.long()).squeeze(-1)
212            else:
213                contrib = net(x_i.float()).squeeze(-1)
214
215            eta = eta + contrib
216
217        # Interaction subnetwork contributions
218        for icfg in self.interaction_configs:
219            key = f"{icfg.feature_i}_x_{icfg.feature_j}"
220            i_idx = self.feature_name_to_idx[icfg.feature_i]
221            j_idx = self.feature_name_to_idx[icfg.feature_j]
222
223            x_i = X[:, i_idx].float()
224            x_j = X[:, j_idx].float()
225
226            contrib = self.interaction_nets[key](x_i, x_j).squeeze(-1)
227            eta = eta + contrib
228
229        # Exposure offset
230        if log_exposure is not None:
231            eta = eta + log_exposure
232
233        # Apply link inverse
234        return self._link_inverse(eta)
235
236    def _link_inverse(self, eta: torch.Tensor) -> torch.Tensor:
237        """Apply inverse link function to linear predictor."""
238        if self.link == "log":
239            return torch.exp(eta)
240        elif self.link == "identity":
241            return eta
242        elif self.link == "logit":
243            return torch.sigmoid(eta)
244        else:
245            raise ValueError(f"Unknown link function: {self.link!r}")
246
247    def linear_predictor(
248        self,
249        X: torch.Tensor,
250        log_exposure: Optional[torch.Tensor] = None,
251    ) -> torch.Tensor:
252        """Return eta (linear predictor) without applying link inverse.
253
254        Useful for inspecting additive contributions before exponentiation.
255        """
256        batch_size = X.shape[0]
257        eta = self.bias.expand(batch_size)
258
259        for i, cfg in enumerate(self.feature_configs):
260            x_i = X[:, i]
261            net = self.feature_nets[cfg.name]
262            if cfg.feature_type == "categorical":
263                contrib = net(x_i.long()).squeeze(-1)
264            else:
265                contrib = net(x_i.float()).squeeze(-1)
266            eta = eta + contrib
267
268        for icfg in self.interaction_configs:
269            key = f"{icfg.feature_i}_x_{icfg.feature_j}"
270            i_idx = self.feature_name_to_idx[icfg.feature_i]
271            j_idx = self.feature_name_to_idx[icfg.feature_j]
272            contrib = self.interaction_nets[key](
273                X[:, i_idx].float(), X[:, j_idx].float()
274            ).squeeze(-1)
275            eta = eta + contrib
276
277        if log_exposure is not None:
278            eta = eta + log_exposure
279
280        return eta
281
282    def feature_contribution(
283        self, X: torch.Tensor, feature_name: str
284    ) -> torch.Tensor:
285        """Return the contribution of a single feature for each observation.
286
287        Useful for explaining individual predictions: the contribution from
288        feature i is exactly f_i(x_i), the output of that subnetwork.
289
290        Parameters
291        ----------
292        X:
293            Feature matrix (batch, n_features).
294        feature_name:
295            Name of the feature to inspect.
296
297        Returns
298        -------
299        torch.Tensor
300            Shape (batch,). Values of f_i(x_i) for each row.
301        """
302        idx = self.feature_name_to_idx[feature_name]
303        cfg = self.feature_configs[idx]
304        net = self.feature_nets[feature_name]
305        x_i = X[:, idx]
306
307        if cfg.feature_type == "categorical":
308            return net(x_i.long()).squeeze(-1)
309        else:
310            return net(x_i.float()).squeeze(-1)
311
312    def project_monotone_weights(self) -> None:
313        """Enforce monotonicity constraints on all relevant subnetworks.
314
315        Call this after optimizer.step() in the training loop.
316        Does nothing for non-monotone and categorical features.
317        """
318        for cfg in self.feature_configs:
319            if cfg.feature_type == "continuous" and cfg.monotonicity != "none":
320                net = self.feature_nets[cfg.name]
321                assert isinstance(net, FeatureNetwork)
322                net.project_weights()
323
324    def feature_importance(self) -> Dict[str, float]:
325        """Estimate feature importance as the L2 norm of output layer weights.
326
327        Larger norm = larger potential contribution from that feature. This
328        is a quick heuristic for feature selection — not a replacement for
329        proper permutation importance or SHAP values on the additive model.
330
331        Returns
332        -------
333        Dict[str, float]
334            Feature name -> importance score, sorted descending.
335        """
336        importances: Dict[str, float] = {}
337
338        for name, net in self.feature_nets.items():
339            total_norm = 0.0
340            for param in net.parameters():
341                total_norm += param.data.norm(2).item() ** 2
342            importances[name] = float(total_norm ** 0.5)
343
344        return dict(sorted(importances.items(), key=lambda x: x[1], reverse=True))
345
346    @property
347    def n_features(self) -> int:
348        """Number of features this model was built for."""
349        return len(self.feature_configs)
350
351    @property
352    def feature_names(self) -> List[str]:
353        """Ordered list of feature names."""
354        return [cfg.name for cfg in self.feature_configs]