Skip to content

ML API

Datasets & loaders

edl_ml.ml.dataset

Torch dataset helpers for the capacitance surrogate.

StandardScalerTensor dataclass

Simple mean/std standardiser stored as tensors.

Fitting is performed on the training set so no test information leaks into the normalisation statistics.

Source code in src/edl_ml/ml/dataset.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@dataclass(slots=True)
class StandardScalerTensor:
    """Simple mean/std standardiser stored as tensors.

    Fitting is performed on the training set so no test information leaks
    into the normalisation statistics.
    """

    mean: torch.Tensor
    std: torch.Tensor

    @classmethod
    def fit(cls, array: NDArray[np.float64]) -> StandardScalerTensor:
        """Fit mean/std on the given 2D array and return a scaler."""
        mean = torch.as_tensor(array.mean(axis=0), dtype=torch.float32)
        std = torch.as_tensor(array.std(axis=0), dtype=torch.float32)
        std = torch.where(std > 1e-8, std, torch.ones_like(std))
        return cls(mean=mean, std=std)

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        """Scale ``x`` using the fitted mean and std."""
        return (x - self.mean) / self.std

    def inverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        """Invert the transformation."""
        return x * self.std + self.mean

fit classmethod 

fit(array: NDArray[float64]) -> StandardScalerTensor

Fit mean/std on the given 2D array and return a scaler.

Source code in src/edl_ml/ml/dataset.py 
30
31
32
33
34
35
36
@classmethod
def fit(cls, array: NDArray[np.float64]) -> StandardScalerTensor:
    """Fit mean/std on the given 2D array and return a scaler."""
    mean = torch.as_tensor(array.mean(axis=0), dtype=torch.float32)
    std = torch.as_tensor(array.std(axis=0), dtype=torch.float32)
    std = torch.where(std > 1e-8, std, torch.ones_like(std))
    return cls(mean=mean, std=std)

inverse_transform 

inverse_transform(x: Tensor) -> torch.Tensor

Invert the transformation.

Source code in src/edl_ml/ml/dataset.py 
42
43
44
def inverse_transform(self, x: torch.Tensor) -> torch.Tensor:
    """Invert the transformation."""
    return x * self.std + self.mean

transform 

transform(x: Tensor) -> torch.Tensor

Scale x using the fitted mean and std.

Source code in src/edl_ml/ml/dataset.py 
38
39
40
def transform(self, x: torch.Tensor) -> torch.Tensor:
    """Scale ``x`` using the fitted mean and std."""
    return (x - self.mean) / self.std

CapacitanceDataset

Bases: Dataset[tuple[Tensor, Tensor]]

PyTorch dataset wrapping a long-format capacitance DataFrame.

Parameters:

Name Type Description Default
df DataFrame

DataFrame produced by :func:edl_ml.data.generate.build_capacitance_dataset.

 required
x_scaler StandardScalerTensor

Fitted scaler for the input features. Required.

 required
y_scaler StandardScalerTensor

Fitted scaler for the target. Required.

 required
Source code in src/edl_ml/ml/dataset.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
class CapacitanceDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
    """PyTorch dataset wrapping a long-format capacitance DataFrame.

    Parameters
    ----------
    df
        DataFrame produced by
        :func:`edl_ml.data.generate.build_capacitance_dataset`.
    x_scaler
        Fitted scaler for the input features. Required.
    y_scaler
        Fitted scaler for the target. Required.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        x_scaler: StandardScalerTensor,
        y_scaler: StandardScalerTensor,
    ) -> None:
        self._x = torch.as_tensor(
            df[list(INPUT_COLUMNS)].to_numpy(dtype=np.float32),
            dtype=torch.float32,
        )
        self._y = torch.as_tensor(
            df[[TARGET_COLUMN]].to_numpy(dtype=np.float32),
            dtype=torch.float32,
        )
        self._x = x_scaler.transform(self._x)
        self._y = y_scaler.transform(self._y)

    def __len__(self) -> int:
        return int(self._x.shape[0])

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self._x[idx], self._y[idx]

LoaderBundle dataclass

Container for the three dataloaders and the fitted scalers.

Attributes:

Name Type Description
train, val, test

DataLoaders for each split.
x_scaler, y_scaler

Scalers fitted on the training split only.
Source code in src/edl_ml/ml/dataset.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@dataclass(frozen=True, slots=True)
class LoaderBundle:
    """Container for the three dataloaders and the fitted scalers.

    Attributes
    ----------
    train, val, test
        DataLoaders for each split.
    x_scaler, y_scaler
        Scalers fitted on the training split only.
    """

    train: DataLoader[tuple[torch.Tensor, torch.Tensor]]
    val: DataLoader[tuple[torch.Tensor, torch.Tensor]]
    test: DataLoader[tuple[torch.Tensor, torch.Tensor]]
    x_scaler: StandardScalerTensor
    y_scaler: StandardScalerTensor

build_loaders 

build_loaders(train_df: DataFrame, val_df: DataFrame, test_df: DataFrame, batch_size: int = 256, num_workers: int = 0, seed: int | None = 0) -> LoaderBundle

Fit scalers on train_df and wrap each split in a DataLoader.

Parameters:

Name Type Description Default
train_df DataFrame

Splits returned by :func:edl_ml.data.generate.split_by_sample.

 required
val_df DataFrame

Splits returned by :func:edl_ml.data.generate.split_by_sample.

 required
test_df DataFrame

Splits returned by :func:edl_ml.data.generate.split_by_sample.

 required
batch_size int

Mini-batch size.

 256
num_workers int

DataLoader workers. Keep at 0 on macOS MPS.

 0
seed int | None

Seed for the training shuffle generator.

 0
Source code in src/edl_ml/ml/dataset.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def build_loaders(
    train_df: pd.DataFrame,
    val_df: pd.DataFrame,
    test_df: pd.DataFrame,
    batch_size: int = 256,
    num_workers: int = 0,
    seed: int | None = 0,
) -> LoaderBundle:
    """Fit scalers on ``train_df`` and wrap each split in a DataLoader.

    Parameters
    ----------
    train_df, val_df, test_df
        Splits returned by
        :func:`edl_ml.data.generate.split_by_sample`.
    batch_size
        Mini-batch size.
    num_workers
        DataLoader workers. Keep at 0 on macOS MPS.
    seed
        Seed for the training shuffle generator.
    """
    x_scaler = StandardScalerTensor.fit(train_df[list(INPUT_COLUMNS)].to_numpy(dtype=np.float32))
    y_scaler = StandardScalerTensor.fit(train_df[[TARGET_COLUMN]].to_numpy(dtype=np.float32))
    gen = torch.Generator()
    if seed is not None:
        gen.manual_seed(seed)

    def _loader(
        df: pd.DataFrame, *, shuffle: bool
    ) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
        ds = CapacitanceDataset(df, x_scaler=x_scaler, y_scaler=y_scaler)
        return DataLoader(
            ds,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            generator=gen if shuffle else None,
            drop_last=False,
        )

    return LoaderBundle(
        train=_loader(train_df, shuffle=True),
        val=_loader(val_df, shuffle=False),
        test=_loader(test_df, shuffle=False),
        x_scaler=x_scaler,
        y_scaler=y_scaler,
    )

Model

edl_ml.ml.model

Torch MLP surrogate for the Gouy-Chapman-Stern solver.

CapacitanceMLP

Bases: Module

Fully connected regressor predicting the scaled capacitance.

The network consumes a feature vector of physical parameters concatenated with the electrode potential and emits a single scalar (the standardised differential capacitance at that potential).

Parameters:

Name Type Description Default
config MLPConfig

Hyperparameters.

 required
Source code in src/edl_ml/ml/model.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
class CapacitanceMLP(nn.Module):
    """Fully connected regressor predicting the scaled capacitance.

    The network consumes a feature vector of physical parameters concatenated
    with the electrode potential and emits a single scalar (the standardised
    differential capacitance at that potential).

    Parameters
    ----------
    config
        Hyperparameters.
    """

    def __init__(self, config: MLPConfig) -> None:
        super().__init__()
        self.config = config
        layers: list[nn.Module] = []
        in_dim = config.input_dim
        for h in config.hidden_dims:
            layers.append(nn.Linear(in_dim, h))
            if config.use_batch_norm:
                layers.append(nn.BatchNorm1d(h))
            layers.append(_make_activation(config.activation))
            if config.dropout > 0.0:
                layers.append(nn.Dropout(config.dropout))
            in_dim = h
        layers.append(nn.Linear(in_dim, 1))
        self.net = nn.Sequential(*layers)
        self._init_weights()

    def _init_weights(self) -> None:
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, a=0.0, nonlinearity="relu")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Return the scaled capacitance prediction for input ``x``."""
        out: torch.Tensor = self.net(x)
        return out

    def count_parameters(self) -> int:
        """Number of trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

count_parameters 

count_parameters() -> int

Number of trainable parameters.

Source code in src/edl_ml/ml/model.py 
 98
 99
100
def count_parameters(self) -> int:
    """Number of trainable parameters."""
    return sum(p.numel() for p in self.parameters() if p.requires_grad)

forward 

forward(x: Tensor) -> torch.Tensor

Return the scaled capacitance prediction for input x.

Source code in src/edl_ml/ml/model.py 
93
94
95
96
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Return the scaled capacitance prediction for input ``x``."""
    out: torch.Tensor = self.net(x)
    return out

MLPConfig dataclass

Hyperparameters of the capacitance MLP.

Parameters:

Name Type Description Default
input_dim int

Number of input features. Equal to len(INPUT_COLUMNS) in the reference setup.

 6
hidden_dims tuple[int, ...]

Sequence of hidden-layer widths.

 (128, 128, 64)
activation str

One of "relu", "silu" or "gelu".

 'silu'
dropout float

Probability for the dropout applied after each hidden layer.

 0.05
use_batch_norm bool

Whether to insert a BatchNorm layer between linear and activation.

 False
Source code in src/edl_ml/ml/model.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@dataclass(frozen=True, slots=True)
class MLPConfig:
    """Hyperparameters of the capacitance MLP.

    Parameters
    ----------
    input_dim
        Number of input features. Equal to ``len(INPUT_COLUMNS)`` in the
        reference setup.
    hidden_dims
        Sequence of hidden-layer widths.
    activation
        One of ``"relu"``, ``"silu"`` or ``"gelu"``.
    dropout
        Probability for the dropout applied after each hidden layer.
    use_batch_norm
        Whether to insert a BatchNorm layer between linear and activation.
    """

    input_dim: int = 6
    hidden_dims: tuple[int, ...] = (128, 128, 64)
    activation: str = "silu"
    dropout: float = 0.05
    use_batch_norm: bool = False

    def __post_init__(self) -> None:
        if self.input_dim < 1:
            raise ValueError("input_dim must be positive")
        if any(h < 1 for h in self.hidden_dims):
            raise ValueError("hidden_dims must be positive integers")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must lie in [0, 1)")
        if self.activation not in {"relu", "silu", "gelu"}:
            raise ValueError(f"unknown activation: {self.activation}")

Training

edl_ml.ml.train

Training loop for the capacitance surrogate with optional MLflow logging.

TrainConfig dataclass

Training hyperparameters.

Parameters:

Name Type Description Default
learning_rate float

Standard optimiser/loop controls.

 0.001
weight_decay float

Standard optimiser/loop controls.

 0.001
max_epochs float

Standard optimiser/loop controls.

 0.001
batch_size float

Standard optimiser/loop controls.

 0.001
patience int

Early-stopping patience in epochs of unimproved validation loss.

 20
grad_clip float

Gradient-norm clip threshold; set to 0 to disable.

 1.0
device str

"auto" selects CUDA > MPS > CPU. Otherwise a torch device string.

 'auto'
seed int

Seed for torch RNGs.

 0
mlflow_experiment str | None

If not None, enable MLflow logging under this experiment.

 None
mlflow_tracking_uri str | None

Optional MLflow tracking URI; defaults to local mlruns/.

 None
Source code in src/edl_ml/ml/train.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
@dataclass(frozen=True, slots=True)
class TrainConfig:
    """Training hyperparameters.

    Parameters
    ----------
    learning_rate, weight_decay, max_epochs, batch_size
        Standard optimiser/loop controls.
    patience
        Early-stopping patience in epochs of unimproved validation loss.
    grad_clip
        Gradient-norm clip threshold; set to 0 to disable.
    device
        ``"auto"`` selects CUDA > MPS > CPU. Otherwise a torch device string.
    seed
        Seed for torch RNGs.
    mlflow_experiment
        If not ``None``, enable MLflow logging under this experiment.
    mlflow_tracking_uri
        Optional MLflow tracking URI; defaults to local ``mlruns/``.
    """

    learning_rate: float = 1e-3
    weight_decay: float = 1e-5
    max_epochs: int = 200
    batch_size: int = 256
    patience: int = 20
    grad_clip: float = 1.0
    device: str = "auto"
    seed: int = 0
    mlflow_experiment: str | None = None
    mlflow_tracking_uri: str | None = None

TrainReport dataclass

Summary of a training run.

Attributes:

Name Type Description
model CapacitanceMLP

The trained model (loaded with the best validation checkpoint).
best_val_loss float

Minimum validation loss observed.
train_losses, val_losses

Per-epoch training and validation mean squared error on the scaled target.
test_metrics dict[str, float]

Dictionary of metrics evaluated on the test set in unscaled units (µF/cm²): mse, rmse, mae, r2, mape.
x_scaler, y_scaler

Scalers used during training; persisted alongside the model for deployment-time inference.
Source code in src/edl_ml/ml/train.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@dataclass(slots=True)
class TrainReport:
    """Summary of a training run.

    Attributes
    ----------
    model
        The trained model (loaded with the best validation checkpoint).
    best_val_loss
        Minimum validation loss observed.
    train_losses, val_losses
        Per-epoch training and validation mean squared error on the scaled
        target.
    test_metrics
        Dictionary of metrics evaluated on the test set in unscaled units
        (µF/cm²): ``mse``, ``rmse``, ``mae``, ``r2``, ``mape``.
    x_scaler, y_scaler
        Scalers used during training; persisted alongside the model for
        deployment-time inference.
    """

    model: CapacitanceMLP
    best_val_loss: float
    train_losses: list[float] = field(default_factory=list)
    val_losses: list[float] = field(default_factory=list)
    test_metrics: dict[str, float] = field(default_factory=dict)
    x_scaler: StandardScalerTensor | None = None
    y_scaler: StandardScalerTensor | None = None

train_model 

train_model(loaders: LoaderBundle, model_config: MLPConfig, train_config: TrainConfig, *, checkpoint_path: Path | str | None = None) -> TrainReport

Train the capacitance MLP and return a :class:TrainReport.

Supports optional MLflow logging. A best-on-validation checkpoint is kept in memory and optionally serialised to checkpoint_path. Early stopping triggers when the validation loss fails to improve for train_config.patience consecutive epochs.

Source code in src/edl_ml/ml/train.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def train_model(
    loaders: LoaderBundle,
    model_config: MLPConfig,
    train_config: TrainConfig,
    *,
    checkpoint_path: Path | str | None = None,
) -> TrainReport:
    """Train the capacitance MLP and return a :class:`TrainReport`.

    Supports optional MLflow logging. A best-on-validation checkpoint is kept
    in memory and optionally serialised to ``checkpoint_path``. Early
    stopping triggers when the validation loss fails to improve for
    ``train_config.patience`` consecutive epochs.
    """
    _set_seed(train_config.seed)
    device = _resolve_device(train_config.device)

    model = CapacitanceMLP(model_config).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=train_config.learning_rate,
        weight_decay=train_config.weight_decay,
    )
    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, train_config.max_epochs))
    loss_fn = nn.MSELoss()

    mlflow_ctx = _maybe_mlflow(train_config, model_config)

    best_val = float("inf")
    best_state: dict[str, torch.Tensor] | None = None
    patience_left = train_config.patience
    report = TrainReport(
        model=model,
        best_val_loss=float("inf"),
        x_scaler=loaders.x_scaler,
        y_scaler=loaders.y_scaler,
    )

    with mlflow_ctx as run:
        if run is not None:
            _log_params(run, model_config, train_config)

        for epoch in range(train_config.max_epochs):
            train_loss = _epoch_pass(
                model,
                loaders.train,
                loss_fn,
                device,
                optimizer=optimizer,
                grad_clip=train_config.grad_clip,
            )
            val_loss = _epoch_pass(model, loaders.val, loss_fn, device)
            scheduler.step()
            report.train_losses.append(train_loss)
            report.val_losses.append(val_loss)
            if run is not None:
                _log_metrics(
                    run,
                    {"train_loss": train_loss, "val_loss": val_loss},
                    step=epoch,
                )
            logger.info(
                "epoch=%d train=%.6f val=%.6f lr=%.2e",
                epoch,
                train_loss,
                val_loss,
                optimizer.param_groups[0]["lr"],
            )
            if val_loss + 1e-9 < best_val:
                best_val = val_loss
                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
                patience_left = train_config.patience
            else:
                patience_left -= 1
                if patience_left <= 0:
                    logger.info("early stopping at epoch %d", epoch)
                    break

        if best_state is not None:
            model.load_state_dict(best_state)
        report.best_val_loss = best_val
        report.test_metrics = _evaluate_unscaled(model, loaders.test, loaders.y_scaler, device)
        if run is not None:
            _log_metrics(run, {f"test_{k}": v for k, v in report.test_metrics.items()})

    if checkpoint_path is not None:
        _save_checkpoint(checkpoint_path, model, loaders, model_config, train_config)
    return report

edl_ml.ml.tune

Optuna hyperparameter search for the capacitance surrogate.

TuneConfig dataclass

Controls for the Optuna study.

Parameters:

Name Type Description Default
n_trials int

Number of hyperparameter trials.

 40
timeout_seconds float | None

Optional wall-clock cap. None means no cap.

 None
max_epochs int

Upper bound on training epochs inside each trial.

 150
patience int

Early-stopping patience inside each trial.

 15
study_name str

Optuna study identifiers. storage=None uses an in-memory study.

 'edl-ml-surrogate'
storage str

Optuna study identifiers. storage=None uses an in-memory study.

 'edl-ml-surrogate'
direction str

"minimize" (on val loss) or "maximize" (e.g. on -MSE).

 'minimize'
seed int

RNG seed controlling the TPE sampler.

 0
Source code in src/edl_ml/ml/tune.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@dataclass(frozen=True, slots=True)
class TuneConfig:
    """Controls for the Optuna study.

    Parameters
    ----------
    n_trials
        Number of hyperparameter trials.
    timeout_seconds
        Optional wall-clock cap. ``None`` means no cap.
    max_epochs
        Upper bound on training epochs inside each trial.
    patience
        Early-stopping patience inside each trial.
    study_name, storage
        Optuna study identifiers. ``storage=None`` uses an in-memory study.
    direction
        ``"minimize"`` (on val loss) or ``"maximize"`` (e.g. on -MSE).
    seed
        RNG seed controlling the TPE sampler.
    """

    n_trials: int = 40
    timeout_seconds: float | None = None
    max_epochs: int = 150
    patience: int = 15
    study_name: str = "edl-ml-surrogate"
    storage: str | None = None
    direction: str = "minimize"
    seed: int = 0

run_optuna_study 

run_optuna_study(loaders: LoaderBundle, input_dim: int, tune_config: TuneConfig, base_train_config: TrainConfig | None = None) -> Any

Launch an Optuna TPE study optimising validation MSE.

Parameters:

Name Type Description Default
loaders LoaderBundle

Pre-built data loaders (train/val/test + scalers).

 required
input_dim int

Number of input features.

 required
tune_config TuneConfig

Study configuration.

 required
base_train_config TrainConfig | None

Template training config; only learning-rate / weight-decay are overridden by the sampler.

 None

Returns:

Type Description
Study

The completed study; best trial accessible via .best_trial.
Source code in src/edl_ml/ml/tune.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def run_optuna_study(
    loaders: LoaderBundle,
    input_dim: int,
    tune_config: TuneConfig,
    base_train_config: TrainConfig | None = None,
) -> Any:
    """Launch an Optuna TPE study optimising validation MSE.

    Parameters
    ----------
    loaders
        Pre-built data loaders (train/val/test + scalers).
    input_dim
        Number of input features.
    tune_config
        Study configuration.
    base_train_config
        Template training config; only learning-rate / weight-decay are
        overridden by the sampler.

    Returns
    -------
    optuna.Study
        The completed study; best trial accessible via ``.best_trial``.
    """
    import optuna

    base = base_train_config or TrainConfig(
        max_epochs=tune_config.max_epochs,
        patience=tune_config.patience,
    )

    def objective(trial: optuna.Trial) -> float:
        mcfg = _suggest_model_config(trial, input_dim)
        tcfg = _suggest_train_config(trial, base)
        report = train_model(loaders, mcfg, tcfg)
        return float(report.best_val_loss)

    sampler = optuna.samplers.TPESampler(seed=tune_config.seed)
    study = optuna.create_study(
        study_name=tune_config.study_name,
        storage=tune_config.storage,
        direction=tune_config.direction,
        load_if_exists=tune_config.storage is not None,
        sampler=sampler,
    )
    study.optimize(
        objective,
        n_trials=tune_config.n_trials,
        timeout=tune_config.timeout_seconds,
        gc_after_trial=True,
    )
    return study

Explanations

edl_ml.ml.explain

SHAP and permutation-based explanations for the capacitance surrogate.

ShapResult dataclass

Output of a SHAP explanation run.

Attributes:

Name Type Description
values NDArray[float64]

SHAP values in unscaled capacitance units, shape (n_samples, n_features).
features NDArray[float64]

Corresponding raw feature values, shape (n_samples, n_features).
feature_names tuple[str, ...]

Names of the features in column order.
base_value float

The expected model output (mean over the background set) in the same units as values.
Source code in src/edl_ml/ml/explain.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
@dataclass(frozen=True, slots=True)
class ShapResult:
    """Output of a SHAP explanation run.

    Attributes
    ----------
    values
        SHAP values in *unscaled* capacitance units, shape
        ``(n_samples, n_features)``.
    features
        Corresponding raw feature values, shape ``(n_samples, n_features)``.
    feature_names
        Names of the features in column order.
    base_value
        The expected model output (mean over the background set) in the same
        units as ``values``.
    """

    values: NDArray[np.float64]
    features: NDArray[np.float64]
    feature_names: tuple[str, ...]
    base_value: float

permutation_feature_importance 

permutation_feature_importance(model: CapacitanceMLP, features: NDArray[float64], targets: NDArray[float64], x_scaler: StandardScalerTensor, y_scaler: StandardScalerTensor, *, n_repeats: int = 20, seed: int = 0) -> NDArray[np.float64]

Return mean permutation feature importance (increase in MAE per feature).

Useful as a lightweight, no-external-dependency alternative to SHAP. Each feature column is independently shuffled n_repeats times and the mean increase in mean-absolute error (µF/cm²) is recorded.

Source code in src/edl_ml/ml/explain.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def permutation_feature_importance(
    model: CapacitanceMLP,
    features: NDArray[np.float64],
    targets: NDArray[np.float64],
    x_scaler: StandardScalerTensor,
    y_scaler: StandardScalerTensor,
    *,
    n_repeats: int = 20,
    seed: int = 0,
) -> NDArray[np.float64]:
    """Return mean permutation feature importance (increase in MAE per feature).

    Useful as a lightweight, no-external-dependency alternative to SHAP. Each
    feature column is independently shuffled ``n_repeats`` times and the mean
    increase in mean-absolute error (µF/cm²) is recorded.
    """
    rng = np.random.default_rng(seed)
    base_pred = _predict_unscaled(model, features, x_scaler, y_scaler)
    base_mae = float(np.mean(np.abs(base_pred - targets)))
    importances = np.zeros(features.shape[1])
    for col in range(features.shape[1]):
        deltas = np.zeros(n_repeats)
        for rep in range(n_repeats):
            shuffled = features.copy()
            rng.shuffle(shuffled[:, col])
            pred = _predict_unscaled(model, shuffled, x_scaler, y_scaler)
            deltas[rep] = float(np.mean(np.abs(pred - targets))) - base_mae
        importances[col] = float(np.mean(deltas))
    return importances

shap_explain 

shap_explain(model: CapacitanceMLP, background: NDArray[float64], samples: NDArray[float64], x_scaler: StandardScalerTensor, y_scaler: StandardScalerTensor, *, nsamples: int = 200) -> ShapResult

Compute KernelSHAP values for the surrogate on raw feature space.

Parameters:

Name Type Description Default
model CapacitanceMLP

Trained :class:CapacitanceMLP.

 required
background NDArray[float64]

Background dataset, shape (m, n_features). KernelSHAP approximates marginal expectations against this reference distribution.

 required
samples NDArray[float64]

Samples to explain, shape (n, n_features).

 required
x_scaler StandardScalerTensor

Scalers used when training model.

 required
y_scaler StandardScalerTensor

Scalers used when training model.

 required
nsamples int

Number of coalitions used by KernelSHAP per sample.

 200

Returns:

Type Description
ShapResult

Results in unscaled output units.
Source code in src/edl_ml/ml/explain.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def shap_explain(
    model: CapacitanceMLP,
    background: NDArray[np.float64],
    samples: NDArray[np.float64],
    x_scaler: StandardScalerTensor,
    y_scaler: StandardScalerTensor,
    *,
    nsamples: int = 200,
) -> ShapResult:
    """Compute KernelSHAP values for the surrogate on raw feature space.

    Parameters
    ----------
    model
        Trained :class:`CapacitanceMLP`.
    background
        Background dataset, shape ``(m, n_features)``. KernelSHAP approximates
        marginal expectations against this reference distribution.
    samples
        Samples to explain, shape ``(n, n_features)``.
    x_scaler, y_scaler
        Scalers used when training ``model``.
    nsamples
        Number of coalitions used by KernelSHAP per sample.

    Returns
    -------
    ShapResult
        Results in unscaled output units.
    """
    import shap

    def f(x: NDArray[np.float64]) -> NDArray[np.float64]:
        return _predict_unscaled(model, x, x_scaler, y_scaler)

    explainer = shap.KernelExplainer(f, background)
    values = np.asarray(explainer.shap_values(samples, nsamples=nsamples))
    base_value = float(np.mean(f(background)))
    return ShapResult(
        values=values,
        features=np.asarray(samples, dtype=float),
        feature_names=INPUT_COLUMNS,
        base_value=base_value,
    )