setting environment params for pylops

alex-rakowski · Nov 22, 2023 · 1cdc541 · 1cdc541
1 parent 155e1e9
commit 1cdc541
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 11 deletions.
diff --git a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py
@@ -19,11 +19,15 @@
 if IN_COLAB is False:
     try:
         import cupy as cp
-        import pylops
     except (ImportError,ModuleNotFoundError,Exception):
         cp = np
-
-
+else:
+    # I need to set some envrionment variables so pylops doesnt try to load cupy
+    import os   
+    os.environ["CUPY_PYLOPS"] = "0"
+    #CUSIGNAL_PYLOPS
+    os.environ["CUSIGNAL_PYLOPS"] = "0"
+import pylops
 from emdfile import Custom, tqdmnd
 from py4DSTEM import DataCube
 from py4DSTEM.process.phase.iterative_base_class import PtychographicReconstruction

diff --git a/py4DSTEM/process/phase/iterative_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_multislice_ptychography.py
@@ -17,10 +17,15 @@
 if IN_COLAB is False:
     try:
         import cupy as cp
-        import pylops
     except (ImportError,ModuleNotFoundError,Exception):
         cp = np
-
+else:
+    # I need to set some envrionment variables so pylops doesnt try to load cupy
+    import os   
+    os.environ["CUPY_PYLOPS"] = "0"
+    #CUSIGNAL_PYLOPS
+    os.environ["CUSIGNAL_PYLOPS"] = "0"
+import pylops
 from emdfile import Custom, tqdmnd
 from py4DSTEM import DataCube
 from py4DSTEM.process.phase.iterative_base_class import PtychographicReconstruction

diff --git a/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py b/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py
@@ -19,10 +19,16 @@
 if IN_COLAB is False:
     try:
         import cupy as cp
-        import pylops
     except (ImportError,ModuleNotFoundError,Exception):
         cp = np
-
+else:
+    # I need to set some envrionment variables so pylops doesnt try to load cupy
+    import os   
+    os.environ["CUPY_PYLOPS"] = "0"
+    #CUSIGNAL_PYLOPS
+    os.environ["CUSIGNAL_PYLOPS"] = "0"
+
+import pylops
 from emdfile import Custom, tqdmnd
 from py4DSTEM import DataCube
 from py4DSTEM.process.phase.iterative_base_class import PtychographicReconstruction

diff --git a/py4DSTEM/process/phase/iterative_overlap_tomography.py b/py4DSTEM/process/phase/iterative_overlap_tomography.py
@@ -19,10 +19,17 @@
 if IN_COLAB is False:
     try:
         import cupy as cp
-        import pylops
+
     except (ImportError,ModuleNotFoundError,Exception):
         cp = np
-
+else:
+    # I need to set some envrionment variables so pylops doesnt try to load cupy
+    import os   
+    os.environ["CUPY_PYLOPS"] = "0"
+    #CUSIGNAL_PYLOPS
+    os.environ["CUSIGNAL_PYLOPS"] = "0"
+
+import pylops
 
 from emdfile import Custom, tqdmnd
 from py4DSTEM import DataCube

diff --git a/py4DSTEM/process/phase/iterative_ptychographic_constraints.py b/py4DSTEM/process/phase/iterative_ptychographic_constraints.py
@@ -12,8 +12,13 @@
 from py4DSTEM.process.utils import get_CoM
 import sys
 IN_COLAB = 'google.colab' in sys.modules
-if IN_COLAB is False:
-    import pylops
+if IN_COLAB is True:
+    import os
+    # I need to set some envrionment variables so pylops doesnt try to load cupy
+    os.environ["CUPY_PYLOPS"] = "0"
+    #CUSIGNAL_PYLOPS
+    os.environ["CUSIGNAL_PYLOPS"] = "0"
+import pylops
 
 
 class PtychographicConstraints: