Data Dataset

dataset

CLASS	DESCRIPTION
DatasetWeighted	Dataset with weights and optional targets for enhanced data loading.

Classes

DataDict

Bases: TypedDict

Type-safe container for converted dataset inputs.

DatasetWeighted(X: ArrayType | list[ArrayType], weights: ArrayType | list[ArrayType] | None = None, target: ArrayType | list[ArrayType] | None = None, reduce: Literal['prod', 'sum'] | Callable[[torch.Tensor], torch.Tensor] = 'prod', labels: list[str] | None = None, dtype: torch.dtype = torch.float32, transform: Callable | None = None, *args, **kwargs)

Bases: Dataset

Dataset with weights and optional targets for enhanced data loading.

Supports multiple input formats including tensors, numpy arrays, and lists of arrays. Provides automatic weight generation, label creation, and comprehensive indexing capabilities for machine learning workflows.

PARAMETER	DESCRIPTION
X	Input data. Can be tensor, numpy array, or list of arrays. TYPE: ArrayType | list
weights	Data weights. Auto-generated as ones if None. TYPE: ArrayType | list | None, optional, by default None DEFAULT: None
target	Target values for supervised learning, if None, dataset is for unsupervised learning. TYPE: ArrayType | list | None, optional, by default None DEFAULT: None
reduce	TYPE: (Literal['prod', 'sum'] | Callable[[Tensor], Tensor],) DEFAULT: 'prod'
optional	Weight reduction method for 2D weights only. For 1D weights, this parameter has no effect and weights are used as-is. Use "prod" for product across features, "sum" for sum, or provide a custom callable that accepts a 2D tensor (n_samples, n_features) and returns a 1D tensor (n_samples,).
by	Weight reduction method for 2D weights only. For 1D weights, this parameter has no effect and weights are used as-is. Use "prod" for product across features, "sum" for sum, or provide a custom callable that accepts a 2D tensor (n_samples, n_features) and returns a 1D tensor (n_samples,).
labels	Feature labels. Auto-generated as ["x_1", "x_2", ...] if None. TYPE: list | None, optional, by default None DEFAULT: None
dtype	Tensor data type for all arrays. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32
transform	Optional transform function applied to data during __getitem__. Should accept and return a torch.Tensor. TYPE: Callable | None, optional, by default None DEFAULT: None

Properties

data : torch.Tensor Input data tensor.

weights : torch.Tensor Weight tensor.

target : torch.Tensor | None Target tensor, if provided.

labels : list[str] Feature labels.

Examples:

Using tensors:

>>> import torch
>>> from spectre.data import DatasetWeighted
>>> X = torch.randn(100, 5)
>>> weights = torch.rand(100)
>>> target = torch.randint(0, 2, (100,))
>>> dataset = DatasetWeighted(X, weights=weights, target=target)
>>> len(dataset)
100
>>> data, w, t = dataset[0]
>>> data.shape
torch.Size([5])
>>> w.shape
torch.Size([])
>>> t.shape
torch.Size([1])

Using list of arrays:

>>> X_list = [torch.randn(50, 5), torch.randn(50, 5)]
>>> dataset = DatasetWeighted(X_list)
>>> len(dataset)
100
>>> data, w = dataset[0]
>>> data.shape
torch.Size([5])
>>> w.shape
torch.Size([])
>>> # Without target
>>> dataset = DatasetWeighted(X, weights=weights)
>>> data, w = dataset[0]
>>> data.shape
torch.Size([5])
>>> w.shape
torch.Size([])

Using multiple weight columns:

>>> weights_2d = torch.rand(100, 3)  # Three weight columns
>>> dataset = DatasetWeighted(X, weights=weights_2d, reduce="sum")
>>> data, w = dataset[0]
>>> data.shape
torch.Size([5])
>>> w.shape
torch.Size([])
>>> w  # Weight is sum across columns
tensor(1.2345)  # Example output

Using custom reduction function:

>>> def geometric_mean(w: torch.Tensor) -> torch.Tensor:
...     return torch.exp(torch.log(w).mean(dim=1))
>>> weights_2d = torch.rand(100, 3)
>>> dataset = DatasetWeighted(X, weights=weights_2d, reduce=geometric_mean)
>>> data, w = dataset[0]
>>> w.shape
torch.Size([])

METHOD	DESCRIPTION
validate	Validate dataset for common data quality issues.
get_statistics	Compute basic statistics for the dataset.
check_nan	Check if dataset contains NaN values in data, weights, or targets.
check_inf	Check if dataset contains infinite values in data, weights, or targets.
has_target	Check if the dataset has target values.

Source code in spectre/data/dataset.py

def __init__(
    self,
    X: ArrayType | list[ArrayType],
    weights: ArrayType | list[ArrayType] | None = None,
    target: ArrayType | list[ArrayType] | None = None,
    reduce: Literal["prod", "sum"]
    | Callable[[torch.Tensor], torch.Tensor] = "prod",
    labels: list[str] | None = None,
    dtype: torch.dtype = torch.float32,
    transform: Callable | None = None,
    *args,
    **kwargs,
) -> None:
    super().__init__(*args, **kwargs)

    if transform is not None and not callable(transform):
        raise TypeError("`transform` must be callable or None.")
    self.transform = transform

    # Convert and validate all inputs
    converted = data_to_tensors(X, weights, target, dtype, reduce)

    # Extract properly typed tensors
    data = converted["data"]
    weights_tensor = converted["weights"]
    target_tensor = converted["target"]

    # Validate and generate labels
    if labels is not None:
        if data.shape[1] != len(labels):
            raise ValueError(
                "Passed `labels` length does not match the number of features."
            )
    else:
        labels = [f"x_{i + 1}" for i in range(data.shape[1])]

    self.data = data
    self.labels = labels
    self.weights = weights_tensor
    self.target = target_tensor

Functions

validate() -> dict[str, bool]

Validate dataset for common data quality issues.

Checks for NaN values, infinite values, and negative weights.

RETURNS	DESCRIPTION
dict[str, bool]	Dictionary with validation results: has_nan_data, has_inf_data, has_nan_weights, has_inf_weights, has_negative_weights, has_nan_target, has_inf_target.

Source code in spectre/data/dataset.py

def validate(self) -> dict[str, bool]:
    """
    Validate dataset for common data quality issues.

    Checks for NaN values, infinite values, and negative weights.

    Returns
    -------
    dict[str, bool]
        Dictionary with validation results: `has_nan_data`, `has_inf_data`,
        `has_nan_weights`, `has_inf_weights`, `has_negative_weights`,
        `has_nan_target`, `has_inf_target`.
    """
    data_nan = torch.isnan(self.data).any()
    data_inf = torch.isinf(self.data).any()
    weights_nan = torch.isnan(self.weights).any()
    weights_inf = torch.isinf(self.weights).any()
    weights_negative = (self.weights < 0).any()
    weights_zero = (self.weights == 0).any()

    results = {
        "has_nan_data": data_nan.item(),
        "has_inf_data": data_inf.item(),
        "has_nan_weights": weights_nan.item(),
        "has_inf_weights": weights_inf.item(),
        "has_negative_weights": weights_negative.item(),
        "has_zero_weights": weights_zero.item(),
    }

    if self.has_target():
        target_nan = torch.isnan(self.target).any()
        target_inf = torch.isinf(self.target).any()
        results["has_nan_target"] = target_nan.item()
        results["has_inf_target"] = target_inf.item()
    else:
        results["has_nan_target"] = False
        results["has_inf_target"] = False

    return results

get_statistics() -> dict[str, torch.Tensor]

Compute basic statistics for the dataset.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary with statistics: mean, std, min, max for data, and weights_sum, weights_mean for weights.

Source code in spectre/data/dataset.py

def get_statistics(self) -> dict[str, torch.Tensor]:
    """
    Compute basic statistics for the dataset.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with statistics: `mean`, `std`, `min`, `max` for data,
        and `weights_sum`, `weights_mean` for weights.
    """
    stats = {
        "mean": self.data.mean(dim=0),
        "std": self.data.std(dim=0),
        "min": self.data.min(dim=0).values,
        "max": self.data.max(dim=0).values,
        "weights_sum": self.weights.sum(),
        "weights_mean": self.weights.mean(),
    }

    if self.has_target():
        stats["target_mean"] = self.target.mean(dim=0)
        stats["target_std"] = self.target.std(dim=0)

    return stats

check_nan() -> bool

Check if dataset contains NaN values in data, weights, or targets.

RETURNS	DESCRIPTION
bool	True if NaN values are found, False otherwise.

Source code in spectre/data/dataset.py

def check_nan(self) -> bool:
    """
    Check if dataset contains NaN values in data, weights, or targets.

    Returns
    -------
    bool
        True if NaN values are found, False otherwise.
    """
    # Combine checks using logical_or for efficiency
    has_nan = torch.logical_or(
        torch.isnan(self.data).any(), torch.isnan(self.weights).any()
    )

    if self.has_target():
        has_nan = torch.logical_or(has_nan, torch.isnan(self.target).any())

    return has_nan.item()

check_inf() -> bool

Check if dataset contains infinite values in data, weights, or targets.

RETURNS	DESCRIPTION
bool	True if infinite values are found, False otherwise.

Source code in spectre/data/dataset.py

def check_inf(self) -> bool:
    """
    Check if dataset contains infinite values in data, weights, or targets.

    Returns
    -------
    bool
        True if infinite values are found, False otherwise.
    """
    # Combine checks using logical_or for efficiency
    has_inf = torch.logical_or(
        torch.isinf(self.data).any(), torch.isinf(self.weights).any()
    )

    if self.has_target():
        has_inf = torch.logical_or(has_inf, torch.isinf(self.target).any())

    return has_inf.item()

has_target() -> bool

Check if the dataset has target values.

RETURNS	DESCRIPTION
bool	True if dataset has targets, False otherwise.

Source code in spectre/data/dataset.py

def has_target(self) -> bool:
    """
    Check if the dataset has target values.

    Returns
    -------
    bool
        True if dataset has targets, False otherwise.
    """
    return self.target is not None

Functions

concatenate(arrays: list[ArrayType], tensors: list[torch.Tensor], dtype: torch.dtype, validate_shape: bool = True) -> torch.Tensor

Concatenate list of arrays into single tensor with validation.

PARAMETER	DESCRIPTION
arrays	List of arrays to concatenate. TYPE: list[ArrayType]
X_tensors	Reference tensors for shape validation. TYPE: list[Tensor]
dtype	Target data type. TYPE: dtype
validate_shape	Whether to validate shapes against X_tensors, by default True. TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
Tensor	Concatenated tensor.

Source code in spectre/data/dataset.py

def concatenate(
    arrays: list[ArrayType],
    tensors: list[torch.Tensor],
    dtype: torch.dtype,
    validate_shape: bool = True,
) -> torch.Tensor:
    """
    Concatenate list of arrays into single tensor with validation.

    Parameters
    ----------
    arrays : list[ArrayType]
        List of arrays to concatenate.

    X_tensors : list[torch.Tensor]
        Reference tensors for shape validation.

    dtype : torch.dtype
        Target data type.

    validate_shape : bool, optional
        Whether to validate shapes against X_tensors, by default True.

    Returns
    -------
    torch.Tensor
        Concatenated tensor.
    """
    if len(arrays) != len(tensors):
        raise ValueError(
            f"Length mismatch: {len(arrays)} arrays vs {len(tensors)} batches."
        )

    # Convert all arrays to tensors
    tensor_list: list[torch.Tensor] = []
    for arr, x_batch in zip(arrays, tensors):
        if validate_shape:
            check_same_shape(x_batch, arr, axis=0)
        tensor = check_to_tensor(arr, dtype)
        tensor_list.append(tensor)

    # Determine dimensions from first tensor
    if tensor_list[0].ndim == 1:
        # Stack 1D tensors into 2D (n_samples, 1)
        return torch.cat(
            [t.unsqueeze(1) if t.ndim == 1 else t for t in tensor_list], dim=0
        )
    else:
        # Concatenate 2D tensors directly
        return torch.cat(tensor_list, dim=0)

data_to_tensors(X: ArrayType | list[ArrayType], weights: ArrayType | list[ArrayType] | None, target: ArrayType | list[ArrayType] | None, dtype: torch.dtype, reduce: Literal['prod', 'sum'] | Callable[[torch.Tensor], torch.Tensor]) -> DataDict

Convert input data to tensors with validation.

Handles list of arrays by concatenating, validates shapes, and ensures all outputs are proper torch.Tensor types for type safety.

PARAMETER	DESCRIPTION
X	Input data. TYPE: ArrayType | list[ArrayType]
weights	Data weights. TYPE: ArrayType | list[ArrayType] | None
target	Target values. TYPE: ArrayType | list[ArrayType] | None
dtype	Tensor data type. TYPE: dtype
reduce	Weight reduction method for 2D weights. Use "prod" for product, "sum" for sum, or provide a custom callable that accepts a 2D tensor (n_samples, n_features) and returns a 1D tensor (n_samples,). TYPE: Literal['prod', 'sum'] | Callable[[Tensor], Tensor]

RETURNS	DESCRIPTION
DataDict	TypedDict with data, weights, and target as torch.Tensor.

Source code in spectre/data/dataset.py

def data_to_tensors(
    X: ArrayType | list[ArrayType],
    weights: ArrayType | list[ArrayType] | None,
    target: ArrayType | list[ArrayType] | None,
    dtype: torch.dtype,
    reduce: Literal["prod", "sum"] | Callable[[torch.Tensor], torch.Tensor],
) -> DataDict:
    """
    Convert input data to tensors with validation.

    Handles list of arrays by concatenating, validates shapes, and ensures
    all outputs are proper torch.Tensor types for type safety.

    Parameters
    ----------
    X : ArrayType | list[ArrayType]
        Input data.

    weights : ArrayType | list[ArrayType] | None
        Data weights.

    target : ArrayType | list[ArrayType] | None
        Target values.

    dtype : torch.dtype
        Tensor data type.

    reduce : Literal["prod", "sum"] | Callable[[torch.Tensor], torch.Tensor]
        Weight reduction method for 2D weights. Use "prod" for product, "sum" for sum,
        or provide a custom callable that accepts a 2D tensor (n_samples, n_features)
        and returns a 1D tensor (n_samples,).

    Returns
    -------
    DataDict
        TypedDict with `data`, `weights`, and `target` as torch.Tensor.
    """
    X_tensor: torch.Tensor
    weights_tensor: torch.Tensor
    target_tensor: torch.Tensor | None

    if isinstance(X, list):
        if len(X) == 0:
            raise RuntimeError("expected a non-empty list of tensors.")

        # Convert all X batches to tensors
        X_tensors: list[torch.Tensor] = []
        for x_batch in X:
            check_array_type(x_batch)
            X_tensors.append(check_to_tensor(x_batch, dtype))

        total_samples = sum(len(x) for x in X_tensors)

        # Concatenate data using helper
        X_tensor = concatenate(X, X_tensors, dtype, validate_shape=False)

        # Handle weights
        if weights is not None:
            if not isinstance(weights, list):
                raise ValueError(f"`X` is a list but weights are {type(weights)}.")
            check_same_len(X_tensors, weights)
            weights_tensor = concatenate(weights, X_tensors, dtype)
        else:
            weights_tensor = torch.ones(total_samples, dtype=dtype)

        # Handle target
        if target is not None:
            if not isinstance(target, list):
                raise ValueError(f"`X` is a list but target is {type(target)}.")
            check_same_len(X_tensors, target)
            target_tensor = concatenate(target, X_tensors, dtype)
        else:
            target_tensor = None

    else:
        # Single array input - validate type first
        if not isinstance(X, (torch.Tensor, np.ndarray)):
            raise ValueError(
                f"Type `{type(X)}` is not supported. Pass `torch.Tensor`, "
                f"`np.ndarray`, or `list`."
            )
        X_tensor = check_to_tensor(X, dtype)

        if len(X_tensor) == 0:
            raise ValueError("Dataset cannot be empty (0 samples).")

        if weights is not None:
            weights_tensor = check_to_tensor(weights, dtype)
            check_same_len(X_tensor, weights_tensor)
        else:
            weights_tensor = torch.ones(len(X_tensor), dtype=dtype)

        if target is not None:
            target_tensor = check_to_tensor(target, dtype)
            check_same_len(X_tensor, target_tensor)
        else:
            target_tensor = None

    # Ensure proper dimensions
    if X_tensor.ndim == 1:
        X_tensor = X_tensor.unsqueeze(1)

    if target_tensor is not None and target_tensor.ndim == 1:
        target_tensor = target_tensor.unsqueeze(1)

    if weights_tensor.ndim == 2:
        # Check string literals first (common case)
        if reduce == "prod":
            weights_tensor = torch.prod(weights_tensor, dim=1)
            # Clamp to prevent numerical issues in downstream calculations
            weights_tensor = torch.clamp(
                weights_tensor, min=torch.finfo(weights_tensor.dtype).min
            )
        elif reduce == "sum":
            weights_tensor = torch.sum(weights_tensor, dim=1)
        elif callable(reduce):
            try:
                reduced = reduce(weights_tensor)
            except Exception as e:
                raise RuntimeError(
                    f"Custom `reduce` callable failed with error: {e}."
                ) from e

            if not isinstance(reduced, torch.Tensor):
                raise TypeError(
                    f"Custom `reduce` callable must return torch.Tensor, "
                    f"got {type(reduced)}."
                )

            if reduced.ndim != 1:
                raise ValueError(
                    f"Custom `reduce` callable must return 1D tensor, "
                    f"got shape {reduced.shape}."
                )

            if len(reduced) != len(weights_tensor):
                raise ValueError(
                    f"Custom `reduce` callable output length ({len(reduced)}) "
                    f"must match input length ({len(weights_tensor)})."
                )

            weights_tensor = reduced
        else:
            raise ValueError(
                f"Invalid `reduce={reduce}`, choose from 'prod', 'sum', or provide "
                f"a callable."
            )

    return {"data": X_tensor, "weights": weights_tensor, "target": target_tensor}