Merge pull request #126 from wjp/tests

Add basic tests to check an astra installation

2 files changed, 100 insertions, 1 deletions

diff --git a/python/astra/__init__.py b/python/astra/__init__.py
index b73fff5..8238a88 100644
--- a/python/astra/__init__.py
+++ b/python/astra/__init__.py
@@ -27,6 +27,7 @@ from . import matlab as m
 from .creators import astra_dict,create_vol_geom, create_proj_geom, create_backprojection, create_sino, create_reconstruction, create_projector,create_sino3d_gpu, create_backprojection3d_gpu
 from .functions import data_op, add_noise_to_sino, clear, move_vol_geom
 from .extrautils import clipCircle
+from .astra import set_gpu_index, get_gpu_info
 from . import data2d
 from . import astra
 from . import data3d
@@ -38,6 +39,7 @@ from . import plugin
 from . import plugins
 from . import log
 from .optomo import OpTomo
+from .tests import test_noCUDA, test_CUDA
 
 __version__ = '1.8'
 
@@ -45,4 +47,4 @@ import os
 
 if 'ASTRA_GPU_INDEX' in os.environ:
     L = [ int(x) for x in os.environ['ASTRA_GPU_INDEX'].split(',') ]
-    astra.set_gpu_index(L)
+    set_gpu_index(L)
diff --git a/python/astra/tests.py b/python/astra/tests.py
new file mode 100644
index 0000000..32afd36
--- /dev/null
+++ b/python/astra/tests.py
@@ -0,0 +1,97 @@
+# -----------------------------------------------------------------------
+# Copyright: 2010-2017, iMinds-Vision Lab, University of Antwerp
+#            2013-2017, CWI, Amsterdam
+#
+# Contact: astra@uantwerpen.be
+# Website: http://www.astra-toolbox.com/
+#
+# This file is part of the ASTRA Toolbox.
+#
+#
+# The ASTRA Toolbox is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# The ASTRA Toolbox is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>.
+#
+# -----------------------------------------------------------------------
+
+from __future__ import print_function, absolute_import
+
+def _basic_par2d_fp(type):
+  import astra
+  import numpy as np
+  vg = astra.create_vol_geom(2, 32)
+  pg = astra.create_proj_geom('parallel', 1, 32, [0])
+  proj_id = astra.create_projector(type, pg, vg)
+  vol = np.random.rand(2, 32)
+  (sino_id, sino) = astra.create_sino(vol, proj_id)
+  astra.data2d.delete(sino_id)
+  astra.projector.delete(proj_id)
+  err = np.max(np.abs(sino[0,:] - np.sum(vol,axis=0)))
+  return err < 1e-6
+
+def _basic_par3d_fp():
+  import astra
+  import numpy as np
+  vg = astra.create_vol_geom(2, 32, 32)
+  pg = astra.create_proj_geom('parallel3d', 1, 1, 32, 32, [0])
+  vol = np.random.rand(32, 2, 32)
+  (sino_id, sino) = astra.create_sino3d_gpu(vol, pg, vg)
+  astra.data3d.delete(sino_id)
+  err = np.max(np.abs(sino[:,0,:] - np.sum(vol,axis=1)))
+  return err < 1e-6
+
+
+def _basic_par2d():
+  print("Testing basic CPU 2D functionality... ", end="")
+  if _basic_par2d_fp('line'):
+    print("Ok")
+    return True
+  else:
+    print("Error")
+    return False
+
+def _basic_par2d_cuda():
+  print("Testing basic CUDA 2D functionality... ", end="")
+  if _basic_par2d_fp('cuda'):
+    print("Ok")
+    return True
+  else:
+    print("Error")
+    return False
+
+def _basic_par3d_cuda():
+  print("Testing basic CUDA 3D functionality... ", end="")
+  if _basic_par3d_fp():
+    print("Ok")
+    return True
+  else:
+    print("Error")
+    return False
+
+def test_noCUDA():
+  """Perform a very basic functionality test, without CUDA"""
+  
+  ok = _basic_par2d()
+  if not ok:
+    raise RuntimeError("Test failed")
+
+def test_CUDA():
+  """Perform a very basic functionality test, including CUDA"""
+  
+  import astra
+  print("Getting GPU info... ", end="")
+  print(astra.get_gpu_info())
+  ok1 = _basic_par2d()
+  ok2 = _basic_par2d_cuda()
+  ok3 = _basic_par3d_cuda()
+  if not (ok1 and ok2 and ok3):
+    raise RuntimeError("Test failed")