From 330188f7550ca0675490499720d0a582dc91d247 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sun, 17 Dec 2023 20:27:34 +0300 Subject: [PATCH] minor: restyling deps.py --- pylops/utils/deps.py | 170 +++++++++++++++++++++---------------------- 1 file changed, 84 insertions(+), 86 deletions(-) diff --git a/pylops/utils/deps.py b/pylops/utils/deps.py index 2edbc86c..3b84488c 100644 --- a/pylops/utils/deps.py +++ b/pylops/utils/deps.py @@ -17,10 +17,73 @@ # error message at import of available package -def devito_import(message): +def cupy_import(message: Optional[str] = None) -> str: + # detect if cupy is available and the user is expecting to be used + cupy_test = ( + util.find_spec("cupy") is not None and int(os.getenv("CUPY_PYLOPS", 1)) == 1 + ) + # if cupy should be importable + if cupy_test: + # try importing it + try: + import_module("cupy") # noqa: F401 + + # if successful set the message to None. + cupy_message = None + # if unable to import but the package is installed + except (ImportError, ModuleNotFoundError) as e: + cupy_message = ( + f"Failed to import cupy, Falling back to CPU (error: {e}). " + "Please ensure your CUDA environment is set up correctly " + "for more details visit 'https://docs.cupy.dev/en/stable/install.html'" + ) + print(UserWarning(cupy_message)) + # if cupy_test is False, it means not installed or environment variable set to 0 + else: + cupy_message = ( + "Cupy package not installed or os.getenv('CUPY_PYLOPS') == 0. " + f"In order to be able to use {message} " + "ensure 'os.getenv('CUPY_PYLOPS') == 1' and run " + "'pip install cupy'; " + "for more details visit 'https://docs.cupy.dev/en/stable/install.html'" + ) + + return cupy_message + + +def cusignal_import(message: Optional[str] = None) -> str: + cusignal_test = ( + util.find_spec("cusignal") is not None + and int(os.getenv("CUSIGNAL_PYLOPS", 1)) == 1 + ) + if cusignal_test: + try: + import_module("cusignal") # noqa: F401 + + cusignal_message = None + except (ImportError, ModuleNotFoundError) as e: + cusignal_message = ( + f"Failed to import cusignal. Falling back to CPU (error: {e}) . " + "Please ensure your CUDA environment is set up correctly; " + "for more details visit 'https://github.com/rapidsai/cusignal#installation'" + ) + print(UserWarning(cusignal_message)) + else: + cusignal_message = ( + "Cusignal not installed or os.getenv('CUSIGNAL_PYLOPS') == 0. " + f"In order to be able to use {message} " + "ensure 'os.getenv('CUSIGNAL_PYLOPS') == 1' and run " + "'conda install cusignal'; " + "for more details visit ''https://github.com/rapidsai/cusignal#installation''" + ) + + return cusignal_message + + +def devito_import(message: Optional[str] = None) -> str: if devito_enabled: try: - import devito # noqa: F401 + import_module("devito") # noqa: F401 devito_message = None except Exception as e: @@ -34,10 +97,10 @@ def devito_import(message): return devito_message -def numba_import(message): +def numba_import(message: Optional[str] = None) -> str: if numba_enabled: try: - import numba # noqa: F401 + import_module("numba") # noqa: F401 numba_message = None except Exception as e: @@ -53,10 +116,10 @@ def numba_import(message): return numba_message -def pyfftw_import(message): +def pyfftw_import(message: Optional[str] = None) -> str: if pyfftw_enabled: try: - import pyfftw # noqa: F401 + import_module("pyfftw") # noqa: F401 pyfftw_message = None except Exception as e: @@ -72,10 +135,10 @@ def pyfftw_import(message): return pyfftw_message -def pywt_import(message): +def pywt_import(message: Optional[str] = None) -> str: if pywt_enabled: try: - import pywt # noqa: F401 + import_module("pywt") # noqa: F401 pywt_message = None except Exception as e: @@ -91,10 +154,10 @@ def pywt_import(message): return pywt_message -def skfmm_import(message): +def skfmm_import(message: Optional[str] = None) -> str: if skfmm_enabled: try: - import skfmm # noqa: F401 + import_module("skfmm") # noqa: F401 skfmm_message = None except Exception as e: @@ -109,10 +172,10 @@ def skfmm_import(message): return skfmm_message -def spgl1_import(message): +def spgl1_import(message: Optional[str] = None) -> str: if spgl1_enabled: try: - import spgl1 # noqa: F401 + import_module("spgl1") # noqa: F401 spgl1_message = None except Exception as e: @@ -126,10 +189,10 @@ def spgl1_import(message): return spgl1_message -def sympy_import(message): +def sympy_import(message: Optional[str] = None) -> str: if sympy_enabled: try: - import sympy # noqa: F401 + import_module("sympy") # noqa: F401 sympy_message = None except Exception as e: @@ -143,78 +206,13 @@ def sympy_import(message): return sympy_message -def cupy_import(message: Optional[str] = None): - # detect if cupy should be importable - cupy_test = ( - util.find_spec("cupy") is not None and int(os.getenv("CUPY_PYLOPS", 1)) == 1 - ) - # if cupy should be importable - if cupy_test: - # try importing it - try: - import_module("cupy") # noqa: F401 - - # if successful set the message to None. - cupy_message = None - # if unable to import but it is installed - except (ImportError, ModuleNotFoundError) as e: - cupy_message = ( - f"Failed to import cupy, Falling back to CPU (error: {e}). " - f"" - "Please ensure your CUDA envrionment is set up correctly " - "for more details visit 'https://docs.cupy.dev/en/stable/install.html'" - ) - print(UserWarning(cupy_message)) - # if cupy_test is False it means not installed or envrionment variable set to 0 - else: - cupy_message = ( - f"cupy package not installed or os.getenv('CUPY_PYLOPS') == 0. In order to be able to use " - f"{message} " - "ensure 'os.getenv('CUPY_PYLOPS') == 1' and run " - "'pip install cupy'. " - "for more details visit 'https://docs.cupy.dev/en/stable/install.html'" - ) - - return cupy_message - - -def cusignal_import(message: Optional[str] = None): - # detect if cupy should be importable - cusignal_test = ( - util.find_spec("cusignal") is not None - and int(os.getenv("CUSIGNAL_PYLOPS", 1)) == 1 - ) - # if cupy should be importable - if cusignal_test: - # try importing it - try: - import_module("cusignal") # noqa: F401 - - # if successful set the message to None. - cusignal_message = None - # if unable to import but it is installed - except (ImportError, ModuleNotFoundError) as e: - cusignal_message = ( - f"Failed to import cusignal. Falling back to CPU (error: {e}) . " - f"" - "Please ensure your CUDA envrionment is set up correctly " - "for more details visit 'https://github.com/rapidsai/cusignal#installation'" - ) - print(UserWarning(cusignal_message)) - # if cupy_test is False it means not installed or envrionment variable set to 0 - else: - cusignal_message = ( - f"cusignal package not installed or os.getenv('CUSIGNAL_PYLOPS') == 0. In order to be able to use " - f"{message} " - "ensure 'os.getenv('CUSIGNAL_PYLOPS') == 1' and run " - "'conda install cusignal'. " - "for more details visit ''https://github.com/rapidsai/cusignal#installation''" - ) - - return cusignal_message - - -# Set package avlaiblity booleans +# Set package availability booleans +# cupy and cusignal: the package is imported to check everything is working correctly, +# if not the package is disabled. We do this here as both libraries are used as drop-in +# replacement for many numpy and scipy routines when cupy arrays are provided. +# all other libraries: we simply check if the package is available and postpone its import +# to check everything is working correctly when a user tries to create an operator that requires +# such a package cupy_enabled: bool = ( True if (cupy_import() is None and int(os.getenv("CUPY_PYLOPS", 1)) == 1) else False )