Parametric Base

base

CLASS	DESCRIPTION
Parametric	Abstract base class for parametric methods using neural networks.

Classes

Parametric(model: ModelType | dict[str, ModelType], device: Literal['cuda', 'cpu'] = 'cuda', store_loss_context: bool = False, 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, *args, **kwargs)

Bases: ABC, LightningModule

Abstract base class for parametric methods using neural networks.

This class provides a unified interface for parametric machine learning techniques that learn embeddings and transformations through neural network optimization. It integrates with PyTorch Lightning for distributed training, automatic logging, and model management, etc.

The class serves as a foundation for methods such as SpectralMap, Autoencoder, and other parametric dimensionality reduction techniques. It handles model parsing, device management, optimizer configuration, and provides standardized training and evaluation interfaces.

Features:

• Automatic model parsing and device placement.
• PyTorch Lightning integration for scalable training.
• Standardized scoring and evaluation interfaces.
• Optimizer and scheduler configuration.
• Training status tracking for model validation.

PARAMETER	DESCRIPTION
model	Neural network model(s) for the parametric method. Single model neural network architecture (ModelType). Model: Pre-constructed Model instance from spectre.core torch.nn.Sequential: PyTorch Sequential container torch.nn.Module: Individual PyTorch module OrderedDict: Dictionary of named layers list or tuple: Sequence of layers to be wrapped in torch.nn.Sequential Model dictionary (dict[str, ModelType]): Dictionary of multiple named models for multi-model architectures (e.g., autoencoders). Keys are model names ("encoder", "decoder", etc.) and values are model specifications in any format supported above. Example: {"encoder": encoder_net, "decoder": decoder_net}. TYPE: ModelType | dict[str, ModelType]
device	Computation device. Falls back to "cpu" if CUDA is not available. TYPE: Literal["cuda", "cpu"], by default "cuda" DEFAULT: 'cuda'
store_loss_context	Whether to store the loss context dictionary for metric extraction in callbacks. When enabled, the loss_context attribute will contain the context dict passed to the loss function, enabling custom extractors to compute metrics from intermediate values like kernel matrices, embeddings, weights, etc. TYPE: bool, by default False DEFAULT: False
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "Adam", "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) When specified, this parameter takes priority over the legacy optimizer_kwargs property. TYPE: type[torch.optim.Optimizer] | str | None, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation. Only used when optimizer_fn is specified. For example: {"lr": 1e-3, "weight_decay": 1e-4}. TYPE: dict | None, 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., "StepLR", "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. Special keys "monitor" and "interval" control PyTorch Lightning scheduler configuration. Defaults: monitor="val_loss",interval="epoch". TYPE: dict | None, by default None DEFAULT: None
*args	Additional arguments passed to parent classes. TYPE: optional DEFAULT: ()
**kwargs	Additional arguments passed to parent classes. TYPE: optional DEFAULT: ()

ATTRIBUTE	DESCRIPTION
model	Primary parsed neural network model on the specified device. For multi-model architectures, this is the first model from the dictionary. TYPE: Model
model_dict	Dictionary of all parsed models on the specified device. For single-model initialization, contains {'model': model}. TYPE: dict[str, Model]
is_trained	Training status flag. Set to True after successful training completion. TYPE: bool
optimizer_kwargs	Configuration dictionary for optimizer and learning rate scheduler. TYPE: dict
loss_context	Loss context dictionary containing intermediate computation results. Only populated when store_loss_context=True. Typically contains keys like 'K' (kernel matrix), 'Z' (embeddings), 'eigenvalues', 'weights', etc. TYPE: dict | None

Notes

This is an abstract base class that requires subclasses to implement:

• loss(X, weights=None, target=None): Loss computation method
• training_step(batch, batch_idx): PyTorch Lightning training step
• forward_step(X): Forward pass computation step
• score_step(X, weights=None, target=None): Scoring computation step

The score method is only available after training (is_trained=True). Use score_step for evaluation without training status checks.

METHOD	DESCRIPTION
add_transform	Add a transformation to the model.
loss	Abstract method for loss computation.
training_step	Abstract method for PyTorch Lightning training step.
forward_step	Abstract method for forward pass computation.
score_step	Abstract method for score calculation without training validation.
forward	Main forward pass method.
score	Score the trained model performance.
configure_optimizers	Configure optimizer and learning rate scheduler for PyTorch Lightning.
validation_step	Validation step during training.
test_step	Test step during evaluation.
predict_step	Prediction step for inference.
predict	Predict method for inference on raw input data. Can be used
fit	Fit the parametric model using PyTorch Lightning.
parse_model	Parse and prepare model for training.
parse_model_dict	Parse and prepare multiple models for training.
save_model	Save model to file with optional TorchScript tracing.
load_model	Load model from file.
from_pretrained	Load pretrained model from checkpoint.
on_train_end	PyTorch Lightning hook called at the end of training.
on_fit_end	PyTorch Lightning hook called at the end of fit (training +
log_training_metrics	Standardized training/validation metric logging.
log_metrics_dict	Log a dictionary of metrics with flexible naming and automatic tensor

Source code in spectre/parametric/base.py

def __init__(
    self,
    model: ModelType | dict[str, ModelType],
    device: Literal["cuda", "cpu"] = "cuda",
    store_loss_context: bool = False,
    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,
    *args,
    **kwargs,
) -> None:
    super().__init__(*args, **kwargs)

    if isinstance(model, dict):
        self.model_dict = self.parse_model_dict(model, device)
        self.model = next(iter(self.model_dict.values()))
    elif isinstance(model, ModelType):
        self.model = self.parse_model(model, device)
        self.model_dict = {"model": self.model}
    else:
        raise TypeError(
            f"Expected model to be `ModelType | dict[str, ModelType]`, "
            f"got {type(model)}."
        )

    self._optimizer_fn = optimizer_fn
    self._optimizer_init_kwargs = optimizer_kwargs or {}
    self._lr_scheduler_fn = lr_scheduler_fn
    self._lr_scheduler_init_kwargs = lr_scheduler_kwargs or {}

    # Backward compatibility
    self._optimizer_kwargs = {}
    self._lr_scheduler_kwargs = {}

    self._is_trained = False
    self.store_loss_context = store_loss_context
    self.loss_context = None

Attributes

is_trained: bool property writable

Check if model has been trained.

in_features: int property

Get number of input features.

out_features: int property

Get number of output features.

optimizer_kwargs: dict property writable

Get optimizer configuration.

Functions

add_transform(transform: Any, name: str) -> None

Add a transformation to the model.

PARAMETER	DESCRIPTION
transform	Transformation to add (e.g., data preprocessing, embedding alignment). TYPE: Any
name	Name of the transformation. TYPE: str

Source code in spectre/parametric/base.py

def add_transform(self, transform: Any, name: str) -> None:
    """
    Add a transformation to the model.

    Parameters
    ----------
    transform : Any
        Transformation to add (e.g., data preprocessing, embedding
        alignment).

    name : str
        Name of the transformation.
    """
    self.model.add_transform(transform, name=name)

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

Abstract method for loss computation.

This method must be implemented by subclasses to define the specific loss function for the parametric technique (e.g., spectral loss, reconstruction loss, regularized loss, etc.).

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor
weights	Sample weights for weighted loss computation. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values for supervised learning. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor | tuple[Tensor, ...]	Scalar loss value for optimization.

Source code in spectre/parametric/base.py

@abstractmethod
def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
    """
    Abstract method for loss computation.

    This method must be implemented by subclasses to define the specific
    loss function for the parametric technique (e.g., spectral loss,
    reconstruction loss, regularized loss, etc.).

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    weights : torch.Tensor | None, by default None
        Sample weights for weighted loss computation.

    target : torch.Tensor | None, by default None
        Target values for supervised learning.

    Returns
    -------
    torch.Tensor | tuple[torch.Tensor, ...]
        Scalar loss value for optimization.
    """
    raise NotImplementedError()

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

Abstract method for PyTorch Lightning training step.

This method defines how a single training batch is processed during training. It should compute the loss and optionally log metrics for monitoring.

PARAMETER	DESCRIPTION
batch	Batch of data containing 'data', 'weights', and optionally 'target'. TYPE: DataBatch
batch_idx	Index of the current batch within the epoch. TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar training loss for backpropagation.

Notes

Implementations should use self.log() for metric tracking and call self.loss() for loss computation.

Source code in spectre/parametric/base.py

@abstractmethod
def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Abstract method for PyTorch Lightning training step.

    This method defines how a single training batch is processed during
    training. It should compute the loss and optionally log metrics for
    monitoring.

    Parameters
    ----------
    batch : DataBatch
        Batch of data containing 'data', 'weights', and optionally
        'target'.

    batch_idx : int
        Index of the current batch within the epoch.

    Returns
    -------
    torch.Tensor
        Scalar training loss for backpropagation.

    Notes
    -----
    Implementations should use `self.log()` for metric tracking and call
    `self.loss()` for loss computation.
    """
    raise NotImplementedError()

forward_step(X: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, ...] abstractmethod

Abstract method for forward pass computation.

This method defines the forward transformation applied to input data. It should implement the core computation of the parametric technique (e.g., encoding, dimensionality reduction, prediction).

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor | tuple[Tensor, ...]	Transformed output tensor.

Notes

This method is used by both forward() and during loss computation. It should be differentiable for gradient-based optimization.

Source code in spectre/parametric/base.py

@abstractmethod
def forward_step(self, X: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, ...]:
    """
    Abstract method for forward pass computation.

    This method defines the forward transformation applied to input data.
    It should implement the core computation of the parametric technique
    (e.g., encoding, dimensionality reduction, prediction).

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    Returns
    -------
    torch.Tensor | tuple[torch.Tensor, ...]
        Transformed output tensor.

    Notes
    -----
    This method is used by both `forward()` and during loss computation. It
    should be differentiable for gradient-based optimization.
    """
    raise NotImplementedError()

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

Abstract method for score calculation without training validation.

This method must be implemented by subclasses to define the specific performance metric for the parametric technique (e.g., reconstruction error, clustering quality, embedding preservation, etc.).

Unlike score(), this method does not require training status validation and can be used for intermediate evaluation during training or for pre-trained models.

PARAMETER	DESCRIPTION
X	Input data for evaluation. TYPE: Tensor
weights	Sample weights for weighted scoring metrics. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values for supervised evaluation metrics. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Performance score. Higher values typically indicate better performance, but the interpretation depends on the specific metric implementation.

Source code in spectre/parametric/base.py

@abstractmethod
def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Abstract method for score calculation without training validation.

    This method must be implemented by subclasses to define the specific
    performance metric for the parametric technique (e.g., reconstruction
    error, clustering quality, embedding preservation, etc.).

    Unlike `score()`, this method does not require training status
    validation and can be used for intermediate evaluation during training
    or for pre-trained models.

    Parameters
    ----------
    X : torch.Tensor
        Input data for evaluation.

    weights : torch.Tensor | None, by default None
        Sample weights for weighted scoring metrics.

    target : torch.Tensor | None, by default None
        Target values for supervised evaluation metrics.

    Returns
    -------
    torch.Tensor
        Performance score. Higher values typically indicate better
        performance, but the interpretation depends on the specific metric
        implementation.
    """
    raise NotImplementedError()

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

Main forward pass method.

Performs the complete forward transformation by calling the abstract forward_step method implemented by subclasses.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Transformed output tensor.

Notes

This method serves as the standard PyTorch torch.nn.Module forward interface and delegates the actual computation to forward_step.

Source code in spectre/parametric/base.py

def forward(self, X: torch.Tensor) -> torch.Tensor:
    """
    Main forward pass method.

    Performs the complete forward transformation by calling the abstract
    `forward_step` method implemented by subclasses.

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    Returns
    -------
    torch.Tensor
        Transformed output tensor.

    Notes
    -----
    This method serves as the standard PyTorch `torch.nn.Module` forward
    interface and delegates the actual computation to `forward_step`.
    """
    X = self.forward_step(X)
    return X

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

Score the trained model performance.

Evaluates the model using the scoring metric defined by the specific parametric technique. This method requires the model to be trained (is_trained=True) and performs evaluation in no-gradient mode.

PARAMETER	DESCRIPTION
X	Input data for evaluation. TYPE: Tensor
weights	Sample weights for weighted scoring. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values for supervised evaluation. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Model performance score.

RAISES	DESCRIPTION
RuntimeError	If the model has not been trained (is_trained=False).

Notes

This method enforces training validation and wraps score_step with torch.no_grad() for efficient evaluation. Use score_step directly for evaluation without training status checks.

Source code in spectre/parametric/base.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Score the trained model performance.

    Evaluates the model using the scoring metric defined by the specific
    parametric technique. This method requires the model to be trained
    (`is_trained=True`) and performs evaluation in no-gradient mode.

    Parameters
    ----------
    X : torch.Tensor
        Input data for evaluation.

    weights : torch.Tensor | None, by default None
        Sample weights for weighted scoring.

    target : torch.Tensor | None, by default None
        Target values for supervised evaluation.

    Returns
    -------
    torch.Tensor
        Model performance score.

    Raises
    ------
    RuntimeError
        If the model has not been trained (`is_trained=False`).

    Notes
    -----
    This method enforces training validation and wraps `score_step` with
    `torch.no_grad()` for efficient evaluation. Use `score_step` directly
    for evaluation without training status checks.
    """
    self.check_trained()

    with torch.no_grad():
        s = self.score_step(X, weights=weights, target=target)
    return s

configure_optimizers() -> torch.optim.Optimizer | dict

Configure optimizer and learning rate scheduler for PyTorch Lightning.

Supports two configuration modes:

1. Using optimizer_fn parameter in __init__
2. Using optimizer_kwargs property

RETURNS	DESCRIPTION
Optimizer or dict	If no scheduler is configured, returns the optimizer directly. If scheduler is configured, returns dict with "optimizer" and "lr_scheduler" keys for PyTorch Lightning.

RAISES	DESCRIPTION
ValueError	If the specified optimizer name is not found in torch.optim.

Examples:

Configure via optimizer_fn parameter:

>>> sm = SpectralMap(
...     model=encoder,
...     optimizer_fn="AdamW",
...     optimizer_kwargs={"lr": 1e-3, "weight_decay": 1e-4},
...     lr_scheduler_fn="ReduceLROnPlateau",
...     lr_scheduler_kwargs={"patience": 5},
... )

Source code in spectre/parametric/base.py

def configure_optimizers(self) -> torch.optim.Optimizer | dict:  # ty: ignore
    """
    Configure optimizer and learning rate scheduler for PyTorch Lightning.

    Supports two configuration modes:

    1. Using `optimizer_fn` parameter in `__init__`
    2. Using `optimizer_kwargs` property

    Returns
    -------
    torch.optim.Optimizer or dict
        If no scheduler is configured, returns the optimizer directly. If
        scheduler is configured, returns dict with "optimizer" and
        "lr_scheduler" keys for PyTorch Lightning.

    Raises
    ------
    ValueError
        If the specified optimizer name is not found in `torch.optim`.

    Examples
    --------
    Configure via `optimizer_fn` parameter:

    >>> sm = SpectralMap(
    ...     model=encoder,
    ...     optimizer_fn="AdamW",
    ...     optimizer_kwargs={"lr": 1e-3, "weight_decay": 1e-4},
    ...     lr_scheduler_fn="ReduceLROnPlateau",
    ...     lr_scheduler_kwargs={"patience": 5},
    ... )
    """
    if self._optimizer_fn is not None:
        optimizer_cls = initialize_optimizer_fn(
            self._optimizer_fn, self._optimizer_init_kwargs
        )
        optimizer = optimizer_cls(self.parameters(), **self._optimizer_init_kwargs)

        # Handle scheduler if provided
        if self._lr_scheduler_fn is not None:
            scheduler = self._build_scheduler(optimizer)
            return {
                "optimizer": optimizer,
                "lr_scheduler": scheduler,
            }

        return optimizer

    else:
        if "name" in self._optimizer_kwargs:
            optimizer_name = self._optimizer_kwargs.get("name", "Adam")
            if not hasattr(torch.optim, optimizer_name):
                raise ValueError(
                    f"Module `torch.optim` does not have {optimizer_name} optimizer."
                )
            # Create copy without 'name' for instantiation
            opt_kwargs = {
                k: v for k, v in self._optimizer_kwargs.items() if k != "name"
            }
        else:
            optimizer_name = "Adam"
            opt_kwargs = self._optimizer_kwargs

        optimizer = getattr(torch.optim, optimizer_name)(
            self.parameters(), **opt_kwargs
        )

        if self._lr_scheduler_kwargs:
            scheduler_cls = self._lr_scheduler_kwargs["scheduler"]
            scheduler_kwargs = {
                k: v
                for k, v in self._lr_scheduler_kwargs.items()
                if k != "scheduler"
            }
            lr_scheduler = scheduler_cls(optimizer, **scheduler_kwargs)
            return [optimizer], [lr_scheduler]
        else:
            return optimizer

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

Validation step during training.

PARAMETER	DESCRIPTION
batch	Input batch data. TYPE: DataBatch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Validation loss.

Source code in spectre/parametric/base.py

def validation_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Validation step during training.

    Parameters
    ----------
    batch : DataBatch
        Input batch data.

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Validation loss.
    """
    return self.training_step(batch, batch_idx)

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

Test step during evaluation.

PARAMETER	DESCRIPTION
batch	Input batch data. TYPE: DataBatch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Test loss.

Source code in spectre/parametric/base.py

def test_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Test step during evaluation.

    Parameters
    ----------
    batch : DataBatch
        Input batch data.

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Test loss.
    """
    return self.training_step(batch, batch_idx)

predict_step(batch: DataBatch, batch_idx: int, dataloader_idx: int | None = None) -> torch.Tensor

Prediction step for inference.

PARAMETER	DESCRIPTION
batch	Input batch data. TYPE: DataBatch
batch_idx	Batch index. TYPE: int
dataloader_idx	Dataloader index for multiple dataloaders. TYPE: int | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Model predictions

Source code in spectre/parametric/base.py

def predict_step(
    self,
    batch: DataBatch,
    batch_idx: int,
    dataloader_idx: int | None = None,
) -> torch.Tensor:
    """
    Prediction step for inference.

    Parameters
    ----------
    batch : DataBatch
        Input batch data.

    batch_idx : int
        Batch index.

    dataloader_idx : int | None, optional, by default None
        Dataloader index for multiple dataloaders.

    Returns
    -------
    torch.Tensor
        Model predictions
    """
    return self.forward(batch.data)

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

Predict method for inference on raw input data. Can be used independently of Pytorch Lightning training.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Model predictions.

Source code in spectre/parametric/base.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict method for inference on raw input data. Can be used
    independently of Pytorch Lightning training.

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    Returns
    -------
    torch.Tensor
        Model predictions.
    """
    self.eval()
    with torch.no_grad():
        return self.forward(X)

fit(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None, trainer: pl.Trainer | None = None, n_epochs: int = 100, batch_size: int = 1000, val: float = 0.2, **trainer_kwargs) -> Parametric

Fit the parametric model using PyTorch Lightning.

This method provides a scikit-learn-style API for training parametric models while using PyTorch Lightning internally. It creates a DataModuleWeighted, configures a Trainer, and handles the training loop.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values for supervised learning. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
trainer	Pre-configured PyTorch Lightning Trainer instance. If provided, n_epochs and trainer_kwargs are ignored. If None, a new Trainer is created using n_epochs and trainer_kwargs. TYPE: pl.Trainer | None, optional, by default None DEFAULT: None
n_epochs	Maximum number of training epochs. Only used if trainer is None. TYPE: int, by default 100 DEFAULT: 100
batch_size	Batch size for training and validation. Only used if trainer is None. TYPE: int, by default 1000 DEFAULT: 1000
val	Fraction of data to use for validation (0.0-1.0). Only used if trainer is None. TYPE: float, by default 0.2 DEFAULT: 0.2
**trainer_kwargs	Additional keyword arguments for PyTorch Lightning Trainer. Only used if trainer is None. callbacks: List of callbacks (EarlyStopping, etc.) logger: Lightning logger accelerator: Hardware accelerator ('gpu', 'cpu') devices: Number or list of devices DEFAULT: {}

RETURNS	DESCRIPTION
Parametric	Return self.

Examples:

Basic usage with default settings:

>>> from spectre.parametric import SpectralMap
>>> sm = SpectralMap(n_states=3, model=network)
>>> sm.fit(X, n_epochs=50)

Using a pre-configured Trainer:

>>> import pytorch_lightning as pl
>>> trainer = pl.Trainer(
...     max_epochs=100,
...     callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss")],
...     logger=pl.loggers.WandbLogger(project="my-project"),
... )
>>> sm.fit(X, trainer=trainer)

With custom Trainer kwargs:

>>> sm.fit(
...     X,
...     n_epochs=100,
...     batch_size=512,
...     val=0.15,
...     callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss", patience=10)],
...     logger=pl.loggers.TensorBoardLogger("logs/"),
...     accelerator="gpu",
...     devices=2,
... )

Notes

• This method calls datamodule.setup() before training
• Training status is set via Lightning hooks on_train_end() and on_fit_end()
• For full control over the training process, use PyTorch Lightning's standard API: create a Trainer and call trainer.fit(model, datamodule)
• This method is compatible with ensemble methods that expect a fit() signature of (X, weights, target)

Source code in spectre/parametric/base.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
    trainer: pl.Trainer | None = None,
    n_epochs: int = 100,
    batch_size: int = 1000,
    val: float = 0.2,
    **trainer_kwargs,
) -> "Parametric":
    """
    Fit the parametric model using PyTorch Lightning.

    This method provides a scikit-learn-style API for training parametric
    models while using PyTorch Lightning internally. It creates a
    [DataModuleWeighted](spectre.data.DataModuleWeighted), configures a
    Trainer, and handles the training loop.

    Parameters
    ----------
    X : torch.Tensor
        Training data of shape (n_samples, in_features).

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

    target : torch.Tensor | None, optional, by default None
        Target values for supervised learning.

    trainer : pl.Trainer | None, optional, by default None
        Pre-configured PyTorch Lightning Trainer instance.

        If provided, `n_epochs` and `trainer_kwargs` are ignored. If None,
        a new Trainer is created using `n_epochs` and `trainer_kwargs`.

    n_epochs : int, by default 100
        Maximum number of training epochs.

        Only used if `trainer` is None.

    batch_size : int, by default 1000
        Batch size for training and validation.

        Only used if `trainer` is None.

    val : float, by default 0.2
        Fraction of data to use for validation (0.0-1.0).

        Only used if `trainer` is None.

    **trainer_kwargs
        Additional keyword arguments for PyTorch Lightning Trainer.

        Only used if `trainer` is None.

        - `callbacks`: List of callbacks (EarlyStopping, etc.)
        - `logger`: Lightning logger
        - `accelerator`: Hardware accelerator ('gpu', 'cpu')
        - `devices`: Number or list of devices

    Returns
    -------
    Parametric
        Return self.

    Examples
    --------
    Basic usage with default settings:

    >>> from spectre.parametric import SpectralMap
    >>> sm = SpectralMap(n_states=3, model=network)
    >>> sm.fit(X, n_epochs=50)

    Using a pre-configured Trainer:

    >>> import pytorch_lightning as pl
    >>> trainer = pl.Trainer(
    ...     max_epochs=100,
    ...     callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss")],
    ...     logger=pl.loggers.WandbLogger(project="my-project"),
    ... )
    >>> sm.fit(X, trainer=trainer)

    With custom Trainer kwargs:

    >>> sm.fit(
    ...     X,
    ...     n_epochs=100,
    ...     batch_size=512,
    ...     val=0.15,
    ...     callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss", patience=10)],
    ...     logger=pl.loggers.TensorBoardLogger("logs/"),
    ...     accelerator="gpu",
    ...     devices=2,
    ... )

    Notes
    -----
    - This method calls `datamodule.setup()` before training
    - Training status is set via Lightning hooks `on_train_end()` and
      `on_fit_end()`
    - For full control over the training process, use PyTorch Lightning's
      standard API: create a Trainer and call `trainer.fit(model,
      datamodule)`
    - This method is compatible with ensemble methods that expect a fit()
      signature of (X, weights, target)
    """
    dataset_dict = {"X": X, "weights": weights, "target": target}

    datamodule = DataModuleWeighted(
        dataset=dataset_dict,
        batch_size=batch_size,
        val=val if val > 0 else None,
    )
    datamodule.setup()

    if trainer is None:
        trainer = pl.Trainer(max_epochs=n_epochs, **trainer_kwargs)

    trainer.fit(self, datamodule)

    return self

parse_model(model: ModelType, device: Literal['cuda', 'cpu'] = 'cuda') -> Model

Parse and prepare model for training.

PARAMETER	DESCRIPTION
model	Neural network model or compatible structure. Can be a Model instance, Sequential module, individual torch.nn.Module, OrderedDict of layers, or list/tuple of layers. TYPE: ModelType
device	Target device for model. TYPE: Literal["cuda", "cpu"], by default "cuda" DEFAULT: 'cuda'

RETURNS	DESCRIPTION
Model	Parsed model on the specified device.

Source code in spectre/parametric/base.py

def parse_model(
    self,
    model: ModelType,
    device: Literal["cuda", "cpu"] = "cuda",
) -> Model:
    """
    Parse and prepare model for training.

    Parameters
    ----------
    model : ModelType
        Neural network model or compatible structure.

        Can be a Model instance, Sequential module, individual
        torch.nn.Module, OrderedDict of layers, or list/tuple of layers.

    device : Literal["cuda", "cpu"], by default "cuda"
        Target device for model.

    Returns
    -------
    Model
        Parsed model on the specified device.
    """
    device = self._validate_device(device)
    model = self._validate_model(model)
    return model.to(device)

parse_model_dict(model_dict: dict[str, ModelType], device: Literal['cuda', 'cpu'] = 'cuda') -> dict[str, Model]

Parse and prepare multiple models for training.

PARAMETER	DESCRIPTION
model_dict	Dictionary of named models where keys are model names (e.g., 'encoder', 'decoder') and values are model specifications compatible with parse_model(). TYPE: dict[str, ModelType]
device	Target device for all models. TYPE: Literal["cuda", "cpu"], by default "cuda" DEFAULT: 'cuda'

RETURNS	DESCRIPTION
dict[str, Model]	Dictionary of parsed Model instances on the specified device.

RAISES	DESCRIPTION
ValueError	If models dict is empty or contains invalid model specifications.
TypeError	If models is not a dictionary.

Examples:

>>> model_dict = {"encoder": encoder_sequential, "decoder": decoder_list}
>>> parsed = self.parse_model_dict(model_dict, device="cuda")
>>> encoder = parsed["encoder"]
>>> decoder = parsed["decoder"]

Source code in spectre/parametric/base.py

def parse_model_dict(
    self,
    model_dict: dict[str, ModelType],
    device: Literal["cuda", "cpu"] = "cuda",
) -> dict[str, Model]:
    """
    Parse and prepare multiple models for training.

    Parameters
    ----------
    model_dict : dict[str, ModelType]
        Dictionary of named models where keys are model names (e.g.,
        'encoder', 'decoder') and values are model specifications
        compatible with `parse_model()`.

    device : Literal["cuda", "cpu"], by default "cuda"
        Target device for all models.

    Returns
    -------
    dict[str, Model]
        Dictionary of parsed Model instances on the specified device.

    Raises
    ------
    ValueError
        If models dict is empty or contains invalid model specifications.
    TypeError
        If models is not a dictionary.

    Examples
    --------
    >>> model_dict = {"encoder": encoder_sequential, "decoder": decoder_list}
    >>> parsed = self.parse_model_dict(model_dict, device="cuda")
    >>> encoder = parsed["encoder"]
    >>> decoder = parsed["decoder"]
    """
    if not model_dict:
        raise ValueError("Model dictionary cannot be empty.")

    parsed_models = {}
    for name, model in model_dict.items():
        if not isinstance(name, str):
            raise ValueError(
                f"Model names must be strings, got {type(name)} for key {name}."
            )
        parsed_models[name] = self.parse_model(model, device)

    return parsed_models

save_model(filename: str, trace: bool = False) -> None

Save model to file with optional TorchScript tracing.

PARAMETER	DESCRIPTION
filename	Output file path. TYPE: str
trace	Whether to use TorchScript tracing for deployment optimization. TYPE: bool, by default False DEFAULT: False

Source code in spectre/parametric/base.py

def save_model(self, filename: str, trace: bool = False) -> None:
    """
    Save model to file with optional TorchScript tracing.

    Parameters
    ----------
    filename : str
        Output file path.

    trace : bool, by default False
        Whether to use TorchScript tracing for deployment optimization.
    """
    if not filename or not filename.strip():
        raise ValueError("Filename cannot be empty.")

    if not hasattr(self, "model") or self.model is None:
        raise ValueError("Model is not initialized.")

    try:
        if not trace:
            torch.save({"model_state_dict": self.state_dict()}, filename)
        else:
            self.eval()
            example_inputs = torch.rand(
                size=(1, self.in_features),
                dtype=torch.float32,
                device=next(self.model.parameters()).device,
            )
            model_traced = torch.jit.trace(self.model, example_inputs)
            model_traced.save(filename)
    except Exception as e:
        raise OSError(f"Failed to save model to '{filename}': {e}.") from e

load_model(filename: str, trace: bool = False) -> None

Load model from file.

PARAMETER	DESCRIPTION
filename	Input file path. TYPE: str
trace	Whether model was saved with TorchScript tracing. TYPE: bool, by default False DEFAULT: False

Source code in spectre/parametric/base.py

def load_model(self, filename: str, trace: bool = False) -> None:
    """
    Load model from file.

    Parameters
    ----------
    filename : str
        Input file path.

    trace : bool, by default False
        Whether model was saved with TorchScript tracing.
    """
    if not filename or not filename.strip():
        raise ValueError("Filename cannot be empty.")

    try:
        if not trace:
            state_dict = torch.load(filename, map_location="cpu")
            if "model_state_dict" not in state_dict:
                raise KeyError(
                    "Expected 'model_state_dict' key in saved file. "
                    "File may be corrupted or saved in different format."
                )
            self.load_state_dict(state_dict["model_state_dict"])
        else:
            self.model = torch.jit.load(filename)
            self.eval()
    except FileNotFoundError:
        raise FileNotFoundError(f'Model file "{filename}" not found.')
    except Exception as e:
        raise RuntimeError(f"Failed to load model from '{filename}': {e}.") from e

from_pretrained(path: str, map_location: str | None = None, **override_kwargs) -> Parametric classmethod

Load pretrained model from checkpoint.

PARAMETER	DESCRIPTION
path	Path to checkpoint file. TYPE: str
map_location	Device mapping (e.g., 'cpu', 'cuda:0'). TYPE: str | None, optional, by default None DEFAULT: None
**override_kwargs	Override configuration parameters. DEFAULT: {}

RETURNS	DESCRIPTION
Parametric	Loaded model instance.

Examples:

>>> model = SpectralMap.from_pretrained(
...     "checkpoints/best.ckpt",
...     map_location="cuda:0",
... )

Source code in spectre/parametric/base.py

@classmethod
def from_pretrained(
    cls,
    path: str,
    map_location: str | None = None,
    **override_kwargs,
) -> "Parametric":
    """
    Load pretrained model from checkpoint.

    Parameters
    ----------
    path : str
        Path to checkpoint file.

    map_location : str | None, optional, by default None
        Device mapping (e.g., 'cpu', 'cuda:0').

    **override_kwargs
        Override configuration parameters.

    Returns
    -------
    Parametric
        Loaded model instance.

    Examples
    --------
    >>> model = SpectralMap.from_pretrained(
    ...     "checkpoints/best.ckpt",
    ...     map_location="cuda:0",
    ... )
    """
    checkpoint = torch.load(path, map_location=map_location)

    # Extract config
    config = checkpoint.get("hyper_parameters", {})
    config.update(override_kwargs)

    # Create model instance
    model = cls(**config)

    # Load state dict
    model.load_state_dict(checkpoint["state_dict"])
    model.is_trained = True

    return model

on_train_end() -> None

PyTorch Lightning hook called at the end of training.

Sets the training status to True, enabling the use of the score method for model evaluation.

Source code in spectre/parametric/base.py

def on_train_end(self) -> None:
    """
    PyTorch Lightning hook called at the end of training.

    Sets the training status to True, enabling the use of the `score`
    method for model evaluation.
    """
    self.is_trained = True
    self.model.is_trained = True

on_fit_end() -> None

PyTorch Lightning hook called at the end of fit (training + validation).

Sets the training status to True, enabling the use of the score method for model evaluation.

Source code in spectre/parametric/base.py

def on_fit_end(self) -> None:
    """
    PyTorch Lightning hook called at the end of fit (training +
    validation).

    Sets the training status to True, enabling the use of the `score`
    method for model evaluation.
    """
    self.is_trained = True
    self.model.is_trained = True

log_training_metrics(loss: torch.Tensor, prog_bar: bool = True, auto_expand_tensors: bool = True, **extra_metrics: torch.Tensor | float | int) -> None

Standardized training/validation metric logging.

Logs the main loss with consistent naming and optional extra metrics. Handles train/validation mode distinction and applies standard logging parameters for PyTorch Lightning. Can expand multi-element tensors into separate metrics.

PARAMETER	DESCRIPTION
loss	Main loss value to log with progress bar. TYPE: Tensor
prog_bar	Whether to show loss in progress bar. TYPE: bool, by default True DEFAULT: True
auto_expand_tensors	Whether to expand multi-element tensors in extra_metrics. TYPE: bool, by default True DEFAULT: True
**extra_metrics	Additional metrics to log with train/val prefix. Multi-element tensors will be expanded if auto_expand_tensors=True. TYPE: Tensor | float | int DEFAULT: {}

Examples:

>>> # Simple loss logging
>>> self.log_training_metrics(loss)
>>>
>>> # With additional scalar metrics
>>> self.log_training_metrics(loss, mse_loss=mse, l1_loss=l1)
>>>
>>> # With tensor metrics (auto-expanded)
>>> eigvals = torch.tensor([1.0, 0.8, 0.6])
>>> self.log_training_metrics(loss, eigval=eigvals)
>>> # Logs: train_loss, train_eigval_0, train_eigval_1, train_eigval_2
>>>
>>> # Disable auto-expansion
>>> self.log_training_metrics(loss, weights=weight, auto_expand_tensors=False)

Source code in spectre/parametric/base.py

def log_training_metrics(
    self,
    loss: torch.Tensor,
    prog_bar: bool = True,
    auto_expand_tensors: bool = True,
    **extra_metrics: torch.Tensor | float | int,
) -> None:
    """
    Standardized training/validation metric logging.

    Logs the main loss with consistent naming and optional extra metrics.
    Handles train/validation mode distinction and applies standard logging
    parameters for PyTorch Lightning. Can expand multi-element tensors into
    separate metrics.

    Parameters
    ----------
    loss : torch.Tensor
        Main loss value to log with progress bar.

    prog_bar : bool, by default True
        Whether to show loss in progress bar.

    auto_expand_tensors : bool, by default True
        Whether to expand multi-element tensors in extra_metrics.

    **extra_metrics : torch.Tensor | float | int
        Additional metrics to log with train/val prefix. Multi-element
        tensors will be expanded if auto_expand_tensors=True.

    Examples
    --------
    >>> # Simple loss logging
    >>> self.log_training_metrics(loss)
    >>>
    >>> # With additional scalar metrics
    >>> self.log_training_metrics(loss, mse_loss=mse, l1_loss=l1)
    >>>
    >>> # With tensor metrics (auto-expanded)
    >>> eigvals = torch.tensor([1.0, 0.8, 0.6])
    >>> self.log_training_metrics(loss, eigval=eigvals)
    >>> # Logs: train_loss, train_eigval_0, train_eigval_1, train_eigval_2
    >>>
    >>> # Disable auto-expansion
    >>> self.log_training_metrics(loss, weights=weight, auto_expand_tensors=False)
    """
    train_or_val = "train" if self.training else "val"

    # Log main loss (always scalar for loss)
    self.log(
        f"{train_or_val}_loss",
        loss,
        prog_bar=prog_bar,
        on_epoch=True,
        on_step=False,
        # Setting `batch_size=-1` silences warning introduced in Lightning 1.5:
        # Trying to infer the `batch_size` from an ambiguous collection. To avoid
        # any miscalculations, use `self.log(..., batch_size=batch_size)`.
        # https://github.com/Lightning-AI/pytorch-lightning/issues/10349
        batch_size=-1,
    )

    # Log any extra metrics with automatic tensor expansion
    all_parsed_metrics = {}
    for name, value in extra_metrics.items():
        parsed = self._parse_tensor_value(name, value, auto_expand_tensors)
        all_parsed_metrics.update(parsed)

    # Log all parsed metrics with train/val prefix
    for metric_name, metric_value in all_parsed_metrics.items():
        self.log(
            f"{train_or_val}_{metric_name}",
            metric_value,
            on_epoch=True,
            on_step=False,
            batch_size=-1,
        )

log_metrics_dict(metrics: dict[str, torch.Tensor | float | int], prefix: str = '', add_train_val_prefix: bool = True, auto_expand_tensors: bool = True, on_step: bool = False, on_epoch: bool = True) -> None

Log a dictionary of metrics with flexible naming and automatic tensor expansion.

PARAMETER	DESCRIPTION
metrics	Dictionary of metric names and values to log. Multi-element tensors will be expanded if auto_expand_tensors=True. TYPE: dict[str, Tensor | float | int]
prefix	Additional prefix to add before metric names. TYPE: str, by default "" DEFAULT: ''
add_train_val_prefix	Whether to add train/val prefix based on current mode. TYPE: bool, by default True DEFAULT: True
auto_expand_tensors	Whether to expand multi-element tensors. TYPE: bool, by default True DEFAULT: True
on_step	Whether to log on each step. TYPE: bool, by default False DEFAULT: False
on_epoch	Whether to log on each epoch. TYPE: bool, by default True DEFAULT: True

Examples:

>>> # Log scalar metrics
>>> loss_components = {"mse": mse_loss, "l1": l1_loss}
>>> self.log_metrics_dict(loss_components, prefix="loss_")
>>>
>>> # Log tensor metrics (auto-expanded)
>>> tensor_metrics = {
...     "weights": torch.tensor([1, 2, 3]),
...     "biases": torch.tensor([[1, 2], [3, 4]]),
... }
>>> self.log_metrics_dict(tensor_metrics)
>>> # Logs: train_weight_0, train_weight_1, train_weight_2, train_biases_0_0
>>>
>>> # Log without train/val prefix or tensor expansion
>>> model_stats = {"n_params": param_count, "weight_matrix": W}
>>> self.log_metrics_dict(
...     model_stats,
...     add_train_val_prefix=False,
...     auto_expand_tensors=False,
... )

Source code in spectre/parametric/base.py

def log_metrics_dict(
    self,
    metrics: dict[str, torch.Tensor | float | int],
    prefix: str = "",
    add_train_val_prefix: bool = True,
    auto_expand_tensors: bool = True,
    on_step: bool = False,
    on_epoch: bool = True,
) -> None:
    """
    Log a dictionary of metrics with flexible naming and automatic tensor
    expansion.

    Parameters
    ----------
    metrics : dict[str, torch.Tensor | float | int]
        Dictionary of metric names and values to log. Multi-element tensors
        will be expanded if `auto_expand_tensors=True`.

    prefix : str, by default ""
        Additional prefix to add before metric names.

    add_train_val_prefix : bool, by default True
        Whether to add train/val prefix based on current mode.

    auto_expand_tensors : bool, by default True
        Whether to expand multi-element tensors.

    on_step : bool, by default False
        Whether to log on each step.

    on_epoch : bool, by default True
        Whether to log on each epoch.

    Examples
    --------
    >>> # Log scalar metrics
    >>> loss_components = {"mse": mse_loss, "l1": l1_loss}
    >>> self.log_metrics_dict(loss_components, prefix="loss_")
    >>>
    >>> # Log tensor metrics (auto-expanded)
    >>> tensor_metrics = {
    ...     "weights": torch.tensor([1, 2, 3]),
    ...     "biases": torch.tensor([[1, 2], [3, 4]]),
    ... }
    >>> self.log_metrics_dict(tensor_metrics)
    >>> # Logs: train_weight_0, train_weight_1, train_weight_2, train_biases_0_0
    >>>
    >>> # Log without train/val prefix or tensor expansion
    >>> model_stats = {"n_params": param_count, "weight_matrix": W}
    >>> self.log_metrics_dict(
    ...     model_stats,
    ...     add_train_val_prefix=False,
    ...     auto_expand_tensors=False,
    ... )
    """
    # Parse all metrics with optional tensor expansion
    all_parsed_metrics = {}
    for name, value in metrics.items():
        parsed = self._parse_tensor_value(name, value, auto_expand_tensors)
        all_parsed_metrics.update(parsed)

    # Build base prefix
    full_prefix = ""
    if add_train_val_prefix:
        full_prefix += "train" if self.training else "val"
        full_prefix += "_"
    if prefix:
        full_prefix += prefix

    # Log all parsed metrics
    metric_dict = {
        f"{full_prefix}{name}": value for name, value in all_parsed_metrics.items()
    }
    self.log_dict(
        metric_dict,
        on_step=on_step,
        on_epoch=on_epoch,
        batch_size=-1,
    )

Functions