pycsamt.ai.training.dataset#

PyTorch Dataset wrapper and normalisation utilities for EM training data.

Normalizer applies z-score normalisation (fit on training data, applied to val/test without refit). EMDataset wraps a ForwardDataset and handles:

z-score normalisation of X and y
log₁₀ transform of thickness values in y
NaN masking (NaN targets are left as torch.nan for masked loss)
Optional on-the-fly noise augmentation

Usage#

>>> from pycsamt.forward.batch import generate_dataset
>>> from pycsamt.ai.training.dataset import EMDataset, Normalizer
>>> ds = generate_dataset(n_samples=1000, seed=0, verbose=False)
>>> train_ds = EMDataset(ds, log_thickness=True)
>>> train_ds.x_norm.mean.shape   # (n_features,)

Classes

`EMDataset`(forward_ds, *[, n_layers, ...])	PyTorch `Dataset`-compatible wrapper for a `ForwardDataset`.
`Normalizer`([eps])	Z-score normaliser that handles NaN values.

class pycsamt.ai.training.dataset.Normalizer(eps=1e-08)[source]#

Bases: object

Z-score normaliser that handles NaN values.

Parameters:: eps (float) – Small constant added to std to avoid division by zero.
Variables:: std (mean,) – Per-feature statistics computed by fit().

fit(X)[source]#

Compute mean and std from X, ignoring NaN.

Parameters:: X (ndarray)
Return type:: Normalizer

transform(X)[source]#

Apply z-score normalisation.

Parameters:: X (ndarray)
Return type:: ndarray

inverse_transform(X_norm)[source]#

Reverse z-score normalisation.

Parameters:: X_norm (ndarray)
Return type:: ndarray

fit_transform(X)[source]#

Parameters:: X (ndarray)
Return type:: ndarray

to_dict()[source]#

Return type:: dict

classmethod from_dict(d)[source]#

Parameters:: d (dict)
Return type:: Normalizer

class pycsamt.ai.training.dataset.EMDataset(forward_ds, *, n_layers=None, log_thickness=True, x_norm=None, y_norm=None, augment_noise=0.0)[source]#

Bases: object

PyTorch Dataset-compatible wrapper for a ForwardDataset.

Handles normalisation, log₁₀ thickness transform, and optional on-the-fly noise augmentation.

Parameters:

forward_ds (ForwardDataset) – Source data produced by generate_dataset().
n_layers (int or None) – If given, only samples with meta.n_layers == n_layers are kept. None keeps all samples (variable n_layers, NaN-padded).
log_thickness (bool) – Apply log₁₀ to thickness values in y before normalising. Strongly recommended when thicknesses span > 2 orders of magnitude.
x_norm (Normalizer or None) – Pre-fitted normaliser for X. If None a new one is fitted on this dataset (use on training split only).
y_norm (Normalizer or None) – Pre-fitted normaliser for y.
augment_noise (float) – If > 0, add Gaussian noise (this level) to X on-the-fly each epoch.

Variables:

X (ndarray, shape (n_valid, n_features))
y (ndarray, shape (n_valid, n_params))
y_norm (x_norm,)
n_features (int)
n_params (int)

inverse_y(y_norm)[source]#

Undo normalisation + log₁₀ thickness transform on predicted y.

Returns raw parameter vector: [rho_0…rho_{n-1}, thick_0…thick_{n-2}] where rho values are in Ω·m (not log) and thicknesses in metres.

Parameters:: y_norm (ndarray)
Return type:: ndarray

inverse_x(x_norm)[source]#

Undo X normalisation.

Parameters:: x_norm (ndarray)
Return type:: ndarray

split(val_frac=0.1, seed=None)[source]#

Split into train and validation EMDataset objects.

The val split uses the same normaliser fitted on the train split.

Returns:

(train_ds, val_ds)

Return type:

tuple of EMDataset

Parameters:

val_frac (float)
seed (int | None)