wavedl 1.2.0__py3-none-any.whl

wavedl/utils/data.py

@@ -0,0 +1,1220 @@
+"""
+Data Loading and Preprocessing Utilities
+=========================================
+
+Provides memory-efficient data loading for large-scale datasets with:
+- Memory-mapped file support for datasets exceeding RAM
+- DDP-safe data preparation with proper synchronization
+- Thread-safe DataLoader worker initialization
+- Multi-format support (NPZ, HDF5, MAT)
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+import gc
+import logging
+import os
+import pickle
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+import h5py
+import numpy as np
+import torch
+from accelerate import Accelerator
+from scipy.sparse import issparse
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from torch.utils.data import DataLoader, Dataset
+from tqdm.auto import tqdm
+
+
+# Optional scipy.io for MATLAB files
+try:
+    import scipy.io
+
+    HAS_SCIPY_IO = True
+except ImportError:
+    HAS_SCIPY_IO = False
+
+
+# ==============================================================================
+# DATA SOURCE ABSTRACTION
+# ==============================================================================
+
+# Supported key names for input/output arrays (priority order, pairwise aligned)
+INPUT_KEYS = ["input_train", "input_test", "X", "data", "inputs", "features", "x"]
+OUTPUT_KEYS = ["output_train", "output_test", "Y", "labels", "outputs", "targets", "y"]
+
+
+class LazyDataHandle:
+    """
+    Context manager wrapper for memory-mapped data handles.
+
+    Provides proper cleanup of file handles returned by load_mmap() methods.
+    Can be used either as a context manager (recommended) or with explicit close().
+
+    Usage:
+        # Context manager (recommended)
+        with source.load_mmap(path) as (inputs, outputs):
+            # Use inputs and outputs
+            pass  # File automatically closed
+
+        # Manual cleanup
+        handle = source.load_mmap(path)
+        inputs, outputs = handle.inputs, handle.outputs
+        # ... use data ...
+        handle.close()
+
+    Attributes:
+        inputs: Input data array/dataset
+        outputs: Output data array/dataset
+    """
+
+    def __init__(self, inputs, outputs, file_handle=None):
+        """
+        Initialize the handle.
+
+        Args:
+            inputs: Input array or lazy dataset
+            outputs: Output array or lazy dataset
+            file_handle: Optional file handle to close on cleanup
+        """
+        self.inputs = inputs
+        self.outputs = outputs
+        self._file = file_handle
+        self._closed = False
+
+    def __enter__(self):
+        """Return (inputs, outputs) tuple for unpacking."""
+        return self.inputs, self.outputs
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Close file handle on context exit."""
+        self.close()
+        return False  # Don't suppress exceptions
+
+    def close(self):
+        """
+        Close the underlying file handle.
+
+        Safe to call multiple times.
+        """
+        if self._closed:
+            return
+        self._closed = True
+
+        # Close the file handle if we have one
+        if self._file is not None:
+            try:
+                self._file.close()
+            except Exception:
+                pass
+            self._file = None
+
+        # Also close any _TransposedH5Dataset wrappers
+        for data in (self.inputs, self.outputs):
+            if hasattr(data, "close"):
+                try:
+                    data.close()
+                except Exception:
+                    pass
+
+    def __repr__(self) -> str:
+        status = "closed" if self._closed else "open"
+        return f"LazyDataHandle(status={status})"
+
+
+class DataSource(ABC):
+    """
+    Abstract base class for data loaders supporting multiple file formats.
+
+    Subclasses must implement the `load()` method to return input/output arrays,
+    and optionally `load_outputs_only()` for memory-efficient target loading.
+    """
+
+    @abstractmethod
+    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Load input and output arrays from a file.
+
+        Args:
+            path: Path to the data file
+
+        Returns:
+            Tuple of (inputs, outputs) as numpy arrays
+        """
+        pass
+
+    @abstractmethod
+    def load_outputs_only(self, path: str) -> np.ndarray:
+        """
+        Load only output/target arrays from a file (memory-efficient).
+
+        This avoids loading large input arrays when only targets are needed,
+        which is critical for HPC environments with memory constraints.
+
+        Args:
+            path: Path to the data file
+
+        Returns:
+            Output/target array
+        """
+        pass
+
+    @staticmethod
+    def detect_format(path: str) -> str:
+        """
+        Auto-detect file format from extension.
+
+        Args:
+            path: Path to data file
+
+        Returns:
+            Format string: 'npz', 'hdf5', or 'mat'
+        """
+        ext = Path(path).suffix.lower()
+        format_map = {
+            ".npz": "npz",
+            ".h5": "hdf5",
+            ".hdf5": "hdf5",
+            ".mat": "mat",
+        }
+        if ext not in format_map:
+            raise ValueError(
+                f"Unsupported file extension: '{ext}'. "
+                f"Supported formats: .npz, .h5, .hdf5, .mat"
+            )
+        return format_map[ext]
+
+    @staticmethod
+    def _find_key(available_keys: list[str], candidates: list[str]) -> str | None:
+        """Find first matching key from candidates in available keys."""
+        for key in candidates:
+            if key in available_keys:
+                return key
+        return None
+
+
+class NPZSource(DataSource):
+    """Load data from NumPy .npz archives."""
+
+    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
+        data = np.load(path, allow_pickle=True)
+        keys = list(data.keys())
+
+        input_key = self._find_key(keys, INPUT_KEYS)
+        output_key = self._find_key(keys, OUTPUT_KEYS)
+
+        if input_key is None or output_key is None:
+            raise KeyError(
+                f"NPZ must contain input and output arrays. "
+                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                f"Found: {keys}"
+            )
+
+        inp = data[input_key]
+        outp = data[output_key]
+
+        # Handle object arrays (e.g., sparse matrices stored as objects)
+        if inp.dtype == object:
+            inp = np.array([x.toarray() if hasattr(x, "toarray") else x for x in inp])
+
+        return inp, outp
+
+    def load_mmap(self, path: str) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Load data using memory-mapped mode for zero-copy access.
+
+        This allows processing large datasets without loading them entirely
+        into RAM. Critical for HPC environments with memory constraints.
+
+        Note: Returns memory-mapped arrays - do NOT modify them.
+        """
+        data = np.load(path, allow_pickle=True, mmap_mode="r")
+        keys = list(data.keys())
+
+        input_key = self._find_key(keys, INPUT_KEYS)
+        output_key = self._find_key(keys, OUTPUT_KEYS)
+
+        if input_key is None or output_key is None:
+            raise KeyError(
+                f"NPZ must contain input and output arrays. "
+                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                f"Found: {keys}"
+            )
+
+        inp = data[input_key]
+        outp = data[output_key]
+
+        return inp, outp
+
+    def load_outputs_only(self, path: str) -> np.ndarray:
+        """Load only targets from NPZ (avoids loading large input arrays)."""
+        data = np.load(path, allow_pickle=True)
+        keys = list(data.keys())
+
+        output_key = self._find_key(keys, OUTPUT_KEYS)
+        if output_key is None:
+            raise KeyError(
+                f"NPZ must contain output array. "
+                f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
+            )
+
+        return data[output_key]
+
+
+class HDF5Source(DataSource):
+    """Load data from HDF5 (.h5, .hdf5) files."""
+
+    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
+        with h5py.File(path, "r") as f:
+            keys = list(f.keys())
+
+            input_key = self._find_key(keys, INPUT_KEYS)
+            output_key = self._find_key(keys, OUTPUT_KEYS)
+
+            if input_key is None or output_key is None:
+                raise KeyError(
+                    f"HDF5 must contain input and output datasets. "
+                    f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                    f"Found: {keys}"
+                )
+
+            # Load into memory (HDF5 datasets are lazy by default)
+            inp = f[input_key][:]
+            outp = f[output_key][:]
+
+        return inp, outp
+
+    def load_mmap(self, path: str) -> LazyDataHandle:
+        """
+        Load HDF5 file with lazy/memory-mapped access.
+
+        Returns a LazyDataHandle that reads from disk on-demand,
+        avoiding loading the entire file into RAM.
+
+        Usage:
+            with source.load_mmap(path) as (inputs, outputs):
+                # Use inputs and outputs
+                pass  # File automatically closed
+        """
+        f = h5py.File(path, "r")  # Keep file open for lazy access
+        keys = list(f.keys())
+
+        input_key = self._find_key(keys, INPUT_KEYS)
+        output_key = self._find_key(keys, OUTPUT_KEYS)
+
+        if input_key is None or output_key is None:
+            f.close()
+            raise KeyError(
+                f"HDF5 must contain input and output datasets. "
+                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                f"Found: {keys}"
+            )
+
+        # Return wrapped handle for proper cleanup
+        return LazyDataHandle(f[input_key], f[output_key], file_handle=f)
+
+    def load_outputs_only(self, path: str) -> np.ndarray:
+        """Load only targets from HDF5 (avoids loading large input arrays)."""
+        with h5py.File(path, "r") as f:
+            keys = list(f.keys())
+
+            output_key = self._find_key(keys, OUTPUT_KEYS)
+            if output_key is None:
+                raise KeyError(
+                    f"HDF5 must contain output dataset. "
+                    f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
+                )
+
+            outp = f[output_key][:]
+
+        return outp
+
+
+class _TransposedH5Dataset:
+    """
+    Lazy transpose wrapper for h5py datasets.
+
+    MATLAB stores arrays in column-major (Fortran) order, while Python/NumPy
+    expects row-major (C) order. This wrapper provides a transposed view
+    without loading the entire dataset into memory.
+
+    Supports:
+        - len(): Returns the transposed first dimension
+        - []: Returns slices with automatic transpose
+        - shape: Returns the transposed shape
+        - dtype: Returns the underlying dtype
+
+    This is critical for MATSource.load_mmap() to return consistent axis
+    ordering with the eager loader (MATSource.load()).
+
+    IMPORTANT: Holds a strong reference to the h5py.File to prevent
+    premature garbage collection while datasets are in use.
+    """
+
+    def __init__(self, h5_dataset, file_handle=None):
+        """
+        Args:
+            h5_dataset: The h5py dataset to wrap
+            file_handle: Optional h5py.File reference to keep alive
+        """
+        self._dataset = h5_dataset
+        self._file = file_handle  # Keep file alive to prevent GC
+        # Transpose shape: MATLAB (cols, rows, ...) -> Python (rows, cols, ...)
+        self.shape = tuple(reversed(h5_dataset.shape))
+        self.dtype = h5_dataset.dtype
+
+        # Precompute transpose axis order for efficiency
+        # For shape (A, B, C) -> reversed (C, B, A), transpose axes are (2, 1, 0)
+        self._transpose_axes = tuple(range(len(h5_dataset.shape) - 1, -1, -1))
+
+    def __len__(self) -> int:
+        return self.shape[0]
+
+    def __getitem__(self, idx):
+        """
+        Fetch data with automatic full transpose.
+
+        Handles integer indexing, slices, and fancy indexing.
+        All operations return data with fully reversed axes to match .T behavior.
+        """
+        if isinstance(idx, int | np.integer):
+            # Single sample: index into last axis of h5py dataset (column-major)
+            # Result needs full transpose of remaining dimensions
+            data = self._dataset[..., idx]
+            if data.ndim == 0:
+                return data
+            elif data.ndim == 1:
+                return data  # 1D doesn't need transpose
+            else:
+                # Full transpose: reverse all axes
+                return np.transpose(data)
+
+        elif isinstance(idx, slice):
+            # Slice indexing: fetch from last axis, then fully transpose
+            start, stop, step = idx.indices(self.shape[0])
+            data = self._dataset[..., start:stop:step]
+
+            # Handle special case: 1D result (e.g., row vector)
+            if data.ndim == 1:
+                return data
+
+            # Full transpose: reverse ALL axes (not just moveaxis)
+            # This matches the behavior of .T on a numpy array
+            return np.transpose(data, axes=self._transpose_axes)
+
+        elif isinstance(idx, list | np.ndarray):
+            # Fancy indexing: load samples one at a time (h5py limitation)
+            # This is slower but necessary for compatibility
+            samples = [self[i] for i in idx]
+            return np.stack(samples, axis=0)
+
+        else:
+            raise TypeError(f"Unsupported index type: {type(idx)}")
+
+    def close(self):
+        """Close the underlying file handle if we own it."""
+        if self._file is not None:
+            try:
+                self._file.close()
+            except Exception:
+                pass
+
+
+class MATSource(DataSource):
+    """
+    Load data from MATLAB .mat files (v7.3+ only, which uses HDF5 format).
+
+    Note: MAT v7.3 files are HDF5 files under the hood, so we use h5py for
+    memory-efficient lazy loading. Save with: save('file.mat', '-v7.3')
+
+    Supports MATLAB sparse matrices (automatically converted to dense).
+
+    For older MAT files (v5/v7), convert to NPZ or save with -v7.3 flag.
+    """
+
+    @staticmethod
+    def _is_sparse_dataset(dataset) -> bool:
+        """Check if an HDF5 dataset/group represents a MATLAB sparse matrix."""
+        # MATLAB v7.3 stores sparse matrices as groups with 'data', 'ir', 'jc' keys
+        if hasattr(dataset, "keys"):
+            keys = set(dataset.keys())
+            return {"data", "ir", "jc"}.issubset(keys)
+        return False
+
+    @staticmethod
+    def _load_sparse_to_dense(group) -> np.ndarray:
+        """Convert MATLAB sparse matrix (CSC format in HDF5) to dense numpy array."""
+        from scipy.sparse import csc_matrix
+
+        data = np.array(group["data"])
+        ir = np.array(group["ir"])  # row indices
+        jc = np.array(group["jc"])  # column pointers
+
+        # Get shape from MATLAB attributes or infer
+        if "MATLAB_sparse" in group.attrs:
+            nrows = group.attrs["MATLAB_sparse"]
+        else:
+            nrows = ir.max() + 1 if len(ir) > 0 else 0
+        ncols = len(jc) - 1
+
+        sparse_mat = csc_matrix((data, ir, jc), shape=(nrows, ncols))
+        return sparse_mat.toarray()
+
+    def _load_dataset(self, f, key: str) -> np.ndarray:
+        """Load a dataset, handling sparse matrices automatically."""
+        dataset = f[key]
+
+        if self._is_sparse_dataset(dataset):
+            # Sparse matrix: convert to dense
+            arr = self._load_sparse_to_dense(dataset)
+        else:
+            # Regular dense array
+            arr = np.array(dataset)
+
+        # Transpose for MATLAB column-major -> Python row-major
+        return arr.T
+
+    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
+        """Load MAT v7.3 file using h5py."""
+        try:
+            with h5py.File(path, "r") as f:
+                keys = list(f.keys())
+
+                input_key = self._find_key(keys, INPUT_KEYS)
+                output_key = self._find_key(keys, OUTPUT_KEYS)
+
+                if input_key is None or output_key is None:
+                    raise KeyError(
+                        f"MAT file must contain input and output arrays. "
+                        f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                        f"Found: {keys}"
+                    )
+
+                # Load with sparse matrix support
+                inp = self._load_dataset(f, input_key)
+                outp = self._load_dataset(f, output_key)
+
+                # Handle 1D outputs that become (1, N) after transpose
+                if outp.ndim == 2 and outp.shape[0] == 1:
+                    outp = outp.T
+
+        except OSError as e:
+            raise ValueError(
+                f"Failed to load MAT file: {path}. "
+                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
+                f"Original error: {e}"
+            )
+
+        return inp, outp
+
+    def load_mmap(self, path: str) -> LazyDataHandle:
+        """
+        Load MAT v7.3 file with lazy/memory-mapped access.
+
+        Returns a LazyDataHandle that reads from disk on-demand,
+        avoiding loading the entire file into RAM.
+
+        Note: For sparse matrices, this will load and convert them.
+        For dense arrays, returns a transposed view wrapper for consistent axis ordering.
+
+        Usage:
+            with source.load_mmap(path) as (inputs, outputs):
+                # Use inputs and outputs
+                pass  # File automatically closed
+        """
+        try:
+            f = h5py.File(path, "r")  # Keep file open for lazy access
+            keys = list(f.keys())
+
+            input_key = self._find_key(keys, INPUT_KEYS)
+            output_key = self._find_key(keys, OUTPUT_KEYS)
+
+            if input_key is None or output_key is None:
+                f.close()
+                raise KeyError(
+                    f"MAT file must contain input and output arrays. "
+                    f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
+                    f"Found: {keys}"
+                )
+
+            # Check for sparse matrices - must load them eagerly
+            inp_dataset = f[input_key]
+            outp_dataset = f[output_key]
+
+            if self._is_sparse_dataset(inp_dataset):
+                inp = self._load_sparse_to_dense(inp_dataset).T
+            else:
+                # Wrap h5py dataset with transpose view for consistent axis order
+                # MATLAB stores column-major, Python expects row-major
+                # Pass file handle to keep it alive
+                inp = _TransposedH5Dataset(inp_dataset, file_handle=f)
+
+            if self._is_sparse_dataset(outp_dataset):
+                outp = self._load_sparse_to_dense(outp_dataset).T
+            else:
+                # Wrap h5py dataset with transpose view (shares same file handle)
+                outp = _TransposedH5Dataset(outp_dataset, file_handle=f)
+
+            # Return wrapped handle for proper cleanup
+            return LazyDataHandle(inp, outp, file_handle=f)
+
+        except OSError as e:
+            raise ValueError(
+                f"Failed to load MAT file: {path}. "
+                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
+                f"Original error: {e}"
+            )
+
+    def load_outputs_only(self, path: str) -> np.ndarray:
+        """Load only targets from MAT v7.3 file (avoids loading large input arrays)."""
+        try:
+            with h5py.File(path, "r") as f:
+                keys = list(f.keys())
+
+                output_key = self._find_key(keys, OUTPUT_KEYS)
+                if output_key is None:
+                    raise KeyError(
+                        f"MAT file must contain output array. "
+                        f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
+                    )
+
+                # Load with sparse matrix support
+                outp = self._load_dataset(f, output_key)
+
+                # Handle 1D outputs
+                if outp.ndim == 2 and outp.shape[0] == 1:
+                    outp = outp.T
+
+        except OSError as e:
+            raise ValueError(
+                f"Failed to load MAT file: {path}. "
+                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
+                f"Original error: {e}"
+            )
+
+        return outp
+
+
+def get_data_source(format: str) -> DataSource:
+    """
+    Factory function to get the appropriate DataSource for a format.
+
+    Args:
+        format: One of 'npz', 'hdf5', 'mat'
+
+    Returns:
+        DataSource instance
+    """
+    sources = {
+        "npz": NPZSource,
+        "hdf5": HDF5Source,
+        "mat": MATSource,
+    }
+
+    if format not in sources:
+        raise ValueError(
+            f"Unsupported format: {format}. Supported: {list(sources.keys())}"
+        )
+
+    return sources[format]()
+
+
+def load_training_data(
+    path: str, format: str = "auto"
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Load training data from file with automatic format detection.
+
+    Supports:
+        - NPZ: NumPy compressed archives (.npz)
+        - HDF5: Hierarchical Data Format (.h5, .hdf5)
+        - MAT: MATLAB files (.mat)
+
+    Flexible key detection supports: input_train/X/data and output_train/y/labels.
+
+    Args:
+        path: Path to data file
+        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)
+
+    Returns:
+        Tuple of (inputs, outputs) arrays
+    """
+    if format == "auto":
+        format = DataSource.detect_format(path)
+
+    source = get_data_source(format)
+    return source.load(path)
+
+
+def load_outputs_only(path: str, format: str = "auto") -> np.ndarray:
+    """
+    Load only output/target arrays from file (memory-efficient).
+
+    This function avoids loading large input arrays when only targets are needed,
+    which is critical for HPC environments with memory constraints during DDP.
+
+    Args:
+        path: Path to data file
+        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)
+
+    Returns:
+        Output/target array
+    """
+    if format == "auto":
+        format = DataSource.detect_format(path)
+
+    source = get_data_source(format)
+    return source.load_outputs_only(path)
+
+
+def load_test_data(
+    path: str,
+    format: str = "auto",
+    input_key: str | None = None,
+    output_key: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """
+    Load test/inference data and return PyTorch tensors ready for model input.
+
+    This is the unified data loading function for inference. It:
+    - Auto-detects file format from extension
+    - Handles custom key names for non-standard datasets
+    - Adds channel dimension if missing (dimension-agnostic)
+    - Returns None for targets if not present in file
+
+    Supports any input dimensionality:
+        - 1D: (N, L) → (N, 1, L)
+        - 2D: (N, H, W) → (N, 1, H, W)
+        - 3D: (N, D, H, W) → (N, 1, D, H, W)
+        - Already has channel: (N, C, ...) → unchanged
+
+    Args:
+        path: Path to data file (NPZ, HDF5, or MAT v7.3)
+        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)
+        input_key: Custom key for input data (overrides auto-detection)
+        output_key: Custom key for output data (overrides auto-detection)
+
+    Returns:
+        Tuple of:
+            - X: Input tensor with channel dimension (N, 1, *spatial_dims)
+            - y: Target tensor (N, T) or None if targets not present
+
+    Example:
+        >>> X, y = load_test_data("test_data.npz")
+        >>> X, y = load_test_data(
+        ...     "data.mat", input_key="waveforms", output_key="params"
+        ... )
+    """
+    if format == "auto":
+        format = DataSource.detect_format(path)
+
+    source = get_data_source(format)
+
+    # Build custom key lists if provided
+    if input_key:
+        custom_input_keys = [input_key] + INPUT_KEYS
+    else:
+        # Prioritize test keys for inference
+        custom_input_keys = ["input_test"] + [
+            k for k in INPUT_KEYS if k != "input_test"
+        ]
+
+    if output_key:
+        custom_output_keys = [output_key] + OUTPUT_KEYS
+    else:
+        custom_output_keys = ["output_test"] + [
+            k for k in OUTPUT_KEYS if k != "output_test"
+        ]
+
+    # Load data using appropriate source
+    try:
+        inp, outp = source.load(path)
+    except KeyError:
+        # Try with just inputs if outputs not found
+        if format == "npz":
+            data = np.load(path, allow_pickle=True)
+            keys = list(data.keys())
+            inp_key = DataSource._find_key(keys, custom_input_keys)
+            if inp_key is None:
+                raise KeyError(
+                    f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
+                )
+            inp = data[inp_key]
+            if inp.dtype == object:
+                inp = np.array(
+                    [x.toarray() if hasattr(x, "toarray") else x for x in inp]
+                )
+            out_key = DataSource._find_key(keys, custom_output_keys)
+            outp = data[out_key] if out_key else None
+        else:
+            raise
+
+    # Handle sparse matrices
+    if issparse(inp):
+        inp = inp.toarray()
+    if outp is not None and issparse(outp):
+        outp = outp.toarray()
+
+    # Convert to tensors
+    X = torch.tensor(np.asarray(inp), dtype=torch.float32)
+
+    if outp is not None:
+        y = torch.tensor(np.asarray(outp), dtype=torch.float32)
+        # Normalize target shape: (N,) → (N, 1)
+        if y.ndim == 1:
+            y = y.unsqueeze(1)
+    else:
+        y = None
+
+    # Add channel dimension if needed (dimension-agnostic)
+    # X.ndim == 2: 1D data (N, L) → (N, 1, L)
+    # X.ndim == 3: 2D data (N, H, W) → (N, 1, H, W)
+    # X.ndim == 4: Check if already has channel dim (C <= 16 heuristic)
+    if X.ndim == 2:
+        X = X.unsqueeze(1)  # 1D signal: (N, L) → (N, 1, L)
+    elif X.ndim == 3:
+        X = X.unsqueeze(1)  # 2D image: (N, H, W) → (N, 1, H, W)
+    elif X.ndim == 4:
+        # Could be 3D volume (N, D, H, W) or 2D with channel (N, C, H, W)
+        # Heuristic: if dim 1 is small (<=16), assume it's already a channel dim
+        if X.shape[1] > 16:
+            X = X.unsqueeze(1)  # 3D volume: (N, D, H, W) → (N, 1, D, H, W)
+    # X.ndim >= 5: assume channel dimension already exists
+
+    return X, y
+
+
+# ==============================================================================
+# WORKER INITIALIZATION
+# ==============================================================================
+def memmap_worker_init_fn(worker_id: int):
+    """
+    Worker initialization function for proper memmap handling in multi-worker DataLoader.
+
+    Each DataLoader worker process runs this function after forking. It:
+    1. Resets the memmap file handle to None, forcing each worker to open its own
+       read-only handle (prevents file descriptor sharing issues and race conditions)
+    2. Seeds numpy's random state per worker to ensure statistical diversity in
+       random augmentations (prevents all workers from applying identical "random"
+       transformations to their batches)
+
+    Args:
+        worker_id: Worker index (0 to num_workers-1), provided by DataLoader
+
+    Usage:
+        DataLoader(dataset, num_workers=8, worker_init_fn=memmap_worker_init_fn)
+    """
+    worker_info = torch.utils.data.get_worker_info()
+    if worker_info is not None:
+        dataset = worker_info.dataset
+        # Force re-initialization of memmap in each worker
+        dataset.data = None
+
+        # Seed numpy RNG per worker using PyTorch's worker seed for reproducibility
+        # This ensures random augmentations (noise, shifts, etc.) are unique per worker
+        np.random.seed(worker_info.seed % (2**32 - 1))
+
+
+# ==============================================================================
+# MEMORY-MAPPED DATASET
+# ==============================================================================
+class MemmapDataset(Dataset):
+    """
+    Zero-copy memory-mapped dataset for large-scale training.
+
+    Uses numpy memory mapping to load data directly from disk, allowing training
+    on datasets that exceed available RAM. The memmap is only opened when first
+    accessed (lazy initialization), and each DataLoader worker maintains its own
+    file handle for thread safety.
+
+    Args:
+        memmap_path: Path to the memory-mapped data file
+        targets: Pre-loaded target tensor (small enough to fit in memory)
+        shape: Full shape of the memmap array (N, C, H, W)
+        indices: Indices into the memmap for this split (train/val)
+
+    Thread Safety:
+        When using with DataLoader num_workers > 0, must use memmap_worker_init_fn
+        as the worker_init_fn to ensure each worker gets its own file handle.
+
+    Example:
+        dataset = MemmapDataset("cache.dat", y_tensor, (10000, 1, 500, 500), train_indices)
+        loader = DataLoader(dataset, num_workers=8, worker_init_fn=memmap_worker_init_fn)
+    """
+
+    def __init__(
+        self,
+        memmap_path: str,
+        targets: torch.Tensor,
+        shape: tuple[int, ...],
+        indices: np.ndarray,
+    ):
+        self.memmap_path = memmap_path
+        self.targets = targets
+        self.shape = shape
+        self.indices = indices
+        self.data: np.memmap | None = None  # Lazy initialization
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.data is None:
+            # Mode 'r' = read-only, prevents accidental data modification
+            self.data = np.memmap(
+                self.memmap_path, dtype="float32", mode="r", shape=self.shape
+            )
+
+        real_idx = self.indices[idx]
+
+        # .copy() detaches from mmap buffer - essential for PyTorch pinned memory
+        x = torch.from_numpy(self.data[real_idx].copy()).contiguous()
+        y = self.targets[real_idx]
+
+        return x, y
+
+    def __repr__(self) -> str:
+        return (
+            f"MemmapDataset(path='{self.memmap_path}', "
+            f"samples={len(self)}, shape={self.shape})"
+        )
+
+
+# ==============================================================================
+# DATA PREPARATION
+# ==============================================================================
+def prepare_data(
+    args: Any,
+    logger: logging.Logger,
+    accelerator: Accelerator,
+    cache_dir: str = ".",
+    val_split: float = 0.2,
+) -> tuple[DataLoader, DataLoader, StandardScaler, tuple[int, ...], int]:
+    """
+    Prepare DataLoaders with DDP synchronization guarantees.
+
+    This function handles:
+    1. Loading raw data and creating memmap cache (Rank 0 only)
+    2. Fitting StandardScaler on training set only (no data leakage)
+    3. Synchronizing all ranks before proceeding
+    4. Creating thread-safe DataLoaders for DDP training
+
+    Supports any input dimensionality:
+        - 1D: (N, L) → returns in_shape = (L,)
+        - 2D: (N, H, W) → returns in_shape = (H, W)
+        - 3D: (N, D, H, W) → returns in_shape = (D, H, W)
+
+    Args:
+        args: Argument namespace with data_path, seed, batch_size, workers
+        logger: Logger instance for status messages
+        accelerator: Accelerator instance for DDP coordination
+        cache_dir: Directory for cache files (default: current directory)
+        val_split: Validation set fraction (default: 0.2)
+
+    Returns:
+        Tuple of:
+            - train_dl: Training DataLoader
+            - val_dl: Validation DataLoader
+            - scaler: Fitted StandardScaler (for inverse transforms)
+            - in_shape: Input spatial dimensions - (L,), (H, W), or (D, H, W)
+            - out_dim: Number of output targets
+
+    Cache Files Created:
+        - train_data_cache.dat: Memory-mapped input data
+        - scaler.pkl: Fitted StandardScaler
+        - data_metadata.pkl: Shape and dimension metadata
+    """
+    CACHE_FILE = os.path.join(cache_dir, "train_data_cache.dat")
+    SCALER_FILE = os.path.join(cache_dir, "scaler.pkl")
+    META_FILE = os.path.join(cache_dir, "data_metadata.pkl")
+
+    # ==========================================================================
+    # PHASE 1: DATA GENERATION (Rank 0 Only)
+    # ==========================================================================
+    # Check cache existence and validity (data path must match)
+    cache_exists = (
+        os.path.exists(CACHE_FILE)
+        and os.path.exists(SCALER_FILE)
+        and os.path.exists(META_FILE)
+    )
+
+    # Validate cache matches current data_path (prevents stale cache corruption)
+    if cache_exists:
+        try:
+            with open(META_FILE, "rb") as f:
+                meta = pickle.load(f)
+            cached_data_path = meta.get("data_path", None)
+            if cached_data_path != os.path.abspath(args.data_path):
+                if accelerator.is_main_process:
+                    logger.warning(
+                        f"⚠️  Cache was created from different data file!\n"
+                        f"   Cached: {cached_data_path}\n"
+                        f"   Current: {os.path.abspath(args.data_path)}\n"
+                        f"   Invalidating cache and regenerating..."
+                    )
+                cache_exists = False
+        except Exception:
+            cache_exists = False
+
+    if not cache_exists:
+        if accelerator.is_main_process:
+            # RANK 0: Create cache (can take a long time for large datasets)
+            # Other ranks will wait at the barrier below
+
+            # Detect format from extension
+            data_format = DataSource.detect_format(args.data_path)
+            logger.info(
+                f"⚡ [Rank 0] Initializing Data Processing from: {args.data_path} (format: {data_format})"
+            )
+
+            # Validate data file exists
+            if not os.path.exists(args.data_path):
+                raise FileNotFoundError(
+                    f"CRITICAL: Data file not found: {args.data_path}"
+                )
+
+            # Load raw data using memory-mapped mode for all formats
+            # This avoids loading the entire dataset into RAM at once
+            try:
+                if data_format == "npz":
+                    source = NPZSource()
+                    inp, outp = source.load_mmap(args.data_path)
+                elif data_format == "hdf5":
+                    source = HDF5Source()
+                    _lazy_handle = source.load_mmap(args.data_path)
+                    inp, outp = _lazy_handle.inputs, _lazy_handle.outputs
+                elif data_format == "mat":
+                    source = MATSource()
+                    _lazy_handle = source.load_mmap(args.data_path)
+                    inp, outp = _lazy_handle.inputs, _lazy_handle.outputs
+                else:
+                    inp, outp = load_training_data(args.data_path, format=data_format)
+                logger.info("   Using memory-mapped loading (low memory mode)")
+            except Exception as e:
+                logger.error(f"Failed to load data file: {e}")
+                raise
+
+            # Detect shape (handle sparse matrices) - DIMENSION AGNOSTIC
+            num_samples = len(inp)
+
+            # Handle 1D targets: (N,) -> treat as single output
+            if outp.ndim == 1:
+                out_dim = 1
+            else:
+                out_dim = outp.shape[1]
+
+            sample_0 = inp[0]
+            if issparse(sample_0) or hasattr(sample_0, "toarray"):
+                sample_0 = sample_0.toarray()
+
+            # Detect dimensionality and validate single-channel assumption
+            raw_shape = sample_0.shape
+
+            # Heuristic for ambiguous multi-channel detection:
+            # Trigger when shape is EXACTLY 3D (could be C,H,W or D,H,W) with:
+            #   - First dim small (<=16) - looks like channels
+            #   - BOTH remaining dims large (>16) - confirming it's an image, not a tiny patch
+            # This distinguishes (3, 256, 256) from (8, 1024) or (128, 128)
+            # Use --single_channel to confirm shallow 3D volumes like (8, 128, 128)
+            is_ambiguous_shape = (
+                len(raw_shape) == 3  # Exactly 3D: could be (C, H, W) or (D, H, W)
+                and raw_shape[0] <= 16  # First dim looks like channels
+                and raw_shape[1] > 16
+                and raw_shape[2] > 16  # Both spatial dims are large
+            )
+
+            # Check for user confirmation via --single_channel flag
+            user_confirmed_single_channel = getattr(args, "single_channel", False)
+
+            if is_ambiguous_shape and not user_confirmed_single_channel:
+                raise ValueError(
+                    f"Ambiguous input shape detected: sample shape {raw_shape}. "
+                    f"This could be either:\n"
+                    f"  - Multi-channel 2D data (C={raw_shape[0]}, H={raw_shape[1]}, W={raw_shape[2]})\n"
+                    f"  - Single-channel 3D volume (D={raw_shape[0]}, H={raw_shape[1]}, W={raw_shape[2]})\n\n"
+                    f"If this is single-channel 3D/shallow volume data, use --single_channel flag.\n"
+                    f"If this is multi-channel 2D data, reshape to (N*C, H, W) with adjusted targets."
+                )
+
+            spatial_shape = raw_shape
+            full_shape = (
+                num_samples,
+                1,
+            ) + spatial_shape  # Add channel dim: (N, 1, ...)
+
+            dim_names = {1: "1D (L)", 2: "2D (H, W)", 3: "3D (D, H, W)"}
+            dim_type = dim_names.get(len(spatial_shape), f"{len(spatial_shape)}D")
+            logger.info(
+                f"   Shape Detected: {full_shape} [{dim_type}] | Output Dim: {out_dim}"
+            )
+
+            # Save metadata (including data path for cache validation)
+            with open(META_FILE, "wb") as f:
+                pickle.dump(
+                    {
+                        "shape": full_shape,
+                        "out_dim": out_dim,
+                        "data_path": os.path.abspath(args.data_path),
+                    },
+                    f,
+                )
+
+            # Create memmap cache
+            if not os.path.exists(CACHE_FILE):
+                logger.info("   Writing Memmap Cache (one-time operation)...")
+                fp = np.memmap(CACHE_FILE, dtype="float32", mode="w+", shape=full_shape)
+
+                chunk_size = 2000
+                pbar = tqdm(
+                    range(0, num_samples, chunk_size),
+                    desc="Caching",
+                    disable=not accelerator.is_main_process,
+                )
+
+                for i in pbar:
+                    end = min(i + chunk_size, num_samples)
+                    batch = inp[i:end]
+
+                    # Handle sparse/dense conversion
+                    if issparse(batch[0]) or hasattr(batch[0], "toarray"):
+                        data_chunk = np.stack(
+                            [x.toarray().astype(np.float32) for x in batch]
+                        )
+                    else:
+                        data_chunk = np.array(batch).astype(np.float32)
+
+                    # Add channel dimension if needed (handles 1D, 2D, and 3D spatial data)
+                    # data_chunk shape: (batch, *spatial) -> need (batch, 1, *spatial)
+                    # full_shape is (N, 1, *spatial), so expected ndim = len(full_shape)
+                    if data_chunk.ndim == len(full_shape) - 1:
+                        # Missing channel dim: (batch, *spatial) -> (batch, 1, *spatial)
+                        data_chunk = np.expand_dims(data_chunk, 1)
+
+                    fp[i:end] = data_chunk
+
+                    # Periodic flush to disk
+                    if i % 10000 == 0:
+                        fp.flush()
+
+                fp.flush()
+                del fp
+                gc.collect()
+
+            # Train/Val split and scaler fitting
+            indices = np.arange(num_samples)
+            tr_idx, val_idx = train_test_split(
+                indices, test_size=val_split, random_state=args.seed
+            )
+
+            if not os.path.exists(SCALER_FILE):
+                logger.info("   Fitting StandardScaler (training set only)...")
+
+                # Convert lazy datasets to numpy for reliable indexing
+                # (h5py and _TransposedH5Dataset may not support fancy indexing)
+                if hasattr(outp, "_dataset") or hasattr(outp, "file"):
+                    # Lazy h5py or _TransposedH5Dataset - load training subset
+                    outp_train = np.array([outp[i] for i in tr_idx])
+                else:
+                    # Already numpy array
+                    outp_train = outp[tr_idx]
+
+                # Ensure 2D for StandardScaler: (N,) -> (N, 1)
+                if outp_train.ndim == 1:
+                    outp_train = outp_train.reshape(-1, 1)
+
+                scaler = StandardScaler()
+                scaler.fit(outp_train)
+                with open(SCALER_FILE, "wb") as f:
+                    pickle.dump(scaler, f)
+
+            # Cleanup: close file handles BEFORE deleting references
+            if "_lazy_handle" in dir() and _lazy_handle is not None:
+                try:
+                    _lazy_handle.close()
+                except Exception:
+                    pass
+            del inp, outp
+            gc.collect()
+
+            logger.info("   ✔ Cache creation complete, synchronizing ranks...")
+        else:
+            # NON-MAIN RANKS: Wait for cache creation
+            # Log that we're waiting (helps with debugging)
+            import time
+
+            wait_start = time.time()
+            while not (
+                os.path.exists(CACHE_FILE)
+                and os.path.exists(SCALER_FILE)
+                and os.path.exists(META_FILE)
+            ):
+                time.sleep(5)  # Check every 5 seconds
+                elapsed = time.time() - wait_start
+                if elapsed > 60 and int(elapsed) % 60 < 5:  # Log every ~minute
+                    logger.info(
+                        f"   [Rank {accelerator.process_index}] Waiting for cache creation... ({int(elapsed)}s)"
+                    )
+            # Small delay to ensure files are fully written
+            time.sleep(2)
+
+    # ==========================================================================
+    # PHASE 2: SYNCHRONIZED LOADING (All Ranks)
+    # ==========================================================================
+    accelerator.wait_for_everyone()
+
+    # Load metadata
+    with open(META_FILE, "rb") as f:
+        meta = pickle.load(f)
+        full_shape = meta["shape"]
+        out_dim = meta["out_dim"]
+
+    # Load and validate scaler
+    with open(SCALER_FILE, "rb") as f:
+        scaler = pickle.load(f)
+
+    if not hasattr(scaler, "scale_") or scaler.scale_ is None:
+        raise RuntimeError("CRITICAL: Scaler is not properly fitted (scale_ is None)")
+
+    # Load targets only (memory-efficient - avoids loading large input arrays)
+    # This is critical for HPC environments with memory constraints during DDP
+    outp = load_outputs_only(args.data_path)
+
+    # Ensure 2D for StandardScaler: (N,) -> (N, 1)
+    if outp.ndim == 1:
+        outp = outp.reshape(-1, 1)
+
+    y_scaled = scaler.transform(outp).astype(np.float32)
+    y_tensor = torch.tensor(y_scaled)
+
+    # Regenerate indices (deterministic with same seed)
+    indices = np.arange(full_shape[0])
+    tr_idx, val_idx = train_test_split(
+        indices, test_size=val_split, random_state=args.seed
+    )
+
+    # Create datasets
+    tr_ds = MemmapDataset(CACHE_FILE, y_tensor, full_shape, tr_idx)
+    val_ds = MemmapDataset(CACHE_FILE, y_tensor, full_shape, val_idx)
+
+    # Create DataLoaders with thread-safe configuration
+    loader_kwargs = {
+        "batch_size": args.batch_size,
+        "num_workers": args.workers,
+        "pin_memory": True,
+        "persistent_workers": (args.workers > 0),
+        "prefetch_factor": 2 if args.workers > 0 else None,
+        "worker_init_fn": memmap_worker_init_fn if args.workers > 0 else None,
+    }
+
+    train_dl = DataLoader(tr_ds, shuffle=True, **loader_kwargs)
+    val_dl = DataLoader(val_ds, shuffle=False, **loader_kwargs)
+
+    # Return spatial shape (H, W) for model initialization
+    in_shape = full_shape[2:]
+
+    return train_dl, val_dl, scaler, in_shape, out_dim