pycsamt.ai#

AI inversion, neural network architectures, training helpers, AI processing, model-zoo utilities, and diagnostic plotting.

Artificial intelligence and machine learning for EM processing and inversion.

Phase 2 additions#

EMInverter1D — full 1-D inversion pipeline (data loading → normalisation → training → prediction)
nets — CNN1D, ResNet1D, FCN1D architectures
training — EMDataset, EMTrainer, metrics
plot — EMStyle, plot_compare, plot_convergence

Quick start (requires PyTorch)#

>>> from pycsamt.forward.batch import generate_dataset
>>> from pycsamt.ai.inversion import EMInverter1D
>>> ds = generate_dataset(n_samples=2_000, seed=0, n_layers=5)
>>> inv = EMInverter1D(arch="resnet", n_layers=5)
>>> inv.fit(ds, epochs=30)
>>> # Predict on new Z objects or ForwardResponse
>>> y_pred = inv.predict(X_test)

class pycsamt.ai.BaseEMNet(arch='cnn1d', n_layers=5, solver='mt1d', device=None)#

Bases: ABC

Abstract base class for EM neural network estimators.

Subclasses must implement fit(), predict(), and _build_network().

All fit / predict calls accept either numpy arrays (the standard path) or EM domain objects (Z tensors, ForwardResponse, SiteCollection) which are coerced to arrays via the _coerce_input() hook.

Parameters:

arch (str) – Architecture identifier (e.g. 'cnn1d', 'resnet').
n_layers (int) – Number of earth layers to invert for.
solver (str) – The EM method this network targets ('mt1d', 'tem1d', 'csamt1d').
device (str or None) – Compute device ('cpu', 'cuda', 'mps'). If None uses CUDA when available.

abstractmethod fit(X, y=None, **kwargs)#

Train the network.

Parameters:

X (ndarray (n_samples, n_features) or path-like or ForwardDataset) – Training data. A string is treated as a path to a .npz file produced by generate_dataset().
y (ndarray (n_samples, n_params) or None) – Targets. If X is a dataset object or path, y is ignored.
**kwargs – Solver-specific keyword arguments (e.g. epochs, batch_size, lr).

Return type:

self

abstractmethod predict(X)#

Predict subsurface model parameters from data.

Parameters:: X (ndarray (n_samples, n_features) or Z object or list of Z) – Input data in array form or as EM domain objects.
Returns:: y_pred – Predicted log₁₀(ρ) and thickness values.
Return type:: ndarray (n_samples, n_params)

score(X, y, *, metric='rmse')#

Evaluate prediction quality.

Parameters:

X (ndarray) – Input features.
y (ndarray) – True model parameters.
metric ({'rmse', 'r2', 'mae', 'relative_rmse'}) – Evaluation metric.

Returns:

score – Lower is better for RMSE/MAE; higher is better for R².

Return type:

float

fit_predict(X, y, **kwargs)#

Fit then predict on the same data (useful for diagnostics).

Parameters:

X (ndarray)
y (ndarray)

Return type:

ndarray

save(path)#

Save the model weights and hyperparameters to path.

Parameters:: path (str or Path) – Destination file path (a .pt extension is used for PyTorch backends; .npz for numpy-only models).
Return type:: None

classmethod load(path)#

Load a model saved with save().

Parameters:: path (str or Path)
Returns:: estimator
Return type:: instance of the calling class

classmethod from_pretrained(name)#

Load a pre-trained model from the pycsamt model zoo.

Note

Pre-trained weights are hosted separately and will be downloaded on first call. This feature is scheduled for Phase 5 of the pycsamt AI roadmap.

Parameters:: name (str) – Model identifier, e.g. 'mt1d-resnet-v1'.
Raises:: NotImplementedError – Until Phase 5 weights are published.
Return type:: BaseEMNet

class pycsamt.ai.BaseEMProcessor#

Bases: ABC

Abstract base class for ML-based EM data processing tasks.

Unlike BaseEMNet, processors transform data to data rather than data to model parameters.

Subclasses in pycsamt.ai.processing include EMDenoiser, EMQCScorer, and AnomalyDetector.

abstractmethod fit(X, **kwargs)#

Fit the processor to training data.

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

abstractmethod transform(X)#

Apply the learned transformation.

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

fit_transform(X, **kwargs)#

Fit then transform in one call.

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

save(path)#

Save to a checkpoint file.

Parameters:: path (str | Path)
Return type:: None

classmethod load(path)#

Parameters:: path (str | Path)
Return type:: BaseEMProcessor

classmethod load_pretrained(name)#

Parameters:: name (str)
Return type:: BaseEMProcessor

class pycsamt.ai.EMCheckpoint(params, weights=None, history=None, meta=None)#

Bases: object

Minimal checkpoint for saving/loading EM network state.

Holds hyperparameters (serialised as JSON) alongside framework- specific weight data (numpy arrays or a torch state dict path).

Parameters:

params (dict) – Hyperparameters (must be JSON-serialisable).
weights (dict or None) – Framework weight dict (numpy arrays).
history (dict or None) – Training history: {'train_loss': [...], 'val_loss': [...]}.
meta (dict or None) – Extra metadata (solver, freq range, training date, …).

save(path)#

Save to a .npz checkpoint file.

Parameters:: path (str | Path)
Return type:: None

classmethod load(path)#

Load from a .npz checkpoint file.

Parameters:: path (str | Path)
Return type:: EMCheckpoint

class pycsamt.ai.EMInverter1D(arch='resnet', n_layers=5, solver='mt1d', *, device=None, log_thickness=True, include_phase=True, augment_noise=0.02, **net_kwargs)#

Bases: BaseEMNet

1-D EM neural-network inverter.

Supports MT, CSAMT (far-field), and TEM step-off data.

Parameters:

arch ({'resnet', 'cnn1d', 'fcn'}) – Network architecture. 'resnet' (Liu 2021 style) gives the best accuracy on typical MT datasets. 'cnn1d' (Puzyrev style) is faster to train. 'fcn' (Moghadas style) handles variable input length.
n_layers (int) – Number of earth layers to invert for.
solver ({'mt1d', 'csamt1d', 'tem1d'}) – EM method this inverter targets (determines default frequency grid for Z-object input coercion).
device (str or None) – Compute device. Auto-detects CUDA/MPS/CPU if None.
log_thickness (bool) – Apply log₁₀ to thickness in training targets. Strongly recommended when thicknesses span > 2 orders of magnitude.
include_phase (bool) – Include impedance phase in the input feature vector for MT/CSAMT. Setting to False halves the input size.
augment_noise (float) – On-the-fly noise level added to training inputs each epoch.
net_kwargs (dict) – Extra keyword arguments forwarded to the network constructor (e.g. channels=(64, 128, 256) for ResNet).

fit(X, y=None, *, epochs=100, batch_size=256, lr=0.001, patience=20, val_frac=0.1, grad_clip=1.0, seed=None, verbose=True)#

Train the inverter on a ForwardDataset or a .npz file path.

Parameters:

X (ForwardDataset or str or Path) – Training data. If y is also given, X and y are treated as raw feature / target numpy arrays.
y (ndarray or None) – Raw targets (only when X is an ndarray).
epochs (int) – Maximum training epochs.
batch_size (int)
lr (float) – Initial learning rate.
patience (int) – Early-stopping patience.
val_frac (float) – Fraction of training data used for validation.
grad_clip (float or None) – Gradient-norm clipping threshold.
seed (int or None) – Random seed for train/val split.
verbose (bool) – Print per-epoch summary.

Return type:

self

predict(X, *, as_log_rho=True)#

Predict model parameters for new data.

Parameters:

X (ndarray (n_samples, n_features) | Z | list of Z | ForwardResponse) – Input data.
as_log_rho (bool) – If True (default), the returned resistivity values are in log₁₀(Ω·m) space matching the training targets. If False, resistivities are back-transformed to Ω·m.

Returns:

y_pred – Parameter vector: [log10(ρ), log10(h)] (or linear if as_log_rho=False).

Return type:

ndarray, shape (n_samples, n_params)

predict_models(X)#

Predict and return a list of LayeredModel.

Resistivity and thickness are back-transformed from log space.

Return type:: list

predict_response(response)#

Convenience: invert a single ForwardResponse and return the predicted LayeredModel.

Return type:: LayeredModel

save(path)#

Save weights + normaliser + hyperparameters to path.

Parameters:: path (str | Path)
Return type:: None

classmethod load(path)#

Load a saved inverter from path.

Parameters:: path (str | Path)
Return type:: EMInverter1D

classmethod from_pretrained(name, *, cache_dir=None)#

Load a pre-trained model from the pycsamt model zoo.

Pre-trained weights are hosted at earthai-tech/pycsamt-models and are downloaded to ~/.pycsamt/model_zoo/ on first call.

Parameters:

name (str) – Model identifier. Call list_pretrained() to see available models.
cache_dir (str or None) – Override the default download directory.

Return type:

EMInverter1D

Raises:

KeyError – If name is not in the model zoo registry.
RuntimeError – If the download fails (weights not yet publicly available).

Examples

>>> from pycsamt.ai.inversion import EMInverter1D
>>> inv = EMInverter1D.from_pretrained("mt1d-resnet-5layer-v1")

class pycsamt.ai.EMInverter2D(n_components=4, n_depth=40, n_stations=20, n_freqs=32, *, arch='unet', unet_depth=None, channels=None, dropout=0.2, device=None, log_rho_out=True, **net_kwargs)#

Bases: BaseEMNet

U-Net–based 2-D MT inversion estimator.

Parameters:

n_components (int, default 4) – Number of input channels (EM data components per frequency and station).
n_depth (int, default 40) – Number of depth cells in the target resistivity section.
n_stations (int, default 20) – Number of stations along the profile.
n_freqs (int, default 32) – Number of frequency channels in the input data.
arch (str, default 'unet') – Network architecture. Only 'unet' is supported in Phase 4.
device (str or None) – Compute device.
log_rho_out (bool, default True) – If True, targets and predictions are in log₁₀(ρ) scale.
**net_kwargs – Additional keyword arguments forwarded to the architecture factory (e.g. channels, dropout).
unet_depth (int | None)
channels (tuple[int, ...] | None)
dropout (float)

fit(X, y=None, *, epochs=100, batch_size=16, lr=0.001, patience=15, val_frac=0.1, grad_clip=1.0, seed=None, verbose=True)#

Train the 2-D inversion network.

Parameters:

X (ndarray (n_profiles, n_components, n_freqs, n_stations)) – Input data panels.
y (ndarray (n_profiles, n_depth, n_stations)) – Target 2-D log₁₀(ρ) sections.
epochs (int) – Standard training hyper-parameters.
batch_size (int) – Standard training hyper-parameters.
lr (float) – Standard training hyper-parameters.
patience (int) – Standard training hyper-parameters.
val_frac (float) – Standard training hyper-parameters.
grad_clip (float | None) – Standard training hyper-parameters.
seed (int | None) – Standard training hyper-parameters.
verbose (bool) – Standard training hyper-parameters.

Return type:

self

predict(X, *, as_log_rho=True)#

Predict 2-D resistivity sections.

Parameters:

X (ndarray (n_profiles, n_components, n_freqs, n_stations))
as_log_rho (bool) – If True (default) output is log₁₀(ρ); otherwise linear ρ.

Returns:

rho_2d

Return type:

ndarray (n_profiles, n_depth, n_stations)

class pycsamt.ai.GCNInverter3D(n_features=40, n_layers=5, hidden=(256, 128, 64), dropout=0.1, device=None, **net_kwargs)#

Bases: BaseEMNet

Graph-convolutional 3-D MT inversion estimator.

Models the spatial context of a multi-station survey as a graph and applies spectral GCN message-passing before regressing per-station 1-D subsurface models. The result is a spatially coherent 3-D resistivity volume without any external graph library dependency.

Parameters:

n_features (int, default 40) – Per-station input feature dimension. A typical choice is 2 × n_freqs (log10(ρ_a) + phase at each frequency).
n_layers (int, default 5) – Number of depth layers per station. Output size per station is 2 * n_layers - 1 (n_layers log10(ρ) + (n_layers-1) log10(h)).
hidden (sequence of int, default (256, 128, 64)) – Width of each GCN message-passing layer.
dropout (float, default 0.1) – Dropout probability applied between GCN layers during training.
device (str or None) – Compute device ('cpu', 'cuda', 'gpu:0', …). None auto-detects.
**net_kwargs – Extra keyword arguments forwarded to GCNNet.

fit(X, y=None, adjacency=None, *, coords=None, radius=5000.0, epochs=100, batch_size=16, lr=0.001, patience=15, val_frac=0.15, grad_clip=1.0, seed=None, verbose=True)#

Train the 3-D GCN inversion network.

Parameters:

X (ndarray (n_samples, n_stations, n_features) or) – (n_stations, n_features) Per-station MT feature matrices.
y (ndarray (n_samples, n_stations, 2*n_layers-1) or) – (n_stations, 2*n_layers-1) Target per-station model parameters (log10 scale recommended).
adjacency (ndarray (n_stations, n_stations), optional) – Pre-computed normalised adjacency matrix. If None, coords and radius must be provided.
coords (ndarray (n_stations, 2), optional) – Station (x, y) positions used to build the adjacency when adjacency is not given.
radius (float) – Maximum inter-station edge distance in the same units as coords; ignored when adjacency is supplied.
epochs (int) – Standard training hyper-parameters.
batch_size (int) – Standard training hyper-parameters.
lr (float) – Standard training hyper-parameters.
patience (int) – Standard training hyper-parameters.
val_frac (float) – Standard training hyper-parameters.
grad_clip (float | None) – Standard training hyper-parameters.
seed (int | None) – Standard training hyper-parameters.
verbose (bool) – Standard training hyper-parameters.

Return type:

self

predict(X, adjacency=None, *, coords=None, radius=5000.0, as_log_rho=True)#

Predict per-station subsurface models.

Parameters:

X (ndarray (n_stations, n_features) or) – (n_samples, n_stations, n_features)
adjacency (ndarray (n_stations, n_stations), optional) – If None, uses the adjacency stored from fit().
coords (ndarray (n_stations, 2), optional) – Build adjacency from coordinates when adjacency is absent.
radius (float) – Edge radius (used only when coords is supplied).
as_log_rho (bool) – Return log10(ρ) when True; linear-scale ρ otherwise.

Returns:

y_pred

Return type:

ndarray (n_stations, n_out) or (n_samples, n_stations, n_out)

predict_with_uncertainty(X, adjacency=None, *, coords=None, radius=5000.0, n_mc=30)#

MC-dropout uncertainty estimate for 3-D predictions.

Runs n_mc stochastic forward passes with dropout active and returns the mean and pointwise standard deviation.

Parameters:

X (ndarray (..., n_stations, n_features))
adjacency (ndarray, optional)
coords (ndarray, optional)
radius (float)
n_mc (int) – Number of Monte-Carlo dropout samples.

Returns:

mean (ndarray — same shape as predict() output)
std (ndarray — same shape)

Return type:

tuple[ndarray, ndarray]

class pycsamt.ai.JointInverter(n_features_list, n_layers=5, *, growth_rate=32, n_dense_layers=6, hidden_dim=256, dropout=0.2, device=None, log_thickness=True, **net_kwargs)#

Bases: BaseEMNet

Multi-modal joint inversion estimator based on DRCNN.

Parameters:

n_features_list (sequence of int) – Feature vector lengths for each modality. E.g. (120, 48) for two modalities.
n_layers (int, default 5) – Number of earth layers to invert for.
growth_rate (int, default 32) – Dense block growth rate.
n_dense_layers (int, default 6) – Sub-layers per dense block.
hidden_dim (int, default 256) – Encoded-feature dimension for each modality and fusion stage.
dropout (float, default 0.2)
device (str or None)
log_thickness (bool, default True) – Apply log₁₀ transform to thickness targets before normalisation.
**net_kwargs – Forwarded to DRCNNNet.

fit(X_list, y, *, epochs=100, batch_size=256, lr=0.001, patience=20, val_frac=0.1, grad_clip=1.0, seed=None, verbose=True)#

Train the joint inverter.

Parameters:

X_list (list of ndarray) – Feature matrices for each modality, each (n_samples, n_features_i). Lengths must match.
y (ndarray (n_samples, 2*n_layers-1)) – Target model parameters: first n_layers columns are log₁₀(ρ) or ρ; last n_layers-1 columns are thicknesses.
epochs (int) – Standard training hyper-parameters.
batch_size (int) – Standard training hyper-parameters.
lr (float) – Standard training hyper-parameters.
patience (int) – Standard training hyper-parameters.
val_frac (float) – Standard training hyper-parameters.
grad_clip (float | None) – Standard training hyper-parameters.
seed (int | None) – Standard training hyper-parameters.
verbose (bool) – Standard training hyper-parameters.

Return type:

self

predict(X_list, *, as_log_rho=True)#

Predict subsurface model parameters.

Parameters:

X_list (list of ndarray, each (n_samples, n_features_i))
as_log_rho (bool) – If True (default), resistivity columns are returned in log₁₀(ρ); otherwise linear.

Returns:

y_pred

Return type:

ndarray (n_samples, 2*n_layers-1)

class pycsamt.ai.EnsembleInverter(base_estimator, n_estimators=5, seeds=None)#

Bases: object

Deep ensemble of EMInverter1D models.

Parameters:

base_estimator (EMInverter1D) – Template estimator (unfitted). Its hyper-parameters are copied for every ensemble member.
n_estimators (int, default 5) – Number of ensemble members.
seeds (list of int or None) – Per-member random seeds. When None, seeds [0, 1, …, n_estimators-1] are used.

Examples

>>> from pycsamt.ai.inversion import EMInverter1D, EnsembleInverter
>>> base = EMInverter1D(arch="resnet", n_layers=5)
>>> ens = EnsembleInverter(base, n_estimators=5)
>>> ens.fit(ds, epochs=50)
EnsembleInverter(n_estimators=5, fitted)
>>> mean, std = ens.predict_with_uncertainty(X_test)

fit(X, y=None, *, epochs=100, batch_size=256, lr=0.001, patience=20, val_frac=0.1, grad_clip=1.0, verbose=True)#

Train all ensemble members.

Parameters:

X (ForwardDataset, path, or ndarray) – Training data (forwarded unchanged to each member’s fit).
y (ndarray or None)
epochs (int) – Forwarded to each member’s fit call; each member receives a unique seed.
batch_size (int) – Forwarded to each member’s fit call; each member receives a unique seed.
lr (float) – Forwarded to each member’s fit call; each member receives a unique seed.
patience (int) – Forwarded to each member’s fit call; each member receives a unique seed.
val_frac (float) – Forwarded to each member’s fit call; each member receives a unique seed.
grad_clip (float | None) – Forwarded to each member’s fit call; each member receives a unique seed.
verbose (bool) – Forwarded to each member’s fit call; each member receives a unique seed.

Return type:

self

predict(X, **kwargs)#

Return the mean ensemble prediction.

Parameters:

X (ndarray (n_samples, n_features))
**kwargs (forwarded to each member’s predict)

Returns:

y_mean

Return type:

ndarray (n_samples, n_params)

predict_with_uncertainty(X, _use_calibrated=True, **kwargs)#

Return (mean, std) ensemble prediction.

If calibrate() has been called, the returned std is the calibrated standard deviation from the PosteriorCalibrator (i.e. the Gaussianization-flow recalibrated sigma). Pass _use_calibrated=False to retrieve the raw inter-member std.

Parameters:

X (ndarray (n_samples, n_features))
_use_calibrated (bool) – When True (default) and a calibrator is attached, return calibrated sigma. Internal flag — not part of the public API.

Returns:

mean (ndarray (n_samples, n_params))
std (ndarray (n_samples, n_params)) – Calibrated or raw inter-member standard deviation.

Return type:

tuple[ndarray, ndarray]

predict_quantiles(X, q=(0.05, 0.25, 0.5, 0.75, 0.95), **kwargs)#

Return per-quantile predictions.

Parameters:

X (ndarray (n_samples, n_features))
q (sequence of float) – Quantiles in (0, 1).

Returns:

quantiles

Return type:

dict {q_value: ndarray (n_samples, n_params)}

calibrate(X_cal, y_cal, *, alpha=0.1)#

Attach a ConformalPredictor and a PosteriorCalibrator fitted on a held-out calibration set.

After calling this method:

predict_intervals() returns conformal prediction bands with guaranteed $\ge 1-\alpha$ marginal coverage (Vovk et al. 2005).
predict_posterior() returns calibrated posterior samples via the Gaussianization normalising flow (Kuleshov et al. 2018).
predict_with_uncertainty() returns the calibrated standard deviation instead of the raw inter-member std.

Parameters:

X_cal (ndarray, shape (n_cal, n_features)) – Calibration inputs (must not overlap with the training set).
y_cal (ndarray, shape (n_cal, n_params)) – True parameter vectors for the calibration inputs.
alpha (float) – Default significance level for conformal intervals.

Return type:

self

predict_intervals(X, alpha=None)#

Conformal prediction intervals with guaranteed marginal coverage.

Requires a prior call to calibrate().

Parameters:

X (ndarray, shape (n_samples, n_features))
alpha (float or None) – Significance level. Defaults to the value used in calibrate().

Returns:

center (ndarray, shape (n_samples, n_params))
lower (ndarray, shape (n_samples, n_params))
upper (ndarray, shape (n_samples, n_params))

Return type:

tuple[ndarray, ndarray, ndarray]

predict_posterior(X, n_samples=500, rng=None)#

Draw calibrated posterior samples via the Gaussianization normalising flow (Kuleshov et al. 2018).

Requires a prior call to calibrate().

Parameters:

X (ndarray, shape (n_pts, n_features))
n_samples (int) – Number of posterior draws per input point.
rng (numpy.random.Generator or None)

Returns:

draws

Return type:

ndarray, shape (n_samples, n_pts, n_params)

coverage_diagnostics(X_test, y_test, alphas=None)#

Per-alpha empirical coverage table (requires prior calibrate()).

Useful for a reliability diagram: nominal coverage 1-alpha on the x-axis, actual coverage on the y-axis. A perfectly calibrated predictor lies on the diagonal.

Parameters:

X_test (ndarray)
y_test (ndarray)
alphas (sequence of float or None)

Returns:

diag

Return type:

dict {alpha: actual_coverage}

score(X, y, *, metric='rmse')#

Score the ensemble mean prediction.

Parameters:

X (ndarray (n_samples, n_features))
y (ndarray (n_samples, n_params))
metric ({'rmse', 'mae', 'r2', 'relative_rmse'})

Returns:

score

Return type:

float

coverage(X, y, *, n_sigma=1.96)#

Fraction of true values within ±n_sigma * std of the ensemble mean.

An uncertainty estimate is well-calibrated when coverage(1.96) ≈ 0.95.

Parameters:

X (ndarray)
y (ndarray)
n_sigma (float)

Returns:

coverage

Return type:

float ∈ [0, 1]

save(path)#

Save all ensemble members.

Creates a directory <path>/ containing member_00.npz, member_01.npz, ….

Parameters:: path (str or Path)
Return type:: None

classmethod load(path, base_class=None)#

Load a saved ensemble.

Parameters:

path (str or Path) – Directory created by save().
base_class (type or None) – Estimator class to use for loading. Defaults to EMInverter1D.

Return type:

EnsembleInverter

plot_uncertainty_profile(X, sample_idx=0, *, y_true=None, n_sigma=1.96, **plot_kwargs)#

Plot a 1-D resistivity profile with uncertainty bands.

Parameters:

X (ndarray (n_samples, n_features))
sample_idx (int) – Which sample to plot.
y_true (ndarray (n_params,) or None) – True model vector for overlay.
n_sigma (float) – Band width in standard deviations.
**plot_kwargs – Forwarded to plot_uncertainty_bands().

Returns:

fig

Return type:

Figure

class pycsamt.ai.ConformalPredictor(base_predictor, alpha=0.1, eps=1e-10)#

Bases: object

Split conformal prediction wrapper for calibrated regression intervals.

Given a fitted base predictor $f(\mathbf{x})$ that also exposes per-output uncertainty estimates $\boldsymbol{\sigma}(\mathbf{x})$ (e.g. EnsembleInverter), this class fits normalised nonconformity scores on a held-out calibration set and uses them to produce prediction bands with provable marginal coverage.

Normalised nonconformity score (per sample, over all output dims):

\[s_i = \max_j \frac{|y_{ij} - f(\mathbf{x}_i)_j|} {\sigma_{ij} + \varepsilon}\]

Calibration quantile (finite-sample correction, Vovk et al. 2005):

\[\hat{q} = \text{Quantile}_{1-\alpha + \frac{1}{n_{\rm cal}+1}} (s_1, \ldots, s_{n_{\rm cal}})\]

Prediction interval for a new point $\mathbf{x}$:

\[\hat{C}_j(\mathbf{x}) = \bigl[f(\mathbf{x})_j - \hat{q}\,\sigma_j(\mathbf{x}),\; f(\mathbf{x})_j + \hat{q}\,\sigma_j(\mathbf{x})\bigr]\]

Marginal coverage guarantee (Vovk et al. 2005):

\[P\!\bigl(y_j \in \hat{C}_j(\mathbf{x})\bigr) \;\ge\; 1 - \alpha \quad \forall j\]

Parameters:

base_predictor (any) – A fitted object exposing predict_with_uncertainty(X) -> (mean, std) where both outputs have shape (n_samples, n_params).
alpha (float) – Default coverage level 1-alpha. Can be overridden per-call.
eps (float) – Floor added to $\boldsymbol{\sigma}$ to prevent division by zero.

Examples

>>> cp = ConformalPredictor(ens_inverter)
>>> cp.calibrate(X_cal, y_cal)
>>> center, lo, hi = cp.predict_intervals(X_new)

calibrate(X_cal, y_cal, alpha=None)#

Fit nonconformity scores on a held-out calibration set.

Parameters:

X_cal (ndarray, shape (n_cal, n_features))
y_cal (ndarray, shape (n_cal, n_params))
alpha (float or None) – Coverage level for this calibration run; defaults to self.alpha.

Return type:

self

predict_intervals(X, alpha=None)#

Return calibrated prediction intervals.

Parameters:

X (ndarray, shape (n_samples, n_features))
alpha (float or None) – Coverage level; defaults to the value used in calibrate().

Returns:

center (ndarray, shape (n_samples, n_params))
lower (ndarray, shape (n_samples, n_params))
upper (ndarray, shape (n_samples, n_params)) – The interval $[\text{center} - \hat{q}\,\sigma,\; \text{center} + \hat{q}\,\sigma]$ satisfies the marginal coverage guarantee.

Return type:

tuple[ndarray, ndarray, ndarray]

coverage(X_test, y_test, alpha=None)#

Empirical marginal coverage on a test set.

A well-calibrated predictor returns a value close to 1 - alpha.

Parameters:

X_test (ndarray)
y_test (ndarray, shape (n_test, n_params))
alpha (float or None)

Returns:

coverage

Return type:

float ∈ [0, 1]

coverage_diagnostics(X_test, y_test, alphas=None)#

Coverage table across a range of significance levels.

Useful for a reliability diagram: plot actual_coverage vs nominal_coverage = 1 - alpha.

Parameters:

X_test (ndarray)
y_test (ndarray)
alphas (sequence of float or None) – Significance levels to evaluate. Defaults to 25 values uniformly spaced in [0.02, 0.50].

Returns:

diag

Return type:

dict {alpha: actual_coverage}

class pycsamt.ai.PosteriorCalibrator(n_bins=100)#

Bases: object

Gaussianization normalising flow for calibrated posterior sampling.

Implements the regression calibration scheme of Kuleshov et al. (2018): a monotone recalibration function $g: [0,1] \to [0,1]$ is learned from calibration-set residuals via isotonic regression (PAVA), then composed with the standard-normal quantile function to produce draws from a calibrated posterior.

Step 1 — fit the recalibration map on $(y_{\rm cal}, \hat{y}_{\rm cal}, \sigma_{\rm cal})$:

\[p_i = \Phi\!\left(\frac{y_i - \hat{y}_i}{\sigma_i}\right), \qquad g = \text{PAVA}\!\left(p_{(i)},\, \frac{i}{n_{\rm cal}}\right)\]

where $p_{(i)}$ are the sorted predicted CDF values and $i/n_{\rm cal}$ are the empirical CDF targets.

Step 2 — calibrated posterior samples:

\[u \sim \mathcal{U}(0,1),\quad p^* = g^{-1}(u),\quad \theta^* = \hat{y} + \sigma \cdot \Phi^{-1}(p^*)\]

The map $g^{-1}$ (inverse PAVA) is computed via linear interpolation. Because the composition $\Phi^{-1} \circ g^{-1} \circ \Phi$ is a monotone bijection $\mathbb{R} \to \mathbb{R}$, this constitutes a one-dimensional normalising flow that transforms the empirical residual distribution into a standard Gaussian.

Parameters:: n_bins (int) – Number of equally spaced CDF bins used as knots for the interpolated inverse map. Higher values give a smoother flow at the cost of slight over-smoothing in data-sparse tails.

Examples

>>> pc = PosteriorCalibrator()
>>> pc.fit(y_cal, y_pred_cal, sigma_cal)
>>> samples = pc.predict_posterior(y_pred_test, sigma_test)
>>> cal_std = pc.calibrated_std(sigma_raw_test)

References

Kuleshov V., Fenner N. & Ermon S. (2018) Accurate uncertainties for deep learning using calibrated regression. ICML, PMLR 80, 2796–2804.

fit(y_cal, y_pred_cal, sigma_cal)#

Fit the Gaussianization recalibration map.

Parameters:

y_cal (ndarray, shape (n_cal, n_params)) – True parameter vectors on the calibration set.
y_pred_cal (ndarray, shape (n_cal, n_params)) – Point predictions (e.g. ensemble mean).
sigma_cal (ndarray, shape (n_cal, n_params)) – Predicted standard deviations (e.g. ensemble std).

Return type:

self

predict_posterior(y_pred, sigma_pred, n_samples=500, rng=None)#

Draw samples from the calibrated posterior via inverse-CDF sampling.

Parameters:

y_pred (ndarray, shape (n_samples_pred, n_params)) – Point predictions (ensemble mean).
sigma_pred (ndarray, shape (n_samples_pred, n_params)) – Raw (uncalibrated) uncertainty estimates.
n_samples (int) – Number of posterior draws per input point.
rng (numpy.random.Generator or None) – Random number generator for reproducibility.

Returns:

draws – Calibrated posterior samples. draws.mean(axis=0) ≈ y_pred and draws.std(axis=0) ≈ calibrated_sigma.

Return type:

ndarray, shape (n_samples, n_samples_pred, n_params)

calibrated_std(sigma_raw)#

Apply the global variance recalibration to raw sigma estimates.

Parameters:: sigma_raw (ndarray)
Returns:: sigma_cal – Recalibrated standard deviation such that coverage is approximately nominal on the calibration distribution.
Return type:: ndarray

calibration_error(y_test, y_pred_test, sigma_test)#

Mean absolute calibration error (MACE, Kuleshov et al. 2018).

For a perfectly calibrated predictor, the fraction of test samples falling below the $p$-quantile of the predicted distribution should equal $p$. MACE measures the average deviation:

\[\text{MACE} = \frac{1}{B} \sum_{b=1}^{B} |\hat{F}(q_b) - q_b|\]

where $B$ is the number of CDF evaluation points.

Parameters:

y_test (ndarray, shape (n_test, n_params))
y_pred_test (ndarray, shape (n_test, n_params))
sigma_test (ndarray, shape (n_test, n_params))

Returns:

mace – 0 = perfectly calibrated, 1 = maximally mis-calibrated.

Return type:

float ∈ [0, 1]

pycsamt.ai.list_pretrained()#

Return a dict of available pre-trained models.

Returns:: models – Mapping name → description.
Return type:: dict

Examples

>>> from pycsamt.ai._zoo import list_pretrained
>>> for name, desc in list_pretrained().items():
...     print(f"{name:35s}  {desc[:60]}")

pycsamt.ai.get_pretrained_info(name)#

Return metadata for a registered pre-trained model.

Parameters:: name (str) – Model identifier (see list_pretrained()).
Returns:: info
Return type:: dict
Raises:: KeyError – If name is not in the registry.

pycsamt.ai.download_checkpoint(name, *, cache_dir=None, force=False, verbose=True)#

Download a pre-trained checkpoint to the local cache.

Parameters:

name (str) – Model name (see list_pretrained()).
cache_dir (str or None) – Override the default cache location (~/.pycsamt/model_zoo/).
force (bool) – Re-download even if the file exists locally.
verbose (bool)

Returns:

path – Local path of the downloaded .npz file.

Return type:

Path

Raises:

KeyError – If name is not in the model zoo.
RuntimeError – If the download fails or the MD5 check fails.

class pycsamt.ai.CNN1DNet(n_features, n_out, *, channels=(32, 64, 128), kernel_size=5, dropout=0.3)#

Bases: object

Factory wrapper — call build() to get an nn.Module.

Use EMInverter1D instead of instantiating this class directly.

Parameters:

n_features (int) – Length of the input feature vector.
n_out (int) – Length of the output parameter vector (2*n_layers - 1).
channels (sequence of int) – Number of filters in each convolutional block.
kernel_size (int) – Convolutional kernel width (same padding applied).
dropout (float) – Dropout probability in the FC head.

build()#

class pycsamt.ai.ResNet1DNet(n_features, n_out, *, channels=(64, 128, 256), dropout=0.3, n_blocks=2)#

Bases: object

Factory wrapper for the 1-D residual CNN.

Parameters:

n_features (int) – Input feature-vector length.
n_out (int) – Output parameter vector length (2*n_layers - 1).
channels (sequence of int) – Filter counts for the three residual stages. Default (64, 128, 256) replicates the I8RFCN paper.
dropout (float) – Dropout before the final linear layer.
n_blocks (int) – Number of residual blocks per stage (default 2).

build()#: Return the nn.Module.

class pycsamt.ai.FCN1DNet(n_features, n_out, *, channels=(32, 64, 128, 64), dropout=0.2)#

Bases: object

Factory wrapper for the fully-convolutional 1-D EM network.

Parameters:

n_features (int) – Input feature vector length.
n_out (int) – Output parameter vector length (2*n_layers - 1).
channels (sequence of int) – Channel widths for each convolutional block. The last entry is used for the bottleneck and global pool.
dropout (float) – Spatial dropout rate.

build()#: Return the nn.Module.

class pycsamt.ai.UNet2DNet(n_in, n_out=1, *, channels=(32, 64, 128, 256, 512), dropout=0.2)#

Bases: object

Factory wrapper for the 2-D U-Net.

All heavy imports are deferred to build().

Parameters:

n_in (int) – Input channels — typically n_components (e.g. 4 for off-diagonal MT impedance: log|Zxy|, phi_xy, log|Zyx|, phi_yx).
n_out (int) – Output channels — 1 for a single $\\log_{10}(\\rho)$ section.
channels (tuple of int, default) – (32, 64, 128, 256, 512) Channel widths at each encoder stage plus the bridge. len(channels) - 1 determines the number of pooling/upsampling stages.
dropout (float, default 0.2) – 2-D spatial dropout probability in each convolutional block.

Examples

>>> net = UNet2DNet(n_in=4, n_out=1)
>>> model = net.build()

build()#

Instantiate and return the nn.Module.

Return type:: Any

class pycsamt.ai.DRCNNNet(n_features_list, n_out, *, growth_rate=32, n_layers=6, hidden_dim=256, dropout=0.2)#

Bases: object

Factory wrapper for the Dense Residual CNN.

Parameters:

n_features_list (sequence of int) – Feature vector length for each input modality. E.g. (120, 48) for 120 MT features and 48 seismic features.
n_out (int) – Output dimension (number of subsurface parameters to predict).
growth_rate (int, default 32) – New channels added by each sub-layer in a dense block.
n_layers (int, default 6) – Number of sub-layers per dense block.
hidden_dim (int, default 256) – Dimension of the encoded representation from each modality and the fusion block output.
dropout (float, default 0.2)

Examples

>>> # MT (120 features) + TEM (48 features)
>>> drcnn = DRCNNNet((120, 48), n_out=9)
>>> model = drcnn.build()

build()#

Instantiate and return the nn.Module.

Return type:: Any

class pycsamt.ai.GCNNet(n_features, n_out, hidden=(256, 128, 64), dropout=0.1)#

Bases: object

Factory that builds a PyTorch or TensorFlow GCN.

Call build() after import to obtain the framework module.

Parameters:

n_features (int) – Dimensionality of per-node input features (e.g. 2 × n_freqs).
n_out (int) – Per-node output size (2n-1 for an n-layer model: n resistivities + n-1 thicknesses).
hidden (sequence of int) – Hidden-layer widths for each GCN message-passing step.
dropout (float) – Dropout probability applied after each hidden GCN layer.

build()#

Return a torch.nn.Module for this architecture.

Return type:: torch.nn.Module

build_tf()#

Return a tf.keras.Model for this architecture.

Return type:: tf.keras.Model

pycsamt.ai.build_adjacency(coords, radius, *, self_loops=True, normalise=True)#

Build a symmetric adjacency matrix from 2-D station coordinates.

Parameters:

coords (ndarray, shape (n_stations, 2)) – Station (x, y) positions in any consistent unit (metres, degrees).
radius (float) – Maximum inter-station distance for an edge to exist. Uses the same unit as coords.
self_loops (bool) – If True (default), add $\tilde{A} = A + I$.
normalise (bool) – If True (default), apply symmetric normalisation $\tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}$.

Returns:

A – Adjacency matrix, optionally normalised.

Return type:

ndarray, shape (n_stations, n_stations), float32

class pycsamt.ai.Normalizer(eps=1e-08)#

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)#

Compute mean and std from X, ignoring NaN.

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

transform(X)#

Apply z-score normalisation.

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

inverse_transform(X_norm)#

Reverse z-score normalisation.

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

fit_transform(X)#

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

to_dict()#

Return type:: dict

classmethod from_dict(d)#

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

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

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)#

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)#

Undo X normalisation.

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

split(val_frac=0.1, seed=None)#

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)

class pycsamt.ai.EMTrainer(model, *, lr=0.001, weight_decay=1e-05, patience=20, min_delta=1e-05, batch_size=256, device='cpu', grad_clip=None, verbose=True)#

Bases: object

Training loop manager for EM 1-D inversion networks.

Parameters:

model (nn.Module) – The PyTorch network to train.
lr (float) – Initial learning rate. Default 1e-3.
weight_decay (float) – L2 regularisation. Default 1e-5.
patience (int) – Early-stopping patience (epochs without val-loss improvement).
min_delta (float) – Minimum val-loss improvement to reset patience counter.
batch_size (int) – Mini-batch size.
device (str) – Compute device ('cpu', 'cuda', 'mps').
grad_clip (float or None) – If set, clip gradient norms to this value.
verbose (bool) – Print per-epoch summary.

Variables:

history (dict) – Keys 'train_loss' and 'val_loss' — lists of per-epoch averages. Also 'lr' and 'epoch_time'.
best_epoch (int) – Epoch index at which the best validation loss was achieved.
best_val_loss (float)

fit(train_ds, val_ds, epochs=100)#

Train the network.

Parameters:

train_ds (EMDataset) – Training split.
val_ds (EMDataset) – Validation split.
epochs (int) – Maximum number of epochs.

Return type:

self

get_weights()#

Return model weights as numpy dict.

Return type:: dict[str, ndarray]

load_weights(weights)#

Restore weights from numpy dict.

Parameters:: weights (dict[str, ndarray])
Return type:: None

pycsamt.ai.rmse(y_true, y_pred)#

Root mean square error, ignoring NaN.

Operates element-wise on y_true and y_pred. Both arrays are assumed to be in log₁₀(Ω·m) or normalised space.

Parameters:

y_true (ndarray)
y_pred (ndarray)

Return type:

float

pycsamt.ai.mae(y_true, y_pred)#

Mean absolute error, ignoring NaN.

Parameters:

y_true (ndarray)
y_pred (ndarray)

Return type:

float

pycsamt.ai.r2(y_true, y_pred)#

Coefficient of determination R², ignoring NaN.

Parameters:

y_true (ndarray)
y_pred (ndarray)

Return type:

float

pycsamt.ai.relative_rmse(y_true, y_pred)#

Normalised RMSE: sqrt(mean((y_true - y_pred)² / y_true²)).

Useful when comparing models with very different resistivity ranges.

Parameters:

y_true (ndarray)
y_pred (ndarray)

Return type:

float

pycsamt.ai.depth_rmse(y_true, y_pred, n_layers, *, depth_weight=True)#

RMSE on the resistivity sub-vector only (first n_layers columns), optionally weighted by layer index so that deeper (harder) layers have less influence on the score.

Parameters:

y_true (ndarray, shape (n_samples, n_params))
y_pred (ndarray, shape (n_samples, n_params))
n_layers (int) – Number of resistivity values in the parameter vector.
depth_weight (bool) – If True, weight by 1 / (1 + layer_index).

Return type:

float

pycsamt.ai.layer_rmse(y_true, y_pred)#

Per-column RMSE vector.

Returns:

rmse_per_col – RMSE for each parameter column independently.

Return type:

ndarray, shape (n_params,)

Parameters:

y_true (ndarray)
y_pred (ndarray)

pycsamt.ai.masked_mse_loss(pred, target)#

MSE loss that ignores NaN (padding) entries in target.

Parameters:

pred (torch.Tensor, shape (batch, n_out))
target (torch.Tensor, shape (batch, n_out))

Returns:

loss

Return type:

torch.Tensor (scalar)

pycsamt.ai.summarise(y_true, y_pred, n_layers=None)#

Compute all scalar metrics and return as a dict.

Returns:

metrics – 'rmse', 'mae', 'r2', 'relative_rmse', 'depth_rmse' (if n_layers given).

Return type:

dict with keys

Parameters:

y_true (ndarray)
y_pred (ndarray)
n_layers (int | None)

class pycsamt.ai.AugmentNoise(sigma=0.05, phase_sigma=None, clip=3.0)#

Bases: _BaseAugmenter

Add multiplicative log-Gaussian noise to feature amplitudes.

For each sample a noise vector

\[\boldsymbol{\varepsilon} \sim \mathcal{N}(\mathbf{0},\, \sigma^2 \mathbf{I})\]

is added to the feature vector. If phase_sigma is given a separate (typically smaller) noise level is used on phase features (assumed to occupy the second half of the feature vector when include_phase=True).

Parameters:

sigma (float, default 0.05) – Noise standard deviation applied to all features.
phase_sigma (float or None) – Separate noise level for phase features. None → use sigma.
clip (float or None) – Clip noise magnitude to [-clip, +clip] standard deviations.

class pycsamt.ai.AugmentStaticShift(shift_range=(0.3, 3.0), n_amp_features=None, per_sample=True)#

Bases: _BaseAugmenter

Apply a random static shift to amplitude features.

Static shift is a common MT artefact: all apparent-resistivity values at one site are scaled by a constant factor due to local near-surface heterogeneity. This augmenter simulates the effect by multiplying the first n_amp_features feature columns by a site-specific log-uniform random factor.

Parameters:

shift_range ((lo, hi), default (0.3, 3.0)) – Range of the multiplicative factor in linear scale.
n_amp_features (int or None) – Number of amplitude columns at the start of the feature vector. None → shift all features (use when features are log₁₀ amplitude only, without phase).
per_sample (bool) – If True draw an independent shift for each sample; otherwise draw one shift per batch.

class pycsamt.ai.AugmentFreqDrop(drop_rate=0.1, contiguous=False, fill_value=0.0)#

Bases: _BaseAugmenter

Randomly zero-out a fraction of frequency channels.

Simulates missing data at individual frequencies (bad data periods, cultural noise bursts, dead-band contamination).

Parameters:

drop_rate (float, default 0.1) – Fraction of feature columns to zero-out per sample.
contiguous (bool, default False) – When True drop a contiguous block of channels (simulates a dead band in frequency).
fill_value (float, default 0.0) – Replacement value for dropped channels.

class pycsamt.ai.AugmentMixup(alpha=0.2)#

Bases: _BaseAugmenter

Mixup augmentation (Zhang et al. 2018).

Randomly pairs samples within the batch and creates convex combinations:

\[\tilde{\mathbf{x}} = \lambda \mathbf{x}_i + (1-\lambda) \mathbf{x}_j, \quad \tilde{\mathbf{y}} = \lambda \mathbf{y}_i + (1-\lambda) \mathbf{y}_j\]

where $\lambda \sim \mathrm{Beta}(\alpha, \alpha)$.

Parameters:: alpha (float, default 0.2) – Beta distribution shape parameter. Small values (≈0.1) produce near-original samples; alpha=1.0 is uniform mixing.

class pycsamt.ai.Compose(augmenters, seed=None)#

Bases: object

Sequentially apply a list of augmenters.

Parameters:

augmenters (list of _BaseAugmenter) – Applied in order.
seed (int or None) – Seed for the shared random generator.

Examples

>>> aug = Compose([
...     AugmentStaticShift(shift_range=(0.5, 2.0)),
...     AugmentNoise(sigma=0.03),
...     AugmentFreqDrop(drop_rate=0.05),
... ])
>>> X_aug, y_aug = aug(X_train, y_train)

class pycsamt.ai.RandomApply(augmenter, p=0.5, seed=None)#

Bases: object

Apply an augmenter with probability p.

Parameters:

augmenter (_BaseAugmenter)
p (float, default 0.5)
seed (int or None)

Examples

>>> aug = RandomApply(AugmentMixup(alpha=0.3), p=0.4)
>>> X_aug, y_aug = aug(X_train, y_train)

class pycsamt.ai.EMStyle(overrides=None)#

Bases: object

Context manager that applies the shared EM plot style.

Parameters:: overrides (dict or None) – Additional rcParams to merge on top of the base style.

Examples

>>> with EMStyle():
...     fig, ax = plt.subplots(figsize=EM_FIGSIZE['double'])
...     ax.semilogy(period, rho_a)

pycsamt.ai.em_context(**overrides)#

Lightweight context manager shorthand.

>>> with em_context(figure_dpi=150):
...     make_plot()

pycsamt.ai.add_colorbar(mappable, ax, label='$\\log_{10}(\\rho)$ (Ω·m)', *, size='4%', pad=0.08)#

Attach a right-side colorbar to ax using make_axes_locatable.

Parameters:

mappable – The ScalarMappable (output of imshow, pcolormesh, etc.)
ax – Host axes.
label (str) – Colorbar label.
size (str) – Width as a percentage of the parent axes.
pad (float) – Padding between axes and colorbar.

Returns:

cbar

Return type:

Colorbar

pycsamt.ai.plot_compare(true_models, pred_models, *, n_cols=5, max_sites=20, depth_max=None, log_scale=True, show_rmse=True, site_labels=None, figsize_per_panel=(2.2, 4.0), title=None, style=True)#

Multi-panel comparison of true and predicted 1-D resistivity profiles.

Parameters:

true_models (list of LayeredModel or ndarray (n_sites, n_params)) – Ground-truth models.
pred_models (list of LayeredModel or ndarray (n_sites, n_params)) – Network-predicted models (same order as true_models).
n_cols (int) – Number of columns in the panel grid.
max_sites (int) – Maximum number of sites to plot (first max_sites are used).
depth_max (float or None) – Maximum depth axis value [m].
log_scale (bool) – Log₁₀ x-axis for resistivity.
show_rmse (bool) – Annotate each panel with per-site log₁₀(ρ) RMSE.
site_labels (list of str or None) – Panel titles (default: 'Site 0', 'Site 1', …).
figsize_per_panel ((w, h)) – Size of each individual panel in inches.
title (str or None) – Figure suptitle.
style (bool) – Apply EMStyle.

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_profile_pair(true_model, pred_model, *, ax=None, depth_max=None, log_scale=True, show_rmse=True, legend=True, style=True)#

Plot a single true/predicted resistivity–depth pair on ax.

Parameters:

true_model (LayeredModel or ndarray (n_params,))
pred_model (LayeredModel or ndarray (n_params,))
ax (Axes or None)
depth_max (float or None)
log_scale (bool)
show_rmse (bool)
legend (bool)
style (bool)

Returns:

ax

Return type:

Axes

pycsamt.ai.plot_convergence(history, *, ax=None, log_scale=True, best_epoch=None, smoothing=0.0, show_lr=True, title='Training convergence', style=True)#

Plot train / validation loss curves from a trainer history dict.

Parameters:

history (dict or list of dict) – A dict with keys 'train_loss' and 'val_loss' (and optionally 'lr'), as returned by history. A list of such dicts (from multiple runs) activates the mean ± 1-σ band mode.
ax (Axes or None) – Target axes. If None, a new figure/axes is created.
log_scale (bool) – Log₁₀ y-axis for the loss.
best_epoch (int or None) – If given, draw a vertical dashed line at this epoch.
smoothing (float in [0, 1)) – Exponential moving average smoothing coefficient. 0 = no smoothing.
show_lr (bool) – Overlay learning rate on a twin y-axis (if 'lr' in history).
title (str) – Axes title.
style (bool) – Apply EMStyle.

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_lr_schedule(lr_history, *, ax=None, title='Learning rate schedule', style=True)#

Standalone learning-rate schedule plot.

Parameters:

lr_history (list of float) – Per-epoch LR values (from trainer.history['lr']).
ax (Axes or None)
title (str)
style (bool)

Returns:

ax

Return type:

Axes

pycsamt.ai.plot_section(rho_2d, *, depths=None, stations=None, depth_max=2000.0, station_spacing=1.0, log_scale=True, vmin=None, vmax=None, cmap=None, title='', xlabel='Station', ylabel='Depth (m)', show_sites=True, figsize=None, ax=None, style=True)#

Plot a 2-D resistivity section.

Parameters:

rho_2d (ndarray, shape (n_depth, n_stations)) – Resistivity or log₁₀(ρ) 2-D model.
depths (ndarray (n_depth,) or None) – Depth values in metres. Default: linear 0 → depth_max.
stations (ndarray (n_stations,) or None) – Station positions (arbitrary units).
depth_max (float) – Maximum depth used when depths is None.
station_spacing (float) – Station interval when stations is None.
log_scale (bool) – Apply log₁₀ transform before plotting. Set False if rho_2d already contains log₁₀(ρ).
vmin (float or None) – Colour scale limits.
vmax (float or None) – Colour scale limits.
cmap (str or None) – Matplotlib colormap name. Defaults to 'RdYlBu_r'.
title (str)
xlabel (str)
ylabel (str)
show_sites (bool) – Draw site-marker triangles at the surface.
figsize ((width, height) or None)
ax (Axes or None)
style (bool) – Apply EMStyle context.

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_section_pair(true_2d, pred_2d, *, depths=None, stations=None, depth_max=2000.0, station_spacing=1.0, log_scale=True, vmin=None, vmax=None, cmap=None, show_difference=True, figsize=None, style=True)#

Side-by-side comparison of true and predicted 2-D resistivity sections.

Optionally shows the absolute difference (show_difference=True).

Parameters:

true_2d (ndarray (n_depth, n_stations))
pred_2d (ndarray (n_depth, n_stations))
depths (see plot_section())
stations (see plot_section())
depth_max (see plot_section())
station_spacing (see plot_section())
log_scale (see plot_section())
vmin (float or None) – Shared colour scale. When None, computed from true_2d.
vmax (float or None) – Shared colour scale. When None, computed from true_2d.
cmap (str or None)
show_difference (bool) – Add a third panel with the signed difference.
figsize ((width, height) or None)
style (bool)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_pseudo_section(rho_a_2d, *, freqs=None, stations=None, station_spacing=1.0, log_freq=True, log_rho=True, vmin=None, vmax=None, cmap=None, component='xy', title='', figsize=None, ax=None, style=True)#

Apparent-resistivity or phase pseudo-section plot.

Parameters:

rho_a_2d (ndarray (n_freqs, n_stations)) – Apparent resistivity or phase values.
freqs (ndarray (n_freqs,) or None) – Frequencies in Hz.
stations (ndarray (n_stations,) or None) – Station positions.
station_spacing (float)
log_freq (bool) – Use log₁₀(period) for the y-axis.
log_rho (bool) – Apply log₁₀ transform to the data.
vmin (float or None)
vmax (float or None)
cmap (str or None)
component (str) – Label for the colour bar (e.g. 'xy', 'yx', 'phase_xy').
title (str)
figsize ((width, height) or None)
ax (Axes or None)
style (bool)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_confusion_matrix(y_true, y_pred, *, class_names=None, normalise=True, cmap='Blues', title='Confusion Matrix', figsize=None, ax=None, style=True)#

Plot a confusion matrix for a classification model.

Parameters:

y_true (int ndarray (n_samples,)) – Ground-truth class labels.
y_pred (int ndarray (n_samples,)) – Predicted class labels.
class_names (list of str or None) – Display labels for each class. Defaults to ['1-D','2-D','3-D'] for three-class problems.
normalise (bool) – Show row-normalised (recall) fractions.
cmap (str) – Matplotlib colormap.
title (see plot_section().)
figsize (see plot_section().)
ax (see plot_section().)
style (see plot_section().)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_residuals(y_true, y_pred, *, param_names=None, n_cols=4, figsize_per_panel=(2.5, 2.5), style=True)#

Scatter plots of predicted vs. true for each model parameter.

Each panel shows the 1:1 line, coloured scatter, and per-parameter R² annotation.

Parameters:

y_true (ndarray (n_samples, n_params))
y_pred (ndarray (n_samples, n_params))
param_names (list of str or None)
n_cols (int)
figsize_per_panel ((width, height))
style (bool)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_layer_errors(y_true, y_pred, n_layers, *, log_rho=True, ax=None, figsize=None, style=True)#

Per-layer mean absolute error bar chart.

Parameters:

y_true (ndarray (n_samples, 2*n_layers-1))
y_pred (ndarray (n_samples, 2*n_layers-1))
n_layers (int)
log_rho (bool) – Label ρ columns as log₁₀(ρ).
ax (Axes or None)
figsize ((width, height) or None)
style (bool)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_uncertainty_bands(x, y_pred, y_upper, y_lower, y_true=None, *, ax=None, xlabel='', ylabel='', title='Prediction with Uncertainty', figsize=None, style=True)#

1-D prediction curve with uncertainty bands.

Suitable for showing per-site model parameter predictions (e.g. resistivity vs. depth) with ±1σ or 10/90 percentile bands.

Parameters:

x (ndarray (n_points,)) – X-axis values (e.g. depth or frequency).
y_pred (ndarray (n_points,)) – Central prediction.
y_upper (ndarray (n_points,)) – Upper and lower uncertainty bounds.
y_lower (ndarray (n_points,)) – Upper and lower uncertainty bounds.
y_true (ndarray (n_points,) or None) – Ground-truth values to overlay.
ax (Axes or None)
xlabel (str)
ylabel (str)
title (str)
figsize ((width, height) or None)
style (bool)

Returns:

fig

Return type:

Figure

pycsamt.ai.plot_feature_importance(importances, feature_names=None, *, top_n=20, horizontal=True, ax=None, figsize=None, title='Feature Importance', style=True)#

Horizontal bar chart of feature importances.

Compatible with scikit-learn feature_importances_ arrays and any non-negative importance measure.

Parameters:

importances (ndarray (n_features,))
feature_names (list of str or None)
top_n (int) – Show only the top n features by importance.
horizontal (bool) – Use horizontal bars (default) for long feature names.
ax (Axes or None)
figsize ((width, height) or None)
title (str)
style (bool)

Returns:

fig

Return type:

Figure

class pycsamt.ai.EMDenoiser(n_freqs=None, n_components=4, *, channels=(64, 128, 64), dropout=0.1, device=None)#

Bases: BaseEMProcessor

Convolutional autoencoder for MT impedance tensor denoising.

The network is trained on clean (or synthetic) impedance data by adding controlled noise and learning to reconstruct the original signal. Supports both PyTorch and TensorFlow backends via pycsamt.backends.

Parameters:

n_freqs (int or None, default None) – Number of frequency channels in the input. When None, inferred from X.shape[2] on the first call to fit().
n_components ({4, 8}, default 4) – Feature channels per frequency: 4 = off-diagonal $Z_{xy}$, $Z_{yx}$ amplitudes and phases; 8 = all four tensor components.
channels (tuple of int, default (64, 128, 64)) – Channel widths for the encoder/decoder stages.
dropout (float, default 0.1) – Dropout probability in the encoder bottleneck.
device (str or None) – "cpu", "cuda", "mps", "/GPU:0", or None for auto-detect via the active backend.

Notes

When neither PyTorch nor TensorFlow is available the denoiser falls back to a per-channel Gaussian smoothing filter (requires scipy).

Call prepare_z_features() to convert a site collection to the (n_samples, n_components, n_freqs) array expected here.

Examples

>>> import numpy as np
>>> X_clean = np.random.randn(200, 4, 32).astype("float32")
>>> den = EMDenoiser()
>>> den.fit(X_clean, epochs=5, verbose=False)
EMDenoiser(n_freqs=32, n_components=4)
>>> X_denoised = den.transform(X_clean)

fit(X, *, noise_level=0.05, epochs=50, batch_size=64, lr=0.001, val_frac=0.1, seed=None, verbose=True)#

Train the denoiser on clean impedance data.

Parameters:

X (ndarray, shape (n_samples, n_components, n_freqs)) – Clean training data.
noise_level (float) – Standard deviation of additive Gaussian noise relative to the per-channel standard deviation of X.
epochs (int) – Maximum number of training epochs.
batch_size (int)
lr (float) – Initial Adam learning rate.
val_frac (float) – Fraction of training samples held out for validation.
seed (int or None)
verbose (bool)

Return type:

self

transform(X)#

Apply the trained denoiser.

Parameters:: X (ndarray, shape (n_samples, n_components, n_freqs)) – Noisy impedance data.
Returns:: X_denoised
Return type:: ndarray, same shape as input

pycsamt.ai.prepare_z_features(sites, *, n_components=4, log_amp=True, freq_ref=None)#

Extract an impedance feature array from a site collection.

Parameters:

sites (SiteCollection, dict, or list of Z objects) – MT impedance data in any format accepted by ensure_sites().
n_components ({4, 8}) – 4 uses off-diagonal components only; 8 includes diagonal.
log_amp (bool) – If True, amplitudes are transformed as $\log_{10}(|Z| + \varepsilon)$ before packing.
freq_ref (ndarray or None) – Common frequency grid to interpolate onto. When None, the grid of the first site is used.

Returns:

X – Feature array ready for EMDenoiser.

Return type:

ndarray, shape (n_sites, n_components, n_freqs)

class pycsamt.ai.EMQCScorer(contamination=0.05, snr_threshold=3.0, skew_threshold=0.3, score_threshold=0.5, use_ml=True, n_estimators=100, random_state=None)#

Bases: BaseEMProcessor

ML-based per-frequency QC scorer for MT impedance data.

Combines hard thresholds on SNR and Swift skew with an IsolationForest anomaly model fitted on extracted signal-quality features. The final quality score for each (site, frequency) cell is $s \in [0, 1]$; observations below score_threshold should be rejected before inversion.

Parameters:

contamination (float, default 0.05) – Expected fraction of contaminated samples, passed directly to sklearn.ensemble.IsolationForest.
snr_threshold (float, default 3.0) – Observations with SNR below this value are hard-flagged as bad regardless of the ML score.
skew_threshold (float, default 0.3) – Observations with Swift skew above this value are hard-flagged.
score_threshold (float, default 0.5) – Quality scores below this value are considered bad.
use_ml (bool, default True) – When False, only the rule-based thresholds are applied and the IsolationForest is not fitted.
n_estimators (int, default 100) – Number of trees in the IsolationForest.
random_state (int or None)

Examples

>>> from pycsamt.ai.processing import EMQCScorer
>>> scorer = EMQCScorer(snr_threshold=5.0, use_ml=False)
>>> scorer.fit(sites)
EMQCScorer(rule_only)
>>> tbl = scorer.score_table(sites)

fit(X, **kwargs)#

Fit the QC model on a training set.

Parameters:: X (SiteCollection, dict, list of Z, or ndarray (n_samples, 5)) – Training data. Site collections are converted automatically. Pass a precomputed feature matrix to skip extraction.
Return type:: self

transform(X)#

Compute quality scores.

Parameters:: X (SiteCollection or ndarray (n_samples, 5))
Returns:: scores – Quality scores in [0, 1]; 1 = good, 0 = bad.
Return type:: ndarray, shape (n_samples,)

score_table(sites)#

Return per-(site, frequency) QC table.

Parameters:: sites (SiteCollection or compatible)
Returns:: df – Columns: station, freq, snr, swift_skew, asym, phase_xy, phase_yx, score, flag (0=bad, 1=good).
Return type:: DataFrame

class pycsamt.ai.AnomalyDetector(n_features=None, latent_dim=32, *, channels=(128, 64), threshold_percentile=95.0, device=None)#

Bases: BaseEMProcessor

Profile-level unsupervised anomaly detector.

Parameters:

n_features (int or None, default None) – Feature vector length per site (typically n_freqs × n_components). When None, inferred from X.shape[1] on the first call to fit().
latent_dim (int, default 32) – Bottleneck dimension of the autoencoder.
channels (tuple of int, default (128, 64)) – Hidden layer widths in the encoder (decoder is mirrored).
threshold_percentile (float, default 95.0) – Sites with reconstruction error above this percentile of the training distribution are flagged as anomalous.
device (str or None) – Compute device. Ignored when using the PCA fallback.

Notes

When neither PyTorch nor TensorFlow is installed the detector silently uses a PCA reconstruction model via scikit-learn. The interface is identical; only the accuracy differs.

Examples

>>> import numpy as np
>>> X = np.random.randn(100, 120).astype("float32")
>>> det = AnomalyDetector(latent_dim=16)
>>> det.fit(X, epochs=10, verbose=False)
AnomalyDetector(n_features=120, latent_dim=16)
>>> scores = det.transform(X)
>>> flags = det.flag_anomalies(X)

fit(X, *, epochs=80, batch_size=32, lr=0.001, val_frac=0.1, seed=None, verbose=True)#

Train the anomaly detector on profile data.

Parameters:

X (ndarray, shape (n_sites, n_features)) – Per-site feature vectors (normal / clean data).
epochs (int) – Training epochs (ignored when using PCA fallback).
batch_size (int)
lr (float)
val_frac (float)
seed (int or None)
verbose (bool)

Return type:

self

transform(X)#

Compute per-site reconstruction error (anomaly score).

Parameters:: X (ndarray, shape (n_samples, n_features))
Returns:: scores – Per-site mean squared reconstruction error.
Return type:: ndarray, shape (n_samples,)

flag_anomalies(X)#

Return a boolean mask of anomalous sites.

Parameters:: X (ndarray, shape (n_samples, n_features))
Returns:: flags
Return type:: bool ndarray, shape (n_samples,)

class pycsamt.ai.DimensionalityClassifier(hidden=(128, 64), dropout=0.2, n_classes=3, lr=0.001, device=None)#

Bases: BaseEMProcessor

MLP classifier for MT data dimensionality (1-D / 2-D / 3-D).

Parameters:

hidden (tuple of int, default (128, 64)) – Hidden layer widths of the shared MLP backbone.
dropout (float, default 0.2) – Dropout probability in each hidden layer.
n_classes (int, default 3) – Number of dimensionality classes (0=1D, 1=2D, 2=3D).
lr (float, default 1e-3) – Default learning rate; can be overridden in fit().
device (str or None)

Notes

Falls back to a random-forest classifier (scikit-learn) when neither PyTorch nor TensorFlow is available.

Examples

>>> from pycsamt.ai.processing import DimensionalityClassifier
>>> clf = DimensionalityClassifier()
>>> clf.fit(X_train, y_train, epochs=30)
DimensionalityClassifier(n_classes=3, torch)
>>> clf.predict(X_test)
array([0, 1, 2, ...])
>>> clf.predict_strike(X_2d)
array([ 35., -12., ...])

classmethod from_features_table(df, *, label_col='dim', strike_col=None, skew_th=3.0, ellipt_th=0.2, **fit_kwargs)#

Construct and train a classifier from a phase_features_table() DataFrame.

Parameters:

df (DataFrame)
label_col (str or None) – Column for pre-computed labels; None → rule-based labels.
strike_col (str or None) – Column for strike direction in degrees.
skew_th (float) – Rule-based thresholds (used when label_col is absent).
ellipt_th (float) – Rule-based thresholds (used when label_col is absent).
**fit_kwargs – Passed to fit() (e.g. epochs=50).

Return type:

DimensionalityClassifier

fit(X, y=None, *, strike=None, epochs=80, batch_size=256, lr=None, val_frac=0.15, seed=None, verbose=True)#

Train the dimensionality classifier.

Parameters:

X (ndarray (n_samples, 5) or DataFrame)
y (int ndarray (n_samples,) or None)
strike (float ndarray (n_samples,) or None)
epochs (int) – Training hyper-parameters.
batch_size (int) – Training hyper-parameters.
lr (float | None) – Training hyper-parameters.
val_frac (float) – Training hyper-parameters.
seed (int | None) – Training hyper-parameters.
verbose (bool) – Training hyper-parameters.

Return type:

self

transform(X)#

Compute class probabilities.

Returns:: proba
Return type:: ndarray, shape (n_samples, n_classes)
Parameters:: X (ndarray | DataFrame)

predict(X)#

Predict dimensionality class (0=1D, 1=2D, 2=3D).

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

predict_strike(X)#

Predict geoelectric strike direction (degrees).

Returns NaN for non-2-D sites and when using the RF fallback.

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

predict_table(sites)#

Classify an entire site collection and return a result DataFrame.

Returns:: df – Columns: station, freq, period, dim (0/1/2), dim_label (str), strike (°), confidence.
Return type:: DataFrame
Parameters:: sites (Any)

AI Modules#

`pycsamt.ai.inversion.inv1d`	End-to-end 1-D EM inversion workflow.
`pycsamt.ai.inversion.inv2d`	EMInverter2D — high-level U-Net–based 2-D MT inversion pipeline.
`pycsamt.ai.inversion.inv3d`	GCNInverter3D — graph-convolutional 3-D MT spatial inversion.
`pycsamt.ai.inversion.ensemble`	EnsembleInverter — deep ensemble for uncertainty-aware EM inversion.
`pycsamt.ai.inversion.joint`	JointInverter — multi-modal / multi-physics joint inversion.
`pycsamt.ai.inversion.calibration`	Calibrated uncertainty quantification for EM neural-network inverters.
`pycsamt.ai.nets.cnn1d`	1-D CNN for EM inversion — Puzyrev (2019/2021) architecture.
`pycsamt.ai.nets.drcnn`	DRCNNNet — Dense Residual CNN for joint / multi-modal EM inversion.
`pycsamt.ai.nets.fcn`	Fully Convolutional Network for EM inversion — Moghadas (2020) style.
`pycsamt.ai.nets.gcn`	Graph Convolutional Network for spatial EM inversion.
`pycsamt.ai.nets.resnet`	1-D Residual CNN for EM inversion — Liu et al. (2021) I8RFCN style.
`pycsamt.ai.nets.unet`	UNet2DNet — 2-D U-Net architecture for EM section inversion.
`pycsamt.ai.processing.anomaly`	AnomalyDetector — unsupervised profile-level anomaly detection.
`pycsamt.ai.processing.classify`	DimensionalityClassifier — MLP-based MT data dimensionality classifier.
`pycsamt.ai.processing.denoise`	EMDenoiser — convolutional autoencoder for MT impedance tensor denoising.
`pycsamt.ai.processing.qc`	EMQCScorer — ML-based per-frequency quality control scorer.
`pycsamt.ai.training.augment`	Data augmentation for EM training datasets.
`pycsamt.ai.training.dataset`	PyTorch Dataset wrapper and normalisation utilities for EM training data.
`pycsamt.ai.training.metrics`	Evaluation metrics for 1-D EM inversion networks.
`pycsamt.ai.training.trainer`	Training loop for EM neural networks.
`pycsamt.ai.plot.compare`	Side-by-side true vs predicted 1-D resistivity profile panels.
`pycsamt.ai.plot.convergence`	Training convergence visualisation.
`pycsamt.ai.plot.diagnostics`	Diagnostic plots for AI/ML model evaluation.
`pycsamt.ai.plot.section`	2-D resistivity section and pseudo-section visualisation.