55@dataclass
56class SharedTrunkConfig:
57    """Hyperparameters for the shared encoder trunk.
58
59    Parameters
60    ----------
61    hidden_dims:
62        Width of each hidden layer in the trunk.  The trunk has
63        ``len(hidden_dims)`` hidden layers.  A good starting point for motor
64        pricing is ``[128, 64]``; reduce to ``[64, 32]`` for smaller datasets.
65    latent_dim:
66        Dimension of the latent representation h that feeds both heads.
67        Typically 16–64.  Higher values give more capacity but slower training.
68    dropout:
69        Dropout probability applied after each hidden layer (except the last).
70        Set to 0.0 to disable.
71    activation:
72        Activation function for all hidden layers.  ``"elu"`` is preferred for
73        count/severity data because it produces smooth gradients near zero.
74        ``"relu"`` and ``"tanh"`` are also supported.
75    use_batch_norm:
76        Whether to apply batch normalisation after each hidden layer.  Helps
77        training stability when input features have very different scales, but
78        can hurt if batch sizes are very small (< 32).
79    """
80
81    hidden_dims: List[int] = field(default_factory=lambda: [128, 64])
82    latent_dim: int = 32
83    dropout: float = 0.1
84    activation: Literal["elu", "relu", "tanh"] = "elu"
85    use_batch_norm: bool = True

 62@dataclass
 63class TrainingConfig:
 64    """Hyperparameters for the training loop.
 65
 66    Parameters
 67    ----------
 68    max_epochs:
 69        Maximum number of training epochs.
 70    batch_size:
 71        Mini-batch size.  Larger batches give more stable gradient estimates
 72        but use more memory.  For most UK motor datasets 512-2048 works well.
 73    lr:
 74        Learning rate for the Adam optimiser.
 75    trunk_lr_multiplier:
 76        If not 1.0, the trunk uses ``lr * trunk_lr_multiplier`` while the heads
 77        use ``lr``.  A value < 1.0 (e.g. 0.3) slows trunk updates, which can
 78        help when the heads need to adapt faster.
 79    weight_decay:
 80        L2 regularisation applied to all parameters.
 81    loss_weight_sev:
 82        Fixed severity loss weight (used when ``auto_balance=False``).
 83    auto_balance:
 84        Equalise Poisson and Gamma loss magnitudes automatically each step.
 85    patience:
 86        Early stopping patience (epochs without improvement on validation loss).
 87        Set to None to disable early stopping.
 88    min_delta:
 89        Minimum improvement in validation loss to count as an improvement.
 90    lr_reduce_factor:
 91        Factor by which the LR scheduler reduces the learning rate on plateau.
 92    lr_patience:
 93        Epochs without improvement before the LR scheduler fires.
 94    verbose:
 95        Whether to print epoch-level training summaries.
 96    device:
 97        PyTorch device string.  ``"auto"`` selects CUDA if available, else CPU.
 98    """
 99
100    max_epochs: int = 100
101    batch_size: int = 512
102    lr: float = 1e-3
103    trunk_lr_multiplier: float = 1.0
104    weight_decay: float = 1e-4
105    loss_weight_sev: float = 1.0
106    auto_balance: bool = True
107    patience: Optional[int] = 15
108    min_delta: float = 1e-4
109    lr_reduce_factor: float = 0.5
110    lr_patience: int = 5
111    verbose: bool = True
112    device: str = "auto"

 55def prepare_features(
 56    df: pd.DataFrame,
 57    numeric_cols: List[str],
 58    categorical_cols: Optional[List[str]] = None,
 59    transformer: Optional[ColumnTransformer] = None,
 60) -> Tuple[np.ndarray, ColumnTransformer]:
 61    """Encode a DataFrame into a numeric matrix.
 62
 63    Numeric columns are standardised (zero mean, unit variance).  Categorical
 64    columns are one-hot encoded (unknown categories at inference are set to
 65    all-zero rows, not an error).
 66
 67    Does not require torch.
 68
 69    Parameters
 70    ----------
 71    df:
 72        Input data.  Must contain all columns in ``numeric_cols`` and
 73        ``categorical_cols``.
 74    numeric_cols:
 75        Names of numeric/continuous columns.
 76    categorical_cols:
 77        Names of categorical columns.  Pass ``None`` or ``[]`` if there are no
 78        categoricals.
 79    transformer:
 80        A fitted ``ColumnTransformer`` from a previous call to this function.
 81        Pass this when encoding held-out or test data to ensure the same
 82        encoding is applied.  When ``None``, a new transformer is fitted on
 83        ``df``.
 84
 85    Returns
 86    -------
 87    X: np.ndarray of shape (n, n_features_out)
 88        Encoded feature matrix.
 89    transformer: ColumnTransformer
 90        Fitted transformer (reuse for test data).
 91
 92    Examples
 93    --------
 94    >>> X_train, ct = prepare_features(df_train, numeric_cols=["age", "value"],
 95    ...                                categorical_cols=["vehicle_class"])
 96    >>> X_test, _  = prepare_features(df_test, numeric_cols=["age", "value"],
 97    ...                                categorical_cols=["vehicle_class"],
 98    ...                                transformer=ct)
 99    """
100    categorical_cols = categorical_cols or []
101
102    if transformer is None:
103        transformers = []
104        if numeric_cols:
105            transformers.append(
106                ("num", StandardScaler(), numeric_cols)
107            )
108        if categorical_cols:
109            transformers.append(
110                (
111                    "cat",
112                    OneHotEncoder(handle_unknown="ignore", sparse_output=False),
113                    categorical_cols,
114                )
115            )
116        transformer = ColumnTransformer(transformers, remainder="drop")
117        transformer.fit(df)
118
119    X = transformer.transform(df).astype(np.float32)
120    return X, transformer

 37def make_dependent_claims(
 38    n_policies: int = 10_000,
 39    gamma: float = -0.15,
 40    base_freq: float = 0.08,
 41    base_sev: float = 3_000.0,
 42    phi: float = 1.5,
 43    n_features: int = 5,
 44    seed: int = 42,
 45    test_fraction: float = 0.2,
 46) -> Tuple[pd.DataFrame, pd.DataFrame]:
 47    """Generate synthetic motor insurance claims with known frequency-severity dependence.
 48
 49    Data generating process
 50    -----------------------
 51    For each policy i with covariates xᵢ and exposure tᵢ:
 52
 53    1. log λᵢ = log(tᵢ) + β₀ + Xᵢ · β_freq       (Poisson frequency)
 54    2. Nᵢ ~ Poisson(λᵢ · tᵢ)
 55    3. log μᵢ = α₀ + Xᵢ · β_sev + γ · Nᵢ         (GGS conditional severity)
 56    4. If Nᵢ > 0: Ȳᵢ ~ Gamma(shape=Nᵢ/φ, mean=μᵢ); else Ȳᵢ = 0
 57
 58    The parameter γ is the true dependence parameter the model should recover.
 59    Negative γ (default −0.15) means higher-claim-count policies have lower
 60    average severity — the pattern typically found in UK motor.
 61
 62    Parameters
 63    ----------
 64    n_policies:
 65        Number of policies in the full dataset (train + test combined).
 66    gamma:
 67        True dependence parameter.  γ=0 gives independence; γ<0 gives negative
 68        frequency-severity correlation (typical for motor).
 69    base_freq:
 70        Baseline claim frequency (λ at x=0, t=1).
 71    base_sev:
 72        Baseline average severity (μ at x=0, N=0), in pounds.
 73    phi:
 74        Gamma dispersion parameter φ.  Higher values give more severity
 75        variability.
 76    n_features:
 77        Number of synthetic covariates.  First half affect frequency, second
 78        half affect severity; they all share the same feature matrix so the
 79        trunk has genuine heterogeneity to exploit.
 80    seed:
 81        Random seed.
 82    test_fraction:
 83        Fraction of data to hold out as test set.
 84
 85    Returns
 86    -------
 87    df_train, df_test : pd.DataFrame
 88        DataFrames with columns:
 89        ``feature_0``, …, ``feature_{n_features-1}``,
 90        ``exposure``, ``n_claims``, ``avg_severity``, ``total_loss``,
 91        ``true_lambda``, ``true_mu``.
 92
 93    Examples
 94    --------
 95    >>> df_train, df_test = make_dependent_claims(n_policies=20_000, gamma=-0.15)
 96    >>> df_train.head()
 97    """
 98    rng = np.random.default_rng(seed)
 99
100    # --- Covariates ---
101    X = rng.standard_normal((n_policies, n_features)).astype(np.float32)
102
103    # --- Exposure: mix of short- and full-year policies ---
104    exposure = rng.choice([0.25, 0.5, 0.75, 1.0, 1.0, 1.0], size=n_policies).astype(np.float32)
105
106    # --- Regression coefficients ---
107    n_freq_feats = max(1, n_features // 2)
108    n_sev_feats = n_features - n_freq_feats
109    beta_freq = rng.uniform(-0.3, 0.3, size=n_freq_feats).astype(np.float32)
110    beta_sev = rng.uniform(-0.2, 0.2, size=n_sev_feats).astype(np.float32)
111
112    # --- Frequency ---
113    log_lambda_base = np.log(base_freq) + X[:, :n_freq_feats] @ beta_freq
114    lambda_ = np.exp(log_lambda_base) * exposure
115    n_claims = rng.poisson(lambda_).astype(np.float32)
116
117    # --- Conditional severity ---
118    log_mu = np.log(base_sev) + X[:, n_freq_feats:] @ beta_sev + gamma * n_claims
119    mu = np.exp(log_mu)
120
121    avg_severity = np.zeros(n_policies, dtype=np.float32)
122    pos = n_claims > 0
123    if pos.sum() > 0:
124        alpha = n_claims[pos] / phi
125        rate = n_claims[pos] / (phi * mu[pos] + 1e-10)
126        avg_severity[pos] = rng.gamma(shape=alpha, scale=1.0 / rate).astype(np.float32)
127
128    total_loss = n_claims * avg_severity
129
130    feature_cols = [f"feature_{i}" for i in range(n_features)]
131    df = pd.DataFrame(X, columns=feature_cols)
132    df["exposure"] = exposure
133    df["n_claims"] = n_claims
134    df["avg_severity"] = avg_severity
135    df["total_loss"] = total_loss
136    df["true_lambda"] = lambda_ / exposure  # per unit exposure
137    df["true_mu"] = mu
138
139    # --- Split ---
140    n_test = max(1, int(n_policies * test_fraction))
141    df_test = df.iloc[-n_test:].copy().reset_index(drop=True)
142    df_train = df.iloc[:-n_test].copy().reset_index(drop=True)
143
144    return df_train, df_test

147def make_independent_claims(
148    n_policies: int = 10_000,
149    base_freq: float = 0.08,
150    base_sev: float = 3_000.0,
151    phi: float = 1.5,
152    n_features: int = 5,
153    seed: int = 42,
154    test_fraction: float = 0.2,
155) -> Tuple[pd.DataFrame, pd.DataFrame]:
156    """Generate synthetic claims with γ=0 (frequency-severity independence).
157
158    Identical to ``make_dependent_claims`` with ``gamma=0``.  Use this as the
159    null comparison to verify the model does not overfit spurious dependence.
160
161    Parameters
162    ----------
163    n_policies, base_freq, base_sev, phi, n_features, seed, test_fraction:
164        Same as ``make_dependent_claims``.
165
166    Returns
167    -------
168    df_train, df_test : pd.DataFrame
169    """
170    return make_dependent_claims(
171        n_policies=n_policies,
172        gamma=0.0,
173        base_freq=base_freq,
174        base_sev=base_sev,
175        phi=phi,
176        n_features=n_features,
177        seed=seed,
178        test_fraction=test_fraction,
179    )