wavedl 1.4.6__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.6"
+__version__ = "1.5.0"
 __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/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:

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.6.dist-info → wavedl-1.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.6
+Version: 1.5.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -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>
 
@@ -766,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
@@ -938,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
 ```
 
@@ -952,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 |
@@ -963,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.6.dist-info → wavedl-1.5.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=ItdZLt3f7sbtAMgiwUtGwwG5Cko4tPLugC_OVhfHMno,1177
+wavedl/__init__.py,sha256=jeYx6dZ0_UD4yXl-I4_Aa63820RzS5-IxJP_KUni8Pw,1177
 wavedl/hpc.py,sha256=-iOjjKkXPcV_quj4vAsMBJN_zWKtD1lMRfIZZBhyGms,8756
-wavedl/hpo.py,sha256=JQvwPgiVHj3sB9Wombn1QO4ammpuo0QAMpRee0LjkuI,14731
+wavedl/hpo.py,sha256=DGCGyt2yhr3WAifAuljhE26gg07CHdaQW4wpDaTKbyo,14968
 wavedl/test.py,sha256=oWGSSC7178loqOxwti-oDXUVogOqbwHL__GfoXSE5Ss,37846
-wavedl/train.py,sha256=9l4aVW1Jd1Sq6yBr8BOoVIKUYmxASDO8XK6BqEkLLWs,50151
+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.6.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.6.dist-info/METADATA,sha256=Hnot8ui2oksCz2UXhj3FHd_Z9MtoP8MJyiMzC6eWq5s,42453
-wavedl-1.4.6.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.6.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.6.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.6.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,,