make read / recon example array api agnostic

ETSIhackers · Nov 15, 2023 · 7755231 · 7755231
1 parent 7b15e92
commit 7755231
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/python/demo_read_reconstruct.py b/python/demo_read_reconstruct.py
@@ -7,9 +7,18 @@
 from pathlib import Path
 
 import array_api_compat.numpy as np
-import numpy.array_api as xp
 
+#----------------------------------------------------------------
+#-- Choose you favorite python backend and device here ----------
+#----------------------------------------------------------------
+
+# import numpy.array_api as xp
+import array_api_compat.torch as xp
 dev = "cpu"
+
+#----------------------------------------------------------------
+#----------------------------------------------------------------
+
 lm_data_dir: str = "../data/sim_LM_acq_1"
 sens_img_file: str = "sensitivity_image.npy"
 prd_file: str = "simulated_lm.prd"

diff --git a/python/prd_io.py b/python/prd_io.py
@@ -133,7 +133,7 @@ def read_prd_to_numpy_arrays(
         # read the detector coordinate look up table
         scanner_lut = xp.asarray(
             [[det.x, det.y, det.z] for det in header.scanner.detectors],
-            dtype=np.float32,
+            dtype=xp.float32,
             device=dev,
         )