Merge pull request #93 from wjp/GPULink

GPULink support

1 files changed, 9 insertions, 6 deletions

diff --git a/python/astra/data3d.py b/python/astra/data3d.py
index 9c2bea4..ecb99d0 100644
--- a/python/astra/data3d.py
+++ b/python/astra/data3d.py
@@ -26,6 +26,8 @@
 from . import data3d_c as d
 import numpy as np
 
+from .pythonutils import GPULink
+
 def create(datatype,geometry,data=None):
     """Create a 3D object.
 
@@ -52,12 +54,13 @@ def link(datatype, geometry, data):
     :returns: :class:`int` -- the ID of the constructed object.
     
     """
-    if not isinstance(data,np.ndarray):
-        raise ValueError("Input should be a numpy array")
-    if not data.dtype==np.float32:
-        raise ValueError("Numpy array should be float32")
-    if not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
-        raise ValueError("Numpy array should be C_CONTIGUOUS and ALIGNED")
+    if not isinstance(data,np.ndarray) and not isinstance(data,GPULink):
+        raise TypeError("Input should be a numpy ndarray or GPULink object")
+    if isinstance(data, np.ndarray):
+        if data.dtype != np.float32:
+            raise ValueError("Numpy array should be float32")
+        if not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
+            raise ValueError("Numpy array should be C_CONTIGUOUS and ALIGNED")
     return d.create(datatype,geometry,data,True)