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

wavedl/models/vit.py

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

wavedl/test.py

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

wavedl/train.py

@@ -931,7 +931,11 @@ def main():
             for x, y in pbar:
                 with accelerator.accumulate(model):
                     pred = model(x)
-                    loss = criterion(pred, y)
+                    # Pass inputs for input-dependent constraints (x_mean, x[...], etc.)
+                    if isinstance(criterion, PhysicsConstrainedLoss):
+                        loss = criterion(pred, y, x)
+                    else:
+                        loss = criterion(pred, y)
 
                     accelerator.backward(loss)
 
@@ -981,7 +985,11 @@ def main():
             with torch.inference_mode():
                 for x, y in val_dl:
                     pred = model(x)
-                    loss = criterion(pred, y)
+                    # Pass inputs for input-dependent constraints
+                    if isinstance(criterion, PhysicsConstrainedLoss):
+                        loss = criterion(pred, y, x)
+                    else:
+                        loss = criterion(pred, y)
 
                     val_loss_sum += loss.detach() * x.size(0)
                     val_samples += x.size(0)
@@ -998,13 +1006,45 @@ def main():
             cpu_preds = torch.cat(local_preds)
             cpu_targets = torch.cat(local_targets)
 
-            # Gather predictions and targets across all ranks
-            # Use accelerator.gather (works with all accelerate versions)
-            gpu_preds = cpu_preds.to(accelerator.device)
-            gpu_targets = cpu_targets.to(accelerator.device)
-            all_preds_gathered = accelerator.gather(gpu_preds).cpu()
-            all_targets_gathered = accelerator.gather(gpu_targets).cpu()
-            gathered = [(all_preds_gathered, all_targets_gathered)]
+            # Gather predictions and targets to rank 0 only (memory-efficient)
+            # Avoids duplicating full validation set on every GPU
+            if torch.distributed.is_initialized():
+                # DDP mode: gather only to rank 0
+                # NCCL backend requires CUDA tensors for collective ops
+                gpu_preds = cpu_preds.to(accelerator.device)
+                gpu_targets = cpu_targets.to(accelerator.device)
+
+                if accelerator.is_main_process:
+                    # Rank 0: allocate gather buffers on GPU
+                    all_preds_list = [
+                        torch.zeros_like(gpu_preds)
+                        for _ in range(accelerator.num_processes)
+                    ]
+                    all_targets_list = [
+                        torch.zeros_like(gpu_targets)
+                        for _ in range(accelerator.num_processes)
+                    ]
+                    torch.distributed.gather(
+                        gpu_preds, gather_list=all_preds_list, dst=0
+                    )
+                    torch.distributed.gather(
+                        gpu_targets, gather_list=all_targets_list, dst=0
+                    )
+                    # Move back to CPU for metric computation
+                    gathered = [
+                        (
+                            torch.cat(all_preds_list).cpu(),
+                            torch.cat(all_targets_list).cpu(),
+                        )
+                    ]
+                else:
+                    # Other ranks: send to rank 0, don't allocate gather buffers
+                    torch.distributed.gather(gpu_preds, gather_list=None, dst=0)
+                    torch.distributed.gather(gpu_targets, gather_list=None, dst=0)
+                    gathered = [(cpu_preds, cpu_targets)]  # Placeholder, not used
+            else:
+                # Single-GPU mode: no gathering needed
+                gathered = [(cpu_preds, cpu_targets)]
 
             # Synchronize validation metrics (scalars only - efficient)
             val_loss_scalar = val_loss_sum.item()

wavedl/utils/cross_validation.py

@@ -128,6 +128,12 @@ def train_fold(
     best_state = None
     history = []
 
+    # Determine if scheduler steps per batch (OneCycleLR) or per epoch
+    # Use isinstance check since class name 'OneCycleLR' != 'onecycle' string in is_epoch_based
+    from torch.optim.lr_scheduler import OneCycleLR
+
+    step_per_batch = isinstance(scheduler, OneCycleLR)
+
     for epoch in range(epochs):
         # Training
         model.train()
@@ -144,6 +150,10 @@ def train_fold(
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             optimizer.step()
 
+            # Per-batch LR scheduling (OneCycleLR)
+            if step_per_batch:
+                scheduler.step()
+
             train_loss += loss.item() * x.size(0)
             train_samples += x.size(0)
 
@@ -186,8 +196,8 @@ def train_fold(
             }
         )
 
-        # LR scheduling
-        if hasattr(scheduler, "step"):
+        # LR scheduling (epoch-based only, not for per-batch schedulers)
+        if not step_per_batch and hasattr(scheduler, "step"):
             if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                 scheduler.step(avg_val_loss)
             else:

wavedl/utils/data.py

@@ -201,9 +201,32 @@ class DataSource(ABC):
 class NPZSource(DataSource):
     """Load data from NumPy .npz archives."""
 
+    @staticmethod
+    def _safe_load(path: str, keys_to_probe: list[str], mmap_mode: str | None = None):
+        """Load NPZ with pickle only if needed (sparse matrix support).
+
+        The error for object arrays happens at ACCESS time, not load time.
+        So we need to probe the keys to detect if pickle is required.
+        """
+        data = np.load(path, allow_pickle=False, mmap_mode=mmap_mode)
+        try:
+            # Probe keys to trigger error if object arrays exist
+            for key in keys_to_probe:
+                if key in data:
+                    _ = data[key]  # This raises ValueError for object arrays
+            return data
+        except ValueError as e:
+            if "allow_pickle=False" in str(e):
+                # Fallback for sparse matrices stored as object arrays
+                data.close() if hasattr(data, "close") else None
+                return np.load(path, allow_pickle=True, mmap_mode=mmap_mode)
+            raise
+
     def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
-        data = np.load(path, allow_pickle=True)
-        keys = list(data.keys())
+        """Load NPZ file (pickle enabled only for sparse matrices)."""
+        # First pass to find keys without loading data
+        with np.load(path, allow_pickle=False) as probe:
+            keys = list(probe.keys())
 
         input_key = self._find_key(keys, INPUT_KEYS)
         output_key = self._find_key(keys, OUTPUT_KEYS)
@@ -215,6 +238,7 @@ class NPZSource(DataSource):
                 f"Found: {keys}"
             )
 
+        data = self._safe_load(path, [input_key, output_key])
         inp = data[input_key]
         outp = data[output_key]
 
@@ -233,8 +257,9 @@ class NPZSource(DataSource):
 
         Note: Returns memory-mapped arrays - do NOT modify them.
         """
-        data = np.load(path, allow_pickle=True, mmap_mode="r")
-        keys = list(data.keys())
+        # First pass to find keys without loading data
+        with np.load(path, allow_pickle=False) as probe:
+            keys = list(probe.keys())
 
         input_key = self._find_key(keys, INPUT_KEYS)
         output_key = self._find_key(keys, OUTPUT_KEYS)
@@ -246,6 +271,7 @@ class NPZSource(DataSource):
                 f"Found: {keys}"
             )
 
+        data = self._safe_load(path, [input_key, output_key], mmap_mode="r")
         inp = data[input_key]
         outp = data[output_key]
 
@@ -253,8 +279,9 @@ class NPZSource(DataSource):
 
     def load_outputs_only(self, path: str) -> np.ndarray:
         """Load only targets from NPZ (avoids loading large input arrays)."""
-        data = np.load(path, allow_pickle=True)
-        keys = list(data.keys())
+        # First pass to find keys without loading data
+        with np.load(path, allow_pickle=False) as probe:
+            keys = list(probe.keys())
 
         output_key = self._find_key(keys, OUTPUT_KEYS)
         if output_key is None:
@@ -263,6 +290,7 @@ class NPZSource(DataSource):
                 f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
             )
 
+        data = self._safe_load(path, [output_key])
         return data[output_key]
 
 
@@ -677,6 +705,7 @@ def load_test_data(
     format: str = "auto",
     input_key: str | None = None,
     output_key: str | None = None,
+    input_channels: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Load test/inference data and return PyTorch tensors ready for model input.
@@ -698,6 +727,9 @@ def load_test_data(
         format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)
         input_key: Custom key for input data (overrides auto-detection)
         output_key: Custom key for output data (overrides auto-detection)
+        input_channels: Explicit number of input channels. If provided, bypasses
+            the heuristic for 4D data. Use input_channels=1 for 3D volumes that
+            look like multi-channel 2D (e.g., depth ≤16).
 
     Returns:
         Tuple of:
@@ -737,19 +769,22 @@ def load_test_data(
     except KeyError:
         # Try with just inputs if outputs not found (inference-only mode)
         if format == "npz":
-            data = np.load(path, allow_pickle=True)
-            keys = list(data.keys())
+            # First pass to find keys
+            with np.load(path, allow_pickle=False) as probe:
+                keys = list(probe.keys())
             inp_key = DataSource._find_key(keys, custom_input_keys)
             if inp_key is None:
                 raise KeyError(
                     f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
                 )
+            out_key = DataSource._find_key(keys, custom_output_keys)
+            keys_to_probe = [inp_key] + ([out_key] if out_key else [])
+            data = NPZSource._safe_load(path, keys_to_probe)
             inp = data[inp_key]
             if inp.dtype == object:
                 inp = np.array(
                     [x.toarray() if hasattr(x, "toarray") else x for x in inp]
                 )
-            out_key = DataSource._find_key(keys, custom_output_keys)
             outp = data[out_key] if out_key else None
         elif format == "hdf5":
             # HDF5: input-only loading for inference
@@ -822,15 +857,20 @@ def load_test_data(
     # Add channel dimension if needed (dimension-agnostic)
     # X.ndim == 2: 1D data (N, L) → (N, 1, L)
     # X.ndim == 3: 2D data (N, H, W) → (N, 1, H, W)
-    # X.ndim == 4: Check if already has channel dim (C <= 16 heuristic)
+    # X.ndim == 4: Check if already has channel dim
     if X.ndim == 2:
         X = X.unsqueeze(1)  # 1D signal: (N, L) → (N, 1, L)
     elif X.ndim == 3:
         X = X.unsqueeze(1)  # 2D image: (N, H, W) → (N, 1, H, W)
     elif X.ndim == 4:
         # Could be 3D volume (N, D, H, W) or 2D with channel (N, C, H, W)
-        # Heuristic: if dim 1 is small (<=16), assume it's already a channel dim
-        if X.shape[1] > 16:
+        if input_channels is not None:
+            # Explicit override: user specifies channel count
+            if input_channels == 1:
+                X = X.unsqueeze(1)  # Add channel: (N, D, H, W) → (N, 1, D, H, W)
+            # else: already has channels, leave as-is
+        elif X.shape[1] > 16:
+            # Heuristic fallback: large dim 1 suggests 3D volume needing channel
             X = X.unsqueeze(1)  # 3D volume: (N, D, H, W) → (N, 1, D, H, W)
     # X.ndim >= 5: assume channel dimension already exists
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.5.0
+Version: 1.5.2
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -99,6 +99,7 @@ The framework handles the engineering challenges of large-scale deep learning
 
 ## ✨ Features
 
+<div align="center">
 <table width="100%">
 <tr>
 <td width="50%" valign="top">
@@ -189,6 +190,7 @@ Deploy models anywhere:
 </td>
 </tr>
 </table>
+</div>
 
 ---
 
@@ -277,6 +279,10 @@ python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
 # Export model to ONNX for deployment (LabVIEW, MATLAB, C++, etc.)
 python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
   --export onnx --export_path <output_file.onnx>
+
+# For 3D volumes with small depth (e.g., 8×128×128), override auto-detection
+python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
+  --input_channels 1
 ```
 
 **Output:**
@@ -372,6 +378,7 @@ WaveDL/
 │       └── utils/                # Utilities
 │           ├── data.py           # Memory-mapped data pipeline
 │           ├── metrics.py        # R², Pearson, visualization
+│           ├── constraints.py    # Physical constraints for training
 │           ├── distributed.py    # DDP synchronization
 │           ├── losses.py         # Loss function factory
 │           ├── optimizers.py     # Optimizer factory

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

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=jeYx6dZ0_UD4yXl-I4_Aa63820RzS5-IxJP_KUni8Pw,1177
+wavedl/__init__.py,sha256=-Xrp2nvZZmC3EqroLASOAApSKxBBqAL8YO0BQQYKRTs,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/test.py,sha256=WIHG3HWT-uF399FQApPpxjggBVFn59cC54HAL4990QU,38550
+wavedl/train.py,sha256=P14r46X-A2wBqZ6dxL8Fv6PADDgt_wdNzbDryvznBkA,54458
 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
@@ -19,20 +19,20 @@ wavedl/models/resnet3d.py,sha256=C7CL4XeSnRlIBuwf5Ei-z183uzIBObrXfkM9Iwuc5e0,874
 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/models/vit.py,sha256=68o9nNjkftvHFArAPupU2ew5e5yCsI2AYaT9TQinVMk,12075
 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/cross_validation.py,sha256=tXiBOY1T7eyO9FwOcxvOkPlhMDdm5rCH1TGDPE-jZak,17961
+wavedl/utils/data.py,sha256=KYw6YJQKPrcqwd_iqJSedws8bAWBOz5fqd3RSb2Wff4,50334
 wavedl/utils/distributed.py,sha256=7wQ3mRjkp_xjPSxDWMnBf5dSkAGUaTzntxbz0BhC5v0,4145
 wavedl/utils/losses.py,sha256=5762M-TBC_hz6uyj1NPbU1vZeFOJQq7fR3-j7OygJRo,7254
 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.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,,
+wavedl-1.5.2.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
+wavedl-1.5.2.dist-info/METADATA,sha256=1iOcg_SkMZW6ju4rDDhyH_dcyOHe4yZqGQEynVT73b0,45512
+wavedl-1.5.2.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+wavedl-1.5.2.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
+wavedl-1.5.2.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
+wavedl-1.5.2.dist-info/RECORD,,