Skip to content

Commit

Permalink
minor: restyling deps.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@mrava87
mrava87 committed Dec 17, 2023
1 parent 49cd053 commit 330188f
Showing 1 changed file with 84 additions and 86 deletions.
170 changes: 84 additions & 86 deletions pylops/utils/deps.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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
)
Expand Down

0 comments on commit 330188f

Please sign in to comment.