wavedl 1.5.7__py3-none-any.whl → 1.6.1__py3-none-any.whl

wavedl/utils/data.py

@@ -474,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]
@@ -965,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:
@@ -1083,11 +1101,26 @@ def load_test_data(
                 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)"
+                    "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

wavedl/utils/losses.py

@@ -1,216 +1,216 @@
-"""
-Loss Functions for Regression Tasks
-====================================
-
-Provides a comprehensive set of loss functions for regression problems,
-with a factory function for easy selection via CLI arguments.
-
-Supported Losses:
-    - mse: Mean Squared Error (default)
-    - mae: Mean Absolute Error (L1)
-    - huber: Huber Loss (smooth blend of MSE and MAE)
-    - smooth_l1: Smooth L1 Loss (PyTorch native Huber variant)
-    - log_cosh: Log-Cosh Loss (smooth approximation to MAE)
-    - weighted_mse: Per-target weighted MSE
-
-Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
-"""
-
-import torch
-import torch.nn as nn
-
-
-# ==============================================================================
-# CUSTOM LOSS FUNCTIONS
-# ==============================================================================
-class LogCoshLoss(nn.Module):
-    """
-    Log-Cosh Loss: A smooth approximation to Mean Absolute Error.
-
-    The loss is defined as: loss = log(cosh(pred - target))
-
-    Properties:
-        - Smooth everywhere (differentiable)
-        - Behaves like L2 for small errors, L1 for large errors
-        - More robust to outliers than MSE
-
-    Example:
-        >>> criterion = LogCoshLoss()
-        >>> loss = criterion(predictions, targets)
-    """
-
-    def __init__(self, reduction: str = "mean"):
-        """
-        Args:
-            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
-        """
-        super().__init__()
-        if reduction not in ("none", "mean", "sum"):
-            raise ValueError(f"Invalid reduction mode: {reduction}")
-        self.reduction = reduction
-
-    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        """
-        Compute Log-Cosh loss.
-
-        Args:
-            pred: Predicted values of shape (N, *)
-            target: Target values of shape (N, *)
-
-        Returns:
-            Loss value (scalar if reduction is 'mean' or 'sum')
-        """
-        diff = pred - target
-        # log(cosh(x)) = x + softplus(-2x) - log(2)
-        # This formulation is numerically stable
-        loss = diff + torch.nn.functional.softplus(-2.0 * diff) - 0.693147  # log(2)
-
-        if self.reduction == "none":
-            return loss
-        elif self.reduction == "sum":
-            return loss.sum()
-        else:  # mean
-            return loss.mean()
-
-
-class WeightedMSELoss(nn.Module):
-    """
-    Weighted Mean Squared Error Loss.
-
-    Applies different weights to each target dimension, allowing
-    prioritization of specific outputs (e.g., prioritize thickness
-    over velocity in NDE applications).
-
-    Example:
-        >>> # 3 targets, prioritize first target
-        >>> criterion = WeightedMSELoss(weights=[2.0, 1.0, 1.0])
-        >>> loss = criterion(predictions, targets)
-    """
-
-    def __init__(
-        self, weights: list[float] | torch.Tensor | None = None, reduction: str = "mean"
-    ):
-        """
-        Args:
-            weights: Per-target weights. If None, equal weights (standard MSE).
-                     Length must match number of output targets.
-            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
-        """
-        super().__init__()
-        if reduction not in ("none", "mean", "sum"):
-            raise ValueError(f"Invalid reduction mode: {reduction}")
-        self.reduction = reduction
-
-        if weights is not None:
-            if isinstance(weights, list):
-                weights = torch.tensor(weights, dtype=torch.float32)
-            self.register_buffer("weights", weights)
-        else:
-            self.weights = None
-
-    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        """
-        Compute weighted MSE loss.
-
-        Args:
-            pred: Predicted values of shape (N, T) where T is number of targets
-            target: Target values of shape (N, T)
-
-        Returns:
-            Loss value (scalar if reduction is 'mean' or 'sum')
-
-        Raises:
-            ValueError: If weight dimension doesn't match target dimension
-        """
-        mse = (pred - target) ** 2
-
-        if self.weights is not None:
-            # Validate weight dimension matches target dimension
-            if self.weights.shape[0] != pred.shape[-1]:
-                raise ValueError(
-                    f"Weight dimension ({self.weights.shape[0]}) must match "
-                    f"output dimension ({pred.shape[-1]}). "
-                    f"Check your --loss_weights argument."
-                )
-            # Use local variable to avoid mutating registered buffer
-            # (mutating self.weights breaks state_dict semantics)
-            weights = self.weights.to(mse.device)
-            # Apply per-target weights with correct broadcasting: (N, T) * (T,) -> (N, T)
-            mse = mse * weights
-
-        if self.reduction == "none":
-            return mse
-        elif self.reduction == "sum":
-            return mse.sum()
-        else:  # mean
-            return mse.mean()
-
-
-# ==============================================================================
-# LOSS REGISTRY
-# ==============================================================================
-_LOSS_REGISTRY = {
-    "mse": nn.MSELoss,
-    "mae": nn.L1Loss,
-    "l1": nn.L1Loss,  # Alias for mae
-    "huber": nn.HuberLoss,
-    "smooth_l1": nn.SmoothL1Loss,
-    "log_cosh": LogCoshLoss,
-    "logcosh": LogCoshLoss,  # Alias
-    "weighted_mse": WeightedMSELoss,
-}
-
-
-def list_losses() -> list[str]:
-    """
-    Return list of available loss function names.
-
-    Returns:
-        List of registered loss function names (excluding aliases)
-    """
-    # Return unique loss names (exclude aliases)
-    primary_names = ["mse", "mae", "huber", "smooth_l1", "log_cosh", "weighted_mse"]
-    return primary_names
-
-
-def get_loss(
-    name: str, weights: list[float] | None = None, delta: float = 1.0, **kwargs
-) -> nn.Module:
-    """
-    Factory function to create loss function by name.
-
-    Args:
-        name: Loss function name (see list_losses())
-        weights: Per-target weights for weighted_mse
-        delta: Delta parameter for Huber loss (default: 1.0)
-        **kwargs: Additional arguments passed to loss constructor
-
-    Returns:
-        Instantiated loss function (nn.Module)
-
-    Raises:
-        ValueError: If loss name is not recognized
-
-    Example:
-        >>> criterion = get_loss("mse")
-        >>> criterion = get_loss("huber", delta=0.5)
-        >>> criterion = get_loss("weighted_mse", weights=[2.0, 1.0, 1.0])
-    """
-    name_lower = name.lower().replace("-", "_")
-
-    if name_lower not in _LOSS_REGISTRY:
-        available = ", ".join(list_losses())
-        raise ValueError(
-            f"Unknown loss function: '{name}'. Available options: {available}"
-        )
-
-    loss_cls = _LOSS_REGISTRY[name_lower]
-
-    # Special handling for specific loss types
-    if name_lower == "huber":
-        return loss_cls(delta=delta, **kwargs)
-    elif name_lower == "weighted_mse":
-        return loss_cls(weights=weights, **kwargs)
-    else:
-        return loss_cls(**kwargs)
+"""
+Loss Functions for Regression Tasks
+====================================
+
+Provides a comprehensive set of loss functions for regression problems,
+with a factory function for easy selection via CLI arguments.
+
+Supported Losses:
+    - mse: Mean Squared Error (default)
+    - mae: Mean Absolute Error (L1)
+    - huber: Huber Loss (smooth blend of MSE and MAE)
+    - smooth_l1: Smooth L1 Loss (PyTorch native Huber variant)
+    - log_cosh: Log-Cosh Loss (smooth approximation to MAE)
+    - weighted_mse: Per-target weighted MSE
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+import torch
+import torch.nn as nn
+
+
+# ==============================================================================
+# CUSTOM LOSS FUNCTIONS
+# ==============================================================================
+class LogCoshLoss(nn.Module):
+    """
+    Log-Cosh Loss: A smooth approximation to Mean Absolute Error.
+
+    The loss is defined as: loss = log(cosh(pred - target))
+
+    Properties:
+        - Smooth everywhere (differentiable)
+        - Behaves like L2 for small errors, L1 for large errors
+        - More robust to outliers than MSE
+
+    Example:
+        >>> criterion = LogCoshLoss()
+        >>> loss = criterion(predictions, targets)
+    """
+
+    def __init__(self, reduction: str = "mean"):
+        """
+        Args:
+            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
+        """
+        super().__init__()
+        if reduction not in ("none", "mean", "sum"):
+            raise ValueError(f"Invalid reduction mode: {reduction}")
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Log-Cosh loss.
+
+        Args:
+            pred: Predicted values of shape (N, *)
+            target: Target values of shape (N, *)
+
+        Returns:
+            Loss value (scalar if reduction is 'mean' or 'sum')
+        """
+        diff = pred - target
+        # log(cosh(x)) = x + softplus(-2x) - log(2)
+        # This formulation is numerically stable
+        loss = diff + torch.nn.functional.softplus(-2.0 * diff) - 0.693147  # log(2)
+
+        if self.reduction == "none":
+            return loss
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:  # mean
+            return loss.mean()
+
+
+class WeightedMSELoss(nn.Module):
+    """
+    Weighted Mean Squared Error Loss.
+
+    Applies different weights to each target dimension, allowing
+    prioritization of specific outputs (e.g., prioritize thickness
+    over velocity in NDE applications).
+
+    Example:
+        >>> # 3 targets, prioritize first target
+        >>> criterion = WeightedMSELoss(weights=[2.0, 1.0, 1.0])
+        >>> loss = criterion(predictions, targets)
+    """
+
+    def __init__(
+        self, weights: list[float] | torch.Tensor | None = None, reduction: str = "mean"
+    ):
+        """
+        Args:
+            weights: Per-target weights. If None, equal weights (standard MSE).
+                     Length must match number of output targets.
+            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
+        """
+        super().__init__()
+        if reduction not in ("none", "mean", "sum"):
+            raise ValueError(f"Invalid reduction mode: {reduction}")
+        self.reduction = reduction
+
+        if weights is not None:
+            if isinstance(weights, list):
+                weights = torch.tensor(weights, dtype=torch.float32)
+            self.register_buffer("weights", weights)
+        else:
+            self.weights = None
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """
+        Compute weighted MSE loss.
+
+        Args:
+            pred: Predicted values of shape (N, T) where T is number of targets
+            target: Target values of shape (N, T)
+
+        Returns:
+            Loss value (scalar if reduction is 'mean' or 'sum')
+
+        Raises:
+            ValueError: If weight dimension doesn't match target dimension
+        """
+        mse = (pred - target) ** 2
+
+        if self.weights is not None:
+            # Validate weight dimension matches target dimension
+            if self.weights.shape[0] != pred.shape[-1]:
+                raise ValueError(
+                    f"Weight dimension ({self.weights.shape[0]}) must match "
+                    f"output dimension ({pred.shape[-1]}). "
+                    f"Check your --loss_weights argument."
+                )
+            # Use local variable to avoid mutating registered buffer
+            # (mutating self.weights breaks state_dict semantics)
+            weights = self.weights.to(mse.device)
+            # Apply per-target weights with correct broadcasting: (N, T) * (T,) -> (N, T)
+            mse = mse * weights
+
+        if self.reduction == "none":
+            return mse
+        elif self.reduction == "sum":
+            return mse.sum()
+        else:  # mean
+            return mse.mean()
+
+
+# ==============================================================================
+# LOSS REGISTRY
+# ==============================================================================
+_LOSS_REGISTRY = {
+    "mse": nn.MSELoss,
+    "mae": nn.L1Loss,
+    "l1": nn.L1Loss,  # Alias for mae
+    "huber": nn.HuberLoss,
+    "smooth_l1": nn.SmoothL1Loss,
+    "log_cosh": LogCoshLoss,
+    "logcosh": LogCoshLoss,  # Alias
+    "weighted_mse": WeightedMSELoss,
+}
+
+
+def list_losses() -> list[str]:
+    """
+    Return list of available loss function names.
+
+    Returns:
+        List of registered loss function names (excluding aliases)
+    """
+    # Return unique loss names (exclude aliases)
+    primary_names = ["mse", "mae", "huber", "smooth_l1", "log_cosh", "weighted_mse"]
+    return primary_names
+
+
+def get_loss(
+    name: str, weights: list[float] | None = None, delta: float = 1.0, **kwargs
+) -> nn.Module:
+    """
+    Factory function to create loss function by name.
+
+    Args:
+        name: Loss function name (see list_losses())
+        weights: Per-target weights for weighted_mse
+        delta: Delta parameter for Huber loss (default: 1.0)
+        **kwargs: Additional arguments passed to loss constructor
+
+    Returns:
+        Instantiated loss function (nn.Module)
+
+    Raises:
+        ValueError: If loss name is not recognized
+
+    Example:
+        >>> criterion = get_loss("mse")
+        >>> criterion = get_loss("huber", delta=0.5)
+        >>> criterion = get_loss("weighted_mse", weights=[2.0, 1.0, 1.0])
+    """
+    name_lower = name.lower().replace("-", "_")
+
+    if name_lower not in _LOSS_REGISTRY:
+        available = ", ".join(list_losses())
+        raise ValueError(
+            f"Unknown loss function: '{name}'. Available options: {available}"
+        )
+
+    loss_cls = _LOSS_REGISTRY[name_lower]
+
+    # Special handling for specific loss types
+    if name_lower == "huber":
+        return loss_cls(delta=delta, **kwargs)
+    elif name_lower == "weighted_mse":
+        return loss_cls(weights=weights, **kwargs)
+    else:
+        return loss_cls(**kwargs)