Parametric Elastic Linear

elastic_linear

CLASS	DESCRIPTION
ElasticLinear	Elastic net linear regression with combined L1 and L2 regularization.

Classes

ElasticLinear(*, model: ModelType | dict[str, ModelType], alpha: float = 1.0, l1_ratio: float = 0.5, loss_kwargs: dict | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, device: Literal['cuda', 'cpu'] = 'cuda')

Bases: Parametric

Elastic net linear regression with combined L1 and L2 regularization.

This class implements elastic net regression, which combines Ridge (L2) and Lasso (L1) regularization techniques. Useful for feature selection and handling multicollinearity in high-dimensional datasets.

The elastic net loss function minimizes:

\(L(w) = \frac{1}{2 N} \| y - X w \|^2_2 + \alpha l_r \| w \|_1 + \frac{\alpha}{2} (1 - l_r) \| w \|^2_2\),

where \(\alpha\) is the regularization strength parameter, \(l_r\) is the elastic net mixing parameter in [0, 1], and \(N\) is the number of samples.

Terms:

• The first term is the mean squared error (MSE) loss between the data \(X\) and the target \(y\)
• The second term is L1 penalty (Lasso) for feature selection
• The third term is L2 penalty (Ridge) for regularization

The model supports PyTorch Lightning training and provides \(R^2\) scoring for model evaluation.

PARAMETER	DESCRIPTION
model	Neural network model or compatible structure. Must contain exactly one linear layer for elastic net regularization to be meaningful. TYPE: ModelType | dict[str, ModelType]
alpha	Regularization strength parameter. Controls the overall amount of regularization applied. Must be positive. Higher values increase regularization strength. TYPE: float, by default 1.0 DEFAULT: 1.0
l1_ratio	Elastic net mixing parameter in [0, 1]. Controls the balance between L1 and L2 regularization: l1_ratio = 0.0: Pure Ridge regression (L2 only) l1_ratio = 1.0: Pure Lasso regression (L1 only) 0.0 < l1_ratio < 1.0: Elastic net (combination of L1 and L2) TYPE: float, by default 0.5 DEFAULT: 0.5
loss_kwargs	Additional keyword arguments to pass to the MSELoss constructor. reduce : Literal["sum", "mean"], optional, by default "mean" Reduction method for loss computation. sign : float, optional, by default 1.0 Sign multiplier for loss (1.0 for minimization). TYPE: dict, optional, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: type[torch.optim.Optimizer] | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
device	Computation device. Will automatically fallback to "cpu" if CUDA is not available. TYPE: Literal["cuda", "cpu"], by default "cuda" DEFAULT: 'cuda'

Examples:

>>> import torch
>>> from spectre.parametric import ElasticLinear
>>>
>>> # Create a simple linear model
>>> linear_layer = torch.nn.Linear(10, 1)
>>> elastic = ElasticLinear(
...     model=linear_layer,
...     alpha=0.1,
...     l1_ratio=0.5,
... )
>>>
>>> # Generate sample data
>>> X = torch.randn(100, 10)
>>> y = torch.randn(100, 1)
>>> weights = torch.ones(100)
>>>
>>> # Training step
>>> batch = (X, weights, y)
>>> loss = elastic.training_step(batch, 0)
>>>
>>> # Evaluate after training (requires is_trained=True)
>>> elastic.is_trained = True
>>> r2_score = elastic.score(X, target=y)

METHOD	DESCRIPTION
score_step	Return the coefficient of determination of the prediction.
training_step	Training step for PyTorch Lightning with elastic net regularization.
forward_step	Forward pass through the linear model.
loss	Compute elastic net loss for given data and target.

ATTRIBUTE	DESCRIPTION
weight	Get the weight matrix of the linear layer. TYPE: Tensor
bias	Get the bias vector of the linear layer. TYPE: Tensor

Source code in spectre/parametric/elastic_linear.py

def __init__(
    self,
    *,
    model: ModelType | dict[str, ModelType],
    alpha: float = 1.0,
    l1_ratio: float = 0.5,
    loss_kwargs: dict | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    device: Literal["cuda", "cpu"] = "cuda",
) -> None:
    super().__init__(
        model=model,
        device=device,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )
    if len(self.model) != 1:
        raise ValueError(
            f"ElasticLinear requires a single linear layer model, got "
            f"{len(self.model)} layers."
        )

    if not isinstance(alpha, (int, float)):
        raise TypeError(f"Parameter `alpha` must be numeric, got {type(alpha)}.")
    check_in_interval(alpha, "(0, inf)")
    self.alpha = alpha

    if not isinstance(l1_ratio, (int, float)):
        raise TypeError(
            f"Parameter `l1_ratio` must be numeric, got {type(l1_ratio)}."
        )
    check_in_interval(l1_ratio, "[0, 1]")
    self.l1_ratio = l1_ratio

    if loss_kwargs is None:
        loss_kwargs = {}
    if not isinstance(loss_kwargs, dict):
        raise TypeError(
            "Additional loss parameters `loss_kwargs` must be a dictionary."
        )

    self.loss_fn = MSELoss(**loss_kwargs)

Attributes

weight: torch.Tensor property

Get the weight matrix of the linear layer.

RETURNS	DESCRIPTION
Tensor	Weight matrix of shape (out_features, in_features).

bias: torch.Tensor property

Get the bias vector of the linear layer.

RETURNS	DESCRIPTION
Tensor	Bias vector of shape (out_features, ).

Functions

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Return the coefficient of determination of the prediction.

\(R^2\) score represents the proportion of variance in the target variable that is predictable from the input features. If the score is 1.0, the model perfectly predicts the target. A score of 0.0 indicates that the model does not explain any of the variance in the target. Negative values indicate that the model performs worse than a horizontal line (mean of the target).

Defined as

\(R^2 = 1 - \frac{\sum_{i=1}^{N} (y_i - \hat{y}_i)^2} {\sum_{i=1}^{N} (y_i - \bar{y})^2}\),

where \(y\) is the true target, \(\hat{y}\) is the predicted target, and \(\bar{y}\) is the mean of the true target.

For multi-dimensional targets, \(R^2\) score is computed for each dimension and then averaged.

If sample weights are provided, a weighted \(R^2\) score is computed.

PARAMETER	DESCRIPTION
X	Dataset of samples. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Target values for X. TYPE: torch.Tensor of shape (n_samples, out_features) | None, by default None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If target values are not provided.

RETURNS	DESCRIPTION
Tensor	Coefficient of determination of the prediction.

Source code in spectre/parametric/elastic_linear.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Return the coefficient of determination of the prediction.

    $R^2$ score represents the proportion of variance in the target
    variable that is predictable from the input features. If the score is 1.0,
    the model perfectly predicts the target. A score of 0.0 indicates that the
    model does not explain any of the variance in the target. Negative values
    indicate that the model performs worse than a horizontal line (mean of the
    target).

    Defined as

    $R^2 = 1 - \\frac{\\sum_{i=1}^{N} (y_i - \\hat{y}_i)^2}
                    {\\sum_{i=1}^{N} (y_i - \\bar{y})^2}$,

    where $y$ is the true target, $\\hat{y}$ is the predicted target,
    and $\\bar{y}$ is the mean of the true target.

    For multi-dimensional targets, $R^2$ score is computed for each
    dimension and then averaged.

    If sample weights are provided, a weighted $R^2$ score is computed.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Dataset of samples.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights.

    target : torch.Tensor of shape (n_samples, out_features) | None, by default None
        Target values for X.

    Raises
    ------
    ValueError
        If target values are not provided.

    Returns
    -------
    torch.Tensor
        Coefficient of determination of the prediction.
    """
    if target is None:
        raise ValueError("Target values are required for scoring.")

    target_pred = self.model(X)

    return r2_loss(target, target_pred, weights=weights, reduce="mean")

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

Training step for PyTorch Lightning with elastic net regularization.

Computes the elastic net loss combining mean squared error with L1 and L2 regularization terms. Automatically logs the loss for monitoring during training.

PARAMETER	DESCRIPTION
batch	Input batch containing (X, weights, target) where: X : torch.Tensor of shape (n_samples, in_features) Input features. weights : torch.Tensor of shape (n_samples,) Sample weights for weighted loss computation. target : torch.Tensor of shape (n_samples, out_features) Target values. TYPE: Batch
batch_idx	Batch index (required by PyTorch Lightning interface but unused). TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar training loss combining MSE with L1/L2 regularization terms.

Source code in spectre/parametric/elastic_linear.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Training step for PyTorch Lightning with elastic net regularization.

    Computes the elastic net loss combining mean squared error with L1 and L2
    regularization terms. Automatically logs the loss for monitoring during
    training.

    Parameters
    ----------
    batch : Batch
        Input batch containing (X, weights, target) where:

        - X : torch.Tensor of shape (n_samples, in_features)
            Input features.
        - weights : torch.Tensor of shape (n_samples,)
            Sample weights for weighted loss computation.
        - target : torch.Tensor of shape (n_samples, out_features)
            Target values.

    batch_idx : int
        Batch index (required by PyTorch Lightning interface but unused).

    Returns
    -------
    torch.Tensor
        Scalar training loss combining MSE with L1/L2 regularization terms.
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss)

    return loss

forward_step(X: torch.Tensor) -> torch.Tensor

Forward pass through the linear model.

Computes predictions by passing input through the single linear layer.

PARAMETER	DESCRIPTION
X	Input features. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, out_features)	Model predictions.

Source code in spectre/parametric/elastic_linear.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass through the linear model.

    Computes predictions by passing input through the single linear layer.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input features.

    Returns
    -------
    torch.Tensor of shape (n_samples, out_features)
        Model predictions.
    """
    X = self.model(X)
    return X

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute elastic net loss for given data and target.

PARAMETER	DESCRIPTION
X	Input features. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. If None, uniform weighting is applied. TYPE: torch.Tensor of shape (n_samples,) or None DEFAULT: None
target	Target values. TYPE: torch.Tensor of shape (n_samples, out_features) or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar loss value combining MSE with L1 and L2 regularization.

Source code in spectre/parametric/elastic_linear.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute elastic net loss for given data and target.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input features.

    weights : torch.Tensor of shape (n_samples,) or None, default None
        Sample weights for weighted loss computation. If None, uniform
        weighting is applied.

    target : torch.Tensor of shape (n_samples, out_features) or None, default None
        Target values.

    Returns
    -------
    torch.Tensor
        Scalar loss value combining MSE with L1 and L2 regularization.
    """
    target_pred = self.forward_step(X)

    context = {"X": target, "Y": target_pred}
    if weights is not None:
        context["weights"] = weights

    loss_mse = 0.5 * self.loss_fn(context)
    loss_l1 = self.alpha * self.l1_ratio * torch.sum(self.weight.abs())
    loss_l2 = 0.5 * self.alpha * (1 - self.l1_ratio) * torch.sum(self.weight.pow(2))

    return loss_mse + loss_l1 + loss_l2

Functions