Data Dataset Mmap

dataset_mmap

CLASS	DESCRIPTION
DatasetWeightedMemoryMapped	Memory-mapped weighted dataset for large datasets with on-demand loading.

Classes

DatasetWeightedMemoryMapped(data_path: str | Path, weights_path: str | Path | None = None, target_path: str | Path | None = None, mode: Literal['r', 'r+', 'c'] = 'r', labels: list | None = None, dtype: torch.dtype = torch.float32, transform: Callable | None = None, cache_size: int = 0)

Bases: Dataset

Memory-mapped weighted dataset for large datasets with on-demand loading.

Uses numpy memory-mapped arrays to handle datasets that do not fit in RAM. Data is loaded on-demand from disk with minimal memory footprint, enabling efficient handling of large-scale datasets. API-compatible with DatasetWeighted for seamless integration with data loaders and modules.

PARAMETER	DESCRIPTION
data_path	Path to memory-mapped data file (.npy format). TYPE: str | Path
weights_path	Path to memory-mapped weights file. If None, generates unit weights. TYPE: str | Path | None, optional, by default None DEFAULT: None
target_path	Path to memory-mapped target file. If None, dataset is unsupervised. TYPE: str | Path | None, optional, by default None DEFAULT: None
mode	File opening mode: "r" (read-only), "r+" (read-write), "c" (copy-on-write). TYPE: Literal["r", "r+", "c"], optional, by default "r" DEFAULT: 'r'
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 conversion. 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
cache_size	Number of most recent samples to cache in memory. 0 disables caching. Note: Cache stores full tensor copies. For large feature dimensions, memory usage can be significant. Use get_cache_memory_usage() to monitor cache memory consumption. TYPE: int, optional, by default 0 DEFAULT: 0

Properties

data : np.memmap Memory-mapped data array.

weights : np.memmap | np.ndarray Memory-mapped weights array or generated unit weights.

target : np.memmap | None Memory-mapped target array, if provided.

labels : list[str] Feature labels.

shape : tuple[int, ...] Shape of the data array.

Examples:

Create and use memory-mapped dataset:

>>> import numpy as np
>>> import torch
>>> from spectre.data import DatasetWeightedMemoryMapped
>>> # Create large dataset on disk
>>> X = np.random.randn(1_000_000, 50)
>>> np.save("large_data.npy", X)
>>> # Load as memory-mapped dataset
>>> dataset = DatasetWeightedMemoryMapped("large_data.npy")
>>> len(dataset)
1000000
>>> batch = dataset[0]
>>> batch.data.shape
torch.Size([50])

With weights and targets:

>>> weights = np.random.rand(1_000_000)
>>> np.save("weights.npy", weights)
>>> targets = np.random.randint(0, 10, (1_000_000, 1))
>>> np.save("targets.npy", targets)
>>> dataset = DatasetWeightedMemoryMapped(
...     data_path="large_data.npy",
...     weights_path="weights.npy",
...     target_path="targets.npy",
... )
>>> batch = dataset[0]
>>> batch.data.shape, batch.weights.shape, batch.target.shape
(torch.Size([50]), torch.Size([]), torch.Size([1]))

With caching for frequently accessed data:

>>> dataset = DatasetWeightedMemoryMapped("large_data.npy", cache_size=1000)
>>> # First access loads from disk
>>> batch1 = dataset[0]
>>> # Second access uses cache (faster)
>>> batch2 = dataset[0]

Efficient validation and statistics for large datasets using sampling:

>>> # For very large datasets, validate a sample instead of the entire dataset
>>> dataset = DatasetWeightedMemoryMapped("large_data.npy")
>>> validation_results = dataset.validate(sample_size=10000)
>>> validation_results["has_nan_data"]
False
>>> # Quick NaN check on sample
>>> has_nan = dataset.check_nan(sample_size=10000)
>>> # Compute statistics on sample
>>> stats = dataset.get_statistics(sample_size=10000)
>>> stats["mean"].shape
torch.Size([50])

Notes

• Memory-mapped files must be in .npy format created by numpy.save()
• Data is not loaded into memory until accessed via __getitem__()
• Use caching for datasets with repeated access patterns
• Copy-on-write mode ("c") is useful for read-mostly with occasional modifications
• For very large datasets, use sample_size parameter in validation methods

METHOD	DESCRIPTION
clear_cache	Clear the sample cache.
has_target	Check if the dataset has target values.
get_cache_memory_usage	Get current cache memory usage.
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.

ATTRIBUTE	DESCRIPTION
shape	Shape of the data array. TYPE: tuple[int, ...]

Source code in spectre/data/dataset_mmap.py

def __init__(
    self,
    data_path: str | Path,
    weights_path: str | Path | None = None,
    target_path: str | Path | None = None,
    mode: Literal["r", "r+", "c"] = "r",
    labels: list | None = None,
    dtype: torch.dtype = torch.float32,
    transform: Callable | None = None,
    cache_size: int = 0,
) -> None:
    if mode not in ["r", "r+", "c"]:
        raise ValueError(f"Invalid mode '{mode}'. Choose from 'r', 'r+', 'c'.")

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

    # Load memory-mapped data
    data_path = Path(data_path)
    if not data_path.exists():
        raise FileNotFoundError(f"Data file not found: {data_path}")

    self.data = np.load(data_path, mmap_mode=mode)

    if self.data.ndim == 1:
        self.data = self.data.reshape(-1, 1)

    n_samples = self.data.shape[0]

    if weights_path is not None:
        weights_path = Path(weights_path)
        if not weights_path.exists():
            raise FileNotFoundError(f"Weights file not found: {weights_path}")
        self.weights = np.load(weights_path, mmap_mode=mode)
        check_same_len(self.data, self.weights)
    else:
        self.weights = np.ones(n_samples, dtype=np.float32)

    if self.weights.ndim > 1 and self.weights.shape[1] > 1:
        raise ValueError(
            "Multi-column weights not supported. Reduce weights to 1D before saving."
        )
    if self.weights.ndim > 1:
        self.weights = self.weights.squeeze()

    if target_path is not None:
        target_path = Path(target_path)
        if not target_path.exists():
            raise FileNotFoundError(f"Target file not found: {target_path}")
        self.target = np.load(target_path, mmap_mode=mode)
        check_same_len(self.data, self.target)

        if self.target.ndim == 1:
            self.target = self.target.reshape(-1, 1)
    else:
        self.target = None

    if labels is not None:
        if self.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(self.data.shape[1])]

    self.labels = labels
    self.dtype = dtype
    self.transform = transform

    # Setup caching
    if not isinstance(cache_size, int) or cache_size < 0:
        raise ValueError("`cache_size` must be a non-negative integer.")
    self.cache_size = cache_size
    self._cache: dict[int, DataBatch] = {}
    self._cache_order: list[int] = []

Attributes

shape: tuple[int, ...] property

Shape of the data array.

Functions

clear_cache() -> None

Clear the sample cache.

Source code in spectre/data/dataset_mmap.py

def clear_cache(self) -> None:
    """Clear the sample cache."""
    self._cache.clear()
    self._cache_order.clear()

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_mmap.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

get_cache_memory_usage() -> dict[str, float]

Get current cache memory usage.

RETURNS	DESCRIPTION
dict[str, float]	Dictionary with "cache_size" (number of cached items) and "estimated_mb" (estimated memory usage in megabytes).

Examples:

>>> dataset = DatasetWeightedMemoryMapped("data.npy", cache_size=1000)
>>> # Access some data
>>> for i in range(100):
...     _ = dataset[i]
>>> usage = dataset.get_cache_memory_usage()
>>> print(f"Cached {usage['cache_size']} items")
>>> print(f"Using ~{usage['estimated_mb']:.2f} MB")

Source code in spectre/data/dataset_mmap.py

def get_cache_memory_usage(self) -> dict[str, float]:
    """
    Get current cache memory usage.

    Returns
    -------
    dict[str, float]
        Dictionary with "cache_size" (number of cached items) and
        "estimated_mb" (estimated memory usage in megabytes).

    Examples
    --------
    >>> dataset = DatasetWeightedMemoryMapped("data.npy", cache_size=1000)
    >>> # Access some data
    >>> for i in range(100):
    ...     _ = dataset[i]
    >>> usage = dataset.get_cache_memory_usage()
    >>> print(f"Cached {usage['cache_size']} items")
    >>> print(f"Using ~{usage['estimated_mb']:.2f} MB")
    """
    if not self._cache or len(self._cache) == 0:
        return {"cache_size": 0, "estimated_mb": 0.0}

    # Estimate from first cached item
    sample = next(iter(self._cache.values()))
    bytes_per_item = (
        sample.data.numel() * sample.data.element_size()
        + sample.weights.numel() * sample.weights.element_size()
    )
    if sample.target is not None:
        bytes_per_item += sample.target.numel() * sample.target.element_size()

    total_mb = (len(self._cache) * bytes_per_item) / (1024 * 1024)
    return {"cache_size": len(self._cache), "estimated_mb": total_mb}

validate(sample_size: int | None = None) -> dict[str, bool]

Validate dataset for common data quality issues.

Checks for NaN values, infinite values, and zero weights in memory-mapped arrays. For large datasets, use sample_size to validate a random subset.

PARAMETER	DESCRIPTION
sample_size	If provided, validate only a random sample of this size instead of the entire dataset. Useful for large datasets where full validation would be too slow. TYPE: int | None, optional, by default None DEFAULT: None

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_mmap.py

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

    Checks for NaN values, infinite values, and zero weights in memory-mapped
    arrays. For large datasets, use `sample_size` to validate a random subset.

    Parameters
    ----------
    sample_size : int | None, optional, by default None
        If provided, validate only a random sample of this size instead of
        the entire dataset. Useful for large datasets where full validation
        would be too slow.

    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`.
    """
    indices = self._get_sample_indices(sample_size)

    if indices is not None:
        data_sample = self.data[indices]
        weights_sample = self.weights[indices]
        target_sample = self.target[indices] if self.has_target() else None
    else:
        data_sample = self.data
        weights_sample = self.weights
        target_sample = self.target

    # Check data for NaN/inf
    data_nan = np.isnan(data_sample).any()
    data_inf = np.isinf(data_sample).any()

    # Check weights for NaN/inf/negative/zero
    weights_nan = np.isnan(weights_sample).any()
    weights_inf = np.isinf(weights_sample).any()
    weights_negative = (weights_sample < 0).any()
    weights_zero = (weights_sample == 0).any()

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

    if self.has_target():
        target_nan = np.isnan(target_sample).any()
        target_inf = np.isinf(target_sample).any()
        results["has_nan_target"] = bool(target_nan)
        results["has_inf_target"] = bool(target_inf)
    else:
        results["has_nan_target"] = False
        results["has_inf_target"] = False

    return results

get_statistics(sample_size: int | None = None) -> dict[str, torch.Tensor]

Compute basic statistics for the dataset.

Computes statistics on memory-mapped arrays. For large datasets, use sample_size to compute statistics on a random subset.

PARAMETER	DESCRIPTION
sample_size	If provided, compute statistics only on a random sample of this size instead of the entire dataset. Useful for large datasets where computing statistics on the full dataset would be too slow. TYPE: int | None, optional, by default None DEFAULT: None

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_mmap.py

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

    Computes statistics on memory-mapped arrays. For large datasets, use
    `sample_size` to compute statistics on a random subset.

    Parameters
    ----------
    sample_size : int | None, optional, by default None
        If provided, compute statistics only on a random sample of this size
        instead of the entire dataset. Useful for large datasets where computing
        statistics on the full dataset would be too slow.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with statistics: `mean`, `std`, `min`, `max` for data,
        and `weights_sum`, `weights_mean` for weights.
    """
    indices = self._get_sample_indices(sample_size)

    if indices is not None:
        data_sample = self.data[indices]
        weights_sample = self.weights[indices]
        target_sample = self.target[indices] if self.has_target() else None
    else:
        data_sample = self.data
        weights_sample = self.weights
        target_sample = self.target

    # Compute stats on numpy arrays, then convert to tensors
    stats = {
        "mean": torch.from_numpy(np.mean(data_sample, axis=0)).to(self.dtype),
        "std": torch.from_numpy(np.std(data_sample, axis=0)).to(self.dtype),
        "min": torch.from_numpy(np.min(data_sample, axis=0)).to(self.dtype),
        "max": torch.from_numpy(np.max(data_sample, axis=0)).to(self.dtype),
        "weights_sum": torch.tensor(np.sum(weights_sample), dtype=self.dtype),
        "weights_mean": torch.tensor(np.mean(weights_sample), dtype=self.dtype),
    }

    if self.has_target():
        stats["target_mean"] = torch.from_numpy(np.mean(target_sample, axis=0)).to(
            self.dtype
        )
        stats["target_std"] = torch.from_numpy(np.std(target_sample, axis=0)).to(
            self.dtype
        )

    return stats

check_nan(sample_size: int | None = None) -> bool

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

For large datasets, use sample_size to check a random subset instead of the entire dataset.

PARAMETER	DESCRIPTION
sample_size	If provided, check only a random sample of this size instead of the entire dataset. TYPE: int | None, optional, by default None DEFAULT: None

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

Source code in spectre/data/dataset_mmap.py

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

    For large datasets, use `sample_size` to check a random subset instead
    of the entire dataset.

    Parameters
    ----------
    sample_size : int | None, optional, by default None
        If provided, check only a random sample of this size instead of
        the entire dataset.

    Returns
    -------
    bool
        True if NaN values are found, False otherwise.
    """
    indices = self._get_sample_indices(sample_size)

    if indices is not None:
        data_sample = self.data[indices]
        weights_sample = self.weights[indices]
        target_sample = self.target[indices] if self.has_target() else None
    else:
        data_sample = self.data
        weights_sample = self.weights
        target_sample = self.target

    has_nan = np.isnan(data_sample).any() or np.isnan(weights_sample).any()

    if self.has_target():
        has_nan = has_nan or np.isnan(target_sample).any()

    return bool(has_nan)

check_inf(sample_size: int | None = None) -> bool

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

For large datasets, use sample_size to check a random subset instead of the entire dataset.

PARAMETER	DESCRIPTION
sample_size	If provided, check only a random sample of this size instead of the entire dataset. TYPE: int | None, optional, by default None DEFAULT: None

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

Source code in spectre/data/dataset_mmap.py

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

    For large datasets, use `sample_size` to check a random subset instead
    of the entire dataset.

    Parameters
    ----------
    sample_size : int | None, optional, by default None
        If provided, check only a random sample of this size instead of
        the entire dataset.

    Returns
    -------
    bool
        True if infinite values are found, False otherwise.
    """
    indices = self._get_sample_indices(sample_size)

    if indices is not None:
        data_sample = self.data[indices]
        weights_sample = self.weights[indices]
        target_sample = self.target[indices] if self.has_target() else None
    else:
        data_sample = self.data
        weights_sample = self.weights
        target_sample = self.target

    has_inf = np.isinf(data_sample).any() or np.isinf(weights_sample).any()

    if self.has_target():
        has_inf = has_inf or np.isinf(target_sample).any()

    return bool(has_inf)

Functions