wavedl 1.6.2__py3-none-any.whl → 1.7.0__py3-none-any.whl

wavedl/utils/cross_validation.py

@@ -36,7 +36,9 @@ from torch.utils.data import DataLoader
 class CVDataset(torch.utils.data.Dataset):
     """Simple in-memory dataset for cross-validation."""
 
-    def __init__(self, X: np.ndarray, y: np.ndarray, expected_spatial_ndim: int = None):
+    def __init__(
+        self, X: np.ndarray, y: np.ndarray, expected_spatial_ndim: int | None = None
+    ):
         """
         Initialize CV dataset with explicit channel dimension handling.
 
@@ -51,6 +53,11 @@ class CVDataset(torch.utils.data.Dataset):
             - If X.ndim == expected_spatial_ndim + 1: Add channel dim (N, *spatial) -> (N, 1, *spatial)
             - If X.ndim == expected_spatial_ndim + 2: Already has channel (N, C, *spatial)
             - If expected_spatial_ndim is None: Use legacy ndim-based inference
+
+        Warning:
+            Legacy mode (expected_spatial_ndim=None) may misinterpret multichannel
+            3D data as single-channel 4D data. Always pass expected_spatial_ndim
+            explicitly for 3D volumes with >1 channel.
         """
         if expected_spatial_ndim is not None:
             # Explicit mode: use expected_spatial_ndim to determine if channel exists
@@ -100,6 +107,7 @@ def train_fold(
     device: torch.device,
     epochs: int,
     patience: int,
+    grad_clip: float,
     scaler: StandardScaler,
     logger: logging.Logger,
 ) -> dict[str, Any]:
@@ -147,7 +155,8 @@ def train_fold(
             pred = model(x)
             loss = criterion(pred, y)
             loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            if grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
             optimizer.step()
 
             # Per-batch LR scheduling (OneCycleLR)
@@ -289,6 +298,7 @@ def run_cross_validation(
     output_dir: str = "./cv_results",
     workers: int = 4,
     seed: int = 2025,
+    grad_clip: float = 1.0,
     logger: logging.Logger | None = None,
 ) -> dict[str, Any]:
     """
@@ -330,8 +340,10 @@ def run_cross_validation(
         )
         logger = logging.getLogger("CV-Trainer")
 
-    # Set seeds
-    np.random.seed(seed)
+    # Set seeds for reproducibility
+    # Note: sklearn KFold uses random_state parameter directly, not global numpy RNG
+    rng = np.random.default_rng(seed)  # Local RNG for any numpy operations
+    _ = rng  # Silence unused variable warning (available for future use)
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
@@ -444,6 +456,7 @@ def run_cross_validation(
             device=device,
             epochs=epochs,
             patience=patience,
+            grad_clip=grad_clip,
             scaler=scaler,
             logger=logger,
         )

wavedl/utils/data.py

@@ -798,6 +798,128 @@ def load_outputs_only(path: str, format: str = "auto") -> np.ndarray:
     return source.load_outputs_only(path)
 
 
+def _load_npz_for_test(
+    path: str,
+    input_keys: list[str],
+    output_keys: list[str],
+    explicit_input_key: str | None,
+    explicit_output_key: str | None,
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Load NPZ file for test/inference with key validation."""
+    with np.load(path, allow_pickle=False) as probe:
+        keys = list(probe.keys())
+
+    inp_key = DataSource._find_key(keys, input_keys)
+    out_key = DataSource._find_key(keys, output_keys)
+
+    # Strict validation for explicit keys
+    if explicit_input_key is not None and explicit_input_key not in keys:
+        raise KeyError(
+            f"Explicit --input_key '{explicit_input_key}' not found. Available: {keys}"
+        )
+    if inp_key is None:
+        raise KeyError(f"Input key not found. Tried: {input_keys}. Found: {keys}")
+    if explicit_output_key is not None and explicit_output_key not in keys:
+        raise KeyError(
+            f"Explicit --output_key '{explicit_output_key}' not found. Available: {keys}"
+        )
+
+    data = NPZSource._load_and_copy(path, [inp_key] + ([out_key] if out_key else []))
+    inp = data[inp_key]
+    if inp.dtype == object:
+        inp = np.array([x.toarray() if hasattr(x, "toarray") else x for x in inp])
+    outp = data[out_key] if out_key else None
+    return inp, outp
+
+
+def _load_hdf5_for_test(
+    path: str,
+    input_keys: list[str],
+    output_keys: list[str],
+    explicit_input_key: str | None,
+    explicit_output_key: str | None,
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Load HDF5 file for test/inference with key validation and OOM guard."""
+    with h5py.File(path, "r") as f:
+        keys = list(f.keys())
+        inp_key = DataSource._find_key(keys, input_keys)
+        out_key = DataSource._find_key(keys, output_keys)
+
+        # Strict validation
+        if explicit_input_key is not None and explicit_input_key not in keys:
+            raise KeyError(
+                f"Explicit --input_key '{explicit_input_key}' not found. Available: {keys}"
+            )
+        if inp_key is None:
+            raise KeyError(f"Input key not found. Tried: {input_keys}. Found: {keys}")
+        if explicit_output_key is not None and explicit_output_key not in keys:
+            raise KeyError(
+                f"Explicit --output_key '{explicit_output_key}' not found. Available: {keys}"
+            )
+
+        # OOM guard
+        n_samples = f[inp_key].shape[0]
+        if n_samples > 100000:
+            raise ValueError(
+                f"Dataset has {n_samples:,} samples. load_test_data() loads "
+                f"everything into RAM which may cause OOM. For large inference "
+                f"sets, use a DataLoader with HDF5Source.load_mmap() instead."
+            )
+
+        inp = f[inp_key][:]
+        outp = f[out_key][:] if out_key else None
+    return inp, outp
+
+
+def _load_mat_for_test(
+    path: str,
+    input_keys: list[str],
+    output_keys: list[str],
+    explicit_input_key: str | None,
+    explicit_output_key: str | None,
+) -> tuple[np.ndarray, np.ndarray | None]:
+    """Load MAT v7.3 file for test/inference with sparse support."""
+    mat_source = MATSource()
+    with h5py.File(path, "r") as f:
+        keys = list(f.keys())
+        inp_key = DataSource._find_key(keys, input_keys)
+        out_key = DataSource._find_key(keys, output_keys)
+
+        # Strict validation
+        if explicit_input_key is not None and explicit_input_key not in keys:
+            raise KeyError(
+                f"Explicit --input_key '{explicit_input_key}' not found. Available: {keys}"
+            )
+        if inp_key is None:
+            raise KeyError(f"Input key not found. Tried: {input_keys}. Found: {keys}")
+        if explicit_output_key is not None and explicit_output_key not in keys:
+            raise KeyError(
+                f"Explicit --output_key '{explicit_output_key}' not found. Available: {keys}"
+            )
+
+        # OOM guard
+        n_samples = f[inp_key].shape[-1]  # MAT is transposed
+        if n_samples > 100000:
+            raise ValueError(
+                f"Dataset has {n_samples:,} samples. load_test_data() loads "
+                f"everything into RAM which may cause OOM. For large inference "
+                f"sets, use a DataLoader with MATSource.load_mmap() instead."
+            )
+
+        inp = mat_source._load_dataset(f, inp_key)
+        outp = None
+        if out_key:
+            outp = mat_source._load_dataset(f, out_key)
+            # Handle MATLAB transpose
+            num_samples = inp.shape[0]
+            if outp.ndim == 2:
+                if (outp.shape[0] == 1 and outp.shape[1] == num_samples) or (
+                    outp.shape[1] == 1 and outp.shape[0] != num_samples
+                ):
+                    outp = outp.T
+    return inp, outp
+
+
 def load_test_data(
     path: str,
     format: str = "auto",
@@ -865,112 +987,17 @@ def load_test_data(
     # We detect keys first to ensure input_test/output_test are used when present
     try:
         if format == "npz":
-            with np.load(path, allow_pickle=False) as probe:
-                keys = list(probe.keys())
-            inp_key = DataSource._find_key(keys, custom_input_keys)
-            out_key = DataSource._find_key(keys, custom_output_keys)
-            # Strict validation: if user explicitly specified input_key, it must exist exactly
-            if input_key is not None and input_key not in keys:
-                raise KeyError(
-                    f"Explicit --input_key '{input_key}' not found. "
-                    f"Available keys: {keys}"
-                )
-            if inp_key is None:
-                raise KeyError(
-                    f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
-                )
-            # Strict validation: if user explicitly specified output_key, it must exist exactly
-            if output_key is not None and output_key not in keys:
-                raise KeyError(
-                    f"Explicit --output_key '{output_key}' not found. "
-                    f"Available keys: {keys}"
-                )
-            data = NPZSource._load_and_copy(
-                path, [inp_key] + ([out_key] if out_key else [])
+            inp, outp = _load_npz_for_test(
+                path, custom_input_keys, custom_output_keys, input_key, output_key
             )
-            inp = data[inp_key]
-            if inp.dtype == object:
-                inp = np.array(
-                    [x.toarray() if hasattr(x, "toarray") else x for x in inp]
-                )
-            outp = data[out_key] if out_key else None
         elif format == "hdf5":
-            with h5py.File(path, "r") as f:
-                keys = list(f.keys())
-                inp_key = DataSource._find_key(keys, custom_input_keys)
-                out_key = DataSource._find_key(keys, custom_output_keys)
-                # Strict validation: if user explicitly specified input_key, it must exist exactly
-                if input_key is not None and input_key not in keys:
-                    raise KeyError(
-                        f"Explicit --input_key '{input_key}' not found. "
-                        f"Available keys: {keys}"
-                    )
-                if inp_key is None:
-                    raise KeyError(
-                        f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
-                    )
-                # Strict validation: if user explicitly specified output_key, it must exist exactly
-                if output_key is not None and output_key not in keys:
-                    raise KeyError(
-                        f"Explicit --output_key '{output_key}' not found. "
-                        f"Available keys: {keys}"
-                    )
-                # OOM guard: warn if dataset is very large
-                n_samples = f[inp_key].shape[0]
-                if n_samples > 100000:
-                    raise ValueError(
-                        f"Dataset has {n_samples:,} samples. load_test_data() loads "
-                        f"everything into RAM which may cause OOM. For large inference "
-                        f"sets, use a DataLoader with HDF5Source.load_mmap() instead."
-                    )
-                inp = f[inp_key][:]
-                outp = f[out_key][:] if out_key else None
+            inp, outp = _load_hdf5_for_test(
+                path, custom_input_keys, custom_output_keys, input_key, output_key
+            )
         elif format == "mat":
-            mat_source = MATSource()
-            with h5py.File(path, "r") as f:
-                keys = list(f.keys())
-                inp_key = DataSource._find_key(keys, custom_input_keys)
-                out_key = DataSource._find_key(keys, custom_output_keys)
-                # Strict validation: if user explicitly specified input_key, it must exist exactly
-                if input_key is not None and input_key not in keys:
-                    raise KeyError(
-                        f"Explicit --input_key '{input_key}' not found. "
-                        f"Available keys: {keys}"
-                    )
-                if inp_key is None:
-                    raise KeyError(
-                        f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
-                    )
-                # Strict validation: if user explicitly specified output_key, it must exist exactly
-                if output_key is not None and output_key not in keys:
-                    raise KeyError(
-                        f"Explicit --output_key '{output_key}' not found. "
-                        f"Available keys: {keys}"
-                    )
-                # OOM guard: warn if dataset is very large (MAT is transposed)
-                n_samples = f[inp_key].shape[-1]
-                if n_samples > 100000:
-                    raise ValueError(
-                        f"Dataset has {n_samples:,} samples. load_test_data() loads "
-                        f"everything into RAM which may cause OOM. For large inference "
-                        f"sets, use a DataLoader with MATSource.load_mmap() instead."
-                    )
-                inp = mat_source._load_dataset(f, inp_key)
-                if out_key:
-                    outp = mat_source._load_dataset(f, out_key)
-                    # Handle transposed outputs from MATLAB
-                    # Case 1: (1, N) - N samples with 1 target → transpose to (N, 1)
-                    # Case 2: (T, 1) - 1 sample with T targets → transpose to (1, T)
-                    num_samples = inp.shape[0]
-                    if outp.ndim == 2:
-                        if outp.shape[0] == 1 and outp.shape[1] == num_samples:
-                            # 1D vector: (1, N) → (N, 1)
-                            outp = outp.T
-                        elif outp.shape[1] == 1 and outp.shape[0] != num_samples:
-                            # Single sample with multiple targets: (T, 1) → (1, T)
-                            outp = outp.T
-                else:
-                    outp = None
+            inp, outp = _load_mat_for_test(
+                path, custom_input_keys, custom_output_keys, input_key, output_key
+            )
         else:
             # Fallback to default source.load() for unknown formats
             inp, outp = source.load(path)
@@ -984,81 +1011,28 @@ def load_test_data(
                 f"Available keys depend on file format. Original error: {e}"
             ) from e
 
-        # Legitimate fallback: no explicit output_key, outputs just not present
+        # Also fail-fast if explicit input_key was provided but not found
+        # This prevents silently loading a different tensor when user mistyped key
+        if input_key is not None:
+            raise KeyError(
+                f"Explicit --input_key '{input_key}' not found in file. "
+                f"Original error: {e}"
+            ) from e
+
+        # Legitimate fallback: no explicit keys, outputs just not present
+        # Re-call helpers with None for explicit keys (validation already done above)
         if format == "npz":
-            # First pass to find keys
-            with np.load(path, allow_pickle=False) as probe:
-                keys = list(probe.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}"
-                )
-            out_key = DataSource._find_key(keys, custom_output_keys)
-            keys_to_probe = [inp_key] + ([out_key] if out_key else [])
-            data = NPZSource._load_and_copy(path, keys_to_probe)
-            inp = data[inp_key]
-            if inp.dtype == object:
-                inp = np.array(
-                    [x.toarray() if hasattr(x, "toarray") else x for x in inp]
-                )
-            outp = data[out_key] if out_key else None
+            inp, outp = _load_npz_for_test(
+                path, custom_input_keys, custom_output_keys, None, None
+            )
         elif format == "hdf5":
-            # HDF5: input-only loading for inference
-            with h5py.File(path, "r") as f:
-                keys = list(f.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}"
-                    )
-                # Check size - load_test_data is eager, large files should use DataLoader
-                n_samples = f[inp_key].shape[0]
-                if n_samples > 100000:
-                    raise ValueError(
-                        f"Dataset has {n_samples:,} samples. load_test_data() loads "
-                        f"everything into RAM which may cause OOM. For large inference "
-                        f"sets, use a DataLoader with HDF5Source.load_mmap() instead."
-                    )
-                inp = f[inp_key][:]
-                out_key = DataSource._find_key(keys, custom_output_keys)
-                outp = f[out_key][:] if out_key else None
+            inp, outp = _load_hdf5_for_test(
+                path, custom_input_keys, custom_output_keys, None, None
+            )
         elif format == "mat":
-            # MAT v7.3: input-only loading with proper sparse handling
-            mat_source = MATSource()
-            with h5py.File(path, "r") as f:
-                keys = list(f.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}"
-                    )
-                # Check size - load_test_data is eager, large files should use DataLoader
-                n_samples = f[inp_key].shape[-1]  # MAT is transposed
-                if n_samples > 100000:
-                    raise ValueError(
-                        f"Dataset has {n_samples:,} samples. load_test_data() loads "
-                        f"everything into RAM which may cause OOM. For large inference "
-                        f"sets, use a DataLoader with MATSource.load_mmap() instead."
-                    )
-                # Use _load_dataset for sparse support and proper transpose
-                inp = mat_source._load_dataset(f, inp_key)
-                out_key = DataSource._find_key(keys, custom_output_keys)
-                if out_key:
-                    outp = mat_source._load_dataset(f, out_key)
-                    # Handle transposed outputs from MATLAB
-                    # Case 1: (1, N) - N samples with 1 target → transpose to (N, 1)
-                    # Case 2: (T, 1) - 1 sample with T targets → transpose to (1, T)
-                    num_samples = inp.shape[0]
-                    if outp.ndim == 2:
-                        if outp.shape[0] == 1 and outp.shape[1] == num_samples:
-                            # 1D vector: (1, N) → (N, 1)
-                            outp = outp.T
-                        elif outp.shape[1] == 1 and outp.shape[0] != num_samples:
-                            # Single sample with multiple targets: (T, 1) → (1, T)
-                            outp = outp.T
-                else:
-                    outp = None
+            inp, outp = _load_mat_for_test(
+                path, custom_input_keys, custom_output_keys, None, None
+            )
         else:
             raise
 
@@ -1524,21 +1498,43 @@ def prepare_data(
 
             logger.info("   ✔ Cache creation complete, synchronizing ranks...")
         else:
-            # NON-MAIN RANKS: Wait for cache creation
-            # Log that we're waiting (helps with debugging)
+            # NON-MAIN RANKS: Wait for cache creation with timeout
+            # Use monotonic clock (immune to system clock changes)
             import time
 
-            wait_start = time.time()
+            wait_start = time.monotonic()
+
+            # Robust env parsing with guards for invalid/non-positive values
+            DEFAULT_CACHE_TIMEOUT = 3600  # 1 hour default
+            try:
+                env_timeout = os.environ.get("WAVEDL_CACHE_TIMEOUT", "")
+                CACHE_TIMEOUT = (
+                    int(env_timeout) if env_timeout else DEFAULT_CACHE_TIMEOUT
+                )
+                if CACHE_TIMEOUT <= 0:
+                    CACHE_TIMEOUT = DEFAULT_CACHE_TIMEOUT
+            except ValueError:
+                CACHE_TIMEOUT = DEFAULT_CACHE_TIMEOUT
+
             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
+                elapsed = time.monotonic() - wait_start
+
+                if elapsed > CACHE_TIMEOUT:
+                    raise RuntimeError(
+                        f"[Rank {accelerator.process_index}] Timeout waiting for cache "
+                        f"files after {CACHE_TIMEOUT}s. Rank 0 may have failed during "
+                        f"cache generation. Check rank 0 logs for errors."
+                    )
+
                 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)"
+                        f"   [Rank {accelerator.process_index}] Waiting for cache "
+                        f"creation... ({int(elapsed)}s / {CACHE_TIMEOUT}s max)"
                     )
             # Small delay to ensure files are fully written
             time.sleep(2)

wavedl/utils/metrics.py

@@ -106,8 +106,20 @@ def configure_matplotlib_style():
     )
 
 
-# Apply style on import
-configure_matplotlib_style()
+# Lazy style initialization flag
+_style_configured = False
+
+
+def _ensure_style_configured():
+    """Apply matplotlib style on first use (lazy initialization).
+
+    This avoids modifying global plt.rcParams at import time, which could
+    unexpectedly override user configurations.
+    """
+    global _style_configured
+    if not _style_configured:
+        configure_matplotlib_style()
+        _style_configured = True
 
 
 # ==============================================================================
@@ -115,11 +127,14 @@ configure_matplotlib_style()
 # ==============================================================================
 class MetricTracker:
     """
-    Tracks running averages of metrics with thread-safe accumulation.
+    Tracks running averages of metrics during training.
 
     Useful for tracking loss, accuracy, or any scalar metric across batches.
     Handles division-by-zero safely by returning 0.0 when count is zero.
 
+    Note:
+        Not thread-safe. Intended for single-threaded use within training loops.
+
     Attributes:
         val: Most recent value
         avg: Running average
@@ -341,6 +356,7 @@ def plot_scientific_scatter(
     Returns:
         Matplotlib Figure object (can be saved or logged to WandB)
     """
+    _ensure_style_configured()
     y_true, y_pred, param_names, num_params = _prepare_plot_data(
         y_true, y_pred, param_names, max_samples
     )
@@ -415,6 +431,7 @@ def plot_error_histogram(
     Returns:
         Matplotlib Figure object
     """
+    _ensure_style_configured()
     y_true, y_pred, param_names, num_params = _prepare_plot_data(
         y_true, y_pred, param_names
     )
@@ -485,6 +502,7 @@ def plot_residuals(
     Returns:
         Matplotlib Figure object
     """
+    _ensure_style_configured()
     y_true, y_pred, param_names, num_params = _prepare_plot_data(
         y_true, y_pred, param_names, max_samples
     )
@@ -552,6 +570,7 @@ def create_training_curves(
     Returns:
         Matplotlib Figure object
     """
+    _ensure_style_configured()
     epochs = [h.get("epoch", i + 1) for i, h in enumerate(history)]
 
     fig, ax1 = plt.subplots(figsize=(FIGURE_WIDTH_INCH * 0.7, FIGURE_WIDTH_INCH * 0.4))
@@ -605,8 +624,8 @@ def create_training_curves(
         if not valid_data:
             return
         vmin, vmax = min(valid_data), max(valid_data)
-        # Get decade range that covers data (ceil for min to avoid going too low)
-        log_min = int(np.ceil(np.log10(vmin)))
+        # Get decade range that covers data (floor for min to include lowest data)
+        log_min = int(np.floor(np.log10(vmin)))
         log_max = int(np.ceil(np.log10(vmax)))
         # Generate ticks at each power of 10
         ticks = [10.0**i for i in range(log_min, log_max + 1)]
@@ -691,6 +710,7 @@ def plot_bland_altman(
     Returns:
         Matplotlib Figure object
     """
+    _ensure_style_configured()
     y_true, y_pred, param_names, num_params = _prepare_plot_data(
         y_true, y_pred, param_names, max_samples
     )
@@ -764,6 +784,7 @@ def plot_qq(
     Returns:
         Matplotlib Figure object
     """
+    _ensure_style_configured()
     from scipy import stats
 
     num_params = y_true.shape[1] if y_true.ndim > 1 else 1

wavedl/utils/schedulers.py

@@ -76,7 +76,7 @@ def get_scheduler(
     # Step/MultiStep parameters
     step_size: int = 30,
     milestones: list[int] | None = None,
-    gamma: float = 0.1,
+    gamma: float = 0.99,
     # Linear warmup parameters
     warmup_epochs: int = 5,
     start_factor: float = 0.1,
@@ -100,7 +100,7 @@ def get_scheduler(
         pct_start: Percentage of cycle spent increasing LR (OneCycleLR)
         step_size: Period for StepLR
         milestones: Epochs to decay LR for MultiStepLR
-        gamma: Decay factor for step/multistep/exponential
+        gamma: Decay factor (0.1 for step/multistep, 0.99 for exponential)
         warmup_epochs: Number of warmup epochs for linear_warmup
         start_factor: Starting LR factor for warmup (LR * start_factor)
         **kwargs: Additional arguments passed to scheduler