Distribution Families #

Returns:

Tuple of (nll_loss, metric).

Parametric heads - standard backbones#

class twiga.distributions.nn.parametric.NormalDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Normal (Gaussian) output head.

Best for: symmetric, unbounded targets - energy demand, temperature.

Parameters predicted:: mu - mean (unconstrained) sigma - std dev (exp of log_scale, strictly positive)

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Dimensionality of the input latent vector.
forecast_horizon (int) – Number of forecast time steps.
out_activation_function (Module | None) – Optional activation applied to mu. Defaults to Identity.

forward(x)#

Predict mu and sigma.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, sigma), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, sigma)#

Construct a torch Distribution from predicted parameters.

Return type:: Normal

get_log_likelihood(mu, sigma, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.LaplaceDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Laplace output head.

Best for: targets with heavier tails than Normal, robust to outliers - electricity prices, wind speed residuals.

Parameters predicted:: mu - location (unconstrained) scale - scale (exp of log_scale, strictly positive)

family: str | None = 'laplace'#

forward(x)#

Predict mu and scale.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, scale), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, scale)#

Construct a torch Distribution from predicted parameters.

Return type:: Laplace

get_log_likelihood(mu, scale, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.LogNormalDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Log-Normal output head.

Best for: strictly positive, right-skewed targets - renewable generation, gas prices, load with zero floor.

Parameters predicted:: mu - mean of log(target) (unconstrained) sigma - std of log(target) (exp of log_scale, strictly positive)

Note

Targets must be strictly positive. Apply a small epsilon shift in the data pipeline if zeros are possible (e.g., target = max(y, 1e-6)).

forward(x)#

Predict mu and sigma of the underlying Normal (in log space).

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, sigma), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, sigma)#

Construct a torch Distribution from predicted parameters.

Return type:: LogNormal

get_log_likelihood(mu, sigma, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.GammaDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48)#

Gamma output head.

Best for: strictly positive targets with flexible skew - solar irradiance, wind power, load at aggregate level.

Parameters predicted:: concentration - shape parameter α (softplus, strictly positive) rate - rate parameter β (softplus, strictly positive)

Note

Mean of Gamma(α, β) = α/β. Targets must be strictly positive.

forward(x)#

Predict concentration and rate.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (concentration, rate), each of shape (B, forecast_horizon, num_target_output).

get_distribution(concentration, rate)#

Construct a torch Distribution from predicted parameters.

Return type:: Gamma

get_log_likelihood(concentration, rate, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.BetaDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48)#

Beta output head.

Best for: bounded [0, 1] targets - capacity factors, state of charge, fill rates, normalised demand ratios.

Parameters predicted:: alpha - first shape parameter (softplus, strictly positive) beta - second shape parameter (softplus, strictly positive)

Note

Targets must lie strictly in (0, 1). Apply a clamp in the data pipeline if boundary values are possible: target = target.clamp(1e-6, 1 - 1e-6)

forward(x)#

Predict alpha and beta.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (alpha, beta), each of shape (B, forecast_horizon, num_target_output).

get_distribution(alpha, beta)#

Construct a torch Distribution from predicted parameters.

Return type:: Beta

get_log_likelihood(alpha, beta, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.StudentTDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None, min_df=2.1)#

Student-T output head.

Best for: targets with very heavy tails and potential outliers - financial returns, spot electricity prices.

Parameters predicted:: mu - location (unconstrained) sigma - scale (exp of log_scale, strictly positive) df - degrees of freedom (softplus, clipped to ≥ 2 for finite variance)

Note

Degrees of freedom are predicted per-sample but averaged across the forecast horizon so the model learns a single df per series.

forecast(z)#

Inference-mode forward (no gradient).

Parameters:: z (Tensor) – Latent tensor of shape (B, hidden_size) from the backbone.
Return type:: dict[str, Tensor]
Returns:: Dict with "loc" (location/mean) and "scale" (std/dispersion). The first forward() output is always loc; the second is scale.

forward(x)#

Predict mu, sigma, and degrees of freedom.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor, Tensor]
Returns:: Tuple of (mu, sigma, df) – - mu, sigma: shape (B, forecast_horizon, num_target_output) - df: shape (B, 1, num_target_output), broadcast-compatible

get_distribution(mu, sigma, df)#

Construct a torch Distribution from predicted parameters.

Return type:: StudentT

get_log_likelihood(mu, sigma, df, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.CauchyDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Cauchy output head.

Best for: targets whose extremes dominate, where even Student-T under-weights the tails. Included because Chapter 5 of the thesis compares Normal, Laplace and Cauchy density heads for net-load, and the comparison is informative precisely because Cauchy usually loses.

Parameters predicted:: loc - location (unconstrained) scale - scale (exp of log_scale, strictly positive)

Warning

The Cauchy distribution has no finite mean or variance. Its forecast() therefore returns the median as loc, which is the correct central estimate here, and any downstream metric that assumes a mean - or that averages predictions across an ensemble - is not meaningful for this head. The reported scale is the half-width at half-maximum, not a standard deviation, so it is not comparable with the scale of the Normal or Laplace heads.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Dimensionality of the input latent vector.
forecast_horizon (int) – Number of forecast time steps.
out_activation_function (Module | None) – Optional activation applied to loc. Defaults to Identity.

References

Rugema, F. Z., et al. (2020). Application of the Cauchy distribution in modelling electricity demand. Cited in Chapter 2 of the thesis as prior use of heavy-tailed densities for power demand.

forward(x)#

Predict the location and scale.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (loc, scale), each of shape (B, forecast_horizon, num_target_output).

get_distribution(loc, scale)#

Build the Cauchy distribution from predicted parameters.

Return type:: Cauchy

get_log_likelihood(loc, scale, targets)#

Return the mean negative log-likelihood.

Return type:: Tensor

Bayesian heads - epistemic uncertainty#

The heads above are deterministic maps onto distribution parameters, so they express aleatoric uncertainty: irreducible noise in the target. The Bayesian heads place a distribution over their own weights, so repeated forward passes differ and express epistemic uncertainty – what the model does not know, which more data could reduce. Both components are reported separately by forecast() as scale_noiseless and scale_noise, and combined in quadrature as scale.

Stochasticity is confined to the decoder. A full-network Bayesian treatment multiplies the parameter count and the inference cost, which is the opposite of what a substation-scale forecaster needs. Both heads still require n_samples forward passes at inference, so they are slower than the quantile and density alternatives.

class twiga.distributions.nn.bayesian.BayesianHead(num_target_output=1, hidden_size=256, forecast_horizon=48, n_samples=100, alpha=0.5, kl_weight=0.01, confidence_level=0.9, pinball_weight=0.0, coverage_weight=0.0, n_train_samples=20, out_activation_function=None, dataset_size=None)#

Shared behaviour for decoder-stochastic Bayesian heads.

Subclasses supply a stochastic decoder and output layer; this class owns the sampling, the aleatoric/epistemic split, and the loss.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Dimensionality of the input latent vector.
forecast_horizon (int) – Number of forecast time steps.
n_samples (int) – Posterior draws used by forecast().
alpha (float) – Weight on the squared-error term of the point loss; the complementary weight goes to the absolute-error term.
kl_weight (float) – Multiplier on the KL term. Ignored when the head has no KL.
confidence_level (float) – Nominal coverage of the interval shaped by the optional auxiliary losses.
pinball_weight (float) – Weight on sampled_pinball_loss(), a proper objective evaluated on n_train_samples draws. Zero disables it.
coverage_weight (float) – Weight on coverage_width_loss(), a heuristic regulariser that is not a proper scoring rule. Zero (the default) disables it; enable it only alongside a non-zero pinball_weight.
n_train_samples (int) – Draws used by the auxiliary losses during training. Kept small because each draw costs a forward pass.
out_activation_function (Module | None) – Optional activation applied to the prediction, for parity with the parametric heads. Defaults to Identity.
dataset_size (int | None) – Number of training examples, used to normalise the KL term by each minibatch’s share of the full dataset (batch_size / dataset_size, the standard Blundell et al. 2015 minibatch-KL weighting – see step()). None (the default at construction time, since the head is built before training data is seen) disables normalisation; set via self.dataset_size once the real value is known (BaseNeuralForecast.fit() does this).

Raises:

ValueError – If n_samples is below 1, alpha is outside [0, 1], or confidence_level is outside (0, 1).

forecast(z)#

Sample the posterior predictive distribution.

Parameters:: z (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: dict[str, Tensor]
Returns:: Dict with loc (the mean over draws, not a single draw), scale (total predictive standard deviation), samples of shape (n_samples, B, H, O), and the epistemic and aleatoric components scale_noiseless and scale_noise.

forward(x)#

Return a single stochastic prediction.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: tuple[Tensor]
Returns:: One-tuple holding the prediction, matching the head contract.

get_distribution(loc, scale)#

Build the Gaussian predictive distribution implied by the draws.

Return type:: Normal

get_log_likelihood(loc, scale, targets)#

Return the mean negative log-likelihood under the predictive Gaussian.

Return type:: Tensor

kl_divergence()#

Sum the KL contributions of all stochastic submodules.

Return type:: Tensor
Returns:: Scalar KL term; zero when no submodule exposes one.

step(z, y, metric_fn, epoch=None)#

Compute the training loss and metric for one batch.

The loss combines a point term (a convex mix of squared and absolute error), a Gaussian likelihood term that fits the homoscedastic noise, and the KL term. Components are returned separately so each can be logged.

The KL term is a sum over every weight/bias in the stochastic decoder (kl_divergence()), so it scales with parameter count rather than batch size the way the mean-reduced nll does – without normalisation it can be orders of magnitude larger than nll and dominate the loss regardless of kl_weight. Following Blundell et al. 2015 (cited in this module’s docstring), each minibatch is charged only its share of the one-time KL cost, batch_size / dataset_size (equivalent to the common 1/num_batches convention for equal-sized batches), so the full KL is applied exactly once per epoch rather than once per batch. Falls back to no normalisation (factor 1.0) when dataset_size hasn’t been set (e.g. the head hasn’t been fit yet).

Parameters:

z (Tensor) – Latent tensor of shape (B, hidden_size).
y (Tensor) – Targets of shape (B, forecast_horizon, num_target_output).
metric_fn (Callable[..., Any]) – Callable returning a scalar metric given (pred, target).
epoch (int | None) – Current epoch; accepted for interface consistency.

Return type:

tuple[dict[str, Tensor], Tensor]

Returns:

Tuple of (loss_components, metric) where loss_components maps "loss" to the total and also carries "nll" and "kl".

class twiga.distributions.nn.bayesian.MFVIDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, n_samples=100, alpha=0.5, kl_weight=0.01, dropout=0.1, expansion_factor=1, activation=None, prior_weight_std=1.0, confidence_level=0.9, pinball_weight=0.0, coverage_weight=0.0, n_train_samples=20, out_activation_function=None, dataset_size=None)#

Bases: BayesianHead

Mean-field variational-inference head.

Places a Gaussian posterior over the decoder weights and trains against the evidence lower bound, so the KL term is exact rather than a surrogate.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Dimensionality of the input latent vector.
forecast_horizon (int) – Number of forecast time steps.
n_samples (int) – Posterior draws used by forecast().
alpha (float) – Weight on the squared-error term of the point loss.
kl_weight (float) – Multiplier on the KL term. Values near 1 recover the strict ELBO; smaller values trade calibration for sharper point forecasts.
dropout (float) – Dropout inside the variational block.
expansion_factor (int) – Hidden width multiplier of the variational block.
activation (Module | None) – Activation module. Defaults to nn.GELU().
prior_weight_std (float) – Prior standard deviation over the weights.
confidence_level (float) – Nominal coverage for the optional auxiliary losses.
pinball_weight (float) – Weight on the sampled pinball loss (proper). Zero disables it.
coverage_weight (float) – Weight on the coverage-width regulariser (not proper). Enable only alongside a non-zero pinball_weight.
n_train_samples (int) – Draws used by the auxiliary losses during training.
out_activation_function (Module | None) – Optional activation applied to the prediction.
dataset_size (int | None) – Number of training examples; see BayesianHead.

class twiga.distributions.nn.bayesian.MCDropoutDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, n_samples=100, alpha=0.5, kl_weight=0.01, dropout=0.1, expansion_factor=1, activation=None, drop_weight=True, confidence_level=0.9, pinball_weight=0.0, coverage_weight=0.0, n_train_samples=20, out_activation_function=None, dataset_size=None)#

Bases: BayesianHead

Monte-Carlo dropout head.

Cheaper than MFVIDistribution - it adds no parameters - but its KL term is only an $\ell_2$ surrogate, so the posterior it approximates is correspondingly cruder.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Dimensionality of the input latent vector.
forecast_horizon (int) – Number of forecast time steps.
n_samples (int) – Posterior draws used by forecast().
alpha (float) – Weight on the squared-error term of the point loss.
kl_weight (float) – Multiplier on the KL surrogate.
dropout (float) – Dropout probability; the dominant control on predictive spread.
expansion_factor (int) – Hidden width multiplier of the dropout block.
activation (Module | None) – Activation module. Defaults to nn.GELU().
drop_weight (bool) – If True, share one dropout mask per output unit across the batch, which corresponds to sampling a single network from the posterior. If False, resample per example.
confidence_level (float) – Nominal coverage for the optional auxiliary losses.
pinball_weight (float) – Weight on the sampled pinball loss (proper). Zero disables it.
coverage_weight (float) – Weight on the coverage-width regulariser (not proper). Enable only alongside a non-zero pinball_weight.
n_train_samples (int) – Draws used by the auxiliary losses during training.
out_activation_function (Module | None) – Optional activation applied to the prediction.
dataset_size (int | None) – Number of training examples; see BayesianHead.

set_dropout_mode(drop_weight=True)#

Propagate the masking scheme to every dropout layer in the head.

Parameters:: drop_weight (bool) – See MCDropoutLinear.set_dropout_mode().
Return type:: None

Bayesian layers#

class twiga.models.nn.core.bayesian.MFVILinear(dim_in, dim_out, prior_weight_std=1.0, prior_bias_std=1.0, init_std=0.05, sqrt_width_scaling=False)#

Bases: Module

Linear layer with a mean-field Gaussian variational posterior.

Each weight is an independent Gaussian $q(w) = \mathcal{N}(\mu, \sigma^2)$ with learned $\mu$ and $\sigma$. A forward pass draws one sample via the reparameterisation trick, so repeated passes explore the posterior.

Parameters:

dim_in (int) – Input dimensionality.
dim_out (int) – Output dimensionality.
prior_weight_std (float) – Standard deviation of the zero-mean Gaussian weight prior.
prior_bias_std (float) – Standard deviation of the zero-mean Gaussian bias prior.
init_std (float) – Initial posterior standard deviation. Small values start the layer close to deterministic, which stabilises early training.
sqrt_width_scaling (bool) – If True, scale the prior standard deviation by 1/sqrt(dim_in) so the prior predictive variance is invariant to layer width.

Raises:

ValueError – If any dimension or standard deviation is not positive.

property bias_std: Tensor#: Posterior standard deviation of the biases, clamped away from zero.

extra_repr()#

Describe the layer in module printouts.

Return type:: str

forward(x)#

Apply the layer using one posterior sample.

Parameters:: x (Tensor) – Input of shape (..., dim_in).
Return type:: Tensor
Returns:: Output of shape (..., dim_out).

kl_divergence()#

Closed-form $\mathrm{KL}[q \,\|\, p]$ for this layer.

Return type:: Tensor
Returns:: Scalar KL divergence summed over parameters.

reset_parameters(init_std=0.05)#

Initialise posterior means as a deterministic layer, stds at init_std.

Parameters:: init_std (float) – Initial posterior standard deviation.
Return type:: None

property weight_std: Tensor#: Posterior standard deviation of the weights, clamped away from zero.

class twiga.models.nn.core.bayesian.MCDropoutLinear(dim_in, dim_out, dropout_prob=0.1, prior_weight_std=1.0, prior_bias_std=1.0, sqrt_width_scaling=False)#

Bases: Module

Linear layer whose dropout stays active at inference.

Repeated forward passes sample from the approximate posterior of Gal & Ghahramani (2016). Unlike torch.nn.Dropout, this layer does not switch off in eval() mode - that is the whole point, and it is why the layer is written out rather than composed from nn.Dropout.

Parameters:

dim_in (int) – Input dimensionality.
dim_out (int) – Output dimensionality.
dropout_prob (float) – Dropout probability in [0, 1).
prior_weight_std (float) – Prior standard deviation used by the KL surrogate.
prior_bias_std (float) – Prior standard deviation for the bias KL surrogate.
sqrt_width_scaling (bool) – If True, scale the prior by 1/sqrt(dim_in).

Raises:

ValueError – If dimensions are not positive or dropout_prob is outside [0, 1).

extra_repr()#

Describe the layer in module printouts.

Return type:: str

forward(x)#

Apply the layer with dropout active in both train and eval mode.

Parameters:: x (Tensor) – Input of shape (..., dim_in).
Return type:: Tensor
Returns:: Output of shape (..., dim_out).

kl_divergence()#

$\ell_2$ surrogate for the KL term under MC dropout.

Return type:: Tensor
Returns:: Scalar regularisation term.

reset_parameters()#

Initialise weights and biases as for a standard linear layer.

Return type:: None

set_dropout_mode(drop_weight=False)#

Select the masking scheme.

Parameters:: drop_weight (bool) – If True, draw one mask per output unit shared across the batch - the correct behaviour when sampling a single network from the posterior. If False, resample per example.
Return type:: None

twiga.models.nn.core.bayesian.FeedForwardBNN(dim, expansion_factor=2, dropout=0.0, activation=None)#

Two-layer variational feed-forward block.

Parameters:

dim (int) – Input and output dimensionality.
expansion_factor (int) – Hidden width multiplier.
dropout (float) – Dropout applied after each activation.
activation (Module | None) – Activation module. Defaults to nn.GELU().

Return type:

Sequential

Returns:

A nn.Sequential of MFVILinear layers, exposing kl_divergence on each submodule.

twiga.models.nn.core.bayesian.FeedForwardMCD(dim, expansion_factor=2, dropout=0.1, activation=None)#

Two-layer MC-dropout feed-forward block.

Parameters:

dim (int) – Input and output dimensionality.
expansion_factor (int) – Hidden width multiplier.
dropout (float) – Dropout probability used inside each layer.
activation (Module | None) – Activation module. Defaults to nn.GELU().

Return type:

Sequential

Returns:

A nn.Sequential of MCDropoutLinear layers.

Auxiliary losses for sample-based heads#

Two optional terms shape the interval during training. They are off by default and reported as separate loss components when enabled.

sampled_pinball_loss reduces the posterior draws to the central interval’s two quantile levels and applies the pinball loss, which is consistent for those levels – it can safely shape the interval.

coverage_width_loss penalises the squared deviation of a smoothed coverage rate from nominal, plus mean width. It is a heuristic regulariser, not a proper scoring rule: it decouples coverage from width, which is precisely the property that makes a composite coverage-width metric unsound. Use it only alongside a consistent objective, never as the sole training signal.

twiga.distributions.nn.custom_loss.sampled_pinball_loss(samples, targets, confidence_level=0.9, reduction='mean')#

Pinball loss on the empirical quantiles of a set of posterior draws.

Turns a sample-based head – a Bayesian one, or any model producing draws – into something trainable against a proper objective. The draws are reduced to the central interval’s two quantile levels $\{\alpha/2,\, 1-\alpha/2\}$, and the pinball loss is applied to those.

Unlike interval_loss() and coverage_width_loss() this is consistent for the quantiles it targets, so it can safely shape the interval during training rather than merely nudging it.

Parameters:

samples (Tensor) – Posterior draws of shape (S, B, H, O).
targets (Tensor) – Observations of shape (B, H, O).
confidence_level (float) – Nominal coverage $1-\alpha$ of the interval.
reduction (str) – "none", "mean" or "sum".

Return type:

Returns:

The reduced pinball loss.

Raises:

ValueError – If confidence_level is outside (0, 1) or samples is not 4-D.

References

Koenker, R., & Bassett, G. (1978). Regression quantiles. Econometrica, 46(1), 33-50.

Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 102(477), 359-378.

twiga.distributions.nn.custom_loss.coverage_width_loss(targets, lower, upper, confidence_level=0.9, width_weight=1.0, reduction='mean')#

Differentiable coverage-and-width regulariser.

Penalises the squared deviation of the soft coverage rate from the nominal level, plus the mean interval width. Coverage is smoothed with a sigmoid so the term is differentiable, since the exact indicator has zero gradient almost everywhere.

Warning

This is a heuristic regulariser, not a proper scoring rule. It decouples coverage from width – exactly the property that makes the ac{CWE} metric unsound – so a model trained on it alone is not driven toward the true quantiles. Use it, if at all, alongside a consistent objective such as sampled_pinball_loss(), and never as the sole training signal. The proper alternative is the interval score, whose 2/alpha penalty ratio is what makes the true quantiles optimal.

Parameters:

targets (Tensor) – Observations of shape (B, H, O).
lower (Tensor) – Lower interval bounds, same shape.
upper (Tensor) – Upper interval bounds, same shape.
confidence_level (float) – Nominal coverage $1-\alpha$.
width_weight (float) – Weight on the mean-width term relative to the coverage term.
reduction (str) – "none", "mean" or "sum".

Return type:

Returns:

The reduced regularisation term.

Raises:

ValueError – If confidence_level is outside (0, 1) or the shapes of targets, lower and upper disagree.

Parametric heads - additive backbones (MLPGAM / MLPGAF)#

These variants receive the pre-summed additive mean directly from the backbone instead of a latent vector, preserving the GAM decomposition.

class twiga.distributions.nn.parametric.AdditiveNormalDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Normal output head preserving additive backbone structure.

For use with MLPGAMNetwork and MLPGAFNetwork, whose encode() already returns the additive mean directly as a flat (B, H*O) vector. The mean is taken as-is (no projection) to honour the additive decomposition; only the scale is learned via a separate linear layer.

Parameters predicted:: mu - additive mean (reshaped directly from latent, no projection) sigma - std dev (exp of log_scale, strictly positive)

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_size (int) – Must equal forecast_horizon * num_target_output.
forecast_horizon (int) – Number of forecast time steps.
out_activation_function (Module | None) – Optional activation applied to mu. Defaults to Identity.

forward(x)#

Predict mu and sigma.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size) - the additive mean from the backbone.
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, sigma), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, sigma)#

Construct a torch Distribution from predicted parameters.

Return type:: Normal

get_log_likelihood(mu, sigma, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.AdditiveLaplaceDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Laplace output head preserving additive backbone structure.

Analogous to AdditiveNormalDistribution but with Laplace tails. Best for: heavy-tailed, outlier-robust targets - electricity prices, wind residuals.

Parameters predicted:: mu - additive location (reshaped directly from latent, no projection) scale - scale (exp of log_scale, strictly positive)

family: str | None = 'laplace'#

forward(x)#

Predict mu and scale.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size) - the additive location from the backbone.
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, scale), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, scale)#

Construct a torch Distribution from predicted parameters.

Return type:: Laplace

get_log_likelihood(mu, scale, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.AdditiveLogNormalDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None)#

Log-Normal output head preserving additive backbone structure.

Best for: strictly positive, right-skewed targets - renewable generation, gas prices.

Parameters predicted:: mu - additive log-space mean (reshaped directly from latent, no projection) sigma - std of log(target) (exp of log_scale, strictly positive)

Note

Targets must be strictly positive.

forward(x)#

Predict mu and sigma of the underlying Normal (in log space).

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size) - the additive log-mean from the backbone.
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple of (mu, sigma), each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, sigma)#

Construct a torch Distribution from predicted parameters.

Return type:: LogNormal

get_log_likelihood(mu, sigma, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

class twiga.distributions.nn.parametric.AdditiveStudentTDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None, min_df=2.1)#

Student-T output head preserving additive backbone structure.

Best for: very heavy tails - spot electricity prices, financial returns.

Parameters predicted:: mu - additive location (reshaped directly from latent, no projection) sigma - scale (exp of log_scale, strictly positive) df - degrees of freedom (softplus, clipped to ≥ min_df for finite variance)

forecast(z)#

Inference-mode forward (no gradient).

Parameters:: z (Tensor) – Latent tensor of shape (B, hidden_size) from the backbone.
Return type:: dict[str, Tensor]
Returns:: Dict with "loc" (location/mean) and "scale" (std/dispersion). The first forward() output is always loc; the second is scale.

forward(x)#

Predict mu, sigma, and degrees of freedom.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size) - the additive location from the backbone.
Return type:: tuple[Tensor, Tensor, Tensor]
Returns:: Tuple of (mu, sigma, df) – - mu, sigma: shape (B, forecast_horizon, num_target_output) - df: shape (B, 1, num_target_output), broadcast-compatible

get_distribution(mu, sigma, df)#

Construct a torch Distribution from predicted parameters.

Return type:: StudentT

get_log_likelihood(mu, sigma, df, targets)#

Return mean negative log-likelihood over the batch.

Return type:: Tensor

Quantile regression distribution#

class twiga.distributions.nn.quantile.QRDistribution(quantiles=None, num_outputs=1, hidden_size=256, horizon=48, kappa=0.5, output_activation=None, conf_level=0.05, loss_fn='pinball', crossing_penalty=10.0, sort_output=False, monotone_output=False, crossing_mode='relu', softplus_beta=1.0)#

Bases: Module

Fixed-grid quantile regression head.

Predicts a fixed set of quantile levels $tau$ directly from a latent representation via a single linear projection. Boundary quantiles conf_level/2 and 1 - conf_level/2 are automatically prepended / appended so the output always spans the requested coverage band.

Non-crossing is encouraged by a penalty term in step(); it is not enforced structurally. For architectural monotonicity guarantees, use FPQRDistribution with qvn_type="monotone".

Parameters:

quantiles (list[float] | None) – Explicit quantile levels the model outputs. None defaults to [0.1, 0.2, …, 0.9]. The list is used exactly as provided — no boundary levels are auto-inserted. Include conf_level/2 and 1 - conf_level/2 explicitly if raw (non-conformal) interval construction is required.
num_outputs (int) – Number of target series D.
hidden_size (int) – Width of the backbone latent vector fed as input.
horizon (int) – Forecast horizon H.
kappa (float) – Huber transition parameter (only used for huber-pinball).
output_activation (Module | None) – Applied elementwise to the raw quantile outputs. None (default) uses Identity. Warning: non-identity activations (e.g. ReLU) constrain the output range and will produce wrong forecasts for targets that can be negative.
conf_level (float) – Significance level α; stored as metadata for downstream consumers (conformal calibration, interval labelling). Does not affect the quantile grid.
loss_fn (str) – "pinball", "huber-pinball", or "winkler". "winkler" applies the interval score to the outer quantile pair and pinball loss to all K levels.
crossing_penalty (float) – Weight λ for the non-crossing hinge penalty.
sort_output (bool) – If True, sort quantile predictions along the quantile axis (dim=1) at inference and training time. Guarantees non-crossing at inference without changing the parameterisation. Has no effect when monotone_output=True.
monotone_output (bool) – If True, use a cumulative-softplus construction in forward() to guarantee strict monotonicity by parameterisation (anchor + positive increments). Supersedes sort_output and makes crossing_penalty redundant.

forecast(input_tensor)#

Forecast quantiles for the given input tensor.

Parameters:

input_tensor (Tensor) – Input tensor of shape (batch_size, hidden_size).

Return type:

Returns:

dict –

Keys loc (B, H, D), quantiles (B, K, H, D),: conf_level (float), quantile_levels (K,).

forward(x)#

Forward pass producing raw (unconstrained) quantile values.

Outputs K quantile predictions directly from the linear head without structural monotonicity enforcement. Non-crossing is encouraged during training via the crossing_penalty term in step():

loss = pinball_loss + λ · non_crossing_loss(Q̂)

where non_crossing_loss penalises any pair (k, k+1) where Q̂(τₖ) > Q̂(τₖ₊₁). The quantile_value_layer is initialised with small uniform weights to reduce crossing violations at the start of training.

Parameters:: x (Tensor) – Latent features of shape (B, hidden_size).
Return type:: Tensor
Returns:: Estimated quantile values Q(τⱼ | x), shape (B, K, H, D) where K = len(self.taus).

step(z, y, metric_fn, _epoch=None)#

Perform a single training/validation step.

Parameters:

z (Tensor) – Latent tensor of shape (B, hidden_size) from the backbone.
y (Tensor) – Target tensor of shape (B, forecast_horizon, num_outputs).
metric_fn (Callable[..., Any]) – Callable returning a scalar metric given (pred, target).
_epoch (int | None) – Unused; accepted for interface consistency with other heads.

Return type:

Returns:

Tuple of (loss, metric).

twiga.distributions.nn.quantile.get_median_quantile(quantile_hats, probs)#

Compute the median (0.5 quantile) from a quantile tensor of shape (B, N, T, C).

Parameters:

quantile_hats (Tensor) – Tensor of shape (B, N, T, C) containing quantile values.
probs (Tensor | list) – Quantile probabilities (e.g., [0.1, 0.3, 0.5, 0.7, 0.9]).

Return type:

Returns:

torch.Tensor – Median values of shape (B, T, C).

Raises:

ValueError – If input shapes are invalid or interpolation is not possible.

FPQR components#

class twiga.distributions.nn.fpquantile.QuantileProposal(n_quantiles=10, z_dim=64, dropout=0.1, conf_level=0.05, n_outputs=1, crossing_penalty=10.0, stop_gradient=True, num_targets=1, shared_across_horizon=False, eps=0.001, anchor_outer=False)#

Bases: Module

A neural network module for proposing quantile probability levels.

Projects the context vector z through a linear layer and softmax to obtain a valid probability distribution over n_quantiles slots, then computes the cumulative sum to produce strictly increasing probability levels taus ∈ [0, 1].

What this module outputs. tau_hats are probability-level midpoints (values in (0, 1)), not forecast values. The value network maps these levels to actual quantile predictions.

Gradient isolation. tau_hats are detached before being returned. The value-network loss must not back-propagate through the proposal parameters; those are trained separately via QuantileProposalLoss.

Entropy term. entropies (Shannon entropy of the floored proposal distribution) should be subtracted from the total loss (loss -= entropies.mean()), encouraging the proposal to spread its quantile grid rather than concentrate on a few levels.

crossing_penalty. Stored as self.crossing_penalty and read by step() to scale non_crossing_loss(). Not applied inside forward().

Variables:

n_quantiles – Number of quantiles to propose.
n_outputs – Number of forecast horizons.
z_dim – Dimensionality of the input features.
net – Linear layer mapping z to quantile logits.
dropout – Dropout layer for regularization.
conf_level – Stored as metadata for downstream consumers (conformal calibration). Does not affect the tau_hats range unless anchor_outer=True (see below).
tau_0 – Zero lower-boundary buffer prepended to taus.
eps – Increment floor; tau_hats are bounded to [eps/2, 1 - eps/2] by construction.
anchor_outer – When True, pins the outermost tau_hat on each end to exactly conf_level / 2 and 1 - conf_level / 2 instead of letting it float – see the Args entry below.

Parameters:

n_quantiles (int) – Number of quantiles to propose (default: 10).
z_dim (int) – Dimensionality of the input features (default: 64).
dropout (float) – Dropout rate for regularization (default: 0.1).
conf_level (float) – Significance level stored as metadata (default: 0.05).
n_outputs (int) – Number of forecast horizons (default: 1).
crossing_penalty (float) – Penalty weight read externally by FPQRDistribution.step (default: 10.0).
shared_across_horizon (bool) – If True, project z to a single set of n_quantiles logits and broadcast the resulting grid identically across every horizon step, reproducing Chapter 5’s single-grid softmax proposal (“a single quantile grid… applied uniformly across all forecast horizons”). Default False preserves the existing behaviour: one linear layer of width n_quantiles * n_outputs, whose per-horizon output columns are independently normalised and therefore vary across horizon steps even without Chapter 6’s explicit horizon-conditioning — this is the mode Chapter 6/7 experiments already depend on and this flag must never change it by default.
eps (float) – Floor on each probability increment. Must satisfy 0 < eps < 1/n_quantiles so the floor projection preserves the simplex. Default 1e-3 gives a meaningful bound on tau_hats ([5e-4, 1 - 5e-4]); use smaller values for a looser bound.
anchor_outer (bool) – If True, the first and last probability increments are overridden to exactly conf_level each (so the outermost tau_hat sits at conf_level/2 and 1 - conf_level/2, the nominal interval boundary), and the remaining n_quantiles - 2 interior increments are the learned (floored) proposal renormalised to fill the remaining 1 - conf_level span – their relative proportions are preserved, only their scale changes. This is what a fixed-grid model like QR gets for free (its outer levels are config constants); here it trades away part of the adaptive grid’s freedom for a WS-relevant guarantee: the interval the model is scored on is never wider than the proposal itself can currently justify being from the nominal level. Requires n_quantiles >= 3 (two anchors plus at least one interior slot) and is only implemented for the softmax proposal (proposal_type="softmax") – HorizonQuantileProposal and KumaraswamyQuantileProposal do not support it. Default False preserves the original fully-adaptive behaviour.

forward(z)#

Compute probability-level proposals from the context vector.

Parameters:

z (Tensor) – Context vector, shape (batch_size, z_dim).

Return type:

tuple[Tensor, Tensor, Tensor]

Returns:

Tuple of three tensors –

taus (B, N+1, H) — strictly increasing cumulative probability boundaries including the fixed zero lower bound.
tau_hats (B, N, H) — midpoint probability levels between consecutive boundaries, bounded to [eps/2, 1 - eps/2] by construction (no explicit clamp). Detached when stop_gradient=True (default), keeping the pinball gradient isolated from proposal parameters. When stop_gradient=False, gradients flow from the pinball loss through the τ-embedding back to the proposal parameters, supplementing the Wasserstein-1 training signal.
entropies (B, H) — Shannon entropy of the floored proposal distribution (the one that determines taus). Subtract from the total loss to encourage diversity.

Raises:

ValueError – If z does not have shape (batch_size, z_dim).

class twiga.distributions.nn.fpquantile.CosinetauEmbedding(num_cosines=32, z_dim=128, num_outputs=48, tau_embed_type='TimeEmb')#

Bases: Module

A PyTorch module for embedding time values using cosine transformations.

Transforms quantile levels $hat{tau}$ into cosine features and projects them to the backbone hidden dimension via a configurable embedding strategy.

Parameters:

num_cosines (int) – Number of cosine functions. Defaults to 32.
z_dim (int) – Output embedding dimension. Defaults to 128.
num_outputs (int) – Number of horizon steps (scaling factor for input size). Defaults to 48.
tau_embed_type (str) –
Projection strategy applied after the cosine features are computed. One of:
- "TimeEmb" (default) — Linear(H·K → d) + ReLU applied per quantile level independently. No inter-quantile mixing.
- "LinearEmb" — pure linear projection without bias or activation. Slightly fewer parameters than TimeEmb; useful when a non-linear gate is provided elsewhere.
- "ConvEmb" — circular Conv1d(kernel=3) over the quantile axis. Mixes cosine features from neighbouring quantile levels $hat{tau}_{n-1}, hat{tau}_n, hat{tau}_{n+1}$, allowing the embedding to encode relative spacing between adjacent levels.

forward(taus)#

Transforms time values into cosine-based embeddings.

Parameters:: taus (Tensor) – Input tensor of shape (batch_size, num_outputs, num_inputs).
Return type:: Tensor
Returns:: Tensor of shape (batch_size, num_outputs, z_dim) containing embedded time values.

class twiga.distributions.nn.fpquantile.FPQRDistribution(n_quantiles=9, num_outputs=1, hidden_size=256, horizon=48, alpha=0.1, proposal_type='softmax', proposal_shared_horizon=False, proposal_eps=0.001, proposal_anchor_outer=False, tau_embed_type='TimeEmb', rope_type='linear', rope_base=None, dropout=0.1, conf_level=0.05, kappa=0.25, num_cosines=32, output_activation=None, activation_function='ReLU', loss_fn='pinball', crossing_penalty=10.0, sort_output=False, qvn_type='standard', crossing_mode='relu', softplus_beta=1.0, qvn_num_heads=4, umnn_quad_points=8, lambda_cal=0.0, cal_temperature=1.0, proposal_stop_gradient=True, proposal_loss_fn='wasserstein', sqda_weight=0.1, entropy_weight=1.0, per_target_tau=False)#

Bases: Module

A neural network for forecasting quantiles using a quantile value network.

This module predicts quantiles for a forecasting task, combining a quantile proposal layer with a linear output layer to produce quantile forecasts.

Parameters:

n_quantiles (int | None) – Number of quantiles to forecast. If None, defaults to 9.
num_outputs (int) – Number of output features. Defaults to 1.
hidden_dim (int) – Size of the hidden layers. Defaults to 256.
horizon (int) – Number of time steps to forecast. Defaults to 48.
dropout (float) – Dropout rate for the quantile proposal layer. Defaults to 0.1.
conf_level (float) – Confidence level for quantile proposal. Defaults to 0.05.
output_activation (Module | None) – Activation function for the output layer. Defaults to nn.Identity().
proposal_type (Literal['softmax', 'horizon', 'kumaraswamy']) – Quantile proposal strategy - 'softmax' (default, entropy-regularised uniform push), 'horizon' (per-step conditioning), or 'kumaraswamy' (tail-adaptive).
proposal_shared_horizon (bool) – When proposal_type='softmax', propose a single quantile grid and broadcast it identically across every horizon step, reproducing Chapter 5’s stated single-grid design (“a single quantile grid… applied uniformly across all forecast horizons”). Default False preserves the existing behaviour, where each horizon step gets an independently-normalised grid even without Chapter 6’s explicit horizon conditioning — this is the mode Chapter 6/7 experiments already depend on. Only valid with proposal_type='softmax'; raises ValueError otherwise, since 'horizon' is per-horizon by definition and 'kumaraswamy' varies along a different axis (shape parameters, not grid sharing).
tau_embed_type (str) – Projection applied to the cosine features (only when num_cosines > 0). 'TimeEmb' (default, linear + ReLU), 'LinearEmb' (pure linear), or 'ConvEmb' (conv over quantile axis).
rope_type (str) – RoPE variant used when num_cosines=0. One of 'linear' (default, integer slot-index positions), 'tau_pos' (mean τ̄_k scaled to [0, N_q] as float positions), 'factored' (split z_dim: quantile-ordering RoPE ⊕ horizon RoPE), or 'direct' (no projection; backbone z rotated by τ position).
rope_base (float | None) – RoPE base frequency passed to TauRoPEEmbedding. None (default) auto-computes max(2, p_max × 4/π) per axis. Only used when num_cosines=0.
qvn_type (Literal['standard', 'monotone', 'monotone_cal', 'attention', 'lattice', 'spline', 'umnn']) –
Quantile value network architecture.
- 'standard' (default) — direct prediction; an optional crossing penalty enforces monotonicity during training.
- 'monotone' — cumulative-softplus increments (MonotonicQVN); non-crossing is guaranteed by construction and no penalty is applied.
- 'attention' — cross-attention between the tau embedding and backbone context (AttentionQVN), reproducing the legacy FPSe2Q-MLP paper’s fusion mechanism (PhD-Thesis-Anthony’s src/fpq/FPQ.py::FPQSeq2Network.QVN). No architectural non-crossing guarantee; uses the same crossing penalty as 'standard'.
- 'lattice' — non-negative weight reparameterisation (LatticeQVN, Cannon, 2018); monotone in tau by construction, no per-slot sequential dependency.
- 'spline' — harmonic-mean-smoothed rate, integrated over FPQR’s own adaptively-proposed tau knots (SplineQVN, Fritsch & Carlson, 1980).
- 'umnn' — integrates an unconstrained derivative network via Gauss-Legendre quadrature (UMNNQVN, Wehenkel & Louppe, NeurIPS 2019).
The last three were added to test alternatives to 'monotone' after it showed an under-fit, undercovering failure mode on the Chapter 5 FPQR design ablation; see experiment/chapter-5-fpqr/README.md.
entropy_weight (float) – Weight on the entropy regulariser that keeps the proposal grid from collapsing. 1.0 reproduces the previous hard-coded behaviour; 0.0 removes the term, which is what an ablation of the proposal-side objective requires. The value does not carry across n_quantiles: the pinball term sums over slots and grows linearly in N, while the entropy it balances is bounded by log(N), so a weight that keeps the grid healthy at N=9 lets it collapse onto the tau_hats clamp at N=31. Re-tune it, scaling roughly linearly with N, whenever n_quantiles changes.
crossing_mode (Literal['relu', 'softplus']) –
Non-crossing penalty shape used when qvn_type='standard'.
- 'relu' (default) — hinge penalty; zero gradient when quantiles are already ordered.
- 'softplus' — smooth hinge; gradients flow even for near-crossing pairs, removing the dead-gradient zone of ReLU.
softplus_beta (float) – Sharpness parameter shared by both modes. When crossing_mode='softplus' it controls the smoothness of the crossing penalty (larger → closer to ReLU). When qvn_type='monotone' it controls the sharpness of the increment gate in MonotonicQVN. Default 1.0.
qvn_num_heads (int) – Number of attention heads for qvn_type='attention'. Must divide hidden_size. Ignored for every other qvn_type. Default 4.
per_target_tau (bool) –
Learn one quantile grid per target instead of sharing a single grid across all of them. Default False.

The proposal is built with n_outputs = horizon, so by default the same probability levels are applied to every target. Where the targets are distinct physical quantities that compromise is real: active power at a substation is skewed and heavy-tailed once photovoltaic generation is present, while voltage is tight and near-symmetric, and one grid serves neither well.

When enabled, a learned per-target offset is added to the backbone context and the same projection is reused per target, so the cost is one D x hidden_size embedding rather than a D-fold wider proposal. The value head is correspondingly narrowed from hidden_size -> horizon * D to hidden_size -> horizon, so the per-target head is smaller overall than the shared-grid one.

The offsets are zero-initialised, so training starts from the shared-grid model and any divergence between targets has to be earned. The flag is inert when num_outputs == 1, which keeps single-target behaviour and results unchanged.

With this enabled, taus, tau_hats and entropies gain a trailing target axis: (B, N+1, H, D), (B, N, H, D) and (B, H, D) respectively.

forecast(input_tensor)#

Forecast quantiles for the given input tensor.

quantile_levels is the per-sample proposal grid averaged across the batch and horizon dimensions so that downstream consumers receive a fixed 1-D array of representative probability levels, matching the interface expected by ForecastResult.

Parameters:

input_tensor (Tensor) – Input tensor of shape (batch_size, z_dim).

Return type:

dict[str, Any]

Returns:

Dict with keys –

"loc": interval-width-weighted sum over every quantile slot, shape (B, horizon, num_outputs).
"quantiles": quantile forecasts, shape (B, n_quantiles, horizon, num_outputs).
"quantile_levels": representative 1-D numpy array of length n_quantiles (mean of tau_hats over batch and horizon).
"conf_level": significance level scalar.

forward(input_tensor)#

Forward pass of the quantile forecasting network.

Parameters:

input_tensor (Tensor) – Input tensor of shape (batch_size, z_dim).

Return type:

tuple[Tensor, ...]

Returns:

torch.Tensor –

Forecasted quantiles of shape: (batch_size, n_quantiles, horizon, num_outputs).

step(z, y, metric_fn, epoch=None)#

Perform a single training/validation step.

Parameters:

z (Tensor) – Latent tensor of shape (B, hidden_size) from the backbone.
y (Tensor) – Target tensor of shape (B, forecast_horizon, num_outputs).
metric_fn (Callable[..., Any]) – Callable returning a scalar metric given (pred, target).
epoch (int | None) – Current epoch (unused; accepted for interface consistency).

Return type:

Returns:

Tuple of (loss, metric).

CRC distribution classes#

class twiga.distributions.nn.residual_conformal.CRCDistribution(num_target_output=1, hidden_size=256, forecast_horizon=48, out_activation_function=None, sigma_loss_fn='hybrid', alpha=0.1, activation='ReLU')#

CRC head for standard backbones (e.g. MLPFNetwork).

The mean is computed by projecting the latent vector z through a linear layer, then reshaping to (B, H, O). The scale is predicted by a sigma layer applied to the detached mean (flattened), ensuring that gradients do not flow back into the backbone during sigma-only training.

Training follows a two-stage protocol:: Stage 1 — step(): optimise μ (backbone + mu_layer), σ is not updated. Stage 2 — step_sigma(): freeze backbone, optimise σ-head only.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_dim – Dimensionality of the backbone latent vector.
forecast_horizon (int) – Number of forecast time steps.
out_activation_function (Module | None) – Optional activation applied to the mean output.
sigma_loss_fn (str) – Calibration objective for σ. One of SIGMA_LOSS_FNS. Default: "hybrid".
alpha (float) – Weight for MSE vs L1 in the mu and hybrid sigma losses.
activation (str) – Unused; kept for API compatibility.

Raises:

ValueError – If sigma_loss_fn is not one of the supported values.

forecast(z)#

Inference-only prediction (no gradients).

For "hybrid_sqrt" the scale is squared to convert from √|r| space back to residual magnitude space.

Return type:: dict[str, Tensor]
Returns:: Dict with "loc" and "scale" in prediction space.

forward(x)#

Predict mean and scale.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_dim).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple (mu, sigma) each of shape (B, forecast_horizon, num_target_output).

get_distribution(mu, sigma)#

Construct a torch distribution for interval generation.

Returns Laplace(μ, σ) for "laplace" and Normal(μ, σ) for all other modes, treating σ as an empirical scale estimate. For "hybrid_sqrt" callers must pass the forecast-adjusted (squared) σ.

Parameters:

mu (Tensor) – Mean tensor of shape (B, H, O).
sigma (Tensor) – Scale tensor in prediction space, same shape.

Return type:

Normal | Laplace

Returns:

Normal or Laplace distribution.

get_log_likelihood(mu, sigma, targets)#

Compute the sigma calibration loss (NLL for probabilistic modes).

Parameters:

mu (Tensor) – Predicted mean.
sigma (Tensor) – Predicted scale in learned space (as returned by forward).
targets (Tensor) – Ground-truth targets.

Return type:

Returns:

Scalar loss tensor.

step(z, y, metric_fn, epoch=None)#

Stage-1 training step: optimise mean parameters.

Parameters:

z (Tensor) – Latent tensor from backbone.
y (Tensor) – Target values.
metric_fn (Callable[..., Tensor]) – Callable (pred, target) → scalar.
epoch (int | None) – Current epoch (unused; kept for API uniformity).

Return type:

Returns:

Tuple (loss, metric).

step_sigma(z, y, metric_fn)#

Stage-2 training step: optimise sigma with frozen backbone.

The mean is recomputed under no_grad so sigma learns to predict residual magnitude without influencing point predictions.

Parameters:

z (Tensor) – Latent tensor from backbone.
y (Tensor) – Target values.
metric_fn (Callable[..., Tensor]) – Callable (pred, target) → scalar.

Return type:

Returns:

Tuple (sigma_loss, metric).

class twiga.distributions.nn.residual_conformal.AdditiveCRCDistribution(num_target_output=1, hidden_size=None, hidden_dim=256, forecast_horizon=48, out_activation_function=None, sigma_loss_fn='hybrid_sqrt', alpha=0.5, sigma_dropout=0.05, activation='SiLU')#

Bases: CRCDistribution

CRC head for additive backbones (MLPGAMNetwork, MLPGAFNetwork).

Inherits all training and inference logic from CRCDistribution. Two architectural differences from the parent:

mu_layer — identity (or optional activation): the backbone’s encode() already returns the additive mean as (B, H*O), so no linear projection is needed.
sigma_layer — ResidualSigmaHead (two-layer MLP + LayerNorm + Dropout): a deeper network better suited to modeling conditional heteroscedasticity from the additive mean signal.

All other methods — step, step_sigma, forecast, get_distribution, get_log_likelihood — are inherited unchanged.

Parameters:

num_target_output (int) – Number of output features per time step.
hidden_dim (int) – Hidden dimension of the sigma MLP. Independent of H*O.
forecast_horizon (int) – Number of forecast time steps.
out_activation_function (Module | None) – Optional activation applied to the mean.
sigma_loss_fn (str) – Calibration objective (see CRCDistribution). Default: "hybrid_sqrt".
alpha (float) – MSE/L1 weight. Default: 0.5.
sigma_dropout (float) – Dropout rate in the sigma MLP.
activation (str) – Activation name for the sigma MLP. Default: "SiLU".

Raises:

ValueError – If sigma_loss_fn is not one of the supported values.

forward(x)#

Predict mean and scale.

Parameters:: x (Tensor) – Flattened additive mean from backbone, shape (B, H*O).
Return type:: tuple[Tensor, Tensor]
Returns:: Tuple (mu, sigma) each of shape (B, forecast_horizon, num_target_output).

Custom loss functions#

class twiga.distributions.nn.copula.GaussianCopula(num_targets, hidden_size, forecast_horizon, rope_embed_dim=None)#

Bases: Module

Conditional Gaussian copula over the target dimension.

Predicts one correlation matrix per lead time from the backbone latent, parameterised by a Cholesky factor whose rows are normalised to unit length so that $R = LL^{\top}$ has a unit diagonal by construction and is a correlation matrix rather than a covariance.

Dependence is modelled across targets within a lead time and not across lead times, matching MultivariateNormalDistribution.

Parameters:

num_targets (int) – Number of jointly forecast targets, $D$.
hidden_size (int) – Width of the backbone latent.
forecast_horizon (int) – Number of steps ahead, $H$.
rope_embed_dim (int | None) – Width of the HorizonRoPEConditioner’s rotated embedding. Defaults to hidden_size.

Raises:

ValueError – If fewer than two targets are given, since a copula over one margin carries no information.

correlation(corr_tril)#

The correlation matrix itself, for reporting and diagnostics.

Return type:: Tensor

forward(x)#

Map the latent to a correlation Cholesky factor.

Parameters:: x (Tensor) – Latent tensor of shape (B, hidden_size).
Return type:: Tensor
Returns:: Lower-triangular factor of shape (B, H, D, D) with unit-norm rows.

log_prob(u, corr_tril)#

Log density of the copula at PIT values u.

With $z = \Phi^{-1}(u)$ and $R = LL^{\top}$,

\[\log c(u) = -\tfrac12 \log|R| - \tfrac12\left(z^{\top}R^{-1}z - z^{\top}z\right).\]

Parameters:

u (Tensor) – PIT values, shape (B, H, D).
corr_tril (Tensor) – Correlation Cholesky, shape (B, H, D, D).

Return type:

Returns:

Log density, shape (B, H).

sample(corr_tril, quantile_values, tau_hats, num_samples)#

Draw coherent joint samples with the quantile margins left intact.

This is the operation an independent head cannot perform, and the one a constraint such as apparent power against a transformer rating needs.

Parameters:

corr_tril (Tensor) – Correlation Cholesky, shape (B, H, D, D).
quantile_values (Tensor) – Quantile forecasts, shape (B, N, H, D).
tau_hats (Tensor) – Probability levels, (B, N, H) or (B, N, H, D).
num_samples (int) – Number of draws.

Return type:

Returns:

Samples of shape (num_samples, B, H, D).

sample_uniform(corr_tril, num_samples)#

Draw dependent uniforms.

Parameters:

corr_tril (Tensor) – Correlation Cholesky, shape (B, H, D, D).
num_samples (int) – Number of draws.

Return type:

Returns:

Uniforms of shape (num_samples, B, H, D).

twiga.distributions.nn.copula_sampling.fit_gaussian_copula(observations, quantiles, quantile_levels)#

Estimate a Gaussian copula correlation from held-out margins.

The correlation is fitted on normal scores of the PIT values, which is the standard construction: transforming through each margin removes the marginal shape and leaves the dependence.

Fitting on held-out data matters here. A correlation estimated on the same points the forecast is scored against would flatter the coupled model for reasons that have nothing to do with its dependence.

Parameters:

observations (ndarray) – Held-out observations, shape (n_batch, n_horizon, n_targets).
quantiles (ndarray) – Quantile forecasts for the same rows, shape (n_batch, n_levels, n_horizon, n_targets).
quantile_levels (ndarray) – Probability levels.

Return type:

ndarray

Returns:

Correlation matrix, shape (n_targets, n_targets).

Raises:

ValueError – If fewer than two targets are present, since a copula over one margin carries no information.

twiga.distributions.nn.copula_sampling.coupled_samples(quantiles, quantile_levels, num_samples, correlation=None, random_state=None, family='gaussian', theta=None, edge_families=('gaussian', 'gaussian', 'gaussian'))#

Draw a joint ensemble from quantile margins, coupled by a copula.

With correlation omitted (family="gaussian") or theta omitted (family="clayton"/"vine"), the targets are sampled independently, which is the comparison a coupled forecast has to beat: identical margins, no dependence.

Parameters:

quantiles (ndarray) – Quantile forecasts, shape (n_batch, n_levels, n_horizon, n_targets).
quantile_levels (ndarray) – Probability levels.
num_samples (int) – Ensemble size to draw.
correlation (ndarray | None) – Gaussian copula correlation, shape (n_targets, n_targets). Used only when family="gaussian"; None samples the margins independently.
random_state (int | None) – Seed, for reproducibility.
family (str) – "gaussian" (default), "clayton", or "vine".
theta (float | tuple[float, float, float] | None) – Clayton dependence parameter (see fit_clayton_theta()) when family="clayton", or the 3-tuple (theta_pq, theta_qv, theta_pv_given_q) (see fit_dvine3()) when family="vine". None samples the margins independently in both cases.
edge_families (tuple[str, str, str]) – For family="vine" only: the three edges’ individual families, matching fit_dvine3()’s parameter of the same name.

Return type:

ndarray

Returns:

Samples of shape (num_samples, n_batch, n_horizon, n_targets).

Raises:

ValueError – If correlation does not match the target dimension, family is not one of "gaussian"/"clayton"/"vine", or family="vine" is used with a target count other than 3.

twiga.distributions.nn.copula_sampling.quantile_pit(observations, quantiles, quantile_levels)#

Probability integral transform of observations through quantile margins.

The quantile function is known only at the learned levels, so the marginal distribution function is taken to be piecewise linear between them.

Parameters:

observations (ndarray) – Observed values, shape (n_batch, n_horizon, n_targets).
quantiles (ndarray) – Quantile forecasts, shape (n_batch, n_levels, n_horizon, n_targets).
quantile_levels (ndarray) – Probability levels, broadcastable to (n_levels,).

Return type:

ndarray

Returns:

PIT values in (0, 1), shape (n_batch, n_horizon, n_targets).

Raises:

ValueError – If the array shapes are inconsistent.

class twiga.distributions.nn.custom_loss.QuantileLoss(kappa=0.0, reduction='mean')#

Bases: Module

Quantile loss module for quantile regression.

Computes the pinball loss between predicted quantiles and target values.

Variables:

kappa – Smoothing parameter for softplus approximation.
reduction – Reduction method: ‘none’, ‘mean’, or ‘sum’.

Parameters:

kappa (float) – Smoothing parameter for softplus approximation (default: 0.0).
reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘mean’).

__init__(kappa=0.0, reduction='mean')#

Initialize the QuantileLoss module.

Parameters:

kappa (float) – Smoothing parameter for softplus approximation (default: 0.0).
reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘mean’).

Raises:

ValueError – If kappa < 0 or reduction is invalid.

forward(inputs, quantiles, targets)#

Compute the quantile loss.

Parameters:

inputs (Tensor) – Predicted quantiles of shape (B, N) where B is batch size and N is number of quantiles.
quantiles (Tensor) – Quantile levels of shape (N,).
targets (Tensor) – Ground truth targets of shape (B, 1) or (B, N).

Return type:

Returns:

The computed quantile loss tensor, reduced according to the specified method.

class twiga.distributions.nn.custom_loss.QuantileHuberLoss(kappa=1.0, eps=1e-08, reduction='mean')#

Bases: Module

Quantile Huber loss module for robust quantile regression.

Computes the Huber loss adjusted for quantiles, less sensitive to outliers than L2 loss.

Variables:

kappa – Threshold for switching between L2 and L1 loss.
eps – Small value to prevent division by zero.
reduction – Reduction method: ‘none’, ‘mean’, or ‘sum’.

Parameters:

kappa (float) – Threshold for switching between L2 and L1 loss (default: 1.0).
eps (float) – Small value to prevent division by zero (default: 1e-8).
reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘mean’).

__init__(kappa=1.0, eps=1e-08, reduction='mean')#

Initialize the QuantileHuberLoss module.

Parameters:

kappa (float) – Threshold for switching between L2 and L1 loss (default: 1.0).
eps (float) – Small value to prevent division by zero (default: 1e-8).
reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘mean’).

Raises:

ValueError – If kappa <= 0 or eps <= 0.

forward(inputs, quantiles, targets)#

Compute the quantile Huber loss.

Parameters:

inputs (Tensor) – Predicted values of shape (N, Q, T, C).
quantiles (Tensor) – Quantile levels of shape (Q,).
targets (Tensor) – True target values of shape (N, Q, T, C).

Return type:

Returns:

The computed quantile Huber loss tensor, reduced according to the specified method.

class twiga.distributions.nn.custom_loss.QuantileProposalLoss(reduction='none')#

Bases: Module

Quantile proposal loss module to ensure non-crossing quantiles.

Computes a loss to enforce that predicted quantiles adhere to specified levels and do not cross.

Variables:: reduction – Reduction method: ‘none’, ‘mean’, or ‘sum’.
Parameters:: reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘none’).

__init__(reduction='none')#

Initialize the QuantileProposalLoss module.

Parameters:: reduction (str) – Reduction method: ‘none’, ‘mean’, or ‘sum’ (default: ‘none’).

forward(quantile, quantile_hats, taus)#

Compute the quantile proposal loss.

Parameters:

quantile (Tensor) – True quantile values of shape (N, Q, T, C).
quantile_hats (Tensor) – Predicted quantile values of shape (N, Q, T, C).
taus (Tensor) – Quantile levels of shape (N, Q, T).

Return type:

Returns:

The computed quantile proposal loss tensor, reduced according to the specified method.

class twiga.distributions.nn.custom_loss.SQDAProposalLoss(qdf_exponent=-1.0, divergence='squared', tail_weight=2.0, tail_alpha=0.2, eps=1e-06, reduction='mean')#

Bases: Module

Spectral Quantile Density-Adaptive (SQDA) proposal loss.

Trains the quantile proposal network by minimising the distance between the proposal’s interval widths $\Delta\tau_n$ and the theoretically optimal density-adaptive target widths $\Delta\tau^*_n \propto q(\hat\tau_n)^{-1}$, where $q(\tau) = Q'(\tau)$ is the quantile density function (QDF).

Theoretical motivation

For a grid of $N$ evaluation points, the $L^p$ approximation error of the step-function quantile estimator satisfies:

\[E_n^{L^p} = q(\hat\tau_n)^p \cdot \frac{(\Delta\tau_n)^{p+1}}{p+1}\]

Minimising $\sum_n E_n^{L^p}$ subject to $\sum_n \Delta\tau_n = 1$ yields (via Lagrange multipliers) the same criterion for all $p$:

\[\Delta\tau^*_n \propto q(\hat\tau_n)^{-1}\]

For heavy-tailed distributions $q(\tau)$ is large at the tails ($\tau \to 0, 1$) and small near the median. The optimal widths are therefore smaller at the tails (denser grid) and wider in the interior — exactly the tail-adaptive property that the existing Wasserstein-1 signal cannot enforce directly.

QDF estimation

The QDF is approximated from the current (detached) quantile predictions $\hat{Q}_n$ via central finite differences:

\[\hat{q}_n \approx \frac{\hat{Q}_{n+1} - \hat{Q}_{n-1}}{\hat\tau_{n+1} - \hat\tau_{n-1}}\]

with one-sided differences at the boundaries.

Gradient flow

Target widths $\Delta\tau^*$ are computed from detached quantities only — no gradient flows to the value network. The loss gradient flows exclusively through $\Delta\tau_{\text{prop}} = \tau_{n+1} - \tau_n$ (boundaries with gradient) to the proposal parameters, preserving the same gradient isolation as the Wasserstein-1 loss.

Parameters:

qdf_exponent (float) – Exponent applied to the QDF when computing target widths. Default -1.0 gives the theoretically optimal $q^{-1}$ target.
divergence (Literal['squared', 'kl', 'symkl']) – Divergence used to measure the gap between proposal and target widths. One of 'squared' (MSE, default), 'kl' ($\text{KL}(\text{target}\,\|\,\text{proposal})$), or 'symkl' (symmetrised KL).
tail_weight (float) – Up-weight applied to tail quantile slots for the 'squared' divergence only (default 2.0). Slots with $\hat\tau < \text{tail\_alpha}$ or $\hat\tau > 1 - \text{tail\_alpha}$ receive this weight; all others receive weight 1. Not applied for 'kl' or 'symkl' because a non-uniform slot weight breaks KL non-negativity.
tail_alpha (float) – Boundary between tail and interior regions (default 0.2).
eps (float) – Small constant for numerical stability (default 1e-6).
reduction (str) – 'none', 'mean', or 'sum' (default 'mean').

Shape conventions (same as QuantileProposalLoss):

quantile: (B, N-1, H, D) — not used (kept for interface compat).
quantile_hats: (B, N, H, D) — detached $\hat{Q}$ at midpoints.
taus: (B, N+1, H) — tau boundaries with gradient from proposal.

forward(quantile, quantile_hats, taus)#

Compute the SQDA proposal loss.

Parameters:

quantile (Tensor) – Unused; shape (B, N-1, H, D). Kept for interface compatibility with QuantileProposalLoss.
quantile_hats (Tensor) – Detached quantile predictions at midpoints, shape (B, N, H, D).
taus (Tensor) – Tau boundaries with gradient from the proposal network, shape (B, N+1, H).

Return type:

Returns:

Scalar loss (or per-element tensor when reduction='none').

ML parametric models#

class twiga.models.ml.gausscatboost_model.GAUSSCATBOOSTConfig(**data)#

Bases: CATBOOSTConfig

Configuration for the Gaussian CatBoost probabilistic model.

Extends CATBOOSTConfig with:

Tighter hyperparameter search bounds for faster convergence.
od_type / od_wait for built-in overfitting detection.
virtual_ensembles_count to control the number of virtual ensemble snapshots used to estimate epistemic uncertainty.

Variables:

name – Fixed to "gausscatboost".
od_type – Overfitting detection strategy ("Iter" or "IncToDec"). Active only when an eval set is provided to GAUSSCATBOOSTModel.fit().
od_wait – Number of iterations without improvement before early stopping triggers.
virtual_ensembles_count – Number of virtual ensemble snapshots used by predict_with_uncertainty(). Not passed to CatBoostRegressor (excluded from model_dump).

domain: Literal['ml']#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: Literal['gausscatboost']#

od_type: Literal['Iter', 'IncToDec']#

od_wait: int#

search_space: BaseSearchSpace#

virtual_ensembles_count: int#

class twiga.models.ml.gausscatboost_model.GAUSSCATBOOSTModel(model_config=None)#

Bases: BaseRegressor

Probabilistic CatBoost model predicting mean and uncertainty.

Uses loss_function="RMSEWithUncertainty" which maximises the log-likelihood of a Normal distribution, jointly learning the mean μ and aleatoric scale σ in a single model.

For multi-horizon forecasting (H output steps), one CatBoostRegressor is trained per output, yielding H models. This avoids MultiOutputRegressor while preserving support for virtual_ensembles_predict.

Uncertainty decomposition via predict_with_uncertainty() returns three components per output step:

mean - predicted mean (same as predict() μ output).
knowledge_uncertainty - epistemic (model) uncertainty estimated from variance across virtual ensemble snapshots.
data_uncertainty - aleatoric uncertainty encoded in the model’s second output (exp of predicted log-σ).

Parameters:: model_config (GAUSSCATBOOSTConfig | None) – Model configuration. Defaults to GAUSSCATBOOSTConfig.

Example:

model = GAUSSCATBOOSTModel()
model.fit(X_train, y_train)
mu, sigma = model.predict(X_test)
unc = model.predict_with_uncertainty(X_test)  # (B, L, H, 3)

fit(X, y, eval_set=None, verbose=False, trial=None)#

Fit one RMSEWithUncertainty model per output step.

Parameters:

X (ndarray) – Shape (B, L, F) - batch × sequence × features.
y (ndarray) – Shape (B, L, H) - batch × sequence × horizons.
eval_set (tuple[ndarray, ndarray] | None) – Optional (X_val, y_val) for early stopping. When provided, od_type / od_wait from the config activate CatBoost’s overfitting detector.
verbose (bool) – Whether to print CatBoost training logs.
trial (Any | None) – Optuna trial for mid-training pruning. When provided alongside eval_set, a CatBoostPruningCallback is injected into the first per-output model so Hyperband can prune unpromising trials early.

Return type:

GAUSSCATBOOSTModel

Returns:

Self for method chaining.

Raises:

ValueError – If X or y are not 3-dimensional.

models: list[catboost.CatBoostRegressor]#

num_targets: int | None#

predict(X)#

Predict mean (μ) and scale (σ) for all output steps.

Parameters:: X (ndarray) – Shape (B, L, F).
Return type:: tuple[ndarray, ndarray]
Returns:: Tuple (mu, sigma) each of shape (B, L, H). sigma = exp(log_sigma) is guaranteed positive.
Raises:: ValueError – If the model has not been fitted or X is not 3-D.

predict_with_uncertainty(X)#

Return mean, epistemic, and aleatoric uncertainty per output.

Uses virtual_ensembles_predict(prediction_type="TotalUncertainty") which - for models trained with RMSEWithUncertainty - returns three values per sample:

Column 0: mean prediction (μ).
Column 1: knowledge (epistemic) uncertainty - variance across virtual ensemble snapshots.
Column 2: data (aleatoric) uncertainty - derived from the predicted log-σ output.

Parameters:: X (ndarray) – Shape (B, L, F).
Return type:: ndarray
Returns:: Array of shape (B, L, H, 3) - last axis is [mean, knowledge_uncertainty, data_uncertainty].
Raises:: ValueError – If the model has not been fitted.

set_fit_request(*, eval_set='$UNCHANGED$', trial='$UNCHANGED$', verbose='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

eval_setstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for eval_set parameter in fit.
trialstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for trial parameter in fit.
verbosestr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for verbose parameter in fit.

Returns#

selfobject: The updated object.

set_predict_request()#

No-op.

Calling this method has no effect.

Returns#

selfobject: The updated object.

set_score_request(*, sample_weight='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns#

selfobject: The updated object.

update(trial)#

Rebuild model from an Optuna trial’s suggested hyperparameters.

Parameters:: trial – Optuna trial object.
Return type:: None

virtual_ensembles_count: int#

class twiga.models.ml.ngboostnormal_model.NGBOOSTNORMALConfig(**data)#

Bases: BaseModelConfig

Configuration for the NGBoost Normal probabilistic forecasting model.

Uses natural gradient boosting with a Gaussian predictive distribution N(μ, σ²). Unlike the two-stage GAUSS* models, NGBoost jointly optimises μ and σ via the natural gradient of the chosen scoring rule, which tends to produce better-calibrated uncertainty estimates.

Variables:

name – Model identifier fixed to "ngboostnormal".
domain – Domain fixed to "ml". Excluded from tuning.
n_estimators – Number of boosting iterations.
learning_rate – Shrinkage applied to each tree.
minibatch_frac – Row-subsample fraction per iteration.
col_sample – Column-subsample fraction per iteration.
seed – Seed for reproducibility.
score – Scoring rule - "LogScore" (MLE) or "CRPScore".
search_space – Hyperparameter search space for Optuna tuning.

Example

>>> from twiga.models.ml import NGBOOSTNORMALConfig
>>> cfg = NGBOOSTNORMALConfig(n_estimators=200, learning_rate=0.05)

col_sample: float#

domain: Literal['ml']#

learning_rate: float#

minibatch_frac: float#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

n_estimators: int#

name: Literal['ngboostnormal']#

score: Literal['LogScore', 'CRPScore']#

search_space: BaseSearchSpace#

seed: int#

class twiga.models.ml.ngboostnormal_model.NGBOOSTNORMALModel(model_config=None)#

Bases: BaseNGBoostRegressor

NGBoost probabilistic model with a Normal (Gaussian) predictive distribution.

Wraps ngboost.NGBRegressor with Dist=Normal, training one regressor per flattened output column to support multi-horizon and multi-target forecasting.

The model jointly learns the conditional mean μ and standard deviation σ via natural gradient boosting on the selected scoring rule (log-likelihood or CRPS). Unlike GAUSSCATBOOSTModel, which uses CatBoost’s native RMSEWithUncertainty loss, NGBoost directly maximises the scoring rule via natural gradients.

Variables:: model_config – Instance of NGBOOSTNORMALConfig.
Parameters:: model_config (NGBOOSTNORMALConfig | None) – Model configuration. Defaults to NGBOOSTNORMALConfig.

Example

>>> import numpy as np
>>> from twiga.models.ml import NGBOOSTNORMALModel
>>> model = NGBOOSTNORMALModel()
>>> X = np.random.randn(100, 10, 5)
>>> y = np.random.randn(100, 48, 1)
>>> model.fit(X, y)
>>> loc, scale = model.predict(X)
>>> loc.shape, scale.shape
((100, 48, 1), (100, 48, 1))

forecast(x)#

Return Normal distribution parameters.

Parameters:

x (ndarray) – Input features of shape (n_samples, seq_len, n_features).

Return type:

Returns:

Dictionary with –

"loc": conditional mean μ, shape (n_samples, horizon, n_targets).
"scale": conditional std-dev σ, shape (n_samples, horizon, n_targets).

predict(X)#

Predict μ and σ for each forecast horizon.

Parameters:: X (ndarray) – Input features of shape (n_samples, seq_len, n_features).
Return type:: tuple[ndarray, ndarray]
Returns:: Tuple (loc, scale) - conditional mean μ and standard deviation σ, each of shape (n_samples, horizon, n_targets).

set_fit_request(*, eval_set='$UNCHANGED$', trial='$UNCHANGED$', verbose='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

eval_setstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for eval_set parameter in fit.
trialstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for trial parameter in fit.
verbosestr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for verbose parameter in fit.

Returns#

selfobject: The updated object.

set_predict_request()#

No-op.

Calling this method has no effect.

Returns#

selfobject: The updated object.

set_score_request(*, sample_weight='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns#

selfobject: The updated object.

class twiga.models.ml.ngboostlognormal_model.NGBOOSTLOGNORMALConfig(**data)#

Bases: BaseModelConfig

Configuration for the NGBoost LogNormal probabilistic forecasting model.

Uses natural gradient boosting with a log-normal predictive distribution. Suitable for strictly positive targets such as solar irradiance, wind speed, electricity price, or non-negative load.

Follows the scipy.stats.lognorm parameter convention:

scale = exp(μ_log) - geometric mean (location-like quantity).
s = σ_log - standard deviation in log-space (shape parameter).

Variables:

name – Model identifier fixed to "ngboostlognormal".
domain – Domain fixed to "ml". Excluded from tuning.
n_estimators – Number of boosting iterations.
learning_rate – Shrinkage applied to each tree.
minibatch_frac – Row-subsample fraction per iteration.
col_sample – Column-subsample fraction per iteration.
seed – Seed for reproducibility.
score – Scoring rule - "LogScore" (MLE) or "CRPScore".
search_space – Hyperparameter search space for Optuna tuning.

Example

>>> from twiga.models.ml import NGBOOSTLOGNORMALConfig
>>> cfg = NGBOOSTLOGNORMALConfig(n_estimators=300, score="CRPScore")

col_sample: float#

domain: Literal['ml']#

learning_rate: float#

minibatch_frac: float#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

n_estimators: int#

name: Literal['ngboostlognormal']#

score: Literal['LogScore', 'CRPScore']#

search_space: BaseSearchSpace#

seed: int#

class twiga.models.ml.ngboostlognormal_model.NGBOOSTLOGNORMALModel(model_config=None)#

Bases: BaseNGBoostRegressor

NGBoost probabilistic model with a LogNormal predictive distribution.

Wraps ngboost.NGBRegressor with Dist=LogNormal, training one regressor per flattened output column for multi-horizon / multi-target support.

Parameter convention (follows scipy.stats.lognorm):

loc (exposed as "loc") - scale = exp(μ_log), the geometric mean.
scale (exposed as "scale") - s = σ_log, the log-space std-dev.

Use this model when the target variable is strictly positive and its logarithm is approximately Gaussian (e.g. solar irradiance, wind power, non-negative energy prices).

Variables:: model_config – Instance of NGBOOSTLOGNORMALConfig.
Parameters:: model_config (NGBOOSTLOGNORMALConfig | None) – Model configuration. Defaults to NGBOOSTLOGNORMALConfig.

Example

>>> import numpy as np
>>> from twiga.models.ml import NGBOOSTLOGNORMALModel
>>> model = NGBOOSTLOGNORMALModel()
>>> X = np.random.randn(100, 10, 5)
>>> y = np.abs(np.random.randn(100, 48, 1)) + 0.1  # strictly positive
>>> model.fit(X, y)
>>> loc, scale = model.predict(X)
>>> loc.shape, scale.shape
((100, 48, 1), (100, 48, 1))

forecast(x)#

Return LogNormal distribution parameters.

Parameters:

x (ndarray) – Input features of shape (n_samples, seq_len, n_features).

Return type:

Returns:

Dictionary with –

"loc": geometric mean exp(μ_log), always > 0, shape (n_samples, horizon, n_targets).
"scale": log-space std-dev σ_log, shape (n_samples, horizon, n_targets).

predict(X)#

Predict geometric mean and log-space std-dev for each forecast horizon.

Parameters:

X (ndarray) – Input features of shape (n_samples, seq_len, n_features).

Return type:

tuple[ndarray, ndarray]

Returns:

Tuple (loc, scale) where –

loc = exp(μ_log) - geometric mean (strictly positive), shape (n_samples, horizon, n_targets).
scale = σ_log - standard deviation in log-space, shape (n_samples, horizon, n_targets).

set_fit_request(*, eval_set='$UNCHANGED$', trial='$UNCHANGED$', verbose='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to fit.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

eval_setstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for eval_set parameter in fit.
trialstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for trial parameter in fit.
verbosestr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for verbose parameter in fit.

Returns#

selfobject: The updated object.

set_predict_request()#

No-op.

Calling this method has no effect.

Returns#

selfobject: The updated object.

set_score_request(*, sample_weight='$UNCHANGED$')#

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.
False: metadata is not requested and the meta-estimator will not pass it to score.
None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.
str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters#

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED: Metadata routing for sample_weight parameter in score.

Returns#

selfobject: The updated object.

class twiga.models.ml.ngboostexponential_model.NGBOOSTEXPONENTIALConfig(**data)#

Bases: BaseModelConfig

Configuration for the NGBoost Exponential probabilistic forecasting model.

Uses natural gradient boosting with an exponential predictive distribution. The exponential distribution has a single parameter scale = 1 / λ, where λ is the rate. Both the mean and the standard deviation equal scale.

Suited for non-negative targets exhibiting exponential decay or inter-arrival times - e.g. rare demand events, sparse consumption spikes, or inter-event durations in energy systems.

Variables:

name – Model identifier fixed to "ngboostexponential".
domain – Domain fixed to "ml". Excluded from tuning.
n_estimators – Number of boosting iterations.
learning_rate – Shrinkage applied to each tree.
minibatch_frac – Row-subsample fraction per iteration.
col_sample – Column-subsample fraction per iteration.
seed – Seed for reproducibility.
score – Scoring rule - "LogScore" (MLE) or "CRPScore".
search_space – Hyperparameter search space for Optuna tuning.

Example

>>> from twiga.models.ml import NGBOOSTEXPONENTIALConfig
>>> cfg = NGBOOSTEXPONENTIALConfig(n_estimators=400, score="CRPScore")

col_sample: float#

domain: Literal['ml']#

learning_rate: float#

minibatch_frac: float#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

n_estimators: int#

name: Literal['ngboostexponential']#

score: Literal['LogScore', 'CRPScore']#

search_space: BaseSearchSpace#

seed: int#

class twiga.models.ml.ngboostexponential_model.NGBOOSTEXPONENTIALModel(model_config=None)#

Bases: BaseNGBoostRegressor

NGBoost probabilistic model with an Exponential predictive distribution.

Wraps ngboost.NGBRegressor with Dist=Exponential, training one regressor per flattened output column for multi-horizon / multi-target support.

The exponential distribution has one parameter, scale = 1 / λ, which equals both the mean and the standard deviation. Both "loc" and "scale" in the returned forecast dict are set to this parameter.

Use this model for non-negative targets with memoryless, decay-like behaviour - inter-arrival durations, sparse demand spikes, or short-term outage durations.

Variables:: model_config – Instance of NGBOOSTEXPONENTIALConfig.
Parameters:: model_config (NGBOOSTEXPONENTIALConfig | None) – Model configuration. Defaults to NGBOOSTEXPONENTIALConfig.

Example

>>> import numpy as np
>>> from twiga.models.ml import NGBOOSTEXPONENTIALModel
>>> model = NGBOOSTEXPONENTIALModel()
>>> X = np.random.randn(100, 10, 5)
>>> y = np.random.exponential(scale=2.0, size=(100, 48, 1))
>>> model.fit(X, y)
>>> loc, scale = model.predict(X)
>>> loc.shape, scale.shape
((100, 48, 1), (100, 48, 1))

forecast(x)#

Return Exponential distribution parameters.

Parameters:

x (ndarray) – Input features of shape (n_samples, seq_len, n_features).

Return type: