Skip to content

BaseModel

The abstract foundation for all distributional regression models in DRN. Provides unified interface, PyTorch Lightning integration, and common functionality.

Class Definition

Bases: LightningModule, ABC

Functions

icdf(x, p, l=None, u=None, max_iter=1000, tolerance=1e-07)

Calculate the inverse CDF (quantiles) of the distribution for the given cumulative probability.

This is a fallback implementation for PyTorch distributions that don't have icdf implemented.

Args: x: Input features p: cumulative probability value at which to evaluate icdf l: lower bound for the quantile search u: upper bound for the quantile search max_iter: maximum number of iterations tolerance: stopping criteria

Returns: A tensor of shape (1, batch_shape) containing the inverse CDF values.

predict(x_raw)

Creates a distributional prediction for the input data. Here, x_raw is raw data, and this method will apply any model-specific preprocessing.

preprocess(x, targets=False)

Convert input data to a PyTorch tensor. Apply any neural network preprocessing (if applicable).

quantiles(x, percentiles, l=None, u=None, max_iter=1000, tolerance=1e-07)

Calculate the quantile values for the given observations and percentiles (cumulative probabilities * 100).

This unified implementation first checks if the distribution has its own quantiles method, then falls back to icdf-based approach.

Overview

BaseModel is the abstract base class that all DRN models inherit from. It provides:

  • Unified Interface - Consistent .fit(), .predict(), and .quantiles() methods
  • PyTorch Lightning Integration - Built-in training loops, early stopping, checkpointing
  • Automatic Data Handling - Conversion between pandas/numpy and PyTorch tensors
  • GPU Support - Automatic device detection and tensor placement
  • Distribution Interface - Common methods for working with distributional predictions

Key Methods

Training Methods

fit(X_train, y_train, X_val=None, y_val=None, **kwargs)

Main training method that handles the complete training workflow:

  • Data preprocessing and tensor conversion
  • DataLoader creation with batching and shuffling
  • Validation setup with early stopping (if validation data provided)
  • Model checkpointing and best weight restoration
  • GPU acceleration when available

Parameters: - X_train, y_train - Training features and targets (pandas/numpy/tensor) - X_val, y_val - Optional validation data for early stopping - batch_size - Training batch size (default: 128) - epochs - Maximum training epochs (default: 10) - patience - Early stopping patience (default: 5) - **trainer_kwargs - Additional PyTorch Lightning trainer arguments

Prediction Methods

predict(x_raw)

Creates distributional predictions from input features:

  • Automatic preprocessing - Handles pandas/numpy to tensor conversion
  • Returns distribution objects - With .mean, .quantiles(), .cdf(), etc.
  • Consistent interface - Same API across all model types

quantiles(x, percentiles, l=None, u=None, **kwargs)

Unified quantile calculation supporting all distribution types:

  • Multiple percentiles - Calculate several quantiles at once
  • Bounds specification - Optional lower/upper bounds for search
  • Automatic fallback - Uses binary search when analytic quantiles unavailable

Abstract Methods

Subclasses must implement:

loss(x, y) -> torch.Tensor

Define the loss function for model training.

_predict(x) -> Distribution

Core prediction logic returning PyTorch distribution objects.