wavedl 1.5.6__py3-none-any.whl → 1.6.0__py3-none-any.whl

wavedl/models/vit.py

@@ -10,9 +10,9 @@ Supports both 1D (signals) and 2D (images) inputs via configurable patch embeddi
     - 2D: Images/spectrograms → patches are grid squares
 
 **Variants**:
-    - vit_tiny: Smallest (~5.7M params, embed_dim=192, depth=12, heads=3)
-    - vit_small: Light (~22M params, embed_dim=384, depth=12, heads=6)
-    - vit_base: Standard (~86M params, embed_dim=768, depth=12, heads=12)
+    - vit_tiny: Smallest (~5.4M backbone params, embed_dim=192, depth=12, heads=3)
+    - vit_small: Light (~21.4M backbone params, embed_dim=384, depth=12, heads=6)
+    - vit_base: Standard (~85.3M backbone params, embed_dim=768, depth=12, heads=12)
 
 References:
     Dosovitskiy, A., et al. (2021). An Image is Worth 16x16 Words:
@@ -365,7 +365,7 @@ class ViTTiny(ViTBase):
     """
     ViT-Tiny: Smallest Vision Transformer variant.
 
-    ~5.7M parameters. Good for: Quick experiments, smaller datasets.
+    ~5.4M backbone parameters. Good for: Quick experiments, smaller datasets.
 
     Args:
         in_shape: (L,) for 1D or (H, W) for 2D
@@ -398,7 +398,7 @@ class ViTSmall(ViTBase):
     """
     ViT-Small: Light Vision Transformer variant.
 
-    ~22M parameters. Good for: Balanced performance.
+    ~21.4M backbone parameters. Good for: Balanced performance.
 
     Args:
         in_shape: (L,) for 1D or (H, W) for 2D
@@ -429,7 +429,7 @@ class ViTBase_(ViTBase):
     """
     ViT-Base: Standard Vision Transformer variant.
 
-    ~86M parameters. Good for: High accuracy, larger datasets.
+    ~85.3M backbone parameters. Good for: High accuracy, larger datasets.
 
     Args:
         in_shape: (L,) for 1D or (H, W) for 2D

wavedl/test.py

@@ -311,7 +311,7 @@ def load_data_for_inference(
 # ==============================================================================
 def load_checkpoint(
     checkpoint_dir: str,
-    in_shape: tuple[int, int],
+    in_shape: tuple[int, ...],
     out_size: int,
     model_name: str | None = None,
 ) -> tuple[nn.Module, any]:
@@ -320,7 +320,7 @@ def load_checkpoint(
 
     Args:
         checkpoint_dir: Path to checkpoint directory
-        in_shape: Input image shape (H, W)
+        in_shape: Input spatial shape - (L,) for 1D, (H, W) for 2D, or (D, H, W) for 3D
         out_size: Number of output parameters
         model_name: Model architecture name (auto-detect if None)
 
@@ -376,7 +376,11 @@ def load_checkpoint(
                 )
 
     logging.info(f"   Building model: {model_name}")
-    model = build_model(model_name, in_shape=in_shape, out_size=out_size)
+    # Use pretrained=False: checkpoint weights will overwrite any pretrained weights,
+    # so downloading ImageNet weights is wasteful and breaks offline/HPC inference.
+    model = build_model(
+        model_name, in_shape=in_shape, out_size=out_size, pretrained=False
+    )
 
     # Load weights (check multiple formats in order of preference)
     weight_path = None

wavedl/train.py

@@ -238,18 +238,13 @@ def parse_args() -> argparse.Namespace:
         default=[],
         help="Python modules to import before training (for custom models)",
     )
-    parser.add_argument(
-        "--pretrained",
-        action="store_true",
-        default=True,
-        help="Use pretrained weights (default: True)",
-    )
     parser.add_argument(
         "--no_pretrained",
         dest="pretrained",
         action="store_false",
-        help="Train from scratch without pretrained weights",
+        help="Train from scratch without pretrained weights (default: use pretrained)",
     )
+    parser.set_defaults(pretrained=True)
 
     # Configuration File
     parser.add_argument(
@@ -367,6 +362,18 @@ def parse_args() -> argparse.Namespace:
         help="DataLoader workers per GPU (-1=auto-detect based on CPU cores)",
     )
     parser.add_argument("--seed", type=int, default=2025, help="Random seed")
+    parser.add_argument(
+        "--deterministic",
+        action="store_true",
+        help="Enable deterministic mode for reproducibility (slower, disables TF32/cuDNN benchmark)",
+    )
+    parser.add_argument(
+        "--cache_validate",
+        type=str,
+        default="sha256",
+        choices=["sha256", "fast", "size"],
+        help="Cache validation mode: sha256 (full hash), fast (partial), size (quick)",
+    )
     parser.add_argument(
         "--single_channel",
         action="store_true",
@@ -512,11 +519,23 @@ def main():
                     # Import as regular module
                     importlib.import_module(module_name)
                     print(f"✓ Imported module: {module_name}")
-            except ImportError as e:
+            except (ImportError, FileNotFoundError, SyntaxError, PermissionError) as e:
                 print(f"✗ Failed to import '{module_name}': {e}", file=sys.stderr)
-                print(
-                    "  Make sure the module is in your Python path or current directory."
-                )
+                if isinstance(e, FileNotFoundError):
+                    print("  File does not exist. Check the path.", file=sys.stderr)
+                elif isinstance(e, SyntaxError):
+                    print(
+                        f"  Syntax error at line {e.lineno}: {e.msg}", file=sys.stderr
+                    )
+                elif isinstance(e, PermissionError):
+                    print(
+                        "  Permission denied. Check file permissions.", file=sys.stderr
+                    )
+                else:
+                    print(
+                        "  Make sure the module is in your Python path or current directory.",
+                        file=sys.stderr,
+                    )
                 sys.exit(1)
 
     # Handle --list_models flag
@@ -648,6 +667,17 @@ def main():
     )
     set_seed(args.seed)
 
+    # Deterministic mode for scientific reproducibility
+    # Disables TF32 and cuDNN benchmark for exact reproducibility (slower)
+    if args.deterministic:
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        if accelerator.is_main_process:
+            print("🔒 Deterministic mode enabled (slower but reproducible)")
+
     # Configure logging (rank 0 only prints to console)
     logging.basicConfig(
         level=logging.INFO if accelerator.is_main_process else logging.ERROR,
@@ -999,12 +1029,14 @@ def main():
 
             for x, y in pbar:
                 with accelerator.accumulate(model):
-                    pred = model(x)
-                    # Pass inputs for input-dependent constraints (x_mean, x[...], etc.)
-                    if isinstance(criterion, PhysicsConstrainedLoss):
-                        loss = criterion(pred, y, x)
-                    else:
-                        loss = criterion(pred, y)
+                    # Use mixed precision for forward pass (respects --precision flag)
+                    with accelerator.autocast():
+                        pred = model(x)
+                        # Pass inputs for input-dependent constraints (x_mean, x[...], etc.)
+                        if isinstance(criterion, PhysicsConstrainedLoss):
+                            loss = criterion(pred, y, x)
+                        else:
+                            loss = criterion(pred, y)
 
                     accelerator.backward(loss)
 
@@ -1053,12 +1085,14 @@ def main():
 
             with torch.inference_mode():
                 for x, y in val_dl:
-                    pred = model(x)
-                    # Pass inputs for input-dependent constraints
-                    if isinstance(criterion, PhysicsConstrainedLoss):
-                        loss = criterion(pred, y, x)
-                    else:
-                        loss = criterion(pred, y)
+                    # Use mixed precision for validation (consistent with training)
+                    with accelerator.autocast():
+                        pred = model(x)
+                        # Pass inputs for input-dependent constraints
+                        if isinstance(criterion, PhysicsConstrainedLoss):
+                            loss = criterion(pred, y, x)
+                        else:
+                            loss = criterion(pred, y)
 
                     val_loss_sum += loss.detach() * x.size(0)
                     val_samples += x.size(0)

wavedl/utils/constraints.py

@@ -207,22 +207,28 @@ class ExpressionConstraint(nn.Module):
             # Parse indices from the slice
             indices = self._parse_subscript_indices(node.slice)
 
+            # Auto-squeeze channel dimension for single-channel inputs
+            # This allows x[i,j] syntax for (B, 1, H, W) inputs instead of x[c,i,j]
+            inputs_for_indexing = inputs
+            if inputs.ndim >= 3 and inputs.shape[1] == 1:
+                inputs_for_indexing = inputs.squeeze(1)  # (B, 1, H, W) → (B, H, W)
+
             # Validate dimensions match
             # inputs shape: (batch, dim1) or (batch, dim1, dim2) or (batch, dim1, dim2, dim3)
-            input_ndim = inputs.ndim - 1  # Exclude batch dimension
+            input_ndim = inputs_for_indexing.ndim - 1  # Exclude batch dimension
             if len(indices) != input_ndim:
                 raise ValueError(
-                    f"Input has {input_ndim}D shape, but got {len(indices)} indices. "
+                    f"Input has {input_ndim}D shape (after channel squeeze), but got {len(indices)} indices. "
                     f"Use x[i] for 1D, x[i,j] for 2D, x[i,j,k] for 3D inputs."
                 )
 
             # Extract the value at the specified indices (for entire batch)
             if len(indices) == 1:
-                return inputs[:, indices[0]]
+                return inputs_for_indexing[:, indices[0]]
             elif len(indices) == 2:
-                return inputs[:, indices[0], indices[1]]
+                return inputs_for_indexing[:, indices[0], indices[1]]
             elif len(indices) == 3:
-                return inputs[:, indices[0], indices[1], indices[2]]
+                return inputs_for_indexing[:, indices[0], indices[1], indices[2]]
             else:
                 raise ValueError("Only 1D, 2D, or 3D input indexing supported.")
         elif isinstance(node, ast.Expression):

wavedl/utils/data.py

@@ -50,9 +50,11 @@ INPUT_KEYS = ["input_train", "input_test", "X", "data", "inputs", "features", "x
 OUTPUT_KEYS = ["output_train", "output_test", "Y", "labels", "outputs", "targets", "y"]
 
 
-def _compute_file_hash(path: str, chunk_size: int = 8 * 1024 * 1024) -> str:
+def _compute_file_hash(
+    path: str, mode: str = "sha256", chunk_size: int = 8 * 1024 * 1024
+) -> str:
     """
-    Compute SHA256 hash of a file for cache validation.
+    Compute hash of a file for cache validation.
 
     Uses chunked reading to handle large files efficiently without loading
     the entire file into memory. This is more reliable than mtime for detecting
@@ -61,16 +63,34 @@ def _compute_file_hash(path: str, chunk_size: int = 8 * 1024 * 1024) -> str:
 
     Args:
         path: Path to file to hash
+        mode: Validation mode:
+            - 'sha256': Full content hash (default, most reliable)
+            - 'fast': Partial hash (first+last 1MB + size, faster for large files)
+            - 'size': File size only (fastest, least reliable)
         chunk_size: Read buffer size (default 8MB for fast I/O)
 
     Returns:
-        Hex string of SHA256 hash
+        Hash string for cache comparison
     """
-    hasher = hashlib.sha256()
-    with open(path, "rb") as f:
-        while chunk := f.read(chunk_size):
-            hasher.update(chunk)
-    return hasher.hexdigest()
+    if mode == "size":
+        return str(os.path.getsize(path))
+    elif mode == "fast":
+        # Hash first 1MB + last 1MB + file size for quick validation
+        file_size = os.path.getsize(path)
+        hasher = hashlib.sha256()
+        hasher.update(str(file_size).encode())
+        with open(path, "rb") as f:
+            hasher.update(f.read(1024 * 1024))  # First 1MB
+            if file_size > 2 * 1024 * 1024:
+                f.seek(-1024 * 1024, 2)
+                hasher.update(f.read())  # Last 1MB
+        return hasher.hexdigest()
+    else:  # sha256 (full)
+        hasher = hashlib.sha256()
+        with open(path, "rb") as f:
+            while chunk := f.read(chunk_size):
+                hasher.update(chunk)
+        return hasher.hexdigest()
 
 
 class LazyDataHandle:
@@ -454,9 +474,18 @@ class _TransposedH5Dataset:
         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))
+    @property
+    def ndim(self) -> int:
+        """Number of dimensions (derived from shape for numpy compatibility)."""
+        return len(self.shape)
+
+    @property
+    def _transpose_axes(self) -> tuple[int, ...]:
+        """Transpose axis order for reversing dimensions.
+
+        For shape (A, B, C) -> reversed (C, B, A), transpose axes are (2, 1, 0).
+        """
+        return tuple(range(len(self._dataset.shape) - 1, -1, -1))
 
     def __len__(self) -> int:
         return self.shape[0]
@@ -840,10 +869,22 @@ def load_test_data(
                 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 [])
             )
@@ -858,10 +899,22 @@ def load_test_data(
                 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:
@@ -878,10 +931,22 @@ def load_test_data(
                 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:
@@ -909,8 +974,17 @@ def load_test_data(
         else:
             # Fallback to default source.load() for unknown formats
             inp, outp = source.load(path)
-    except KeyError:
-        # Try with just inputs if outputs not found (inference-only mode)
+    except KeyError as e:
+        # IMPORTANT: Only fall back to inference-only mode if outputs are
+        # genuinely missing (auto-detection failed). If user explicitly
+        # provided --output_key, they expect it to exist - don't silently drop.
+        if output_key is not None:
+            raise KeyError(
+                f"Explicit --output_key '{output_key}' not found in file. "
+                f"Available keys depend on file format. Original error: {e}"
+            ) from e
+
+        # Legitimate fallback: no explicit output_key, outputs just not present
         if format == "npz":
             # First pass to find keys
             with np.load(path, allow_pickle=False) as probe:
@@ -1020,9 +1094,33 @@ def load_test_data(
             if input_channels == 1:
                 X = X.unsqueeze(1)  # Add channel: (N, D, H, W) → (N, 1, D, H, W)
             # else: already has channels, leave as-is
-        elif X.shape[1] > 16:
-            # Heuristic fallback: large dim 1 suggests 3D volume needing channel
-            X = X.unsqueeze(1)  # 3D volume: (N, D, H, W) → (N, 1, D, H, W)
+        else:
+            # Detect channels-last format: (N, H, W, C) where C is small (1-4)
+            # and spatial dims are large (>16). This catches common mistakes.
+            if X.shape[-1] <= 4 and X.shape[1] > 16 and X.shape[2] > 16:
+                raise ValueError(
+                    f"Input appears to be channels-last format: {tuple(X.shape)}. "
+                    "WaveDL expects channels-first (N, C, H, W). "
+                    "Convert your data using: X = X.permute(0, 3, 1, 2). "
+                    "If this is actually a 3D volume with small depth, "
+                    "use --input_channels 1 to add a channel dimension."
+                )
+            elif X.shape[1] > 16:
+                # Heuristic fallback: large dim 1 suggests 3D volume needing channel
+                X = X.unsqueeze(1)  # 3D volume: (N, D, H, W) → (N, 1, D, H, W)
+            else:
+                # Ambiguous case: shallow 3D volume (D <= 16) or multi-channel 2D
+                # Default to treating as multi-channel 2D (no modification needed)
+                # Log a warning so users know about the --input_channels option
+                import warnings
+
+                warnings.warn(
+                    f"Ambiguous 4D input shape: {tuple(X.shape)}. "
+                    f"Assuming {X.shape[1]} channels (multi-channel 2D). "
+                    f"For 3D volumes with depth={X.shape[1]}, use --input_channels 1.",
+                    UserWarning,
+                    stacklevel=2,
+                )
     # X.ndim >= 5: assume channel dimension already exists
 
     return X, y
@@ -1207,7 +1305,9 @@ def prepare_data(
                 cache_exists = False
             # Content hash check (robust against cloud sync mtime changes)
             elif cached_content_hash is not None:
-                current_hash = _compute_file_hash(args.data_path)
+                current_hash = _compute_file_hash(
+                    args.data_path, mode=getattr(args, "cache_validate", "sha256")
+                )
                 if cached_content_hash != current_hash:
                     if accelerator.is_main_process:
                         logger.warning(
@@ -1330,7 +1430,9 @@ def prepare_data(
 
             # Save metadata (including data path, size, content hash for cache validation)
             file_stats = os.stat(args.data_path)
-            content_hash = _compute_file_hash(args.data_path)
+            content_hash = _compute_file_hash(
+                args.data_path, mode=getattr(args, "cache_validate", "sha256")
+            )
             with open(META_FILE, "wb") as f:
                 pickle.dump(
                     {