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

wavedl/hpc.py

@@ -26,7 +26,6 @@ import os
 import shutil
 import subprocess
 import sys
-import tempfile
 from pathlib import Path
 
 
@@ -60,26 +59,40 @@ def setup_hpc_environment() -> None:
     Handles restricted home directories (e.g., Compute Canada) and
     offline logging configurations.
     """
-    # Use SLURM_TMPDIR if available, otherwise system temp
-    tmpdir = os.environ.get(
-        "SLURM_TMPDIR", os.environ.get("TMPDIR", tempfile.gettempdir())
-    )
-
-    # Configure directories for systems with restricted home directories
-    os.environ.setdefault("MPLCONFIGDIR", f"{tmpdir}/matplotlib")
-    os.environ.setdefault(
-        "FONTCONFIG_PATH", os.environ.get("FONTCONFIG_PATH", "/etc/fonts")
-    )
-    os.environ.setdefault("XDG_CACHE_HOME", f"{tmpdir}/.cache")
-
-    # Ensure matplotlib config dir exists
-    Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
+    # Check if home is writable
+    home = os.path.expanduser("~")
+    home_writable = os.access(home, os.W_OK)
+
+    # Use SLURM_TMPDIR if available, otherwise CWD for HPC, or system temp
+    if home_writable:
+        # Local machine - let libraries use defaults
+        cache_base = None
+    else:
+        # HPC with restricted home - use CWD for persistent caches
+        cache_base = os.getcwd()
+
+    # Only set environment variables if home is not writable
+    if cache_base:
+        os.environ.setdefault("TORCH_HOME", f"{cache_base}/.torch_cache")
+        os.environ.setdefault("MPLCONFIGDIR", f"{cache_base}/.matplotlib")
+        os.environ.setdefault("FONTCONFIG_CACHE", f"{cache_base}/.fontconfig")
+        os.environ.setdefault("XDG_CACHE_HOME", f"{cache_base}/.cache")
+
+        # Ensure directories exist
+        for env_var in [
+            "TORCH_HOME",
+            "MPLCONFIGDIR",
+            "FONTCONFIG_CACHE",
+            "XDG_CACHE_HOME",
+        ]:
+            Path(os.environ[env_var]).mkdir(parents=True, exist_ok=True)
 
     # WandB configuration (offline by default for HPC)
     os.environ.setdefault("WANDB_MODE", "offline")
-    os.environ.setdefault("WANDB_DIR", f"{tmpdir}/wandb")
-    os.environ.setdefault("WANDB_CACHE_DIR", f"{tmpdir}/wandb_cache")
-    os.environ.setdefault("WANDB_CONFIG_DIR", f"{tmpdir}/wandb_config")
+    if cache_base:
+        os.environ.setdefault("WANDB_DIR", f"{cache_base}/.wandb")
+        os.environ.setdefault("WANDB_CACHE_DIR", f"{cache_base}/.wandb_cache")
+        os.environ.setdefault("WANDB_CONFIG_DIR", f"{cache_base}/.wandb_config")
 
     # Suppress non-critical warnings
     os.environ.setdefault(

wavedl/test.py

@@ -29,29 +29,52 @@ Author: Ductho Le (ductho.le@outlook.com)
 # ==============================================================================
 # ENVIRONMENT CONFIGURATION (must be before matplotlib import)
 # ==============================================================================
+# Auto-configure writable cache directories when home is not writable.
+# Uses current working directory as fallback - works on HPC and local machines.
 import os
 
 
-os.environ.setdefault("MPLCONFIGDIR", os.getenv("TMPDIR", "/tmp") + "/matplotlib")
-os.environ.setdefault("FONTCONFIG_PATH", "/etc/fonts")
-
-import argparse
-import logging
-import pickle
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn as nn
-from sklearn.metrics import mean_absolute_error, r2_score
-from torch.utils.data import DataLoader, TensorDataset
-from tqdm.auto import tqdm
+def _setup_cache_dir(env_var: str, subdir: str) -> None:
+    """Set cache directory to CWD if home is not writable."""
+    if env_var in os.environ:
+        return  # User already set, respect their choice
+
+    # Check if home is writable
+    home = os.path.expanduser("~")
+    if os.access(home, os.W_OK):
+        return  # Home is writable, let library use defaults
+
+    # Home not writable - use current working directory
+    cache_path = os.path.join(os.getcwd(), f".{subdir}")
+    os.makedirs(cache_path, exist_ok=True)
+    os.environ[env_var] = cache_path
+
+
+# Configure cache directories (before any library imports)
+_setup_cache_dir("TORCH_HOME", "torch_cache")
+_setup_cache_dir("MPLCONFIGDIR", "matplotlib")
+_setup_cache_dir("FONTCONFIG_CACHE", "fontconfig")
+_setup_cache_dir("XDG_DATA_HOME", "local/share")
+_setup_cache_dir("XDG_STATE_HOME", "local/state")
+_setup_cache_dir("XDG_CACHE_HOME", "cache")
+
+import argparse  # noqa: E402
+import logging  # noqa: E402
+import pickle  # noqa: E402
+from pathlib import Path  # noqa: E402
+
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import torch  # noqa: E402
+import torch.nn as nn  # noqa: E402
+from sklearn.metrics import mean_absolute_error, r2_score  # noqa: E402
+from torch.utils.data import DataLoader, TensorDataset  # noqa: E402
+from tqdm.auto import tqdm  # noqa: E402
 
 # Local imports
-from wavedl.models import build_model, list_models
-from wavedl.utils import (
+from wavedl.models import build_model, list_models  # noqa: E402
+from wavedl.utils import (  # noqa: E402
     FIGURE_DPI,
     calc_pearson,
     load_test_data,

wavedl/train.py

@@ -40,45 +40,34 @@ from __future__ import annotations
 # =============================================================================
 # HPC Environment Setup (MUST be before any library imports)
 # =============================================================================
-# Set writable cache directories for matplotlib and fontconfig ONLY when
-# the default paths are not writable (common on HPC clusters).
+# Auto-configure writable cache directories when home is not writable.
+# Uses current working directory as fallback - works on HPC and local machines.
 import os
-import tempfile
 
 
-def _setup_cache_dir(env_var: str, default_subpath: str) -> None:
-    """Set cache directory only if default path is not writable."""
+def _setup_cache_dir(env_var: str, subdir: str) -> None:
+    """Set cache directory to CWD if home is not writable."""
     if env_var in os.environ:
         return  # User already set, respect their choice
 
-    # Check if default home config path is writable
+    # Check if home is writable
     home = os.path.expanduser("~")
-    default_path = os.path.join(home, ".config", default_subpath)
-    default_parent = os.path.dirname(default_path)
+    if os.access(home, os.W_OK):
+        return  # Home is writable, let library use defaults
 
-    # If default path or its parent is writable, let the library use defaults
-    if (
-        os.access(default_path, os.W_OK)
-        or (os.path.exists(default_parent) and os.access(default_parent, os.W_OK))
-        or os.access(os.path.join(home, ".config"), os.W_OK)
-    ):
-        return
-
-    # Default not writable - find alternative location
-    for cache_base in [
-        os.environ.get("SCRATCH"),
-        os.environ.get("SLURM_TMPDIR"),
-        tempfile.gettempdir(),
-    ]:
-        if cache_base and os.access(cache_base, os.W_OK):
-            cache_path = os.path.join(cache_base, f".{default_subpath}")
-            os.makedirs(cache_path, exist_ok=True)
-            os.environ[env_var] = cache_path
-            return
+    # Home not writable - use current working directory
+    cache_path = os.path.join(os.getcwd(), f".{subdir}")
+    os.makedirs(cache_path, exist_ok=True)
+    os.environ[env_var] = cache_path
 
 
+# Configure cache directories (before any library imports)
+_setup_cache_dir("TORCH_HOME", "torch_cache")
 _setup_cache_dir("MPLCONFIGDIR", "matplotlib")
 _setup_cache_dir("FONTCONFIG_CACHE", "fontconfig")
+_setup_cache_dir("XDG_DATA_HOME", "local/share")
+_setup_cache_dir("XDG_STATE_HOME", "local/state")
+_setup_cache_dir("XDG_CACHE_HOME", "cache")
 
 # =============================================================================
 # Standard imports (after environment setup)
@@ -1065,7 +1054,17 @@ def main():
                             f,
                         )
 
-                    unwrapped = accelerator.unwrap_model(model)
+                    # Unwrap model for saving (handle torch.compile compatibility)
+                    try:
+                        unwrapped = accelerator.unwrap_model(model)
+                    except KeyError:
+                        # torch.compile model may not have _orig_mod in expected location
+                        # Fall back to getting the module directly
+                        unwrapped = model.module if hasattr(model, "module") else model
+                        # If still compiled, try to get the underlying model
+                        if hasattr(unwrapped, "_orig_mod"):
+                            unwrapped = unwrapped._orig_mod
+
                     torch.save(
                         unwrapped.state_dict(),
                         os.path.join(args.output_dir, "best_model_weights.pth"),

{wavedl-1.4.3.dist-info → wavedl-1.4.4.dist-info}/METADATA

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

{wavedl-1.4.3.dist-info → wavedl-1.4.4.dist-info}/RECORD

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=7uvgr9r21qLGu6RIZDUfQBhg1vXNAfmIWl_P6BMj_KQ,1177
-wavedl/hpc.py,sha256=de_GKERX8GS10sXRX9yXiGzMnk1jjq8JPzRw7QDs6d4,7967
+wavedl/__init__.py,sha256=n0XNSrp0aGEE6HQpzbF1wiKK-jORKgcc0Q2Op4MtQGk,1177
+wavedl/hpc.py,sha256=0h8IZzOT0EzmEv3fU9cKyRVE9V1ivtBzbjuBCaxYadc,8445
 wavedl/hpo.py,sha256=YJXsnSGEBSVUqp_2ah7zu3_VClAUqZrdkuzDaSqQUjU,12952
-wavedl/test.py,sha256=jZmRJaivYYTMMTaccCi0yQjHOfp0a9YWR1wAPeKFH-k,36246
-wavedl/train.py,sha256=TjLABBPCqu9r7FEWlxJlKsT7uAMo6hiDxRfii2SKXe4,49052
+wavedl/test.py,sha256=Wajcze8gFEyJ9VyN_Bq-YadZ_VZtVaX_HicvUmW6MXM,37365
+wavedl/train.py,sha256=_pW7prvlNqfUGrGweHO2QelS87UiAYKvyJwqMAIj6yI,49292
 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
@@ -29,9 +29,9 @@ wavedl/utils/losses.py,sha256=5762M-TBC_hz6uyj1NPbU1vZeFOJQq7fR3-j7OygJRo,7254
 wavedl/utils/metrics.py,sha256=mkCpqZwl_XUpNvA5Ekjf7y-HqApafR7eR6EuA8cBdM8,37287
 wavedl/utils/optimizers.py,sha256=PyIkJ_hRhFi_Fio81Gy5YQNhcME0JUUEl8OTSyu-0RA,6323
 wavedl/utils/schedulers.py,sha256=e6Sf0yj8VOqkdwkUHLMyUfGfHKTX4NMr-zfgxWqCTYI,7659
-wavedl-1.4.3.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.3.dist-info/METADATA,sha256=vs3nt8R2O5lD7q-si9M5ChyrriIKm0fFzDwi4HVIYxw,40386
-wavedl-1.4.3.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.3.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.3.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.3.dist-info/RECORD,,
+wavedl-1.4.4.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
+wavedl-1.4.4.dist-info/METADATA,sha256=4tgSvkwzJmZP3PLCrqx-FYV_w6VT6Mi4XIsB0Dvb6_0,40386
+wavedl-1.4.4.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+wavedl-1.4.4.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
+wavedl-1.4.4.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
+wavedl-1.4.4.dist-info/RECORD,,