wavedl 1.4.6__py3-none-any.whl → 1.5.1__py3-none-any.whl

wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.4.6"
+__version__ = "1.5.1"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl/hpo.py

@@ -89,7 +89,8 @@ def create_objective(args):
         # Suggest hyperparameters
         model = trial.suggest_categorical("model", models)
         lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
-        batch_size = trial.suggest_categorical("batch_size", [64, 128, 256, 512])
+        batch_sizes = args.batch_sizes or [16, 32, 64, 128]
+        batch_size = trial.suggest_categorical("batch_size", batch_sizes)
         optimizer = trial.suggest_categorical("optimizer", optimizers)
         scheduler = trial.suggest_categorical("scheduler", schedulers)
         loss = trial.suggest_categorical("loss", losses)
@@ -317,6 +318,13 @@ Examples:
         default=None,
         help=f"Losses to search (default: {DEFAULT_LOSSES})",
     )
+    parser.add_argument(
+        "--batch_sizes",
+        type=int,
+        nargs="+",
+        default=None,
+        help="Batch sizes to search (default: 16 32 64 128)",
+    )
 
     # Training settings for each trial
     parser.add_argument(

wavedl/models/vit.py

@@ -54,6 +54,16 @@ class PatchEmbed(nn.Module):
         if self.dim == 1:
             # 1D: segment patches
             L = in_shape[0]
+            if L % patch_size != 0:
+                import warnings
+
+                warnings.warn(
+                    f"Input length {L} not divisible by patch_size {patch_size}. "
+                    f"Last {L % patch_size} elements will be dropped. "
+                    f"Consider padding input to {((L // patch_size) + 1) * patch_size}.",
+                    UserWarning,
+                    stacklevel=2,
+                )
             self.num_patches = L // patch_size
             self.proj = nn.Conv1d(
                 1, embed_dim, kernel_size=patch_size, stride=patch_size
@@ -61,6 +71,17 @@ class PatchEmbed(nn.Module):
         elif self.dim == 2:
             # 2D: grid patches
             H, W = in_shape
+            if H % patch_size != 0 or W % patch_size != 0:
+                import warnings
+
+                warnings.warn(
+                    f"Input shape ({H}, {W}) not divisible by patch_size {patch_size}. "
+                    f"Border pixels will be dropped (H: {H % patch_size}, W: {W % patch_size}). "
+                    f"Consider padding to ({((H // patch_size) + 1) * patch_size}, "
+                    f"{((W // patch_size) + 1) * patch_size}).",
+                    UserWarning,
+                    stacklevel=2,
+                )
             self.num_patches = (H // patch_size) * (W // patch_size)
             self.proj = nn.Conv2d(
                 1, embed_dim, kernel_size=patch_size, stride=patch_size

wavedl/test.py

@@ -166,6 +166,13 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Parameter names for output (e.g., 'h' 'v11' 'v12')",
     )
+    parser.add_argument(
+        "--input_channels",
+        type=int,
+        default=None,
+        help="Explicit number of input channels. Bypasses auto-detection heuristics "
+        "for ambiguous 4D shapes (e.g., 3D volumes with small depth).",
+    )
 
     # Inference options
     parser.add_argument(
@@ -235,6 +242,7 @@ def load_data_for_inference(
     format: str = "auto",
     input_key: str | None = None,
     output_key: str | None = None,
+    input_channels: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Load test data for inference using the unified data loading pipeline.
@@ -278,7 +286,11 @@ def load_data_for_inference(
 
     # Use the unified loader from utils.data
     X, y = load_test_data(
-        file_path, format=format, input_key=input_key, output_key=output_key
+        file_path,
+        format=format,
+        input_key=input_key,
+        output_key=output_key,
+        input_channels=input_channels,
     )
 
     # Log results
@@ -452,7 +464,12 @@ def run_inference(
         predictions: Numpy array (N, out_size) - still in normalized space
     """
     if device is None:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            device = torch.device("mps")
+        else:
+            device = torch.device("cpu")
 
     model = model.to(device)
     model.eval()
@@ -463,7 +480,7 @@ def run_inference(
         batch_size=batch_size,
         shuffle=False,
         num_workers=num_workers,
-        pin_memory=device.type == "cuda",
+        pin_memory=device.type in ("cuda", "mps"),
     )
 
     predictions = []
@@ -919,8 +936,13 @@ def main():
     )
     logger = logging.getLogger("Tester")
 
-    # Device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Device (CUDA > MPS > CPU)
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
     logger.info(f"Using device: {device}")
 
     # Load test data
@@ -929,6 +951,7 @@ def main():
         format=args.format,
         input_key=args.input_key,
         output_key=args.output_key,
+        input_channels=args.input_channels,
     )
     in_shape = tuple(X_test.shape[2:])
 

wavedl/train.py

@@ -375,6 +375,36 @@ def parse_args() -> argparse.Namespace:
         help=argparse.SUPPRESS,  # Hidden: use --precision instead
     )
 
+    # Physical Constraints
+    parser.add_argument(
+        "--constraint",
+        type=str,
+        nargs="+",
+        default=[],
+        help="Soft constraint expressions: 'y0 - y1*y2' (penalize violations)",
+    )
+
+    parser.add_argument(
+        "--constraint_file",
+        type=str,
+        default=None,
+        help="Python file with constraint(pred, inputs) function",
+    )
+    parser.add_argument(
+        "--constraint_weight",
+        type=float,
+        nargs="+",
+        default=[0.1],
+        help="Weight(s) for soft constraints (one per constraint, or single shared weight)",
+    )
+    parser.add_argument(
+        "--constraint_reduction",
+        type=str,
+        default="mse",
+        choices=["mse", "mae"],
+        help="Reduction mode for constraint penalties",
+    )
+
     # Logging
     parser.add_argument(
         "--wandb", action="store_true", help="Enable Weights & Biases logging"
@@ -553,7 +583,7 @@ def main():
         return
 
     # ==========================================================================
-    # 1. SYSTEM INITIALIZATION
+    # SYSTEM INITIALIZATION
     # ==========================================================================
     # Initialize Accelerator for DDP and mixed precision
     accelerator = Accelerator(
@@ -609,7 +639,7 @@ def main():
             )
 
     # ==========================================================================
-    # 2. DATA & MODEL LOADING
+    # DATA & MODEL LOADING
     # ==========================================================================
     train_dl, val_dl, scaler, in_shape, out_dim = prepare_data(
         args, logger, accelerator, cache_dir=args.output_dir
@@ -663,7 +693,7 @@ def main():
                 )
 
     # ==========================================================================
-    # 2.5. OPTIMIZER, SCHEDULER & LOSS CONFIGURATION
+    # OPTIMIZER, SCHEDULER & LOSS CONFIGURATION
     # ==========================================================================
     # Parse comma-separated arguments with validation
     try:
@@ -707,6 +737,43 @@ def main():
     # Move criterion to device (important for WeightedMSELoss buffer)
     criterion = criterion.to(accelerator.device)
 
+    # ==========================================================================
+    # PHYSICAL CONSTRAINTS INTEGRATION
+    # ==========================================================================
+    from wavedl.utils.constraints import (
+        PhysicsConstrainedLoss,
+        build_constraints,
+    )
+
+    # Build soft constraints
+    soft_constraints = build_constraints(
+        expressions=args.constraint,
+        file_path=args.constraint_file,
+        reduction=args.constraint_reduction,
+    )
+
+    # Wrap criterion with PhysicsConstrainedLoss if we have soft constraints
+    if soft_constraints:
+        # Pass output scaler so constraints can be evaluated in physical space
+        output_mean = scaler.mean_ if hasattr(scaler, "mean_") else None
+        output_std = scaler.scale_ if hasattr(scaler, "scale_") else None
+        criterion = PhysicsConstrainedLoss(
+            criterion,
+            soft_constraints,
+            weights=args.constraint_weight,
+            output_mean=output_mean,
+            output_std=output_std,
+        )
+        if accelerator.is_main_process:
+            logger.info(
+                f"   🔬 Physical constraints: {len(soft_constraints)} constraint(s) "
+                f"with weight(s) {args.constraint_weight}"
+            )
+            if output_mean is not None:
+                logger.info(
+                    "   📐 Constraints evaluated in physical space (denormalized)"
+                )
+
     # Track if scheduler should step per batch (OneCycleLR) or per epoch
     scheduler_step_per_batch = not is_epoch_based(args.scheduler)
 
@@ -762,7 +829,7 @@ def main():
         )
 
     # ==========================================================================
-    # 3. AUTO-RESUME / RESUME FROM CHECKPOINT
+    # AUTO-RESUME / RESUME FROM CHECKPOINT
     # ==========================================================================
     start_epoch = 0
     best_val_loss = float("inf")
@@ -818,7 +885,7 @@ def main():
             raise FileNotFoundError(f"Checkpoint not found: {args.resume}")
 
     # ==========================================================================
-    # 4. PHYSICAL METRIC SETUP
+    # PHYSICAL METRIC SETUP
     # ==========================================================================
     # Physical MAE = normalized MAE * scaler.scale_
     phys_scale = torch.tensor(
@@ -826,7 +893,7 @@ def main():
     )
 
     # ==========================================================================
-    # 5. TRAINING LOOP
+    # TRAINING LOOP
     # ==========================================================================
     # Dynamic console header
     if accelerator.is_main_process:
@@ -864,7 +931,11 @@ def main():
             for x, y in pbar:
                 with accelerator.accumulate(model):
                     pred = model(x)
-                    loss = criterion(pred, y)
+                    # 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)
 
@@ -914,7 +985,11 @@ def main():
             with torch.inference_mode():
                 for x, y in val_dl:
                     pred = model(x)
-                    loss = criterion(pred, y)
+                    # 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)
@@ -931,13 +1006,45 @@ def main():
             cpu_preds = torch.cat(local_preds)
             cpu_targets = torch.cat(local_targets)
 
-            # Gather predictions and targets across all ranks
-            # Use accelerator.gather (works with all accelerate versions)
-            gpu_preds = cpu_preds.to(accelerator.device)
-            gpu_targets = cpu_targets.to(accelerator.device)
-            all_preds_gathered = accelerator.gather(gpu_preds).cpu()
-            all_targets_gathered = accelerator.gather(gpu_targets).cpu()
-            gathered = [(all_preds_gathered, all_targets_gathered)]
+            # Gather predictions and targets to rank 0 only (memory-efficient)
+            # Avoids duplicating full validation set on every GPU
+            if torch.distributed.is_initialized():
+                # DDP mode: gather only to rank 0
+                # NCCL backend requires CUDA tensors for collective ops
+                gpu_preds = cpu_preds.to(accelerator.device)
+                gpu_targets = cpu_targets.to(accelerator.device)
+
+                if accelerator.is_main_process:
+                    # Rank 0: allocate gather buffers on GPU
+                    all_preds_list = [
+                        torch.zeros_like(gpu_preds)
+                        for _ in range(accelerator.num_processes)
+                    ]
+                    all_targets_list = [
+                        torch.zeros_like(gpu_targets)
+                        for _ in range(accelerator.num_processes)
+                    ]
+                    torch.distributed.gather(
+                        gpu_preds, gather_list=all_preds_list, dst=0
+                    )
+                    torch.distributed.gather(
+                        gpu_targets, gather_list=all_targets_list, dst=0
+                    )
+                    # Move back to CPU for metric computation
+                    gathered = [
+                        (
+                            torch.cat(all_preds_list).cpu(),
+                            torch.cat(all_targets_list).cpu(),
+                        )
+                    ]
+                else:
+                    # Other ranks: send to rank 0, don't allocate gather buffers
+                    torch.distributed.gather(gpu_preds, gather_list=None, dst=0)
+                    torch.distributed.gather(gpu_targets, gather_list=None, dst=0)
+                    gathered = [(cpu_preds, cpu_targets)]  # Placeholder, not used
+            else:
+                # Single-GPU mode: no gathering needed
+                gathered = [(cpu_preds, cpu_targets)]
 
             # Synchronize validation metrics (scalars only - efficient)
             val_loss_scalar = val_loss_sum.item()

wavedl/utils/__init__.py

@@ -15,6 +15,12 @@ from .config import (
     save_config,
     validate_config,
 )
+from .constraints import (
+    ExpressionConstraint,
+    FileConstraint,
+    PhysicsConstrainedLoss,
+    build_constraints,
+)
 from .cross_validation import (
     CVDataset,
     run_cross_validation,
@@ -91,8 +97,11 @@ __all__ = [
     "FIGURE_WIDTH_INCH",
     "FONT_SIZE_TEXT",
     "FONT_SIZE_TICKS",
+    # Constraints
     "CVDataset",
     "DataSource",
+    "ExpressionConstraint",
+    "FileConstraint",
     "HDF5Source",
     "LogCoshLoss",
     "MATSource",
@@ -101,10 +110,12 @@ __all__ = [
     # Metrics
     "MetricTracker",
     "NPZSource",
+    "PhysicsConstrainedLoss",
     "WeightedMSELoss",
     # Distributed
     "broadcast_early_stop",
     "broadcast_value",
+    "build_constraints",
     "calc_pearson",
     "calc_per_target_r2",
     "configure_matplotlib_style",

wavedl/utils/config.py

@@ -306,6 +306,16 @@ def validate_config(
         # Config
         "config",
         "list_models",
+        # Physical Constraints
+        "constraint",
+        "bounds",
+        "constraint_file",
+        "constraint_weight",
+        "constraint_reduction",
+        "positive",
+        "output_bounds",
+        "output_transform",
+        "output_formula",
         # Metadata (internal)
         "_metadata",
     }