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

wavedl/models/__init__.py

@@ -77,6 +77,15 @@ from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
+# Optional RATENet (unpublished, may be gitignored)
+try:
+    from .ratenet import RATENet, RATENetLite, RATENetTiny, RATENetV2
+
+    _HAS_RATENET = True
+except ImportError:
+    _HAS_RATENET = False
+
+
 # Optional timm-based models (imported conditionally)
 try:
     from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
@@ -111,6 +120,7 @@ __all__ = [
     "MC3_18",
     "MODEL_REGISTRY",
     "TCN",
+    # Classes (uppercase first, alphabetically)
     "BaseModel",
     "ConvNeXtBase_",
     "ConvNeXtSmall",
@@ -152,6 +162,7 @@ __all__ = [
     "VimBase",
     "VimSmall",
     "VimTiny",
+    # Functions (lowercase, alphabetically)
     "build_model",
     "get_model",
     "list_models",
@@ -186,3 +197,14 @@ if _HAS_TIMM_MODELS:
             "UniRepLKNetTiny",
         ]
     )
+
+# Add RATENet models to __all__ if available (unpublished)
+if _HAS_RATENET:
+    __all__.extend(
+        [
+            "RATENet",
+            "RATENetLite",
+            "RATENetTiny",
+            "RATENetV2",
+        ]
+    )

wavedl/test.py

@@ -398,6 +398,14 @@ def load_checkpoint(
 
     if HAS_SAFETENSORS and weight_path.suffix == ".safetensors":
         state_dict = load_safetensors(str(weight_path))
+    elif weight_path.suffix == ".safetensors":
+        # Safetensors file exists but library not installed
+        raise ImportError(
+            f"Checkpoint uses safetensors format ({weight_path.name}) but "
+            f"'safetensors' package is not installed. Install it with:\n"
+            f"    pip install safetensors\n"
+            f"Or convert the checkpoint to PyTorch format (model.bin)."
+        )
     else:
         state_dict = torch.load(weight_path, map_location="cpu", weights_only=True)
 

wavedl/train.py

@@ -1266,10 +1266,10 @@ def main():
                         os.path.join(args.output_dir, "best_model_weights.pth"),
                     )
 
-                    # Copy scaler to checkpoint for portability
+                    # Copy scaler to checkpoint for portability (always overwrite to stay current)
                     scaler_src = os.path.join(args.output_dir, "scaler.pkl")
                     scaler_dst = os.path.join(ckpt_dir, "scaler.pkl")
-                    if os.path.exists(scaler_src) and not os.path.exists(scaler_dst):
+                    if os.path.exists(scaler_src):
                         shutil.copy2(scaler_src, scaler_dst)
 
                     logger.info(

wavedl/utils/data.py

@@ -984,7 +984,15 @@ def load_test_data(
                 f"Available keys depend on file format. Original error: {e}"
             ) from e
 
-        # Legitimate fallback: no explicit output_key, outputs just not present
+        # Also fail-fast if explicit input_key was provided but not found
+        # This prevents silently loading a different tensor when user mistyped key
+        if input_key is not None:
+            raise KeyError(
+                f"Explicit --input_key '{input_key}' not found in file. "
+                f"Original error: {e}"
+            ) from e
+
+        # Legitimate fallback: no explicit keys, outputs just not present
         if format == "npz":
             # First pass to find keys
             with np.load(path, allow_pickle=False) as probe:
@@ -1524,21 +1532,43 @@ def prepare_data(
 
             logger.info("   ✔ Cache creation complete, synchronizing ranks...")
         else:
-            # NON-MAIN RANKS: Wait for cache creation
-            # Log that we're waiting (helps with debugging)
+            # NON-MAIN RANKS: Wait for cache creation with timeout
+            # Use monotonic clock (immune to system clock changes)
             import time
 
-            wait_start = time.time()
+            wait_start = time.monotonic()
+
+            # Robust env parsing with guards for invalid/non-positive values
+            DEFAULT_CACHE_TIMEOUT = 3600  # 1 hour default
+            try:
+                env_timeout = os.environ.get("WAVEDL_CACHE_TIMEOUT", "")
+                CACHE_TIMEOUT = (
+                    int(env_timeout) if env_timeout else DEFAULT_CACHE_TIMEOUT
+                )
+                if CACHE_TIMEOUT <= 0:
+                    CACHE_TIMEOUT = DEFAULT_CACHE_TIMEOUT
+            except ValueError:
+                CACHE_TIMEOUT = DEFAULT_CACHE_TIMEOUT
+
             while not (
                 os.path.exists(CACHE_FILE)
                 and os.path.exists(SCALER_FILE)
                 and os.path.exists(META_FILE)
             ):
                 time.sleep(5)  # Check every 5 seconds
-                elapsed = time.time() - wait_start
+                elapsed = time.monotonic() - wait_start
+
+                if elapsed > CACHE_TIMEOUT:
+                    raise RuntimeError(
+                        f"[Rank {accelerator.process_index}] Timeout waiting for cache "
+                        f"files after {CACHE_TIMEOUT}s. Rank 0 may have failed during "
+                        f"cache generation. Check rank 0 logs for errors."
+                    )
+
                 if elapsed > 60 and int(elapsed) % 60 < 5:  # Log every ~minute
                     logger.info(
-                        f"   [Rank {accelerator.process_index}] Waiting for cache creation... ({int(elapsed)}s)"
+                        f"   [Rank {accelerator.process_index}] Waiting for cache "
+                        f"creation... ({int(elapsed)}s / {CACHE_TIMEOUT}s max)"
                     )
             # Small delay to ensure files are fully written
             time.sleep(2)

{wavedl-1.6.2.dist-info → wavedl-1.6.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.6.2
+Version: 1.6.3
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -214,11 +214,11 @@ This installs everything you need: training, inference, HPO, ONNX export.
 ```bash
 git clone https://github.com/ductho-le/WaveDL.git
 cd WaveDL
-pip install -e .
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For contributor setup (pre-commit hooks), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
@@ -273,8 +273,6 @@ accelerate launch --num_machines 2 --main_process_ip <ip> -m wavedl.train --mode
 
 ### Testing & Inference
 
-After training, use `wavedl-test` to evaluate your model on test data:
-
 ```bash
 # Basic inference
 wavedl-test --checkpoint <checkpoint_folder> --data_path <test_data>

{wavedl-1.6.2.dist-info → wavedl-1.6.3.dist-info}/RECORD

@@ -1,9 +1,9 @@
-wavedl/__init__.py,sha256=hFGU_j86Beexkcrn_V3fotGQ4ncwLGvz2lCOejEJ-f0,1177
+wavedl/__init__.py,sha256=CdjeIFEZh4ccwPiNsUxpac5l3pxwb9e1SmCNiLiYG1I,1177
 wavedl/hpo.py,sha256=nEiy-2O_5EhxF5hU8X5TviSAiXfVrTQx0-VE6baW7JQ,14633
 wavedl/launcher.py,sha256=_CFlgpKgHrtZebl1yQbJZJEcob06Y9-fqnRYzwW7UJQ,11776
-wavedl/test.py,sha256=1UUy9phCqrr3h_lN6mGJ7Sj73skDg4KyLk2Yuq9DiKU,38797
-wavedl/train.py,sha256=vBufy6gHShawgj8O6dvVER9TPhORa1s7L6pQtTe-N5M,57824
-wavedl/models/__init__.py,sha256=8OiT2seq1qBiUzKaSkmh_VOLJlLTT9Cn-mjhMHKGFpI,5203
+wavedl/test.py,sha256=5MzBtEH2lWWYG23Fz-VpMFAWR5SfZbFomBbu8ptsZRU,39208
+wavedl/train.py,sha256=DizXhi9BFL8heLmO8ENiNm2QubAMm9mdpDiaBlULeKM,57824
+wavedl/models/__init__.py,sha256=hyR__h_D8PsUQCBSM5tj94yYK00uG8ABjEmj_RR8SGE,5719
 wavedl/models/_pretrained_utils.py,sha256=VPdU1DwJB93ZBf_GFIgb8-6BbAt18Phs4yorwlhLw70,12404
 wavedl/models/_template.py,sha256=J_D8taSPmV8lBaucN_vU-WiG98iFr7CJrZVNNX_Tdts,4600
 wavedl/models/base.py,sha256=bDoHYFli-aR8amcFYXbF98QYaKSCEwZWpvOhN21ODro,9075
@@ -32,15 +32,15 @@ wavedl/utils/__init__.py,sha256=s5R9bRmJ8GNcJrD3OSAOXzwZJIXZbdYrAkZnus11sVQ,3300
 wavedl/utils/config.py,sha256=MXkaVc1_zo8sDro8mjtK1MV65t2z8b1Z6fviwSorNiY,10534
 wavedl/utils/constraints.py,sha256=V9Gyi8-uIMbLUWb2cOaHZD0SliWLxVrHZHFyo4HWK7g,18031
 wavedl/utils/cross_validation.py,sha256=HfInyZ8gUROc_AyihYKzzUE0vnoPt_mFvAI2OPK4P54,17945
-wavedl/utils/data.py,sha256=5ph2Pi8PKvuaSoJaXbFIL9WsX8pTN0A6P8FdmxvXdv4,63469
+wavedl/utils/data.py,sha256=HXod6i6g76oFAjLz7xepBPQEFHRgQ7E1M-YSKwUya-I,64799
 wavedl/utils/distributed.py,sha256=7wQ3mRjkp_xjPSxDWMnBf5dSkAGUaTzntxbz0BhC5v0,4145
 wavedl/utils/losses.py,sha256=KWpU5S5noFzp3bLbcH9RNpkFPajy6fyTIh5cNjI-BYA,7038
 wavedl/utils/metrics.py,sha256=YoqiXWOsUB9Y4_alj8CmHcTgnV4MFcH5PH4XlIC13HY,40304
 wavedl/utils/optimizers.py,sha256=ZoETDSOK1fWUT2dx69PyYebeM8Vcqf9zOIKUERWk5HY,6107
 wavedl/utils/schedulers.py,sha256=K6YCiyiMM9rb0cCRXTp89noXeXcAyUEiePr27O5Cozs,7408
-wavedl-1.6.2.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.6.2.dist-info/METADATA,sha256=2mTyuip32AneUURV3K8oAjZQ2rA_13AB16R-VyRN5s8,47659
-wavedl-1.6.2.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.6.2.dist-info/entry_points.txt,sha256=NuAvdiG93EYYpqv-_1wf6PN0WqBfABanDKalNKe2GOs,148
-wavedl-1.6.2.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.6.2.dist-info/RECORD,,
+wavedl-1.6.3.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
+wavedl-1.6.3.dist-info/METADATA,sha256=BtPoAiMwvE58b-dmR0TfPPLLBqYzaAiJ2GBUd0FBSY0,47613
+wavedl-1.6.3.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+wavedl-1.6.3.dist-info/entry_points.txt,sha256=NuAvdiG93EYYpqv-_1wf6PN0WqBfABanDKalNKe2GOs,148
+wavedl-1.6.3.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
+wavedl-1.6.3.dist-info/RECORD,,