Utils Callbacks

callbacks

CLASS	DESCRIPTION
BaseCallback	Abstract base class for PyTorch Lightning callbacks with enhanced functionality.
TrainingCallback	Training callback for logging and saving training metrics.
ParameterAnnealingCallback	PyTorch Lightning callback for annealing temperature during training.

FUNCTION	DESCRIPTION
model_complexity_extractor	Extract model complexity metrics.
memory_usage_extractor	Extract GPU memory usage.
gradient_norm_extractor	Extract gradient norms.
kernel_params_extractor	Universal kernel parameter extractor.
eigenvalue_extractor	Factory function to create eigenvalue extractor for parametric models.
create_basic_training_callback	Create a basic training callback with standard logging.
create_training_callback	Create a training callback with custom extractors.
create_callback_from_config	Create a callback instance from configuration dictionary.
get_default_callbacks	Get a default set of callbacks.
create_callback_from_template	Create callbacks from predefined templates.
get_callback_config_schema	Get configuration schema for all available callbacks.

Classes

BaseCallback(enabled: bool = True, log_every_n_epochs: int = 1)

Bases: Callback, ABC

Abstract base class for PyTorch Lightning callbacks with enhanced functionality.

Provides common functionality for all callbacks including configuration management, robust metric extraction, and device-aware computations.

PARAMETER	DESCRIPTION
enabled	Whether the callback is active during training. TYPE: bool DEFAULT: True
log_every_n_epochs	Frequency of logging in epochs. Must be positive. TYPE: int, by default 1 DEFAULT: 1

Notes

Subclasses must implement the abstract method get_callback_type() to identify the callback type for configuration and factory systems.

METHOD	DESCRIPTION
get_callback_type	Return the callback type identifier.
should_log_this_epoch	Check if logging should occur this epoch.
safe_tensor_extract	Safely extract scalar value from tensors or other types.
from_config	Create callback instance from configuration dictionary.
to_config	Export callback configuration to dictionary.
validate_config	Validate callback configuration.

Source code in spectre/utils/callbacks.py

def __init__(self, enabled: bool = True, log_every_n_epochs: int = 1):
    super().__init__()
    self.enabled = enabled
    if not isinstance(log_every_n_epochs, int) or log_every_n_epochs < 1:
        raise ValueError(
            f"log_every_n_epochs must be a positive integer, got {log_every_n_epochs}"
        )
    self.log_every_n_epochs = log_every_n_epochs

Functions

get_callback_type() -> str abstractmethod

Return the callback type identifier.

Used by factory methods and configuration systems to identify the specific callback class.

RETURNS	DESCRIPTION
str	Unique identifier for this callback type.

Source code in spectre/utils/callbacks.py

@abstractmethod
def get_callback_type(self) -> str:
    """
    Return the callback type identifier.

    Used by factory methods and configuration systems to identify
    the specific callback class.

    Returns
    -------
    str
        Unique identifier for this callback type.
    """
    pass

should_log_this_epoch(trainer) -> bool

Check if logging should occur this epoch.

PARAMETER	DESCRIPTION
trainer	The PyTorch Lightning trainer instance. TYPE: Trainer

RETURNS	DESCRIPTION
bool	True if logging should occur this epoch.

Source code in spectre/utils/callbacks.py

def should_log_this_epoch(self, trainer) -> bool:
    """
    Check if logging should occur this epoch.

    Parameters
    ----------
    trainer : pytorch_lightning.Trainer
        The PyTorch Lightning trainer instance.

    Returns
    -------
    bool
        True if logging should occur this epoch.
    """
    if not self.enabled or trainer.sanity_checking:
        return False
    return trainer.current_epoch % self.log_every_n_epochs == 0

safe_tensor_extract(value: Any) -> float | None

Safely extract scalar value from tensors or other types.

Handles PyTorch tensors, NumPy arrays, and scalar values with proper error handling and device management.

PARAMETER	DESCRIPTION
value	Value to extract scalar from. TYPE: Any

RETURNS	DESCRIPTION
float | None	Extracted scalar value as float, or None if extraction fails.

Source code in spectre/utils/callbacks.py

def safe_tensor_extract(self, value: Any) -> float | None:
    """
    Safely extract scalar value from tensors or other types.

    Handles PyTorch tensors, NumPy arrays, and scalar values with
    proper error handling and device management.

    Parameters
    ----------
    value : Any
        Value to extract scalar from.

    Returns
    -------
    float | None
        Extracted scalar value as float, or None if extraction fails.
    """
    try:
        if value is None:
            return None
        elif hasattr(value, "item"):
            return float(value.item())
        elif hasattr(value, "detach"):
            return float(value.detach().cpu().item())
        elif isinstance(value, (int, float, np.integer, np.floating)):
            return float(value)
        else:
            return float(value)
    except (RuntimeError, ValueError, TypeError):
        return None

from_config(config: dict[str, Any]) -> BaseCallback classmethod

Create callback instance from configuration dictionary.

This base implementation only handles common parameters. Subclasses should override this method to properly handle their specific parameters.

PARAMETER	DESCRIPTION
config	Configuration dictionary. TYPE: dict[str, Any]

RETURNS	DESCRIPTION
BaseCallback	Configured callback instance.

Source code in spectre/utils/callbacks.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> "BaseCallback":
    """
    Create callback instance from configuration dictionary.

    This base implementation only handles common parameters. Subclasses
    should override this method to properly handle their specific parameters.

    Parameters
    ----------
    config : dict[str, Any]
        Configuration dictionary.

    Returns
    -------
    BaseCallback
        Configured callback instance.
    """
    # Extract common parameters
    enabled = config.get("enabled", True)
    log_every_n_epochs = config.get("log_every_n_epochs", 1)

    # Create instance with common parameters only
    return cls(enabled=enabled, log_every_n_epochs=log_every_n_epochs)

to_config() -> dict[str, Any]

Export callback configuration to dictionary.

RETURNS	DESCRIPTION
dict[str, Any]	Configuration dictionary.

Source code in spectre/utils/callbacks.py

def to_config(self) -> dict[str, Any]:
    """
    Export callback configuration to dictionary.

    Returns
    -------
    dict[str, Any]
        Configuration dictionary.
    """
    return {
        "type": self.get_callback_type(),
        "enabled": self.enabled,
        "log_every_n_epochs": self.log_every_n_epochs,
    }

validate_config(config: dict[str, Any]) -> None

Validate callback configuration.

PARAMETER	DESCRIPTION
config	Configuration to validate. TYPE: dict[str, Any]

RAISES	DESCRIPTION
ValueError	If configuration is invalid.

Source code in spectre/utils/callbacks.py

def validate_config(self, config: dict[str, Any]) -> None:
    """
    Validate callback configuration.

    Parameters
    ----------
    config : dict[str, Any]
        Configuration to validate.

    Raises
    ------
    ValueError
        If configuration is invalid.
    """
    if "log_every_n_epochs" in config:
        log_freq = config["log_every_n_epochs"]
        if not isinstance(log_freq, int) or log_freq < 1:
            raise ValueError(
                f"log_every_n_epochs must be a positive integer, got {log_freq}"
            )

    if "enabled" in config and not isinstance(config["enabled"], bool):
        raise ValueError("enabled must be a boolean")

TrainingCallback(save_path: str | Path | None = None, save_format: str = 'pickle', metric_prefix: str | None = None, custom_extractors: dict[str, Callable] | None = None, static_extractors: dict[str, Callable] | None = None, track_lr: bool = False, track_epoch_time: bool = False, max_history: int | None = None, **kwargs)

Bases: BaseCallback

Training callback for logging and saving training metrics.

Supports flexible metric collection, custom value extraction, and multiple save formats.

PARAMETER	DESCRIPTION
save_path	Path to save metrics file. TYPE: str | Path, optional, by default None DEFAULT: None
save_format	Format to save metrics ('pickle', 'json'). TYPE: str, optional, by default 'pickle' DEFAULT: 'pickle'
metric_prefix	Prefix to filter metrics by name. TYPE: str, optional, by default None DEFAULT: None
custom_extractors	Custom functions to extract time-varying values from trainer/module, called every epoch. TYPE: dict[str, Callable], optional, by default None DEFAULT: None
static_extractors	Custom functions to extract constant values from trainer/module, called only once at training start (e.g., model_complexity_extractor). TYPE: dict[str, Callable], optional, by default None DEFAULT: None
track_lr	Whether to track learning rate. TYPE: bool, optional, by default False DEFAULT: False
track_epoch_time	Whether to track epoch duration. TYPE: bool, optional, by default False DEFAULT: False
max_history	Maximum number of epochs to keep in history. If None, keeps all history. Useful for long training runs to prevent unbounded memory growth. TYPE: int | None, optional, by default None DEFAULT: None

Examples:

Track both static and dynamic metrics:

>>> from spectre.utils.callbacks import (
...     TrainingCallback,
...     model_complexity_extractor,
...     gradient_norm_extractor,
... )
>>> callback = TrainingCallback(
...     static_extractors={
...         "model_complexity_extractor": model_complexity_extractor,
...     },
...     custom_extractors={"gradient_norm_extractor": gradient_norm_extractor},
...     track_lr=True,
... )
>>> # After training, access metrics:
>>> callback.static_metrics  # {'total_params': 1000, 'trainable_params': 900}
>>> callback.metrics["gradient_norm_extractor"]  # [0.5, 0.3, 0.2, ...]

METHOD	DESCRIPTION
from_config	Create TrainingCallback instance from configuration dictionary.
on_train_start	Log static metrics at the start of training.
on_train_epoch_start	Track epoch start time if requested.
on_train_epoch_end	Log metrics at the end of each training epoch.
on_validation_epoch_end	Log validation metrics.
add_custom_extractor	Add a custom metric extractor function.
get_metric	Get a specific metric by name.
get_latest_metrics	Get the most recent values for all metrics.
clear_metrics	Clear all stored metrics.
save	Save metrics to file.
load	Load metrics from file.

Source code in spectre/utils/callbacks.py

def __init__(
    self,
    save_path: str | Path | None = None,
    save_format: str = "pickle",
    metric_prefix: str | None = None,
    custom_extractors: dict[str, Callable] | None = None,
    static_extractors: dict[str, Callable] | None = None,
    track_lr: bool = False,
    track_epoch_time: bool = False,
    max_history: int | None = None,
    **kwargs,
):
    super().__init__(**kwargs)

    self.save_path = Path(save_path) if save_path else None
    self.save_format = save_format.lower()
    self.metric_prefix = metric_prefix
    self.custom_extractors = custom_extractors or {}
    self.static_extractors = static_extractors or {}
    self.track_lr = track_lr
    self.track_epoch_time = track_epoch_time
    self.max_history = max_history

    # Initialize metrics storage
    self.metrics = defaultdict(list)
    self.metrics["epoch"] = []

    # Static metrics (logged only once)
    self.static_metrics = {}

    # For epoch timing
    self._epoch_start_time = None

    # Validate save format
    if self.save_format not in ["pickle", "json"]:
        raise ValueError(f"Unsupported save format: {self.save_format}")

Functions

from_config(config: dict[str, Any]) -> TrainingCallback classmethod

Create TrainingCallback instance from configuration dictionary.

PARAMETER	DESCRIPTION
config	Configuration dictionary. TYPE: dict[str, Any]

RETURNS	DESCRIPTION
TrainingCallback	Configured callback instance.

Source code in spectre/utils/callbacks.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> "TrainingCallback":
    """
    Create TrainingCallback instance from configuration dictionary.

    Parameters
    ----------
    config : dict[str, Any]
        Configuration dictionary.

    Returns
    -------
    TrainingCallback
        Configured callback instance.
    """
    # Whitelist of allowed parameters
    allowed_params = {
        "enabled",
        "log_every_n_epochs",
        "save_path",
        "save_format",
        "metric_prefix",
        "custom_extractors",
        "static_extractors",
        "track_lr",
        "track_epoch_time",
        "max_history",
    }

    # Extract only whitelisted parameters
    kwargs = {k: v for k, v in config.items() if k in allowed_params}

    return cls(**kwargs)

on_train_start(trainer, pl_module) -> None

Log static metrics at the start of training.

Source code in spectre/utils/callbacks.py

def on_train_start(self, trainer, pl_module) -> None:
    """Log static metrics at the start of training."""
    if not self.static_extractors:
        return

    for key, extractor_fn in self.static_extractors.items():
        try:
            value = extractor_fn(trainer, pl_module)

            # Handle dict returns by flattening
            if isinstance(value, dict):
                for sub_key, sub_value in value.items():
                    scalar_value = self.safe_tensor_extract(sub_value)
                    if scalar_value is not None:
                        self.static_metrics[sub_key] = scalar_value
            else:
                scalar_value = self.safe_tensor_extract(value)
                if scalar_value is not None:
                    self.static_metrics[key] = scalar_value
        except Exception as e:
            warnings.warn(
                f"Failed to extract static metric '{key}': {e}",
                RuntimeWarning,
                stacklevel=2,
            )

on_train_epoch_start(trainer, pl_module) -> None

Track epoch start time if requested.

Source code in spectre/utils/callbacks.py

def on_train_epoch_start(self, trainer, pl_module) -> None:
    """Track epoch start time if requested."""
    if self.track_epoch_time:
        self._epoch_start_time = time.time()

on_train_epoch_end(trainer, pl_module) -> None

Log metrics at the end of each training epoch.

Source code in spectre/utils/callbacks.py

def on_train_epoch_end(self, trainer, pl_module) -> None:
    """Log metrics at the end of each training epoch."""
    if trainer.sanity_checking:
        return

    current_epoch = trainer.current_epoch
    self.metrics["epoch"].append(current_epoch)

    # Log standard metrics from trainer
    self._log_trainer_metrics(trainer)

    # Log learning rate if requested
    if self.track_lr:
        self._log_learning_rate(trainer)

    # Log epoch duration if requested
    if self.track_epoch_time:
        self._log_epoch_duration()

    # Log custom extracted values
    self._log_custom_metrics(trainer, pl_module)

    # Trim all metrics to max_history length
    self._trim_metrics()

    # Auto-save if path is provided
    if self.save_path:
        self.save(self.save_path)

on_validation_epoch_end(trainer, pl_module) -> None

Log validation metrics.

Source code in spectre/utils/callbacks.py

def on_validation_epoch_end(self, trainer, pl_module) -> None:
    """Log validation metrics."""
    if trainer.sanity_checking:
        return

    # Only log validation metrics if we're also logging training metrics
    if (
        not self.metrics["epoch"]
        or trainer.current_epoch != self.metrics["epoch"][-1]
    ):
        return

    # Log validation-specific metrics
    validation_metrics = {
        k: v for k, v in trainer.callback_metrics.items() if "val" in k.lower()
    }

    for key, val in validation_metrics.items():
        if self._should_log_metric(key):
            self._safe_append_metric(key, val)

add_custom_extractor(name: str, extractor_fn: Callable) -> None

Add a custom metric extractor function.

PARAMETER	DESCRIPTION
name	Name of the metric to log. TYPE: str
extractor_fn	Function that takes (trainer, pl_module) and returns a value. TYPE: Callable

Source code in spectre/utils/callbacks.py

def add_custom_extractor(self, name: str, extractor_fn: Callable) -> None:
    """
    Add a custom metric extractor function.

    Parameters
    ----------
    name : str
        Name of the metric to log.

    extractor_fn : Callable
        Function that takes (trainer, pl_module) and returns a value.
    """
    self.custom_extractors[name] = extractor_fn

get_metric(name: str) -> list[Any]

Get a specific metric by name.

Source code in spectre/utils/callbacks.py

def get_metric(self, name: str) -> list[Any]:
    """Get a specific metric by name."""
    return self.metrics.get(name, [])

get_latest_metrics() -> dict[str, Any]

Get the most recent values for all metrics.

Source code in spectre/utils/callbacks.py

def get_latest_metrics(self) -> dict[str, Any]:
    """Get the most recent values for all metrics."""
    return {
        key: values[-1] if values else None for key, values in self.metrics.items()
    }

clear_metrics() -> None

Clear all stored metrics.

Source code in spectre/utils/callbacks.py

def clear_metrics(self) -> None:
    """Clear all stored metrics."""
    self.metrics.clear()
    self.metrics["epoch"] = []

save(filename: str | Path, format: str | None = None) -> None

Save metrics to file.

PARAMETER	DESCRIPTION
filename	Output file path. TYPE: str | Path
format	Save format ("pickle" or "json"). If None, uses save_format attribute, or auto-detects from file extension. TYPE: str | None, optional, by default None DEFAULT: None

Notes

Saves both time-varying metrics (in metrics) and static metrics (in static_metrics) to the same file with separate keys.

For JSON format, tensors and arrays are automatically converted to Python native types (scalars or lists) for proper JSON serialization.

Format auto-detection rules:

• .json extension → JSON format
• .pkl, .pickle extension → pickle format
• Otherwise → uses save_format attribute (default: pickle)

Examples:

>>> callback = TrainingCallback()
>>> callback.save("metrics.json")  # Auto-detects JSON format
>>> callback.save("metrics.pkl")  # Auto-detects pickle format
>>> callback.save("data", format="json")  # Explicit format

Source code in spectre/utils/callbacks.py

def save(self, filename: str | Path, format: str | None = None) -> None:
    """
    Save metrics to file.

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

    format : str | None, optional, by default None
        Save format ("pickle" or "json"). If None, uses `save_format`
        attribute, or auto-detects from file extension.

    Notes
    -----
    Saves both time-varying metrics (in `metrics`) and static metrics
    (in `static_metrics`) to the same file with separate keys.

    For JSON format, tensors and arrays are automatically converted to
    Python native types (scalars or lists) for proper JSON serialization.

    Format auto-detection rules:

    - `.json` extension → JSON format
    - `.pkl`, `.pickle` extension → pickle format
    - Otherwise → uses `save_format` attribute (default: pickle)

    Examples
    --------
    >>> callback = TrainingCallback()
    >>> callback.save("metrics.json")  # Auto-detects JSON format
    >>> callback.save("metrics.pkl")  # Auto-detects pickle format
    >>> callback.save("data", format="json")  # Explicit format
    """
    filename = Path(filename)

    # Determine format: explicit > file extension > default
    if format is None:
        ext = filename.suffix.lower()
        if ext == ".json":
            format = "json"
        elif ext in [".pkl", ".pickle"]:
            format = "pickle"
        else:
            format = self.save_format

    # Ensure parent directory exists
    filename.parent.mkdir(parents=True, exist_ok=True)

    # Convert defaultdict to regular dict and include static metrics
    metrics_dict = {
        "metrics": dict(self.metrics),
        "static_metrics": self.static_metrics,
    }

    if format == "pickle":
        with open(filename, "wb") as f:
            pickle.dump(metrics_dict, f, pickle.HIGHEST_PROTOCOL)
    elif format == "json":
        # Convert tensors/arrays to JSON-serializable types
        serializable_dict = self._convert_to_json_serializable(metrics_dict)
        with open(filename, "w") as f:
            json.dump(serializable_dict, f, indent=2)
    else:
        raise ValueError(f"Unsupported save format: {format}")

load(filename: str | Path, format: str | None = None) -> None

Load metrics from file.

PARAMETER	DESCRIPTION
filename	Input file path. TYPE: str | Path
format	Load format ("pickle" or "json"). If None, auto-detects from file extension or uses save_format attribute. TYPE: str | None, optional, by default None DEFAULT: None

Source code in spectre/utils/callbacks.py

def load(self, filename: str | Path, format: str | None = None) -> None:
    """
    Load metrics from file.

    Parameters
    ----------
    filename : str | Path
        Input file path.
    format : str | None, optional, by default None
        Load format ("pickle" or "json"). If None, auto-detects from
        file extension or uses `save_format` attribute.
    """
    filename = Path(filename)

    if not filename.exists():
        raise FileNotFoundError(f"Metrics file not found: {filename}")

    # Determine format: explicit > file extension > default
    if format is None:
        ext = filename.suffix.lower()
        if ext == ".json":
            format = "json"
        elif ext in [".pkl", ".pickle"]:
            format = "pickle"
        else:
            format = self.save_format

    if format == "pickle":
        with open(filename, "rb") as f:
            loaded_data = pickle.load(f)
    elif format == "json":
        with open(filename, "r") as f:
            loaded_data = json.load(f)
    else:
        raise ValueError(f"Unsupported load format: {format}")

    # Handle both old format (direct dict) and new format (with static_metrics)
    if isinstance(loaded_data, dict) and "metrics" in loaded_data:
        # New format with static_metrics
        self.metrics.clear()
        self.metrics.update(defaultdict(list, loaded_data["metrics"]))
        self.static_metrics = loaded_data.get("static_metrics", {})
    else:
        # Old format (backward compatibility)
        self.metrics.clear()
        self.metrics.update(defaultdict(list, loaded_data))

ParameterAnnealingCallback(target_attr: str = 'temperature', start_temperature: float = 1.0, end_temperature: float = 0.01, total_epochs: int | None = None, schedule: str = 'cosine', warmup_epochs: int = 0, **kwargs)

Bases: BaseCallback

PyTorch Lightning callback for annealing temperature during training.

Gradually reduces temperature parameter in loss modules that support temperature-based weighting. This enables exploration-to-exploitation transition in soft-argmax computations.

Supports multiple annealing schedules:

• Linear: T(t) = T_start + (T_end - T_start) * (t / T_total)
• Exponential: T(t) = T_start * (T_end / T_start)^(t / T_total)
• Cosine: T(t) = T_end + 0.5 * (T_start - T_end) * (1 + cos(π * t / T_total))

PARAMETER	DESCRIPTION
target_attr	Name of the temperature attribute to update in the loss module. The loss module must have a settable attribute with this name. TYPE: str, by default "temperature" DEFAULT: 'temperature'
start_temperature	Initial temperature value at epoch 0 (or after warmup). Higher values create broader/smoother distributions. TYPE: float, by default 1.0 DEFAULT: 1.0
end_temperature	Final temperature value at the last epoch. Lower values create sharper/more focused distributions. TYPE: float, by default 0.01 DEFAULT: 0.01
total_epochs	Total number of epochs for annealing schedule. If None, automatically inferred from trainer.max_epochs. TYPE: int | None, by default None DEFAULT: None
schedule	Annealing schedule type: "linear": Uniform temperature reduction "exponential": Faster reduction early, slower later "cosine": Smooth S-curve reduction (recommended) TYPE: (linear, exponential, cosine) DEFAULT: "linear"
warmup_epochs	Number of initial epochs to keep temperature constant at start_temperature. Useful for stable initialization. TYPE: int, by default 0 DEFAULT: 0

Examples:

Linear annealing with AdaptiveEigenvalueLoss:

>>> from spectre.loss import AdaptiveEigenvalueLoss
>>> from spectre.utils.callbacks import ParameterAnnealingCallback
>>> from spectre.parametric import SpectralMap
>>> loss = AdaptiveEigenvalueLoss(temperature=1.0, min_states=2, max_states=10)
>>> callback = ParameterAnnealingCallback(
...     start_temperature=1.0,
...     end_temperature=0.05,
...     total_epochs=100,
...     schedule="linear",
... )
>>> sm = SpectralMap(loss_fn=loss, n_eigen=10)
>>> sm.fit(X, n_epochs=100, callbacks=[callback])

Cosine annealing with warmup for stable training:

>>> callback = ParameterAnnealingCallback(
...     start_temperature=2.0,
...     end_temperature=0.01,
...     total_epochs=50,
...     schedule="cosine",
...     warmup_epochs=5,
... )

Multiple loss parameters can be annealed separately:

>>> temp_callback = ParameterAnnealingCallback(
...     target_attr="temperature",
...     start_temperature=1.0,
...     end_temperature=0.05,
... )
>>> alpha_callback = ParameterAnnealingCallback(
...     target_attr="adaptive_weight",
...     start_temperature=0.3,
...     end_temperature=1.0,
... )
>>> sm.fit(X, callbacks=[temp_callback, alpha_callback])

Notes

In soft-argmax with weights = softmax(values / temperature):

• High temperature (T >> 1): Uniform weighting across all values
• Low temperature (T -> 0): Winner-take-all, focuses on maximum value
• Annealing: Starts with exploration (broad), ends with exploitation (focused)

METHOD	DESCRIPTION
get_callback_type	Return callback type identifier.
from_config	Create ParameterAnnealingCallback instance from configuration dictionary.
on_train_start	Check for temperature mismatch at training start.
on_train_epoch_start	Update temperature at the start of each training epoch.
to_config	Export callback configuration.

Source code in spectre/utils/callbacks.py

def __init__(
    self,
    target_attr: str = "temperature",
    start_temperature: float = 1.0,
    end_temperature: float = 0.01,
    total_epochs: int | None = None,
    schedule: str = "cosine",
    warmup_epochs: int = 0,
    **kwargs,
):
    super().__init__(**kwargs)

    # Validation
    if not isinstance(start_temperature, (int, float)) or start_temperature <= 0:
        raise ValueError(
            f"start_temperature must be numeric > 0, got {start_temperature}"
        )
    if not isinstance(end_temperature, (int, float)) or end_temperature < 0:
        raise ValueError(
            f"end_temperature must be numeric >= 0, got {end_temperature}"
        )
    if total_epochs is not None and (
        not isinstance(total_epochs, int) or total_epochs <= 0
    ):
        raise ValueError(
            f"total_epochs must be positive int or None, got {total_epochs}"
        )
    if not isinstance(warmup_epochs, int) or warmup_epochs < 0:
        raise ValueError(
            f"warmup_epochs must be non-negative int, got {warmup_epochs}"
        )
    if schedule not in ["linear", "exponential", "cosine"]:
        raise ValueError(
            f"schedule must be 'linear', 'exponential', or 'cosine', got {schedule}"
        )
    if (
        schedule in ["linear", "exponential"]
        and start_temperature == end_temperature
    ):
        raise ValueError(
            f"start_temperature and end_temperature must differ for {schedule} schedule, "
            f"got {start_temperature} == {end_temperature}"
        )
    if schedule == "exponential" and end_temperature == 0:
        raise ValueError(
            "end_temperature cannot be 0 for exponential schedule. "
            "Use a small positive value (e.g., 1e-6) or choose 'linear' or 'cosine' schedule."
        )

    self.target_attr = target_attr
    self.start_temperature = float(start_temperature)
    self.end_temperature = float(end_temperature)
    self.total_epochs = total_epochs
    self.schedule = schedule
    self.warmup_epochs = warmup_epochs

Functions

get_callback_type() -> str

Return callback type identifier.

Source code in spectre/utils/callbacks.py

def get_callback_type(self) -> str:
    """Return callback type identifier."""
    return "parameter_annealing_callback"

from_config(config: dict[str, Any]) -> ParameterAnnealingCallback classmethod

Create ParameterAnnealingCallback instance from configuration dictionary.

PARAMETER	DESCRIPTION
config	Configuration dictionary. TYPE: dict[str, Any]

RETURNS	DESCRIPTION
ParameterAnnealingCallback	Configured callback instance.

Source code in spectre/utils/callbacks.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> "ParameterAnnealingCallback":
    """
    Create ParameterAnnealingCallback instance from configuration dictionary.

    Parameters
    ----------
    config : dict[str, Any]
        Configuration dictionary.

    Returns
    -------
    ParameterAnnealingCallback
        Configured callback instance.
    """
    # Whitelist of allowed parameters
    allowed_params = {
        "enabled",
        "log_every_n_epochs",
        "target_attr",
        "start_temperature",
        "end_temperature",
        "total_epochs",
        "schedule",
        "warmup_epochs",
    }

    # Extract only whitelisted parameters
    kwargs = {k: v for k, v in config.items() if k in allowed_params}

    return cls(**kwargs)

on_train_start(trainer, pl_module) -> None

Check for temperature mismatch at training start.

Issues a warning if the callback's start_temperature differs from the current value of the target attribute in the loss module.

Source code in spectre/utils/callbacks.py

def on_train_start(self, trainer, pl_module) -> None:
    """
    Check for temperature mismatch at training start.

    Issues a warning if the callback's `start_temperature` differs from the
    current value of the target attribute in the loss module.
    """
    if not self.enabled or trainer.sanity_checking:
        return

    # Try to get current temperature from loss module
    current_temp = None

    # Try loss_fn first (most common for SpectralMap)
    if hasattr(pl_module, "loss_fn"):
        loss_module = pl_module.loss_fn

        # Check if it's a CompositeLoss
        if hasattr(loss_module, "losses") and isinstance(loss_module.losses, dict):
            # Check all child losses for temperature values
            temps = []
            for child_loss in loss_module.losses.values():
                if hasattr(child_loss, self.target_attr):
                    temps.append(getattr(child_loss, self.target_attr))

            # Use first found temperature for warning check
            if temps:
                current_temp = temps[0]
        elif hasattr(loss_module, self.target_attr):
            current_temp = getattr(loss_module, self.target_attr)

    # Try criterion attribute (some models use this)
    if current_temp is None and hasattr(pl_module, "criterion"):
        loss_module = pl_module.criterion
        if hasattr(loss_module, self.target_attr):
            current_temp = getattr(loss_module, self.target_attr)

    # Try direct attribute on pl_module
    if current_temp is None and hasattr(pl_module, self.target_attr):
        current_temp = getattr(pl_module, self.target_attr)

    # Issue warning if mismatch detected
    if current_temp is not None:
        # Use proper tolerance for float comparison
        tolerance = torch.finfo(torch.float32).eps * 10
        if abs(current_temp - self.start_temperature) > tolerance:
            warn_message = (
                f"Temperature mismatch detected: loss module has {self.target_attr}={current_temp}, "
                f"but callback will start annealing from {self.start_temperature}. The callback's "
                f"start_temperature will override the loss module's initial value."
            )

            warnings.warn(warn_message, UserWarning, stacklevel=2)

on_train_epoch_start(trainer, pl_module) -> None

Update temperature at the start of each training epoch.

Computes new temperature based on current epoch, schedule, and warmup settings, then updates the target attribute in the loss module.

Source code in spectre/utils/callbacks.py

def on_train_epoch_start(self, trainer, pl_module) -> None:
    """
    Update temperature at the start of each training epoch.

    Computes new temperature based on current epoch, schedule, and warmup
    settings, then updates the target attribute in the loss module.
    """
    if not self.enabled or trainer.sanity_checking:
        return

    # Infer total_epochs from trainer if not specified
    total_epochs = self.total_epochs
    if total_epochs is None:
        if trainer.max_epochs is None:
            # If max_epochs not set, cannot anneal
            return
        total_epochs = trainer.max_epochs

    current_epoch = trainer.current_epoch

    # Compute new temperature
    if current_epoch < self.warmup_epochs:
        # Warmup phase: keep at start temperature
        new_temp = self.start_temperature
    else:
        # Annealing phase: compute based on schedule
        annealing_epochs = total_epochs - self.warmup_epochs
        if annealing_epochs <= 0:
            # No annealing epochs (warmup >= total), keep at start
            new_temp = self.start_temperature
        else:
            progress = (current_epoch - self.warmup_epochs) / annealing_epochs
            progress = min(progress, 1.0)  # Clamp to [0, 1]

            if self.schedule == "linear":
                # Linear interpolation
                new_temp = (
                    self.start_temperature
                    + (self.end_temperature - self.start_temperature) * progress
                )
            elif self.schedule == "exponential":
                # Exponential decay
                ratio = self.end_temperature / self.start_temperature
                new_temp = self.start_temperature * (ratio**progress)
            elif self.schedule == "cosine":
                # Cosine annealing
                new_temp = self.end_temperature + 0.5 * (
                    self.start_temperature - self.end_temperature
                ) * (1 + math.cos(math.pi * progress))

    # Update the loss module's temperature attribute
    self._update_temperature(pl_module, new_temp)

    # Log the temperature
    if trainer.logger:
        trainer.logger.log_metrics(
            {f"annealing/{self.target_attr}": new_temp},
            step=trainer.global_step,
        )

to_config() -> dict[str, Any]

Export callback configuration.

Source code in spectre/utils/callbacks.py

def to_config(self) -> dict[str, Any]:
    """Export callback configuration."""
    config = super().to_config()
    config.update(
        {
            "target_attr": self.target_attr,
            "start_temperature": self.start_temperature,
            "end_temperature": self.end_temperature,
            "total_epochs": self.total_epochs,
            "schedule": self.schedule,
            "warmup_epochs": self.warmup_epochs,
        }
    )
    return config

Functions

model_complexity_extractor(trainer, pl_module) -> dict[str, int]

Extract model complexity metrics.

Source code in spectre/utils/callbacks.py

def model_complexity_extractor(trainer, pl_module) -> dict[str, int]:
    """Extract model complexity metrics."""
    total_params = sum(p.numel() for p in pl_module.parameters())
    trainable_params = sum(p.numel() for p in pl_module.parameters() if p.requires_grad)
    return {"total_params": total_params, "trainable_params": trainable_params}

memory_usage_extractor(trainer, pl_module) -> dict[str, int]

Extract GPU memory usage.

RETURNS	DESCRIPTION
dict[str, int]	Dictionary with memory usage in bytes. Returns zeros if CUDA is unavailable or if model is not on GPU.

Source code in spectre/utils/callbacks.py

def memory_usage_extractor(trainer, pl_module) -> dict[str, int]:
    """
    Extract GPU memory usage.

    Returns
    -------
    dict[str, int]
        Dictionary with memory usage in bytes. Returns zeros if CUDA is unavailable
        or if model is not on GPU.
    """
    if not torch.cuda.is_available():
        return {"gpu_memory_allocated": 0, "gpu_memory_reserved": 0}

    # Check if model is actually on CUDA
    if pl_module is None:
        return {"gpu_memory_allocated": 0, "gpu_memory_reserved": 0}

    try:
        model_device = next(pl_module.parameters(), None)
    except (AttributeError, StopIteration):
        return {"gpu_memory_allocated": 0, "gpu_memory_reserved": 0}

    if model_device is None or not model_device.is_cuda:
        return {"gpu_memory_allocated": 0, "gpu_memory_reserved": 0}

    device_idx = model_device.get_device()
    allocated = torch.cuda.memory_allocated(device_idx)
    reserved = torch.cuda.memory_reserved(device_idx)
    max_allocated = torch.cuda.max_memory_allocated(device_idx)

    return {
        "gpu_memory_allocated": allocated,
        "gpu_memory_reserved": reserved,
        "gpu_memory_max": max_allocated,
    }

gradient_norm_extractor(trainer, pl_module) -> float

Extract gradient norms.

Source code in spectre/utils/callbacks.py

def gradient_norm_extractor(trainer, pl_module) -> float:
    """Extract gradient norms."""
    total_norm = 0
    param_count = 0
    for p in pl_module.parameters():
        if p.grad is not None:
            param_norm = p.grad.detach().data.norm(2)
            total_norm += param_norm.item() ** 2
            param_count += 1
    return (total_norm**0.5) if param_count > 0 else 0.0

kernel_params_extractor(trainer, pl_module) -> dict

Universal kernel parameter extractor.

Works with any kernel type through unified get_param_summary() interface. Automatically handles scalar vs vector/matrix parameters.

PARAMETER	DESCRIPTION
trainer	PyTorch Lightning trainer instance. TYPE: Trainer
pl_module	Module containing kernel_fn attribute. TYPE: LightningModule

RETURNS	DESCRIPTION
dict	Dictionary of kernel parameters and metadata: For scalar parameters: {param_name: float_value} For vector/matrix parameters: {param_name_mean, param_name_std, param_name_min, param_name_max: float} kernel_type: str, name of kernel class kernel_learnable: bool, whether parameters are learnable

Examples:

With GaussianKernel (scalar bandwidth):

>>> from spectre.kernel import GaussianKernel
>>> kernel = GaussianKernel(bw_method=torch.tensor(1.5))
>>> # Returns: {'bandwidth': 1.5, 'kernel_type': 'GaussianKernel',
>>> #           'kernel_learnable': False}

With TKernel (vector alpha, scalar beta):

>>> from spectre.kernel import TKernel
>>> kernel = TKernel(alpha=torch.ones(10, 1), beta=torch.tensor(2.0))
>>> # Returns: {'alpha_mean': 1.0, 'alpha_std': 0.0, 'alpha_min': 1.0,
>>> #           'alpha_max': 1.0, 'beta': 2.0, 'kernel_type': 'TKernel',
>>> #           'kernel_learnable': False}

Notes

• Requires pl_module to have kernel_fn attribute
• Kernel must implement get_param_summary() method
• Returns empty dict if kernel_fn not found

Source code in spectre/utils/callbacks.py

def kernel_params_extractor(trainer, pl_module) -> dict:
    """
    Universal kernel parameter extractor.

    Works with any kernel type through unified `get_param_summary()` interface.
    Automatically handles scalar vs vector/matrix parameters.

    Parameters
    ----------
    trainer : pytorch_lightning.Trainer
        PyTorch Lightning trainer instance.

    pl_module : pytorch_lightning.LightningModule
        Module containing `kernel_fn` attribute.

    Returns
    -------
    dict
        Dictionary of kernel parameters and metadata:

        - For scalar parameters: `{param_name: float_value}`
        - For vector/matrix parameters: `{param_name_mean, param_name_std,
          param_name_min, param_name_max: float}`
        - `kernel_type`: str, name of kernel class
        - `kernel_learnable`: bool, whether parameters are learnable

    Examples
    --------
    With GaussianKernel (scalar bandwidth):

    >>> from spectre.kernel import GaussianKernel
    >>> kernel = GaussianKernel(bw_method=torch.tensor(1.5))
    >>> # Returns: {'bandwidth': 1.5, 'kernel_type': 'GaussianKernel',
    >>> #           'kernel_learnable': False}

    With TKernel (vector alpha, scalar beta):

    >>> from spectre.kernel import TKernel
    >>> kernel = TKernel(alpha=torch.ones(10, 1), beta=torch.tensor(2.0))
    >>> # Returns: {'alpha_mean': 1.0, 'alpha_std': 0.0, 'alpha_min': 1.0,
    >>> #           'alpha_max': 1.0, 'beta': 2.0, 'kernel_type': 'TKernel',
    >>> #           'kernel_learnable': False}

    Notes
    -----
    - Requires `pl_module` to have `kernel_fn` attribute
    - Kernel must implement `get_param_summary()` method
    - Returns empty dict if kernel_fn not found
    """
    if not hasattr(pl_module, "kernel_fn"):
        return {}

    kernel = pl_module.kernel_fn

    # Get parameter summary (handles scalars vs vectors automatically)
    summary = kernel.get_param_summary()

    # Flatten nested dicts for logging
    metrics = {}
    for param_name, value in summary.items():
        if isinstance(value, dict):
            # Vector/matrix parameter - flatten stats
            for stat_name, stat_value in value.items():
                if stat_name != "shape":  # Skip shape in metrics
                    metrics[f"{param_name}_{stat_name}"] = stat_value
        else:
            # Scalar parameter
            metrics[param_name] = value

    return metrics

eigenvalue_extractor(n_eigen: int | None = None) -> Callable

Factory function to create eigenvalue extractor for parametric models.

Creates an extractor that retrieves learned eigenvalues from models that compute spectral decompositions during training. Supports eigenvalues attributes.

PARAMETER	DESCRIPTION
n_eigen	Number of eigenvalues to extract. If None, extracts all available eigenvalues. If specified, extracts only the first n_eigen values. TYPE: int | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Callable	Extractor function with signature (trainer, pl_module) -> dict.

Examples:

Extract all eigenvalues:

>>> from spectre.parametric import SpectralMap
>>> from spectre.utils.callbacks import TrainingCallback, eigenvalue_extractor
>>> sm = SpectralMap(...)
>>> callback = TrainingCallback(
...     custom_extractors={"eigenvalues": eigenvalue_extractor(n_eigen=10)},
... )
>>> # After training, metrics will contain all 10 eigenvalues:
>>> # {'eigenvalue_0': 0.95, 'eigenvalue_1': 0.82, ..., 'eigenvalue_9': 0.10}

Notes

• Returns empty dict if module has no eigenvalue attributes
• If n_eigen exceeds available eigenvalues, returns all available

Source code in spectre/utils/callbacks.py

def eigenvalue_extractor(n_eigen: int | None = None) -> Callable:
    """
    Factory function to create eigenvalue extractor for parametric models.

    Creates an extractor that retrieves learned eigenvalues from models
    that compute spectral decompositions during training. Supports
    `eigenvalues` attributes.

    Parameters
    ----------
    n_eigen : int | None, by default None
        Number of eigenvalues to extract. If None, extracts all available
        eigenvalues. If specified, extracts only the first `n_eigen` values.

    Returns
    -------
    Callable
        Extractor function with signature `(trainer, pl_module) -> dict`.

    Examples
    --------
    Extract all eigenvalues:

    >>> from spectre.parametric import SpectralMap
    >>> from spectre.utils.callbacks import TrainingCallback, eigenvalue_extractor
    >>> sm = SpectralMap(...)
    >>> callback = TrainingCallback(
    ...     custom_extractors={"eigenvalues": eigenvalue_extractor(n_eigen=10)},
    ... )
    >>> # After training, metrics will contain all 10 eigenvalues:
    >>> # {'eigenvalue_0': 0.95, 'eigenvalue_1': 0.82, ..., 'eigenvalue_9': 0.10}

    Notes
    -----
    - Returns empty dict if module has no eigenvalue attributes
    - If `n_eigen` exceeds available eigenvalues, returns all available
    """

    def _extract(trainer, pl_module) -> dict:
        # Check for eigenvalues attribute
        eigvals_tensor = None
        if hasattr(pl_module, "eigenvalues"):
            eigvals_tensor = pl_module.eigenvalues
        else:
            return {}

        # Handle None or empty tensors
        if eigvals_tensor is None or not isinstance(eigvals_tensor, torch.Tensor):
            return {}

        # Convert to CPU and extract individual values
        try:
            eigvals_cpu = eigvals_tensor.detach().cpu()

            # Limit to n_eigen if specified
            if n_eigen is not None:
                eigvals_cpu = eigvals_cpu[:n_eigen]

            return {f"eigenvalue_{i}": val.item() for i, val in enumerate(eigvals_cpu)}
        except (RuntimeError, ValueError, AttributeError):
            return {}

    return _extract

kernel_metrics_extractor(metrics: list[str] | None = None, effective_rank_method: str = 'threshold', effective_rank_threshold: float = 0.99) -> Callable

Factory function to create kernel metrics extractor for parametric models.

Creates an extractor that computes kernel quality metrics from the stored loss context. Supports effective rank, kernel alignment, and reconstruction error metrics.

Requires parametric models to be initialized with store_loss_context=True.

PARAMETER	DESCRIPTION
metrics	List of metrics to compute. If None, defaults to ["effective_rank"]. Available metrics: "effective_rank": Effective rank of kernel matrix "kernel_trace": Trace of kernel matrix (sum of diagonal) "kernel_frobenius": Frobenius norm of kernel matrix TYPE: list[str] | None, by default None DEFAULT: None
effective_rank_method	Method for computing effective rank. Options: "threshold": Count eigenvalues above cumulative threshold "entropy": Shannon entropy of normalized eigenvalues TYPE: str, by default "threshold" DEFAULT: 'threshold'
effective_rank_threshold	Threshold for cumulative eigenvalue sum (used with method="threshold"). TYPE: float, by default 0.99 DEFAULT: 0.99

RETURNS	DESCRIPTION
Callable	Extractor function with signature (trainer, pl_module) -> dict.

Examples:

Extract effective rank with default settings:

>>> from spectre.parametric import SpectralMap
>>> from spectre.utils.callbacks import TrainingCallback, kernel_metrics_extractor
>>> sm = SpectralMap(..., store_loss_context=True)
>>> callback = TrainingCallback(
...     custom_extractors={"kernel_metrics_extractor": kernel_metrics_extractor()},
... )
>>> # After training, metrics will contain: {'effective_rank': 15.2, ...}

Extract multiple kernel metrics:

>>> callback = TrainingCallback(
...     custom_extractors={
...         "kernel_metrics_extractor": kernel_metrics_extractor(
...             metrics=["effective_rank", "kernel_trace", "kernel_frobenius"],
...             effective_rank_method="entropy",
...         )
...     },
... )

Notes

• Returns empty dict if module has no loss_context attribute
• Requires store_loss_context=True when initializing parametric models

Source code in spectre/utils/callbacks.py

def kernel_metrics_extractor(
    metrics: list[str] | None = None,
    effective_rank_method: str = "threshold",
    effective_rank_threshold: float = 0.99,
) -> Callable:
    """
    Factory function to create kernel metrics extractor for parametric models.

    Creates an extractor that computes kernel quality metrics from the stored
    loss context. Supports effective rank, kernel alignment, and reconstruction
    error metrics.

    Requires parametric models to be initialized with `store_loss_context=True`.

    Parameters
    ----------
    metrics : list[str] | None, by default None
        List of metrics to compute. If None, defaults to ["effective_rank"].

        Available metrics:

        - "effective_rank": Effective rank of kernel matrix
        - "kernel_trace": Trace of kernel matrix (sum of diagonal)
        - "kernel_frobenius": Frobenius norm of kernel matrix

    effective_rank_method : str, by default "threshold"
        Method for computing effective rank.

        Options:

        - "threshold": Count eigenvalues above cumulative threshold
        - "entropy": Shannon entropy of normalized eigenvalues

    effective_rank_threshold : float, by default 0.99
        Threshold for cumulative eigenvalue sum (used with method="threshold").

    Returns
    -------
    Callable
        Extractor function with signature `(trainer, pl_module) -> dict`.

    Examples
    --------
    Extract effective rank with default settings:

    >>> from spectre.parametric import SpectralMap
    >>> from spectre.utils.callbacks import TrainingCallback, kernel_metrics_extractor
    >>> sm = SpectralMap(..., store_loss_context=True)
    >>> callback = TrainingCallback(
    ...     custom_extractors={"kernel_metrics_extractor": kernel_metrics_extractor()},
    ... )
    >>> # After training, metrics will contain: {'effective_rank': 15.2, ...}

    Extract multiple kernel metrics:

    >>> callback = TrainingCallback(
    ...     custom_extractors={
    ...         "kernel_metrics_extractor": kernel_metrics_extractor(
    ...             metrics=["effective_rank", "kernel_trace", "kernel_frobenius"],
    ...             effective_rank_method="entropy",
    ...         )
    ...     },
    ... )

    Notes
    -----
    - Returns empty dict if module has no `loss_context` attribute
    - Requires `store_loss_context=True` when initializing parametric models
    """
    if metrics is None:
        metrics = ["effective_rank"]

    def _extract(trainer, pl_module) -> dict:
        if not hasattr(pl_module, "loss_context"):
            return {}

        context = pl_module.loss_context
        if context is None or not isinstance(context, dict):
            return {}

        # Extract kernel matrix from context
        K = context.get("K")
        if K is None or not isinstance(K, torch.Tensor):
            return {}

        # Move to CPU and detach for metric computation
        try:
            K = K.detach().cpu()
        except (RuntimeError, AttributeError):
            return {}

        results = {}

        # Compute requested metrics
        for metric_name in metrics:
            try:
                if metric_name == "effective_rank":
                    eigenvalues = context.get("eigenvalues")
                    # Fall back to module attribute
                    if eigenvalues is None and hasattr(pl_module, "eigenvalues"):
                        eigenvalues = pl_module.eigenvalues

                    if eigenvalues is not None and isinstance(
                        eigenvalues, torch.Tensor
                    ):
                        eigenvalues = eigenvalues.detach().cpu()
                    else:
                        eigenvalues = None

                    eff_rank = effective_rank(
                        eigenvalues=eigenvalues,
                        K=K,
                        method=effective_rank_method,
                        threshold=effective_rank_threshold,
                    )
                    results["effective_rank"] = float(eff_rank)

                elif metric_name == "kernel_trace":
                    # Trace of kernel matrix
                    trace = torch.trace(K).item()
                    results["kernel_trace"] = float(trace)

                elif metric_name == "kernel_frobenius":
                    # Frobenius norm of kernel matrix
                    frob_norm = torch.linalg.matrix_norm(K, ord="fro").item()
                    results["kernel_frobenius"] = float(frob_norm)

            except Exception:
                # Skip this metric if computation fails
                continue

        return results

    return _extract

generic_metrics_extractor(metrics: dict[str, Callable] | list[Callable], parameter_mappings: dict[str, dict[str, str]] | None = None, default_params: dict[str, Any] | None = None, metric_params: dict[str, dict[str, Any]] | None = None) -> Callable

Factory function to create generic metric extractor from loss_context.

Automatically inspects metric function signatures and computes metrics only when all required parameters are available in the loss context. Supports flexible parameter mapping to handle naming mismatches between context keys and function parameters.

Requires parametric models to be initialized with store_loss_context=True.

PARAMETER	DESCRIPTION
metrics	Metric functions to compute. If dict, keys are metric names and values are callable metric functions. If list, each element should be a callable with a __name__ attribute for identification. TYPE: dict[str, Callable] | list[Callable]
parameter_mappings	Map context keys to function parameter names for each metric. Format: {metric_name: {context_key: param_name}} Example: { "kernel_alignment": {"K": "K1", "K_ref": "K2"}, "effective_rank": {"K": "K"} } TYPE: dict[str, dict[str, str]] | None, by default None DEFAULT: None
default_params	Default values for parameters across all metrics. These values are used when parameters are not found in context and have no default in signature. TYPE: dict[str, Any] | None, by default None DEFAULT: None
metric_params	Per-metric default parameters that override default_params. Format: {metric_name: {param_name: value}} TYPE: dict[str, dict[str, Any]] | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Callable	Extractor function with signature (trainer, pl_module) -> dict.

Examples:

Simple usage with automatic parameter matching:

>>> from spectre.metrics.kernel import effective_rank, kernel_alignment
>>> extractor = generic_metrics_extractor(
...     metrics=[effective_rank],
... )

Advanced usage with parameter mappings and defaults:

>>> extractor = generic_metrics_extractor(
...     metrics={
...         "eff_rank": effective_rank,
...         "k_align": kernel_alignment,
...     },
...     parameter_mappings={
...         "k_align": {"K": "K1"},  # Map context K to param K1
...     },
...     default_params={"K2": torch.eye(10)},  # Provide K2 for alignment
...     metric_params={"eff_rank": {"method": "entropy"}},
... )

Using with TrainingCallback:

>>> from spectre.parametric import SpectralMap
>>> sm = SpectralMap(..., store_loss_context=True)
>>> callback = TrainingCallback(custom_extractors={"metrics": extractor})
>>> sm.fit(X, callbacks=[callback])

Notes

• Returns empty dict if module has no loss_context attribute
• Requires store_loss_context=True when initializing parametric models
• Metrics are computed independently; failure in one doesn't affect others
• Only metrics with all required parameters available are computed

Source code in spectre/utils/callbacks.py

def generic_metrics_extractor(
    metrics: dict[str, Callable] | list[Callable],
    parameter_mappings: dict[str, dict[str, str]] | None = None,
    default_params: dict[str, Any] | None = None,
    metric_params: dict[str, dict[str, Any]] | None = None,
) -> Callable:
    """
    Factory function to create generic metric extractor from loss_context.

    Automatically inspects metric function signatures and computes metrics only
    when all required parameters are available in the loss context. Supports
    flexible parameter mapping to handle naming mismatches between context keys
    and function parameters.

    Requires parametric models to be initialized with `store_loss_context=True`.

    Parameters
    ----------
    metrics : dict[str, Callable] | list[Callable]
        Metric functions to compute. If dict, keys are metric names and values
        are callable metric functions. If list, each element should be a callable
        with a `__name__` attribute for identification.

    parameter_mappings : dict[str, dict[str, str]] | None, by default None
        Map context keys to function parameter names for each metric.
        Format: {metric_name: {context_key: param_name}}

        Example:
            {
                "kernel_alignment": {"K": "K1", "K_ref": "K2"},
                "effective_rank": {"K": "K"}
            }

    default_params : dict[str, Any] | None, by default None
        Default values for parameters across all metrics. These values are used
        when parameters are not found in context and have no default in signature.

    metric_params : dict[str, dict[str, Any]] | None, by default None
        Per-metric default parameters that override `default_params`.
        Format: {metric_name: {param_name: value}}

    Returns
    -------
    Callable
        Extractor function with signature `(trainer, pl_module) -> dict`.

    Examples
    --------
    Simple usage with automatic parameter matching:

    >>> from spectre.metrics.kernel import effective_rank, kernel_alignment
    >>> extractor = generic_metrics_extractor(
    ...     metrics=[effective_rank],
    ... )

    Advanced usage with parameter mappings and defaults:

    >>> extractor = generic_metrics_extractor(
    ...     metrics={
    ...         "eff_rank": effective_rank,
    ...         "k_align": kernel_alignment,
    ...     },
    ...     parameter_mappings={
    ...         "k_align": {"K": "K1"},  # Map context K to param K1
    ...     },
    ...     default_params={"K2": torch.eye(10)},  # Provide K2 for alignment
    ...     metric_params={"eff_rank": {"method": "entropy"}},
    ... )

    Using with TrainingCallback:

    >>> from spectre.parametric import SpectralMap
    >>> sm = SpectralMap(..., store_loss_context=True)
    >>> callback = TrainingCallback(custom_extractors={"metrics": extractor})
    >>> sm.fit(X, callbacks=[callback])

    Notes
    -----
    - Returns empty dict if module has no `loss_context` attribute
    - Requires `store_loss_context=True` when initializing parametric models
    - Metrics are computed independently; failure in one doesn't affect others
    - Only metrics with all required parameters available are computed
    """
    # Convert list to dict if needed
    if isinstance(metrics, list):
        metrics = {fn.__name__: fn for fn in metrics}

    parameter_mappings = parameter_mappings or {}
    default_params = default_params or {}
    metric_params = metric_params or {}

    # Analyze signatures for each metric
    signatures = {}
    for name, fn in metrics.items():
        try:
            signatures[name] = inspect.signature(fn)
        except (ValueError, TypeError):
            # Skip metrics we can't introspect
            continue

    def _extract(trainer, pl_module) -> dict:
        # Check for loss_context
        if not hasattr(pl_module, "loss_context"):
            return {}

        context = pl_module.loss_context
        if context is None or not isinstance(context, dict):
            return {}

        results = {}

        # Compute each metric independently
        for metric_name, metric_fn in metrics.items():
            if metric_name not in signatures:
                continue

            try:
                sig = signatures[metric_name]

                # Build parameter mapping for this metric
                param_map = parameter_mappings.get(metric_name, {})

                # Collect available parameters
                kwargs = {}

                for param_name, param in sig.parameters.items():
                    # Skip self, *args, **kwargs
                    if param.kind in (
                        inspect.Parameter.VAR_POSITIONAL,
                        inspect.Parameter.VAR_KEYWORD,
                    ):
                        continue

                    value = None
                    found = False

                    # Try mapped context key
                    for context_key, mapped_param in param_map.items():
                        if mapped_param == param_name and context_key in context:
                            value = context[context_key]
                            found = True
                            break

                    # Try direct context key match
                    if not found and param_name in context:
                        value = context[param_name]
                        found = True

                    # Try metric-specific defaults
                    if not found and metric_name in metric_params:
                        if param_name in metric_params[metric_name]:
                            value = metric_params[metric_name][param_name]
                            found = True

                    # Try global defaults
                    if not found and param_name in default_params:
                        value = default_params[param_name]
                        found = True

                    # Try function default if available
                    if not found and param.default != inspect.Parameter.empty:
                        value = param.default
                        found = True

                    # Required parameter not found
                    if not found:
                        # Skip this metric - missing required parameter
                        break

                    # Detach and move tensors to CPU
                    if isinstance(value, torch.Tensor):
                        try:
                            value = value.detach().cpu()
                        except (RuntimeError, AttributeError):
                            pass

                    kwargs[param_name] = value
                else:
                    # All parameters found, compute metric
                    result = metric_fn(**kwargs)

                    # Convert result to float if it's a tensor
                    if isinstance(result, torch.Tensor):
                        result = result.item()

                    results[metric_name] = float(result)

            except Exception:
                # Skip this metric if computation fails
                continue

        return results

    return _extract

create_basic_training_callback(save_path: str = 'metrics.pkl') -> TrainingCallback

Create a basic training callback with standard logging.

Source code in spectre/utils/callbacks.py

def create_basic_training_callback(save_path: str = "metrics.pkl") -> TrainingCallback:
    """Create a basic training callback with standard logging."""
    return TrainingCallback(save_path=save_path, track_lr=True, track_epoch_time=True)

create_training_callback(save_path: str = 'advanced_metrics.json') -> TrainingCallback

Create a training callback with custom extractors.

Source code in spectre/utils/callbacks.py

def create_training_callback(
    save_path: str = "advanced_metrics.json",
) -> TrainingCallback:
    """Create a training callback with custom extractors."""
    callback = TrainingCallback(
        save_path=save_path, save_format="json", track_lr=True, track_epoch_time=True
    )

    # Add custom extractors
    callback.add_custom_extractor(
        "model_params_extractor", lambda _t, m: sum(p.numel() for p in m.parameters())
    )
    callback.add_custom_extractor("gradient_norm_extractor", gradient_norm_extractor)

    if hasattr(__import__("torch"), "cuda") and __import__("torch").cuda.is_available():
        callback.add_custom_extractor(
            "gpu_memory_mb",
            lambda _t, _m: __import__("torch").cuda.memory_allocated() / 1024**2,
        )

    return callback

register_callback(name: str, callback_class: BaseCallback) -> None

Register a new callback type in the factory.

PARAMETER	DESCRIPTION
name	Unique identifier for the callback type. TYPE: str
callback_class	Callback class that inherits from BaseCallback. TYPE: BaseCallback

Examples:

Register a custom callback:

>>> class MyCallback(BaseCallback):
...     def get_callback_type(self) -> str:
...         return "my_custom"
>>> register_callback("my_custom", MyCallback)

Source code in spectre/utils/callbacks.py

def register_callback(name: str, callback_class: BaseCallback) -> None:
    """
    Register a new callback type in the factory.

    Parameters
    ----------
    name : str
        Unique identifier for the callback type.

    callback_class : BaseCallback
        Callback class that inherits from BaseCallback.

    Examples
    --------
    Register a custom callback:

    >>> class MyCallback(BaseCallback):
    ...     def get_callback_type(self) -> str:
    ...         return "my_custom"
    >>> register_callback("my_custom", MyCallback)
    """
    if not issubclass(callback_class, BaseCallback):
        raise ValueError("Callback class must inherit from BaseCallback")
    CALLBACK_REGISTRY[name] = callback_class

list_available_callbacks() -> list[str]

List all available callback types.

RETURNS	DESCRIPTION
list[str]	List of registered callback type names.

Source code in spectre/utils/callbacks.py

def list_available_callbacks() -> list[str]:
    """
    List all available callback types.

    Returns
    -------
    list[str]
        List of registered callback type names.
    """
    return list(CALLBACK_REGISTRY.keys())

create_callback_from_config(config: dict[str, Any]) -> BaseCallback

Create a callback instance from configuration dictionary.

PARAMETER	DESCRIPTION
config	Configuration dictionary containing 'type' field and callback parameters. TYPE: dict[str, Any]

RETURNS	DESCRIPTION
BaseCallback	Configured callback instance.

RAISES	DESCRIPTION
ValueError	If callback type is not recognized or configuration is invalid.

Examples:

Create a spectral gap callback:

>>> config = {"type": "spectral_gap", "n_eigen": 15, "log_every_n_epochs": 2}
>>> callback = create_callback_from_config(config)

Create metrics callback with custom extractors:

>>> config = {
...     "type": "training_callback",
...     "save_path": "metrics.json",
...     "save_format": "json",
...     "track_lr": True,
... }
>>> callback = create_callback_from_config(config)

Source code in spectre/utils/callbacks.py

def create_callback_from_config(config: dict[str, Any]) -> BaseCallback:
    """
    Create a callback instance from configuration dictionary.

    Parameters
    ----------
    config : dict[str, Any]
        Configuration dictionary containing 'type' field and callback parameters.

    Returns
    -------
    BaseCallback
        Configured callback instance.

    Raises
    ------
    ValueError
        If callback type is not recognized or configuration is invalid.

    Examples
    --------
    Create a spectral gap callback:

    >>> config = {"type": "spectral_gap", "n_eigen": 15, "log_every_n_epochs": 2}
    >>> callback = create_callback_from_config(config)

    Create metrics callback with custom extractors:

    >>> config = {
    ...     "type": "training_callback",
    ...     "save_path": "metrics.json",
    ...     "save_format": "json",
    ...     "track_lr": True,
    ... }
    >>> callback = create_callback_from_config(config)
    """
    callback_type = config.get("type")
    if callback_type is None:
        raise ValueError("Configuration must contain 'type' field")

    if callback_type not in CALLBACK_REGISTRY:
        available = ", ".join(CALLBACK_REGISTRY.keys())
        raise ValueError(
            f"Unknown callback type: {callback_type}. Available types: {available}"
        )

    callback_class = CALLBACK_REGISTRY[callback_type]

    # Create a copy of config without the 'type' field
    callback_config = {k: v for k, v in config.items() if k != "type"}

    # Validate configuration before creating instance
    dummy_instance = callback_class()
    dummy_instance.validate_config(callback_config)

    # Create configured instance
    return callback_class.from_config(callback_config)

get_default_callbacks(save_path: str | None = None) -> list[BaseCallback]

Get a default set of callbacks.

PARAMETER	DESCRIPTION
save_path	Path to save metrics. If provided, enables metrics saving. TYPE: str DEFAULT: None

RETURNS	DESCRIPTION
list[BaseCallback]	List of default callbacks (empty list if no save_path).

Examples:

Get basic callbacks:

>>> callbacks = get_default_callbacks()
[]

Get callbacks with metrics saving:

>>> callbacks = get_default_callbacks(save_path="metrics.json")
[TrainingCallback(...)]

Source code in spectre/utils/callbacks.py

def get_default_callbacks(save_path: str | None = None) -> list[BaseCallback]:
    """
    Get a default set of callbacks.

    Parameters
    ----------
    save_path : str, optional
        Path to save metrics. If provided, enables metrics saving.

    Returns
    -------
    list[BaseCallback]
        List of default callbacks (empty list if no save_path).

    Examples
    --------
    Get basic callbacks:

    >>> callbacks = get_default_callbacks()
    []

    Get callbacks with metrics saving:

    >>> callbacks = get_default_callbacks(save_path="metrics.json")
    [TrainingCallback(...)]
    """
    if save_path:
        return [TrainingCallback(save_path=save_path, track_lr=True)]
    return []

create_callback_from_template(template_name: str, **kwargs) -> BaseCallback | list[BaseCallback]

Create callbacks from predefined templates.

PARAMETER	DESCRIPTION
template_name	Name of the callback template. TYPE: str
**kwargs	Template-specific parameters. DEFAULT: {}

RETURNS	DESCRIPTION
BaseCallback | list[BaseCallback]	Callback instance(s) from template.

Examples:

Create from template:

>>> cb = create_callback_from_template(
...     "training_callback",
...     save_path="metrics.json",
... )

Source code in spectre/utils/callbacks.py

def create_callback_from_template(
    template_name: str, **kwargs
) -> BaseCallback | list[BaseCallback]:
    """
    Create callbacks from predefined templates.

    Parameters
    ----------
    template_name : str
        Name of the callback template.

    **kwargs
        Template-specific parameters.

    Returns
    -------
    BaseCallback | list[BaseCallback]
        Callback instance(s) from template.

    Examples
    --------
    Create from template:

    >>> cb = create_callback_from_template(
    ...     "training_callback",
    ...     save_path="metrics.json",
    ... )
    """
    templates = {
        "training_callback": lambda **kw: TrainingCallback(
            track_lr=True, track_epoch_time=True, **kw
        ),
    }

    if template_name not in templates:
        available = ", ".join(templates.keys())
        raise ValueError(f"Unknown template: {template_name}. Available: {available}")

    return templates[template_name](**kwargs)

get_callback_config_schema() -> dict[str, dict[str, Any]]

Get configuration schema for all available callbacks.

RETURNS	DESCRIPTION
dict[str, dict[str, Any]]	Schema dictionary with callback types as keys and parameter specifications as values.

Source code in spectre/utils/callbacks.py

def get_callback_config_schema() -> dict[str, dict[str, Any]]:
    """
    Get configuration schema for all available callbacks.

    Returns
    -------
    dict[str, dict[str, Any]]
        Schema dictionary with callback types as keys and parameter
        specifications as values.
    """
    return {
        "training_callback": {
            "description": "Training callback with custom extractors and multiple formats",
            "parameters": {
                "enabled": {"type": "bool", "default": True},
                "log_every_n_epochs": {
                    "type": "int",
                    "default": 1,
                    "range": "[1, inf)",
                },
                "save_path": {"type": "str", "optional": True},
                "save_format": {
                    "type": "str",
                    "choices": ["pickle", "json"],
                    "default": "pickle",
                },
                "metric_prefix": {"type": "str", "optional": True},
                "track_lr": {"type": "bool", "default": True},
                "track_epoch_time": {"type": "bool", "default": False},
            },
        },
        "parameter_annealing_callback": {
            "description": "Anneal parameter during training for adaptive loss functions",
            "parameters": {
                "enabled": {"type": "bool", "default": True},
                "log_every_n_epochs": {
                    "type": "int",
                    "default": 1,
                    "range": "[1, inf)",
                },
                "target_attr": {
                    "type": "str",
                    "default": "temperature",
                    "description": "Attribute name to update in loss module",
                },
                "start_temperature": {
                    "type": "float",
                    "default": 1.0,
                    "range": "(0, inf)",
                },
                "end_temperature": {
                    "type": "float",
                    "default": 0.01,
                    "range": "(0, inf)",
                },
                "total_epochs": {
                    "type": "int",
                    "optional": True,
                    "default": None,
                    "range": "[1, inf)",
                },
                "schedule": {
                    "type": "str",
                    "choices": ["linear", "exponential", "cosine"],
                    "default": "cosine",
                },
                "warmup_epochs": {
                    "type": "int",
                    "default": 0,
                    "range": "[0, inf)",
                },
            },
        },
    }