oscura 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

oscura/loaders/mmap_loader.py

@@ -0,0 +1,583 @@
+"""Memory-mapped file loader for huge waveform files.
+
+This module provides efficient memory-mapped loading for GB+ files that cannot
+fit in RAM. Unlike eager loading, memory-mapped arrays don't load the entire
+file into memory but access it in chunks on-demand via the OS page cache.
+
+Key features:
+- Zero-copy data access via numpy.memmap
+- Chunked iteration for processing huge files
+- Integration with existing TraceKit loader infrastructure
+- Support for common binary formats (raw, NPY, structured)
+- Automatic fallback to regular loading for small files
+
+Example:
+    >>> from oscura.loaders.mmap_loader import load_mmap
+    >>> # Load 10 GB file without loading all data to RAM
+    >>> trace = load_mmap("huge_trace.npy", sample_rate=1e9)
+    >>> print(f"Length: {len(trace)} samples")
+    >>>
+    >>> # Process in chunks to avoid OOM
+    >>> for chunk in trace.iter_chunks(chunk_size=1_000_000):
+    ...     result = analyze_chunk(chunk)
+
+References:
+    Performance optimization for huge files (>1 GB)
+    API-017: Lazy Loading for Huge Files
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from oscura.core.exceptions import LoaderError
+from oscura.core.types import TraceMetadata
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+    from numpy.typing import DTypeLike, NDArray
+
+
+# File size threshold for automatic mmap suggestion (1 GB)
+MMAP_THRESHOLD = 1024 * 1024 * 1024
+
+
+class MmapWaveformTrace:
+    """Memory-mapped waveform trace for huge files.
+
+    Provides lazy access to waveform data via memory mapping. Data is not
+    loaded into RAM but accessed directly from disk through the OS page cache.
+
+    This allows working with files larger than available RAM without OOM errors.
+
+    Attributes:
+        file_path: Path to the memory-mapped file.
+        sample_rate: Sample rate in Hz.
+        length: Number of samples in the trace.
+        dtype: NumPy dtype of the samples.
+        metadata: Additional trace metadata.
+
+    Example:
+        >>> trace = MmapWaveformTrace(
+        ...     file_path="huge_trace.bin",
+        ...     sample_rate=1e9,
+        ...     length=10_000_000_000,
+        ...     dtype=np.float32
+        ... )
+        >>> # Access subset without loading entire file
+        >>> subset = trace[1000:2000]
+        >>> # Process in chunks
+        >>> for chunk in trace.iter_chunks(chunk_size=1_000_000):
+        ...     process(chunk)
+    """
+
+    def __init__(
+        self,
+        file_path: str | Path,
+        sample_rate: float,
+        length: int,
+        *,
+        dtype: DTypeLike = np.float64,
+        offset: int = 0,
+        metadata: dict[str, Any] | None = None,
+        mode: str = "r",
+    ) -> None:
+        """Initialize memory-mapped trace.
+
+        Args:
+            file_path: Path to binary data file.
+            sample_rate: Sample rate in Hz.
+            length: Number of samples.
+            dtype: Data type of samples.
+            offset: Byte offset to start of data in file.
+            metadata: Additional metadata dictionary.
+            mode: File access mode ('r' for read-only, 'r+' for read-write).
+
+        Raises:
+            LoaderError: If file not found or invalid parameters.
+
+        Example:
+            >>> trace = MmapWaveformTrace(
+            ...     file_path="trace.f32",
+            ...     sample_rate=1e9,
+            ...     length=1_000_000_000,
+            ...     dtype=np.float32
+            ... )
+        """
+        self._file_path = Path(file_path)
+        self._sample_rate = float(sample_rate)
+        self._length = int(length)
+        self._dtype = np.dtype(dtype)
+        self._offset = int(offset)
+        self._metadata = metadata or {}
+        self._mode = mode
+
+        # Memory-mapped array - created on first access
+        self._memmap: np.memmap[Any, np.dtype[Any]] | None = None
+
+        # Validate inputs
+        if self._sample_rate <= 0:
+            raise LoaderError(f"sample_rate must be positive, got {self._sample_rate}")
+        if self._length < 0:
+            raise LoaderError(f"length must be non-negative, got {self._length}")
+        if self._offset < 0:
+            raise LoaderError(f"offset must be non-negative, got {self._offset}")
+
+        # Verify file exists
+        if not self._file_path.exists():
+            raise LoaderError(f"File not found: {self._file_path}")
+
+        # Verify file size
+        expected_size = self._offset + self._length * self._dtype.itemsize
+        actual_size = self._file_path.stat().st_size
+        if actual_size < expected_size:
+            raise LoaderError(
+                f"File too small for requested data. "
+                f"Expected at least {expected_size} bytes, got {actual_size} bytes",
+                file_path=str(self._file_path),
+            )
+
+    @property
+    def sample_rate(self) -> float:
+        """Sample rate in Hz."""
+        return self._sample_rate
+
+    @property
+    def length(self) -> int:
+        """Number of samples."""
+        return self._length
+
+    @property
+    def duration(self) -> float:
+        """Duration in seconds."""
+        return self._length / self._sample_rate
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        """Metadata dictionary."""
+        return self._metadata
+
+    @property
+    def dtype(self) -> np.dtype[Any]:
+        """Data type of samples."""
+        return self._dtype
+
+    @property
+    def file_path(self) -> Path:
+        """Path to memory-mapped file."""
+        return self._file_path
+
+    @property
+    def data(self) -> np.memmap[Any, np.dtype[Any]]:
+        """Memory-mapped data array.
+
+        Returns a numpy.memmap object that behaves like a numpy array
+        but doesn't load data into memory until accessed.
+
+        Returns:
+            Memory-mapped numpy array.
+
+        Example:
+            >>> trace = load_mmap("huge.npy", sample_rate=1e9)
+            >>> data = trace.data  # No data loaded yet
+            >>> subset = data[1000:2000]  # Only this range loaded
+        """
+        if self._memmap is None:
+            self._memmap = np.memmap(  # type: ignore[call-overload]
+                str(self._file_path),
+                dtype=self._dtype,
+                mode=self._mode,
+                offset=self._offset,
+                shape=(self._length,),
+            )
+        return self._memmap
+
+    @property
+    def time_vector(self) -> NDArray[np.float64]:
+        """Time vector in seconds.
+
+        Note: For huge traces, this can consume significant memory.
+        Consider using time values on-demand instead.
+
+        Returns:
+            Array of time values corresponding to samples.
+
+        Example:
+            >>> # For huge traces, avoid materializing full time vector
+            >>> # Instead compute on-demand:
+            >>> t_start = 0
+            >>> t_end = trace.length / trace.sample_rate
+        """
+        return np.arange(self._length, dtype=np.float64) / self._sample_rate
+
+    def __getitem__(self, key: int | slice) -> float | NDArray[np.float64]:
+        """Slice the memory-mapped trace.
+
+        Supports both integer indexing and slicing. Only the requested
+        portion is loaded from disk.
+
+        Args:
+            key: Index or slice.
+
+        Returns:
+            Single sample (float) or array slice.
+
+        Raises:
+            TypeError: If key is not int or slice.
+
+        Example:
+            >>> sample = trace[1000]  # Load single sample
+            >>> chunk = trace[1000:2000]  # Load 1000 samples
+            >>> every_10th = trace[::10]  # Load decimated data
+        """
+        if isinstance(key, (int, slice)):
+            return self.data[key]
+        else:
+            raise TypeError(f"Indices must be int or slice, not {type(key).__name__}")
+
+    def __len__(self) -> int:
+        """Number of samples."""
+        return self._length
+
+    def iter_chunks(
+        self, chunk_size: int = 1_000_000, overlap: int = 0
+    ) -> Iterator[NDArray[np.float64]]:
+        """Iterate over trace in chunks.
+
+        Yields consecutive chunks of data, optionally with overlap between
+        chunks. This is efficient for processing huge files that don't fit
+        in memory.
+
+        Args:
+            chunk_size: Number of samples per chunk.
+            overlap: Number of samples to overlap between chunks.
+
+        Yields:
+            Numpy arrays of chunk_size (or smaller for last chunk).
+
+        Raises:
+            ValueError: If chunk_size or overlap invalid.
+
+        Example:
+            >>> # Process 10 GB file in 1M sample chunks
+            >>> for chunk in trace.iter_chunks(chunk_size=1_000_000):
+            ...     result = compute_fft(chunk)
+            ...
+            >>> # With 50% overlap for windowed processing
+            >>> for chunk in trace.iter_chunks(chunk_size=2048, overlap=1024):
+            ...     spectrum = analyze_spectrum(chunk)
+        """
+        if chunk_size <= 0:
+            raise ValueError(f"chunk_size must be positive, got {chunk_size}")
+        if overlap < 0:
+            raise ValueError(f"overlap must be non-negative, got {overlap}")
+        if overlap >= chunk_size:
+            raise ValueError(f"overlap ({overlap}) must be less than chunk_size ({chunk_size})")
+
+        data = self.data
+        step = chunk_size - overlap
+
+        for start in range(0, self._length, step):
+            end = min(start + chunk_size, self._length)
+            # Convert memmap slice to regular array to avoid keeping file handle open
+            yield np.asarray(data[start:end], dtype=np.float64)
+
+    def to_eager(self) -> Any:
+        """Convert to eager WaveformTrace by loading all data.
+
+        WARNING: This loads the entire file into memory. Only use this
+        if you're sure the data fits in RAM.
+
+        Returns:
+            WaveformTrace with data loaded in memory.
+
+        Example:
+            >>> # Only convert to eager if file is small enough
+            >>> if trace.length < 10_000_000:
+            ...     eager_trace = trace.to_eager()
+        """
+        from oscura.core.types import WaveformTrace
+
+        # Load all data into memory
+        data = np.asarray(self.data, dtype=np.float64)
+
+        metadata = TraceMetadata(
+            sample_rate=self._sample_rate,
+            source_file=str(self._file_path),
+            **self._metadata,  # type: ignore[arg-type]
+        )
+
+        return WaveformTrace(data=data, metadata=metadata)
+
+    def close(self) -> None:
+        """Close memory-mapped file handle.
+
+        Should be called when done with the trace to free resources.
+        The trace cannot be used after closing.
+
+        Example:
+            >>> trace = load_mmap("huge.npy", sample_rate=1e9)
+            >>> # ... use trace ...
+            >>> trace.close()
+        """
+        if self._memmap is not None:
+            # Delete reference to allow garbage collection
+            del self._memmap
+            self._memmap = None
+
+    def __del__(self) -> None:
+        """Cleanup memory map on deletion."""
+        self.close()
+
+    def __repr__(self) -> str:
+        """String representation."""
+        size_mb = (self._length * self._dtype.itemsize) / (1024 * 1024)
+        return (
+            f"MmapWaveformTrace("
+            f"file={self._file_path.name}, "
+            f"sample_rate={self._sample_rate:.2e} Hz, "
+            f"length={self._length:,} samples, "
+            f"size={size_mb:.1f} MB, "
+            f"dtype={self._dtype})"
+        )
+
+    def __enter__(self) -> MmapWaveformTrace:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Context manager exit - close the file."""
+        self.close()
+
+
+def load_mmap(
+    file_path: str | PathLike[str],
+    sample_rate: float | None = None,
+    *,
+    dtype: DTypeLike | None = None,
+    offset: int = 0,
+    length: int | None = None,
+    mode: str = "r",
+    **metadata: Any,
+) -> MmapWaveformTrace:
+    """Load waveform file with memory mapping.
+
+    Creates a memory-mapped trace that doesn't load data into RAM.
+    Supports .npy files (auto-detects format) and raw binary files.
+
+    Args:
+        file_path: Path to waveform file (.npy or raw binary).
+        sample_rate: Sample rate in Hz (required for raw files, optional for .npy).
+        dtype: Data type (required for raw files, auto-detected for .npy).
+        offset: Byte offset to data start (auto-computed for .npy).
+        length: Number of samples (auto-computed if possible).
+        mode: File access mode ('r' for read-only, 'r+' for read-write).
+        **metadata: Additional metadata to store.
+
+    Returns:
+        MmapWaveformTrace for memory-mapped access.
+
+    Raises:
+        LoaderError: If file not found or parameters invalid.
+
+    Example:
+        >>> # Load NumPy file (auto-detects format)
+        >>> trace = load_mmap("huge_trace.npy", sample_rate=1e9)
+        >>>
+        >>> # Load raw binary file
+        >>> trace = load_mmap(
+        ...     "data.f32",
+        ...     sample_rate=1e9,
+        ...     dtype=np.float32,
+        ...     length=1_000_000_000
+        ... )
+        >>>
+        >>> # Use context manager
+        >>> with load_mmap("huge.npy", sample_rate=1e9) as trace:
+        ...     for chunk in trace.iter_chunks(chunk_size=1_000_000):
+        ...         process(chunk)
+
+    References:
+        API-017: Lazy Loading for Huge Files
+    """
+    file_path = Path(file_path)
+
+    if not file_path.exists():
+        raise LoaderError(f"File not found: {file_path}")
+
+    suffix = file_path.suffix.lower()
+
+    # Handle .npy files with automatic format detection
+    if suffix == ".npy":
+        return _load_npy_mmap(file_path, sample_rate, mode, metadata)
+
+    # Handle .npz files (not directly memory-mappable, but can extract)
+    elif suffix == ".npz":
+        raise LoaderError(
+            "NPZ files cannot be directly memory-mapped. "
+            "Extract the array first using np.load() and save as .npy",
+            file_path=str(file_path),
+            fix_hint="Use: np.save('array.npy', np.load('file.npz')['array'])",
+        )
+
+    # Handle raw binary files
+    else:
+        if dtype is None:
+            raise LoaderError(
+                "dtype is required for raw binary files",
+                file_path=str(file_path),
+                fix_hint="Specify dtype, e.g., dtype=np.float32",
+            )
+        if sample_rate is None:
+            raise LoaderError(
+                "sample_rate is required for raw binary files",
+                file_path=str(file_path),
+            )
+
+        # Compute length from file size if not provided
+        dtype_np = np.dtype(dtype)
+        if length is None:
+            file_size = file_path.stat().st_size - offset
+            length = file_size // dtype_np.itemsize
+
+        return MmapWaveformTrace(
+            file_path=file_path,
+            sample_rate=sample_rate,
+            length=length,
+            dtype=dtype_np,
+            offset=offset,
+            metadata=metadata,
+            mode=mode,
+        )
+
+
+def _load_npy_mmap(
+    file_path: Path,
+    sample_rate: float | None,
+    mode: str,
+    metadata: dict[str, Any],
+) -> MmapWaveformTrace:
+    """Load NumPy .npy file with memory mapping.
+
+    Reads the .npy header to extract dtype, shape, and data offset,
+    then creates a memory-mapped array.
+
+    Args:
+        file_path: Path to .npy file.
+        sample_rate: Sample rate in Hz (required).
+        mode: File access mode.
+        metadata: Additional metadata.
+
+    Returns:
+        MmapWaveformTrace for the .npy file.
+
+    Raises:
+        LoaderError: If sample_rate not provided or file invalid.
+    """
+    if sample_rate is None:
+        raise LoaderError(
+            "sample_rate is required for .npy files",
+            file_path=str(file_path),
+            fix_hint="Specify sample_rate, e.g., sample_rate=1e9",
+        )
+
+    try:
+        # Read NumPy header without loading data
+        with open(file_path, "rb") as f:
+            import numpy.lib.format as npf
+
+            # Read header
+            version = npf.read_magic(f)  # type: ignore[no-untyped-call]
+
+            if version == (1, 0):
+                shape, fortran_order, dtype = npf.read_array_header_1_0(f)  # type: ignore[no-untyped-call]
+            elif version == (2, 0):
+                shape, fortran_order, dtype = npf.read_array_header_2_0(f)  # type: ignore[no-untyped-call]
+            else:
+                raise LoaderError(
+                    f"Unsupported NPY version: {version}",
+                    file_path=str(file_path),
+                )
+
+            # Get data offset
+            offset = f.tell()
+
+        # Validate shape
+        if not isinstance(shape, tuple):
+            raise LoaderError(
+                f"Invalid .npy shape: {shape}",
+                file_path=str(file_path),
+            )
+
+        if len(shape) != 1:
+            raise LoaderError(
+                f"Expected 1D array, got shape {shape}",
+                file_path=str(file_path),
+                fix_hint="Reshape to 1D or extract specific column",
+            )
+
+        length = shape[0]
+
+        if fortran_order:
+            raise LoaderError(
+                "Fortran-ordered arrays not supported for memory mapping",
+                file_path=str(file_path),
+                fix_hint="Resave array in C order: np.save('file.npy', arr, allow_pickle=False)",
+            )
+
+        return MmapWaveformTrace(
+            file_path=file_path,
+            sample_rate=sample_rate,
+            length=length,
+            dtype=dtype,
+            offset=offset,
+            metadata=metadata,
+            mode=mode,
+        )
+
+    except Exception as e:
+        if isinstance(e, LoaderError):
+            raise
+        raise LoaderError(
+            f"Failed to load .npy file: {e}",
+            file_path=str(file_path),
+        ) from e
+
+
+def should_use_mmap(file_path: str | PathLike[str], threshold: int = MMAP_THRESHOLD) -> bool:
+    """Check if file should use memory mapping.
+
+    Recommends memory mapping for files larger than threshold (default 1 GB).
+
+    Args:
+        file_path: Path to file.
+        threshold: Size threshold in bytes (default: 1 GB).
+
+    Returns:
+        True if file size >= threshold, False otherwise.
+
+    Example:
+        >>> if should_use_mmap("huge_trace.npy"):
+        ...     trace = load_mmap("huge_trace.npy", sample_rate=1e9)
+        ... else:
+        ...     trace = load("huge_trace.npy", sample_rate=1e9)
+    """
+    file_path = Path(file_path)
+    if not file_path.exists():
+        return False
+
+    file_size = file_path.stat().st_size
+    return file_size >= threshold
+
+
+__all__ = [
+    "MMAP_THRESHOLD",
+    "MmapWaveformTrace",
+    "load_mmap",
+    "should_use_mmap",
+]