wavedl 1.4.2__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.2"
+__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

@@ -37,9 +37,43 @@ Author: Ductho Le (ductho.le@outlook.com)
 
 from __future__ import annotations
 
+# =============================================================================
+# HPC Environment Setup (MUST be before any library imports)
+# =============================================================================
+# Auto-configure writable cache directories when home is not writable.
+# Uses current working directory as fallback - works on HPC and local machines.
+import os
+
+
+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")
+
+# =============================================================================
+# Standard imports (after environment setup)
+# =============================================================================
 import argparse
 import logging
-import os
 import pickle
 import shutil
 import sys
@@ -47,6 +81,10 @@ import time
 import warnings
 from typing import Any
 
+
+# Suppress Pydantic warnings from accelerate's internal Field() usage
+warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -582,9 +620,9 @@ def main():
     # Torch 2.0 compilation (requires compatible Triton on GPU)
     if args.compile:
         try:
-            # Test if Triton is available AND compatible with this PyTorch version
-            # PyTorch needs triton_key from triton.compiler.compiler
-            from triton.compiler.compiler import triton_key
+            # Test if Triton is available - just import the package
+            # Different Triton versions have different internal APIs, so just check base import
+            import triton
 
             model = torch.compile(model)
             if accelerator.is_main_process:
@@ -875,9 +913,13 @@ def main():
             cpu_preds = torch.cat(local_preds)
             cpu_targets = torch.cat(local_targets)
 
-            # Gather to rank 0 only via gather_object (avoids all-gather to every rank)
-            # gather_object returns list of objects from each rank: [(preds0, targs0), (preds1, targs1), ...]
-            gathered = accelerator.gather_object((cpu_preds, cpu_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)]
 
             # Synchronize validation metrics (scalars only - efficient)
             val_loss_scalar = val_loss_sum.item()
@@ -1012,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.2.dist-info → wavedl-1.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.2
+Version: 1.4.4
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -57,6 +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)
 [![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)
 

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

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=K52yq0nkj2B3W0ZpR3tb7RUcHMIiANAE2d1WTuGVZLI,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=Gh02hlfjcote6w1sgUndJzKU_FhkJckTAhlq1aL8FY8,46842
+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.2.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.2.dist-info/METADATA,sha256=xHMuRcGdF8Tdju6aBK601mqLYx0uSIdmtjoQrHpPYmI,40245
-wavedl-1.4.2.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.2.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.2.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.2.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,,