wavedl 1.4.5__py3-none-any.whl → 1.5.0__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.5"
+__version__ = "1.5.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl/hpc.py

@@ -174,7 +174,9 @@ Environment Variables:
     return args, remaining
 
 
-def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
+def print_summary(
+    exit_code: int, wandb_enabled: bool, wandb_mode: str, wandb_dir: str
+) -> None:
     """Print post-training summary and instructions."""
     print()
     print("=" * 40)
@@ -183,7 +185,8 @@ def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
         print("✅ Training completed successfully!")
         print("=" * 40)
 
-        if wandb_mode == "offline":
+        # Only show WandB sync instructions if user enabled wandb
+        if wandb_enabled and wandb_mode == "offline":
             print()
             print("📊 WandB Sync Instructions:")
             print("   From the login node, run:")
@@ -237,6 +240,10 @@ def main() -> int:
         f"--dynamo_backend={args.dynamo_backend}",
     ]
 
+    # Explicitly set multi_gpu to suppress accelerate auto-detection warning
+    if num_gpus > 1:
+        cmd.append("--multi_gpu")
+
     # Add multi-node networking args if specified (required for some clusters)
     if args.main_process_ip:
         cmd.append(f"--main_process_ip={args.main_process_ip}")
@@ -263,8 +270,10 @@ def main() -> int:
         exit_code = 130
 
     # Print summary
+    wandb_enabled = "--wandb" in train_args
     print_summary(
         exit_code,
+        wandb_enabled,
         os.environ.get("WANDB_MODE", "offline"),
         os.environ.get("WANDB_DIR", "/tmp/wandb"),
     )

wavedl/hpo.py

@@ -31,7 +31,7 @@ try:
     import optuna
     from optuna.trial import TrialState
 except ImportError:
-    print("Error: Optuna not installed. Run: pip install -e '.[hpo]'")
+    print("Error: Optuna not installed. Run: pip install wavedl")
     sys.exit(1)
 
 
@@ -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)
@@ -147,6 +148,32 @@ def create_objective(args):
             cmd.extend(["--output_dir", tmpdir])
             history_file = Path(tmpdir) / "training_history.csv"
 
+            # GPU isolation for parallel trials: assign each trial to a specific GPU
+            # This prevents multiple trials from competing for all GPUs
+            env = None
+            if args.n_jobs > 1:
+                import os
+
+                # Detect available GPUs
+                n_gpus = 1
+                try:
+                    import subprocess as sp
+
+                    result_gpu = sp.run(
+                        ["nvidia-smi", "--list-gpus"],
+                        capture_output=True,
+                        text=True,
+                    )
+                    if result_gpu.returncode == 0:
+                        n_gpus = len(result_gpu.stdout.strip().split("\n"))
+                except Exception:
+                    pass
+
+                # Assign trial to a specific GPU (round-robin)
+                gpu_id = trial.number % n_gpus
+                env = os.environ.copy()
+                env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
             # Run training
             try:
                 result = subprocess.run(
@@ -155,6 +182,7 @@ def create_objective(args):
                     text=True,
                     timeout=args.timeout,
                     cwd=Path(__file__).parent,
+                    env=env,
                 )
 
                 # Read best val_loss from training_history.csv (reliable machine-readable)
@@ -248,7 +276,10 @@ Examples:
         "--n_trials", type=int, default=50, help="Number of HPO trials (default: 50)"
     )
     parser.add_argument(
-        "--n_jobs", type=int, default=1, help="Parallel trials (default: 1)"
+        "--n_jobs",
+        type=int,
+        default=-1,
+        help="Parallel trials (-1 = auto-detect GPUs, default: -1)",
     )
     parser.add_argument(
         "--quick",
@@ -287,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(
@@ -315,11 +353,30 @@ Examples:
 
     args = parser.parse_args()
 
+    # Convert to absolute path (child processes may run in different cwd)
+    args.data_path = str(Path(args.data_path).resolve())
+
     # Validate data path
     if not Path(args.data_path).exists():
         print(f"Error: Data file not found: {args.data_path}")
         sys.exit(1)
 
+    # Auto-detect GPUs for n_jobs if not specified
+    if args.n_jobs == -1:
+        try:
+            result_gpu = subprocess.run(
+                ["nvidia-smi", "--list-gpus"],
+                capture_output=True,
+                text=True,
+            )
+            if result_gpu.returncode == 0:
+                args.n_jobs = max(1, len(result_gpu.stdout.strip().split("\n")))
+            else:
+                args.n_jobs = 1
+        except Exception:
+            args.n_jobs = 1
+        print(f"Auto-detected {args.n_jobs} GPU(s) for parallel trials")
+
     # Create study
     print("=" * 60)
     print("WaveDL Hyperparameter Optimization")

wavedl/test.py

@@ -366,13 +366,19 @@ def load_checkpoint(
     logging.info(f"   Building model: {model_name}")
     model = build_model(model_name, in_shape=in_shape, out_size=out_size)
 
-    # Load weights (prefer safetensors)
-    weight_path = checkpoint_dir / "model.safetensors"
-    if not weight_path.exists():
-        weight_path = checkpoint_dir / "pytorch_model.bin"
-
-    if not weight_path.exists():
-        raise FileNotFoundError(f"No model weights found in {checkpoint_dir}")
+    # Load weights (check multiple formats in order of preference)
+    weight_path = None
+    for fname in ["model.safetensors", "model.bin", "pytorch_model.bin"]:
+        candidate = checkpoint_dir / fname
+        if candidate.exists():
+            weight_path = candidate
+            break
+
+    if weight_path is None:
+        raise FileNotFoundError(
+            f"No model weights found in {checkpoint_dir}. "
+            f"Expected one of: model.safetensors, model.bin, pytorch_model.bin"
+        )
 
     if HAS_SAFETENSORS and weight_path.suffix == ".safetensors":
         state_dict = load_safetensors(str(weight_path))

wavedl/train.py

@@ -148,6 +148,24 @@ torch.set_float32_matmul_precision("high")  # Use TF32 for float32 ops
 torch.backends.cudnn.benchmark = True
 
 
+# ==============================================================================
+# LOGGING UTILITIES
+# ==============================================================================
+from contextlib import contextmanager
+
+
+@contextmanager
+def suppress_accelerate_logging():
+    """Temporarily suppress accelerate's verbose checkpoint save messages."""
+    accelerate_logger = logging.getLogger("accelerate.checkpointing")
+    original_level = accelerate_logger.level
+    accelerate_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        accelerate_logger.setLevel(original_level)
+
+
 # ==============================================================================
 # ARGUMENT PARSING
 # ==============================================================================
@@ -357,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"
@@ -535,7 +583,7 @@ def main():
         return
 
     # ==========================================================================
-    # 1. SYSTEM INITIALIZATION
+    # SYSTEM INITIALIZATION
     # ==========================================================================
     # Initialize Accelerator for DDP and mixed precision
     accelerator = Accelerator(
@@ -591,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
@@ -645,7 +693,7 @@ def main():
                 )
 
     # ==========================================================================
-    # 2.5. OPTIMIZER, SCHEDULER & LOSS CONFIGURATION
+    # OPTIMIZER, SCHEDULER & LOSS CONFIGURATION
     # ==========================================================================
     # Parse comma-separated arguments with validation
     try:
@@ -689,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)
 
@@ -744,7 +829,7 @@ def main():
         )
 
     # ==========================================================================
-    # 3. AUTO-RESUME / RESUME FROM CHECKPOINT
+    # AUTO-RESUME / RESUME FROM CHECKPOINT
     # ==========================================================================
     start_epoch = 0
     best_val_loss = float("inf")
@@ -800,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(
@@ -808,7 +893,7 @@ def main():
     )
 
     # ==========================================================================
-    # 5. TRAINING LOOP
+    # TRAINING LOOP
     # ==========================================================================
     # Dynamic console header
     if accelerator.is_main_process:
@@ -1033,7 +1118,8 @@ def main():
             # Step 3: Save checkpoint with all ranks participating
             if is_best_epoch:
                 ckpt_dir = os.path.join(args.output_dir, "best_checkpoint")
-                accelerator.save_state(ckpt_dir)  # All ranks must call this
+                with suppress_accelerate_logging():
+                    accelerator.save_state(ckpt_dir, safe_serialization=False)
 
                 # Step 4: Rank 0 handles metadata and updates tracking variables
                 if accelerator.is_main_process:
@@ -1096,7 +1182,8 @@ def main():
             if periodic_checkpoint_needed:
                 ckpt_name = f"epoch_{epoch + 1}_checkpoint"
                 ckpt_dir = os.path.join(args.output_dir, ckpt_name)
-                accelerator.save_state(ckpt_dir)  # All ranks participate
+                with suppress_accelerate_logging():
+                    accelerator.save_state(ckpt_dir, safe_serialization=False)
 
                 if accelerator.is_main_process:
                     with open(os.path.join(ckpt_dir, "training_meta.pkl"), "wb") as f:
@@ -1147,7 +1234,11 @@ def main():
 
     except KeyboardInterrupt:
         logger.warning("Training interrupted. Saving emergency checkpoint...")
-        accelerator.save_state(os.path.join(args.output_dir, "interrupted_checkpoint"))
+        with suppress_accelerate_logging():
+            accelerator.save_state(
+                os.path.join(args.output_dir, "interrupted_checkpoint"),
+                safe_serialization=False,
+            )
 
     except Exception as e:
         logger.error(f"Critical error: {e}", exc_info=True)

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",
     }

wavedl/utils/constraints.py

@@ -0,0 +1,470 @@
+"""
+Physical Constraints for Training
+=================================
+
+Soft constraint enforcement via penalty-based loss terms.
+
+Usage:
+    # Expression constraints
+    wavedl-train --constraint "y0 > 0" --constraint_weight 0.1
+
+    # Complex constraints via Python file
+    wavedl-train --constraint_file my_constraint.py
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 2.0.0
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+import sys
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+# ==============================================================================
+# SAFE EXPRESSION PARSING
+# ==============================================================================
+SAFE_FUNCTIONS: dict[str, Callable] = {
+    "sin": torch.sin,
+    "cos": torch.cos,
+    "tan": torch.tan,
+    "exp": torch.exp,
+    "log": torch.log,
+    "sqrt": torch.sqrt,
+    "abs": torch.abs,
+    "relu": F.relu,
+    "sigmoid": torch.sigmoid,
+    "softplus": F.softplus,
+    "tanh": torch.tanh,
+    "min": torch.minimum,
+    "max": torch.maximum,
+    "pow": torch.pow,
+    "clamp": torch.clamp,
+}
+
+INPUT_AGGREGATES: dict[str, Callable[[torch.Tensor], torch.Tensor]] = {
+    "x_mean": lambda x: x.mean(dim=tuple(range(1, x.ndim))),
+    "x_sum": lambda x: x.sum(dim=tuple(range(1, x.ndim))),
+    "x_max": lambda x: x.amax(dim=tuple(range(1, x.ndim))),
+    "x_min": lambda x: x.amin(dim=tuple(range(1, x.ndim))),
+    "x_std": lambda x: x.std(dim=tuple(range(1, x.ndim))),
+    "x_energy": lambda x: (x**2).sum(dim=tuple(range(1, x.ndim))),
+}
+
+
+# ==============================================================================
+# SOFT CONSTRAINTS
+# ==============================================================================
+class ExpressionConstraint(nn.Module):
+    """
+    Soft constraint via string expression.
+
+    Parses mathematical expressions using Python's AST for safe evaluation.
+    Supports output variables (y0, y1, ...), input aggregates (x_mean, ...),
+    and whitelisted math functions.
+
+    Example:
+        >>> constraint = ExpressionConstraint("y0 - y1 * y2")
+        >>> penalty = constraint(predictions, inputs)
+
+        >>> constraint = ExpressionConstraint("sin(y0) + cos(y1)")
+        >>> penalty = constraint(predictions, inputs)
+    """
+
+    def __init__(self, expression: str, reduction: str = "mse"):
+        """
+        Args:
+            expression: Mathematical expression to evaluate (should equal 0)
+            reduction: How to reduce violations - 'mse' or 'mae'
+        """
+        super().__init__()
+        self.expression = expression
+        self.reduction = reduction
+        self._tree = ast.parse(expression, mode="eval")
+        self._validate(self._tree)
+
+    def _validate(self, tree: ast.Expression) -> None:
+        """Validate that expression only uses safe functions."""
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+                if node.func.id not in SAFE_FUNCTIONS:
+                    raise ValueError(
+                        f"Unsafe function '{node.func.id}' in constraint. "
+                        f"Allowed: {list(SAFE_FUNCTIONS.keys())}"
+                    )
+
+    def _eval(
+        self, node: ast.AST, pred: torch.Tensor, inputs: torch.Tensor | None
+    ) -> torch.Tensor:
+        """Recursively evaluate AST node."""
+        if isinstance(node, ast.Constant):
+            return torch.tensor(node.value, device=pred.device, dtype=pred.dtype)
+        elif isinstance(node, ast.Name):
+            name = node.id
+            # Output variable: y0, y1, ...
+            if name.startswith("y") and name[1:].isdigit():
+                idx = int(name[1:])
+                if idx >= pred.shape[1]:
+                    raise ValueError(
+                        f"Output index {idx} out of range. "
+                        f"Model has {pred.shape[1]} outputs."
+                    )
+                return pred[:, idx]
+            # Input aggregate: x_mean, x_sum, ...
+            elif name in INPUT_AGGREGATES:
+                if inputs is None:
+                    raise ValueError(
+                        f"Constraint uses '{name}' but inputs not provided."
+                    )
+                return INPUT_AGGREGATES[name](inputs)
+            else:
+                raise ValueError(
+                    f"Unknown variable '{name}'. "
+                    f"Use y0, y1, ... for outputs or {list(INPUT_AGGREGATES.keys())} for inputs."
+                )
+        elif isinstance(node, ast.BinOp):
+            left = self._eval(node.left, pred, inputs)
+            right = self._eval(node.right, pred, inputs)
+            ops = {
+                ast.Add: torch.add,
+                ast.Sub: torch.sub,
+                ast.Mult: torch.mul,
+                ast.Div: torch.div,
+                ast.Pow: torch.pow,
+                ast.Mod: torch.remainder,
+            }
+            if type(node.op) not in ops:
+                raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
+            return ops[type(node.op)](left, right)
+        elif isinstance(node, ast.UnaryOp):
+            operand = self._eval(node.operand, pred, inputs)
+            if isinstance(node.op, ast.USub):
+                return -operand
+            elif isinstance(node.op, ast.UAdd):
+                return operand
+            else:
+                raise ValueError(
+                    f"Unsupported unary operator: {type(node.op).__name__}"
+                )
+        elif isinstance(node, ast.Call):
+            if not isinstance(node.func, ast.Name):
+                raise ValueError("Only direct function calls supported (e.g., sin(x))")
+            func_name = node.func.id
+            if func_name not in SAFE_FUNCTIONS:
+                raise ValueError(f"Unsafe function: {func_name}")
+            args = [self._eval(arg, pred, inputs) for arg in node.args]
+            return SAFE_FUNCTIONS[func_name](*args)
+        elif isinstance(node, ast.Compare):
+            # Comparison operators: y0 > 0, y0 < 1, y0 >= 0, y0 <= 1
+            # Returns penalty (violation amount) when constraint is not satisfied
+            if len(node.ops) != 1 or len(node.comparators) != 1:
+                raise ValueError(
+                    "Only single comparisons supported (e.g., 'y0 > 0', not 'y0 > 0 > y1')"
+                )
+            left = self._eval(node.left, pred, inputs)
+            right = self._eval(node.comparators[0], pred, inputs)
+            op = node.ops[0]
+
+            # Return violation amount (0 if satisfied, positive if violated)
+            if isinstance(
+                op, (ast.Gt, ast.GtE)
+            ):  # y0 > right → penalize if y0 <= right
+                return F.relu(right - left)
+            elif isinstance(
+                op, (ast.Lt, ast.LtE)
+            ):  # y0 < right → penalize if y0 >= right
+                return F.relu(left - right)
+            elif isinstance(op, ast.Eq):  # y0 == right → penalize difference
+                return torch.abs(left - right)
+            elif isinstance(op, ast.NotEq):  # y0 != right → not useful as constraint
+                raise ValueError(
+                    "'!=' is not a valid constraint. Use '==' for equality constraints."
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported comparison operator: {type(op).__name__}"
+                )
+        elif isinstance(node, ast.Subscript):
+            # Input indexing: x[0], x[0,5], x[0,5,10]
+            if not isinstance(node.value, ast.Name) or node.value.id != "x":
+                raise ValueError(
+                    "Subscript indexing only supported for 'x' (inputs). "
+                    "Use x[i], x[i,j], or x[i,j,k]."
+                )
+            if inputs is None:
+                raise ValueError("Constraint uses 'x[...]' but inputs not provided.")
+
+            # Parse indices from the slice
+            indices = self._parse_subscript_indices(node.slice)
+
+            # Validate dimensions match
+            # inputs shape: (batch, dim1) or (batch, dim1, dim2) or (batch, dim1, dim2, dim3)
+            input_ndim = inputs.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"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]]
+            elif len(indices) == 2:
+                return inputs[:, indices[0], indices[1]]
+            elif len(indices) == 3:
+                return inputs[:, indices[0], indices[1], indices[2]]
+            else:
+                raise ValueError("Only 1D, 2D, or 3D input indexing supported.")
+        elif isinstance(node, ast.Expression):
+            return self._eval(node.body, pred, inputs)
+        else:
+            raise ValueError(f"Unsupported AST node type: {type(node).__name__}")
+
+    def _parse_subscript_indices(self, slice_node: ast.AST) -> list[int]:
+        """Parse subscript indices from AST slice node."""
+        if isinstance(slice_node, ast.Constant):
+            # Single index: x[0]
+            return [int(slice_node.value)]
+        elif isinstance(slice_node, ast.Tuple):
+            # Multiple indices: x[0,5] or x[0,5,10]
+            indices = []
+            for elt in slice_node.elts:
+                if not isinstance(elt, ast.Constant):
+                    raise ValueError(
+                        "Only constant indices supported in x[...]. "
+                        "Use x[0,5] not x[i,j]."
+                    )
+                indices.append(int(elt.value))
+            return indices
+        else:
+            raise ValueError(
+                f"Unsupported subscript type: {type(slice_node).__name__}. "
+                "Use x[0], x[0,5], or x[0,5,10]."
+            )
+
+    def forward(
+        self, pred: torch.Tensor, inputs: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """
+        Compute constraint violation penalty.
+
+        Args:
+            pred: Model predictions of shape (N, num_outputs)
+            inputs: Model inputs of shape (N, ...) for input-dependent constraints
+
+        Returns:
+            Scalar penalty value
+        """
+        violation = self._eval(self._tree, pred, inputs)
+        if self.reduction == "mse":
+            return (violation**2).mean()
+        else:  # mae
+            return violation.abs().mean()
+
+    def __repr__(self) -> str:
+        return (
+            f"ExpressionConstraint('{self.expression}', reduction='{self.reduction}')"
+        )
+
+
+class FileConstraint(nn.Module):
+    """
+    Load constraint function from Python file.
+
+    The file must define a function `constraint(pred, inputs=None)` that
+    returns per-sample violation values.
+
+    Example file (my_constraint.py):
+        import torch
+
+        def constraint(pred, inputs=None):
+            # Monotonicity: y0 < y1 < y2
+            diffs = pred[:, 1:] - pred[:, :-1]
+            return torch.relu(-diffs).sum(dim=1)
+
+    Usage:
+        >>> constraint = FileConstraint("my_constraint.py")
+        >>> penalty = constraint(predictions, inputs)
+    """
+
+    def __init__(self, file_path: str, reduction: str = "mse"):
+        """
+        Args:
+            file_path: Path to Python file containing constraint function
+            reduction: How to reduce violations - 'mse' or 'mae'
+        """
+        super().__init__()
+        self.file_path = file_path
+        self.reduction = reduction
+
+        # Load module from file
+        spec = importlib.util.spec_from_file_location("constraint_module", file_path)
+        if spec is None or spec.loader is None:
+            raise ValueError(f"Could not load constraint file: {file_path}")
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["constraint_module"] = module
+        spec.loader.exec_module(module)
+
+        if not hasattr(module, "constraint"):
+            raise ValueError(
+                f"Constraint file must define 'constraint(pred, inputs)' function: {file_path}"
+            )
+
+        self._constraint_fn = module.constraint
+
+    def forward(
+        self, pred: torch.Tensor, inputs: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Evaluate constraint from loaded function."""
+        violation = self._constraint_fn(pred, inputs)
+        if self.reduction == "mse":
+            return (violation**2).mean()
+        else:
+            return violation.abs().mean()
+
+    def __repr__(self) -> str:
+        return f"FileConstraint('{self.file_path}')"
+
+
+# ==============================================================================
+# COMBINED LOSS WRAPPER
+# ==============================================================================
+class PhysicsConstrainedLoss(nn.Module):
+    """
+    Combine base loss with constraint penalties.
+
+    Total Loss = Base Loss + Σ(weight_i × constraint_i)
+
+    Constraints are evaluated in **physical space** (denormalized) while
+    the base loss is computed in normalized space for stable training.
+
+    Example:
+        >>> base_loss = nn.MSELoss()
+        >>> constraints = [ExpressionConstraint("y0 - y1*y2")]
+        >>> criterion = PhysicsConstrainedLoss(
+        ...     base_loss,
+        ...     constraints,
+        ...     weights=[0.1],
+        ...     output_mean=[10, 5, 50],
+        ...     output_std=[2, 1, 10],
+        ... )
+        >>> loss = criterion(pred, target, inputs)
+    """
+
+    def __init__(
+        self,
+        base_loss: nn.Module,
+        constraints: list[nn.Module] | None = None,
+        weights: list[float] | None = None,
+        output_mean: torch.Tensor | list[float] | None = None,
+        output_std: torch.Tensor | list[float] | None = None,
+    ):
+        """
+        Args:
+            base_loss: Base loss function (e.g., MSELoss)
+            constraints: List of constraint modules
+            weights: Weight for each constraint. If shorter than constraints,
+                     last weight is repeated.
+            output_mean: Mean of each output (for denormalization). Shape: (num_outputs,)
+            output_std: Std of each output (for denormalization). Shape: (num_outputs,)
+        """
+        super().__init__()
+        self.base_loss = base_loss
+        self.constraints = nn.ModuleList(constraints or [])
+        self.weights = weights or [0.1]
+
+        # Store scaler as buffers (moves to correct device automatically)
+        if output_mean is not None:
+            if not isinstance(output_mean, torch.Tensor):
+                output_mean = torch.tensor(output_mean, dtype=torch.float32)
+            self.register_buffer("output_mean", output_mean)
+        else:
+            self.register_buffer("output_mean", None)
+
+        if output_std is not None:
+            if not isinstance(output_std, torch.Tensor):
+                output_std = torch.tensor(output_std, dtype=torch.float32)
+            self.register_buffer("output_std", output_std)
+        else:
+            self.register_buffer("output_std", None)
+
+    def _denormalize(self, pred: torch.Tensor) -> torch.Tensor:
+        """Convert normalized predictions to physical values."""
+        if self.output_mean is None or self.output_std is None:
+            return pred
+        return pred * self.output_std + self.output_mean
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        inputs: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Compute combined loss.
+
+        Args:
+            pred: Model predictions (normalized)
+            target: Ground truth targets (normalized)
+            inputs: Model inputs (for input-dependent constraints)
+
+        Returns:
+            Combined loss value
+        """
+        # Base loss in normalized space (stable gradients)
+        loss = self.base_loss(pred, target)
+
+        # Denormalize for constraint evaluation (physical units)
+        pred_physical = self._denormalize(pred)
+
+        for i, constraint in enumerate(self.constraints):
+            weight = self.weights[i] if i < len(self.weights) else self.weights[-1]
+            penalty = constraint(pred_physical, inputs)
+            loss = loss + weight * penalty
+
+        return loss
+
+    def __repr__(self) -> str:
+        has_scaler = self.output_mean is not None
+        return f"PhysicsConstrainedLoss(base={self.base_loss}, constraints={len(self.constraints)}, denormalize={has_scaler})"
+
+
+# ==============================================================================
+# FACTORY FUNCTIONS
+# ==============================================================================
+def build_constraints(
+    expressions: list[str] | None = None,
+    file_path: str | None = None,
+    reduction: str = "mse",
+) -> list[nn.Module]:
+    """
+    Build soft constraint modules from CLI arguments.
+
+    Args:
+        expressions: Expression constraints (e.g., ["y0 - y1*y2", "y0 > 0"])
+        file_path: Path to Python constraint file
+        reduction: Reduction mode for penalties
+
+    Returns:
+        List of constraint modules
+    """
+    constraints: list[nn.Module] = []
+
+    for expr in expressions or []:
+        constraints.append(ExpressionConstraint(expr, reduction))
+
+    if file_path:
+        constraints.append(FileConstraint(file_path, reduction))
+
+    return constraints

wavedl/utils/metrics.py

@@ -560,11 +560,56 @@ def create_training_curves(
         )
         lines.append(line)
 
+    def set_lr_ticks(ax: plt.Axes, data: list[float], n_ticks: int = 4) -> None:
+        """Set n uniformly spaced ticks on LR axis with 10^n format labels."""
+        valid_data = [v for v in data if v is not None and not np.isnan(v) and v > 0]
+        if not valid_data:
+            return
+        vmin, vmax = min(valid_data), max(valid_data)
+        # Snap to clean decade boundaries
+        log_min = np.floor(np.log10(vmin))
+        log_max = np.ceil(np.log10(vmax))
+        # Generate n uniformly spaced ticks as powers of 10
+        log_ticks = np.linspace(log_min, log_max, n_ticks)
+        # Round to nearest integer power of 10 for clean numbers
+        log_ticks = np.round(log_ticks)
+        ticks = 10.0**log_ticks
+        # Remove duplicates while preserving order
+        ticks = list(dict.fromkeys(ticks))
+        ax.set_yticks(ticks)
+        # Format all tick labels as 10^n
+        labels = [f"$10^{{{int(np.log10(t))}}}$" for t in ticks]
+        ax.set_yticklabels(labels)
+        ax.minorticks_off()
+
+    def set_loss_ticks(ax: plt.Axes, data: list[float]) -> None:
+        """Set ticks at powers of 10 that cover the data range."""
+        valid_data = [v for v in data if v is not None and not np.isnan(v) and v > 0]
+        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)))
+        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)]
+        ax.set_yticks(ticks)
+        # Format labels as 10^n
+        labels = [f"$10^{{{i}}}$" for i in range(log_min, log_max + 1)]
+        ax.set_yticklabels(labels)
+        ax.minorticks_off()
+
     ax1.set_xlabel("Epoch")
     ax1.set_ylabel("Loss")
     ax1.set_yscale("log")  # Log scale for loss
     ax1.grid(True, alpha=0.3)
 
+    # Collect all loss values and set clean power of 10 ticks
+    all_loss_values = []
+    for metric in metrics:
+        all_loss_values.extend([h.get(metric, np.nan) for h in history])
+    set_loss_ticks(ax1, all_loss_values)
+
     # Check if learning rate data exists
     has_lr = show_lr and any("lr" in h for h in history)
 
@@ -581,9 +626,11 @@ def create_training_curves(
             alpha=0.7,
             label="Learning Rate",
         )
-        ax2.set_ylabel("Learning Rate", color=COLORS["neutral"])
-        ax2.tick_params(axis="y", labelcolor=COLORS["neutral"])
+        ax2.set_ylabel("Learning Rate")
         ax2.set_yscale("log")  # Log scale for LR
+        set_lr_ticks(ax2, lr_values, n_ticks=4)
+        # Ensure right spine (axis line) is visible
+        ax2.spines["right"].set_visible(True)
         lines.append(line_lr)
 
     # Combined legend

{wavedl-1.4.5.dist-info → wavedl-1.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.5
+Version: 1.5.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -49,7 +49,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 
 ### A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 
-[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
+[![Python 3.11+](https://img.shields.io/badge/Python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
 [![PyTorch 2.x](https://img.shields.io/badge/PyTorch-2.x-ee4c2c.svg?style=plastic&logo=pytorch&logoColor=white)](https://pytorch.org/)
 [![Accelerate](https://img.shields.io/badge/Accelerate-Enabled-yellow.svg?style=plastic&logo=huggingface&logoColor=white)](https://huggingface.co/docs/accelerate/)
 <br>
@@ -57,7 +57,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 [![Lint](https://img.shields.io/github/actions/workflow/status/ductho-le/WaveDL/lint.yml?branch=main&style=plastic&logo=ruff&logoColor=white&label=Lint)](https://github.com/ductho-le/WaveDL/actions/workflows/lint.yml)
 [![Try it on Colab](https://img.shields.io/badge/Try_it_on_Colab-8E44AD?style=plastic&logo=googlecolab&logoColor=white)](https://colab.research.google.com/github/ductho-le/WaveDL/blob/main/notebooks/demo.ipynb)
 <br>
-[![Downloads](https://img.shields.io/pepy/dt/wavedl?style=plastic&logo=pypi&logoColor=white&color=9ACD32)](https://pepy.tech/project/wavedl)
+[![Downloads](https://img.shields.io/badge/dynamic/json?url=https://pypistats.org/api/packages/wavedl/recent?period=month%26mirrors=false&query=data.last_month&style=plastic&logo=pypi&logoColor=white&color=9ACD32&label=Downloads&suffix=/month)](https://pypistats.org/packages/wavedl)
 [![License: MIT](https://img.shields.io/badge/License-MIT-orange.svg?style=plastic)](LICENSE)
 [![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.18012338-008080.svg?style=plastic)](https://doi.org/10.5281/zenodo.18012338)
 
@@ -113,14 +113,12 @@ Train on datasets larger than RAM:
 </td>
 <td width="50%" valign="top">
 
-**🧠 One-Line Model Registration**
+**🧠 Models? We've Got Options**
 
-Plug in any architecture:
-```python
-@register_model("my_net")
-class MyNet(BaseModel): ...
-```
-Design your model. Register with one line.
+38 architectures, ready to go:
+- CNNs, ResNets, ViTs, EfficientNets...
+- All adapted for regression
+- [Add your own](#adding-custom-models) in one line
 
 </td>
 </tr>
@@ -137,12 +135,12 @@ Multi-GPU training without the pain:
 </td>
 <td width="50%" valign="top">
 
-**📊 Publish-Ready Output**
+**🔬 Physics-Constrained Training**
 
-Results go straight to your paper:
-- 11 diagnostic plots with LaTeX styling
-- Multi-format export (PNG, PDF, SVG, ...)
-- MAE in physical units per parameter
+Make your model respect the laws:
+- Enforce bounds, positivity, equations
+- Simple expression syntax or Python
+- [Custom constraints](#physical-constraints) for various laws
 
 </td>
 </tr>
@@ -383,7 +381,7 @@ WaveDL/
 ├── configs/                      # YAML config templates
 ├── examples/                     # Ready-to-run examples
 ├── notebooks/                    # Jupyter notebooks
-├── unit_tests/                   # Pytest test suite (704 tests)
+├── unit_tests/                   # Pytest test suite (725 tests)
 │
 ├── pyproject.toml                # Package config, dependencies
 ├── CHANGELOG.md                  # Version history
@@ -727,6 +725,104 @@ seed: 2025
 
 </details>
 
+<details>
+<summary><b>Physical Constraints</b> — Enforce Physics During Training</summary>
+
+Add penalty terms to the loss function to enforce physical laws:
+
+```
+Total Loss = Data Loss + weight × penalty(violation)
+```
+
+### Expression Constraints
+
+```bash
+# Positivity
+--constraint "y0 > 0"
+
+# Bounds
+--constraint "y0 >= 0" "y0 <= 1"
+
+# Equations (penalize deviations from zero)
+--constraint "y2 - y0 * y1"
+
+# Input-dependent constraints
+--constraint "y0 - 2*x[0]"
+
+# Multiple constraints with different weights
+--constraint "y0 > 0" "y1 - y2" --constraint_weight 0.1 1.0
+```
+
+### Custom Python Constraints
+
+For complex physics (matrix operations, implicit equations):
+
+```python
+# my_constraint.py
+import torch
+
+def constraint(pred, inputs=None):
+    """
+    Args:
+        pred:   (batch, num_outputs)
+        inputs: (batch, features) or (batch, C, H, W) or (batch, C, D, H, W)
+    Returns:
+        (batch,) — violation per sample (0 = satisfied)
+    """
+    # Outputs (same for all data types)
+    y0, y1, y2 = pred[:, 0], pred[:, 1], pred[:, 2]
+
+    # Inputs — Tabular: (batch, features)
+    # x0 = inputs[:, 0]                    # Feature 0
+    # x_sum = inputs.sum(dim=1)            # Sum all features
+
+    # Inputs — Images: (batch, C, H, W)
+    # pixel = inputs[:, 0, 3, 5]           # Pixel at (3,5), channel 0
+    # img_mean = inputs.mean(dim=(1,2,3))  # Mean over C,H,W
+
+    # Inputs — 3D Volumes: (batch, C, D, H, W)
+    # voxel = inputs[:, 0, 2, 3, 5]        # Voxel at (2,3,5), channel 0
+
+    # Example constraints:
+    # return y2 - y0 * y1                                    # Wave equation
+    # return y0 - 2 * inputs[:, 0]                           # Output = 2×input
+    # return inputs[:, 0, 3, 5] * y0 + inputs[:, 0, 6, 7] * y1  # Mixed
+
+    return y0 - y1 * y2
+```
+
+```bash
+--constraint_file my_constraint.py --constraint_weight 1.0
+```
+
+---
+
+### Reference
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--constraint` | — | Expression(s): `"y0 > 0"`, `"y0 - y1*y2"` |
+| `--constraint_file` | — | Python file with `constraint(pred, inputs)` |
+| `--constraint_weight` | `0.1` | Penalty weight(s) |
+| `--constraint_reduction` | `mse` | `mse` (squared) or `mae` (linear) |
+
+#### Expression Syntax
+
+| Variable | Meaning |
+|----------|---------|
+| `y0`, `y1`, ... | Model outputs |
+| `x[0]`, `x[1]`, ... | Input values (1D tabular) |
+| `x[i,j]`, `x[i,j,k]` | Input values (2D/3D: images, volumes) |
+| `x_mean`, `x_sum`, `x_max`, `x_min`, `x_std` | Input aggregates |
+
+**Operators:** `+`, `-`, `*`, `/`, `**`, `>`, `<`, `>=`, `<=`, `==`
+
+**Functions:** `sin`, `cos`, `exp`, `log`, `sqrt`, `sigmoid`, `softplus`, `tanh`, `relu`, `abs`
+
+</details>
+
+
+
 <details>
 <summary><b>Hyperparameter Search (HPO)</b></summary>
 
@@ -734,18 +830,20 @@ Automatically find the best training configuration using [Optuna](https://optuna
 
 **Run HPO:**
 
-You specify which models to search and how many trials to run:
 ```bash
-# Search 3 models with 100 trials
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 100
+# Basic HPO (auto-detects GPUs for parallel trials)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 100
 
-# Search 1 model (faster)
-python -m wavedl.hpo --data_path train.npz --models cnn --n_trials 50
+# Search multiple models
+wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
 
-# Search all your candidate models
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 resnet50 vit_small densenet121 --n_trials 200
+# Quick mode (fewer parameters, faster)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
 ```
 
+> [!TIP]
+> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+
 **Train with best parameters**
 
 After HPO completes, it prints the optimal command:
@@ -764,7 +862,7 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | Schedulers | [all 8](#learning-rate-schedulers) | `--schedulers X Y` |
 | Losses | [all 6](#loss-functions) | `--losses X Y` |
 | Learning rate | 1e-5 → 1e-2 | (always searched) |
-| Batch size | 64, 128, 256, 512 | (always searched) |
+| Batch size | 16, 32, 64, 128 | (always searched) |
 
 **Quick Mode** (`--quick`):
 - Uses minimal defaults: cnn + adamw + plateau + mse
@@ -784,7 +882,7 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
-| `--n_jobs` | `1` | Parallel trials (multi-GPU) |
+| `--n_jobs` | `-1` | Parallel trials (-1 = auto-detect GPUs) |
 | `--max_epochs` | `50` | Max epochs per trial |
 | `--output` | `hpo_results.json` | Output file |
 
@@ -936,12 +1034,12 @@ The `examples/` folder contains a **complete, ready-to-run example** for **mater
 ```bash
 # Run inference on the example data
 python -m wavedl.test --checkpoint ./examples/elastic_cnn_example/best_checkpoint \
-  --data_path ./examples/elastic_cnn_example/Test_data_100.mat \
+  --data_path ./examples/elastic_cnn_example/Test_data_500.mat \
   --plot --save_predictions --output_dir ./examples/elastic_cnn_example/test_results
 
 # Export to ONNX (already included as model.onnx)
 python -m wavedl.test --checkpoint ./examples/elastic_cnn_example/best_checkpoint \
-  --data_path ./examples/elastic_cnn_example/Test_data_100.mat \
+  --data_path ./examples/elastic_cnn_example/Test_data_500.mat \
   --export onnx --export_path ./examples/elastic_cnn_example/model.onnx
 ```
 
@@ -950,7 +1048,7 @@ python -m wavedl.test --checkpoint ./examples/elastic_cnn_example/best_checkpoin
 | File | Description |
 |------|-------------|
 | `best_checkpoint/` | Pre-trained CNN checkpoint |
-| `Test_data_100.mat` | 100 sample test set (500×500 dispersion curves → *h*, √(*E*/ρ), *ν*) |
+| `Test_data_500.mat` | 500 sample test set (500×500 dispersion curves → *h*, √(*E*/ρ), *ν*) |
 | `model.onnx` | ONNX export with embedded de-normalization |
 | `training_history.csv` | Epoch-by-epoch training metrics (loss, R², LR, etc.) |
 | `training_curves.png` | Training/validation loss and learning rate plot |
@@ -961,7 +1059,7 @@ python -m wavedl.test --checkpoint ./examples/elastic_cnn_example/best_checkpoin
 
 <p align="center">
   <img src="examples/elastic_cnn_example/training_curves.png" alt="Training curves" width="600"><br>
-  <em>Training and validation loss over 162 epochs with learning rate schedule</em>
+  <em>Training and validation loss over 227 epochs with <code>onecycle</code> learning rate schedule</em>
 </p>
 
 **Inference Results:**

{wavedl-1.4.5.dist-info → wavedl-1.5.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=2ro7SYQ3wCmq-ejiAm5sd6BeXf6sZgixC9U2vS7Ckbs,1177
-wavedl/hpc.py,sha256=0h8IZzOT0EzmEv3fU9cKyRVE9V1ivtBzbjuBCaxYadc,8445
-wavedl/hpo.py,sha256=YJXsnSGEBSVUqp_2ah7zu3_VClAUqZrdkuzDaSqQUjU,12952
-wavedl/test.py,sha256=81al6vQBDAJ3CpSEtxZn6xzR1c4-jo28R7tX_84KROc,37642
-wavedl/train.py,sha256=_pW7prvlNqfUGrGweHO2QelS87UiAYKvyJwqMAIj6yI,49292
+wavedl/__init__.py,sha256=jeYx6dZ0_UD4yXl-I4_Aa63820RzS5-IxJP_KUni8Pw,1177
+wavedl/hpc.py,sha256=-iOjjKkXPcV_quj4vAsMBJN_zWKtD1lMRfIZZBhyGms,8756
+wavedl/hpo.py,sha256=DGCGyt2yhr3WAifAuljhE26gg07CHdaQW4wpDaTKbyo,14968
+wavedl/test.py,sha256=oWGSSC7178loqOxwti-oDXUVogOqbwHL__GfoXSE5Ss,37846
+wavedl/train.py,sha256=Xy0-W9Kxqdc7WaAyt_HDSZwhZKrfjVbZd2KzgSid6rQ,52457
 wavedl/models/__init__.py,sha256=lfSohEnAUztO14nuwayMJhPjpgySzRN3jGiyAUuBmAU,3206
 wavedl/models/_template.py,sha256=J_D8taSPmV8lBaucN_vU-WiG98iFr7CJrZVNNX_Tdts,4600
 wavedl/models/base.py,sha256=T9iDF9IQM2MYucG_ggQd31rieUkB2fob-nkHyNIl2ak,7337
@@ -20,18 +20,19 @@ wavedl/models/swin.py,sha256=p-okfq3Qm4_neJTxCcMzoHoVzC0BHW3BMnbpr_Ri2U0,13224
 wavedl/models/tcn.py,sha256=RtY13QpFHqz72b4ultv2lStCIDxfvjySVe5JaTx_GaM,12601
 wavedl/models/unet.py,sha256=LqIXhasdBygwP7SZNNmiW1bHMPaJTVBpaeHtPgEHkdU,7790
 wavedl/models/vit.py,sha256=0C3GZk11VsYFTl14d86Wtl1Zk1T5rYJjvkaEfEN4N3k,11100
-wavedl/utils/__init__.py,sha256=YMgzuwndjr64kt9k0_6_9PMJYTVdiaH5veSMff_ZycA,3051
-wavedl/utils/config.py,sha256=fMoucikIQHn85mVhGMa7TnXTuFDcEEPjfXk2EjbkJR0,10591
+wavedl/utils/__init__.py,sha256=s5R9bRmJ8GNcJrD3OSAOXzwZJIXZbdYrAkZnus11sVQ,3300
+wavedl/utils/config.py,sha256=jGW-K7AYB6zrD2BfVm2XPnSY9rbfL_EkM4bwxhBLuwM,10859
+wavedl/utils/constraints.py,sha256=Pof5hzeTSGsPY_E6Sc8iMQDaXc_zfEasQI2tCszk_gw,17614
 wavedl/utils/cross_validation.py,sha256=117ac9KDzaIaqhtP8ZRs15Xpqmq5fLpX2-vqkNvtMaU,17487
 wavedl/utils/data.py,sha256=_OaWvU5oFVJW0NwM5WyDD0Kb1hy5MgvJIFpzvJGux9w,48214
 wavedl/utils/distributed.py,sha256=7wQ3mRjkp_xjPSxDWMnBf5dSkAGUaTzntxbz0BhC5v0,4145
 wavedl/utils/losses.py,sha256=5762M-TBC_hz6uyj1NPbU1vZeFOJQq7fR3-j7OygJRo,7254
-wavedl/utils/metrics.py,sha256=mkCpqZwl_XUpNvA5Ekjf7y-HqApafR7eR6EuA8cBdM8,37287
+wavedl/utils/metrics.py,sha256=EJmJvF7gACQsUoKYldlladN_SbnRiuE-Smj0eSnbraQ,39394
 wavedl/utils/optimizers.py,sha256=PyIkJ_hRhFi_Fio81Gy5YQNhcME0JUUEl8OTSyu-0RA,6323
 wavedl/utils/schedulers.py,sha256=e6Sf0yj8VOqkdwkUHLMyUfGfHKTX4NMr-zfgxWqCTYI,7659
-wavedl-1.4.5.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.5.dist-info/METADATA,sha256=4ltxFDaqPqh4XUAW_K8nkFmvqBzPcL2cxmghH11GMWg,42191
-wavedl-1.4.5.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.5.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.5.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.5.dist-info/RECORD,,
+wavedl-1.5.0.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
+wavedl-1.5.0.dist-info/METADATA,sha256=jlQPfPLtwIXzWHKwgN8xQjzco_xB7L3j7p_dl3TJA8E,45224
+wavedl-1.5.0.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+wavedl-1.5.0.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
+wavedl-1.5.0.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
+wavedl-1.5.0.dist-info/RECORD,,