wavedl 1.5.1__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.1"
+__version__ = "1.5.2"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl/utils/data.py

@@ -202,18 +202,31 @@ class NPZSource(DataSource):
     """Load data from NumPy .npz archives."""
 
     @staticmethod
-    def _safe_load(path: str, mmap_mode: str | None = None):
-        """Load NPZ with pickle only if needed (sparse matrix support)."""
+    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:
-            return np.load(path, allow_pickle=False, mmap_mode=mmap_mode)
-        except ValueError:
-            # Fallback for sparse matrices stored as object arrays
-            return np.load(path, allow_pickle=True, mmap_mode=mmap_mode)
+            # 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]:
         """Load NPZ file (pickle enabled only for sparse matrices)."""
-        data = self._safe_load(path)
-        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)
@@ -225,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]
 
@@ -243,8 +257,9 @@ class NPZSource(DataSource):
 
         Note: Returns memory-mapped arrays - do NOT modify them.
         """
-        data = self._safe_load(path, 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)
@@ -256,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]
 
@@ -263,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 = self._safe_load(path)
-        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:
@@ -273,6 +290,7 @@ class NPZSource(DataSource):
                 f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
             )
 
+        data = self._safe_load(path, [output_key])
         return data[output_key]
 
 
@@ -751,19 +769,22 @@ def load_test_data(
     except KeyError:
         # Try with just inputs if outputs not found (inference-only mode)
         if format == "npz":
-            data = NPZSource._safe_load(path)
-            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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.5.1
+Version: 1.5.2
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT

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

@@ -1,4 +1,4 @@
-wavedl/__init__.py,sha256=9RY06pDbdsRrjTpxRYGYMTMZl7jB6-Tm-elfcPtvY3Y,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=WIHG3HWT-uF399FQApPpxjggBVFn59cC54HAL4990QU,38550
@@ -24,15 +24,15 @@ 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=tXiBOY1T7eyO9FwOcxvOkPlhMDdm5rCH1TGDPE-jZak,17961
-wavedl/utils/data.py,sha256=11N7Y6w5PYUSBcGANIB1P8JOapqPrVHjTj9CuUeTLac,49172
+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.1.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.5.1.dist-info/METADATA,sha256=pzrIe4ZaDrRRGPIFlgQayG9L_k2q2g8TBP_puOQoTio,45512
-wavedl-1.5.1.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.5.1.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.5.1.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.5.1.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,,