pywavelet 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

pywavelet/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.6'
-__version_tuple__ = version_tuple = (0, 2, 6)
+__version__ = version = '0.2.7'
+__version_tuple__ = version_tuple = (0, 2, 7)

pywavelet/backend.py

@@ -1,5 +1,9 @@
 import importlib
 import os
+from rich.table import Table, Text
+from rich.console import Console
+
+
 
 from .logger import logger
 
@@ -8,45 +12,89 @@ CUPY = "cupy"
 NUMPY = "numpy"
 
 
+def cuda_is_available() -> bool:
+    """Check if CUDA is available."""
+    # Check if CuPy is available and CUDA is accessible
+    cupy_available = importlib.util.find_spec("cupy") is not None
+    if cupy_available:
+        import cupy
+
+        try:
+            cupy.cuda.runtime.getDeviceCount()  # Check if any CUDA device is available
+            cuda_available = True
+        except cupy.cuda.runtime.CUDARuntimeError:
+            cuda_available = False
+    else:
+        cuda_available = False
+    return cuda_available
+
+
+def jax_is_available() -> bool:
+    """Check if JAX is available."""
+    return importlib.util.find_spec(JAX) is not None
+
+
+def get_available_backends_table():
+    """Print the available backends as a rich table."""
+
+    jax_avail = jax_is_available()
+    cupy_avail = cuda_is_available()
+    table = Table("Backend", "Available", title="Available backends")
+    true_check = "[green]✓[/green]"
+    false_check = "[red]✗[/red]"
+    table.add_row(JAX, true_check if jax_avail else false_check)
+    table.add_row(CUPY, true_check if cupy_avail else false_check)
+    table.add_row(NUMPY, true_check)
+    console = Console(width=150)
+    with console.capture() as capture:
+        console.print(table)
+    return Text.from_ansi(capture.get())
+
+
 def get_backend_from_env():
     """Select and return the appropriate backend module."""
     backend = os.getenv("PYWAVELET_BACKEND", NUMPY).lower()
 
-    if backend == JAX:
-        if importlib.util.find_spec(JAX):
-            import jax.numpy as xp
-            from jax.numpy.fft import fft, ifft, irfft, rfft, rfftfreq
-            from jax.scipy.special import betainc
-
-            logger.info("Using JAX backend")
-            return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
-        else:
-            logger.warning(
-                "JAX backend requested but not installed. Falling back to NumPy."
-            )
-
-    elif backend == CUPY:
-        if importlib.util.find_spec(CUPY):
-            import cupy as xp
-            from cupy.fft import fft, ifft, irfft, rfft, rfftfreq
-            from cupyx.scipy.special import betainc
-
-            logger.info("Using CuPy backend")
-            return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
-        else:
-            logger.warning(
-                "CuPy backend requested but not installed. Falling back to NumPy."
-            )
-
-    # Default to NumPy
-    import numpy as xp
-    from numpy.fft import fft, ifft, irfft, rfft, rfftfreq
-    from scipy.special import betainc
-
-    logger.info("Using NumPy+Numba backend")
+    if backend == JAX and jax_is_available():
+
+        import jax.numpy as xp
+        from jax.numpy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from jax.scipy.special import betainc
+
+        logger.info("Using JAX backend")
+
+    elif backend == CUPY and cuda_is_available():
+
+        import cupy as xp
+        from cupy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from cupyx.scipy.special import betainc
+
+        logger.info("Using CuPy backend")
+
+    elif backend == NUMPY:
+        import numpy as xp
+        from numpy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from scipy.special import betainc
+
+        logger.info("Using NumPy backend")
+
+
+    else:
+        logger.error(
+            f"Backend {backend} is not available. "
+        )
+        print(get_available_backends_table())
+        logger.warning(
+            f"Setting backend to NumPy. "
+        )
+        os.environ["PYWAVELET_BACKEND"] = NUMPY
+        return get_backend_from_env()
+
     return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
 
 
+cuda_available = cuda_is_available()
+
 # Get the chosen backend
 xp, fft, ifft, irfft, rfft, rfftfreq, betainc, current_backend = (
     get_backend_from_env()

pywavelet/transforms/__init__.py

@@ -1,5 +1,11 @@
+from ..logger import logger
+
 from ..backend import current_backend
 
+
+logger.debug(f"Using {current_backend} backend")
+
+
 if current_backend == "jax":
     from .jax import (
         from_freq_to_wavelet,

pywavelet/transforms/jax/__init__.py

@@ -4,6 +4,30 @@ from .inverse import from_wavelet_to_freq, from_wavelet_to_time
 
 logger.warning("JAX SUBPACKAGE NOT FULLY TESTED")
 
+
+def _log_jax_info():
+    """Log JAX backend and precision information.
+
+    backend : str
+        JAX backend. ["cpu", "gpu", "tpu"]
+    precision : str
+        JAX precision. ["32bit", "64bit"]
+    """
+    import jax
+
+    _backend = jax.default_backend()
+    _precision = "64bit" if jax.config.jax_enable_x64 else "32bit"
+
+    logger.info(f"Jax running on {_backend} [{_precision} precision].")
+    if _precision == "32bit":
+        logger.warning(
+            "Jax is not running in 64bit precision. "
+            "To change, use jax.config.update('jax_enable_x64', True)."
+        )
+
+
+_log_jax_info()
+
 __all__ = [
     "from_wavelet_to_time",
     "from_wavelet_to_freq",

{pywavelet-0.2.6.dist-info → pywavelet-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pywavelet
-Version: 0.2.6
+Version: 0.2.7
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/

{pywavelet-0.2.6.dist-info → pywavelet-0.2.7.dist-info}/RECORD

@@ -1,9 +1,9 @@
 pywavelet/__init__.py,sha256=K7pQ8W2w9d5qwI4KzPdTpRn5-YaUfMpjnJmg7oQnYSM,508
-pywavelet/_version.py,sha256=nObnONsicQ3YX6SG5MVBxmIp5dmRacXDauSqZijWQbY,511
-pywavelet/backend.py,sha256=1AjwqoIlan6vNFZcon_LIVsiPH8HrWQwU3RON7dnjUE,1585
+pywavelet/_version.py,sha256=Xk20v7uvkFqkpy9aLJzVngs1eKQn0FYUP2oyA1MEQUU,511
+pywavelet/backend.py,sha256=Ixl9nVP7Lobbyff5WF-pqXo0zuM7RVbyC1j35dSO3jI,2791
 pywavelet/logger.py,sha256=DyKC-pJ_N9GlVeXL00E1D8hUd8GceBg-pnn7g1YPKcM,391
 pywavelet/utils.py,sha256=FqQ6V41WGHMbLC4wv_1xnwHjOPDVSWnG78sAeqbYtYU,1994
-pywavelet/transforms/__init__.py,sha256=t4cHI8Rd5UnLwqCunr4sCQRmsKhHOnZ5VqkDphhi-VM,784
+pywavelet/transforms/__init__.py,sha256=c5dnTdzm-Se3idp2LP-dHvk6fmv_ynDCixIDfLCjEPw,865
 pywavelet/transforms/phi_computer.py,sha256=jVxeWtfx5P1H-_HdMsK7xHuINZAjH9bj7cA8CJ98isw,3667
 pywavelet/transforms/cupy/__init__.py,sha256=8BBE6msB071WdstA860a7g64C0aHT2PZsqfEgP6nmkA,336
 pywavelet/transforms/cupy/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
@@ -13,7 +13,7 @@ pywavelet/transforms/cupy/forward/main.py,sha256=g2Pl-j4LBg7GLlzzCSoCGuEd6NNCckJ
 pywavelet/transforms/cupy/inverse/__init__.py,sha256=J4KIzPzbHNg_8fV_c1MpPq3slSqHQV0j3VFrjfd1Nog,121
 pywavelet/transforms/cupy/inverse/main.py,sha256=5pTtGNNdwlSGDQV4sqGyzUPnmqFUgFOFUFfpqjZx07Q,1608
 pywavelet/transforms/cupy/inverse/to_freq.py,sha256=gpqu5Y65ZvuET5jANp6UAuAamg2PRkpAlaAjWPh7uBk,1835
-pywavelet/transforms/jax/__init__.py,sha256=D_f-JgFAzOIJ-EuQZhTMziD4MT6lVWS3XV9s51Cu7Kg,335
+pywavelet/transforms/jax/__init__.py,sha256=_JG9EaHKD95U-Nzm5zQ1RHQ7XuBjZarMa7VpZ-y7rgY,941
 pywavelet/transforms/jax/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
 pywavelet/transforms/jax/forward/from_freq.py,sha256=XYtRziPD7MCbeKf4HAucQrMzko4T0zmNV7jg5bziVwA,3910
 pywavelet/transforms/jax/forward/from_time.py,sha256=4RZ8-ah0qOMP20i3-xThVWddxa1QTCvZKnGpNAJbb0g,1765
@@ -37,7 +37,7 @@ pywavelet/types/plotting.py,sha256=qjv5IeuSEc9WWkfJYvz1eQRgTKTspWxj4lwB5N69SbU,1
 pywavelet/types/timeseries.py,sha256=sataMW4BPFqi23h_NBZ_U9-Svuo9pLXVRmUJI6KTXG0,9430
 pywavelet/types/wavelet.py,sha256=lDhpy9bEb_I-YDQbI3elaWuU8l9E2P6wDcuAQONv8lA,13591
 pywavelet/types/wavelet_bins.py,sha256=gBjhWwfjcbbSnbGZVMNUeFFVUo2DVxJS4abDUVCL7ts,1458
-pywavelet-0.2.6.dist-info/METADATA,sha256=WScrhO_gC_5wKwY39T0aI81YBj71xn9O6-MI0GBRucQ,2571
-pywavelet-0.2.6.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-pywavelet-0.2.6.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
-pywavelet-0.2.6.dist-info/RECORD,,
+pywavelet-0.2.7.dist-info/METADATA,sha256=5oaGFUO6ASLeTlweptCKP5D9LUGkLQbmJNfWMVYgwxc,2571
+pywavelet-0.2.7.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
+pywavelet-0.2.7.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
+pywavelet-0.2.7.dist-info/RECORD,,

{pywavelet-0.2.6.dist-info → pywavelet-0.2.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (77.0.3)
+Generator: setuptools (78.0.2)
 Root-Is-Purelib: true
 Tag: py3-none-any