wavedl 1.4.5__tar.gz → 1.5.0__tar.gz

{wavedl-1.4.5/src/wavedl.egg-info → wavedl-1.5.0}/PKG-INFO

@@ -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 → wavedl-1.5.0}/README.md

@@ -4,7 +4,7 @@
 
 ### 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>
@@ -12,7 +12,7 @@
 [![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)
 
@@ -68,14 +68,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>
@@ -92,12 +90,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>
@@ -338,7 +336,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
@@ -682,6 +680,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>
 
@@ -689,18 +785,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:
@@ -719,7 +817,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
@@ -739,7 +837,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 |
 
@@ -891,12 +989,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
 ```
 
@@ -905,7 +1003,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 |
@@ -916,7 +1014,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 → wavedl-1.5.0}/src/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-1.4.5 → wavedl-1.5.0}/src/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-1.4.5 → wavedl-1.5.0}/src/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")