Merge branch 'direct_fpbp'

8 files changed, 159 insertions, 48 deletions

diff --git a/python/astra/PyIncludes.pxd b/python/astra/PyIncludes.pxd
index b9a61a9..f964118 100644
--- a/python/astra/PyIncludes.pxd
+++ b/python/astra/PyIncludes.pxd
@@ -236,9 +236,17 @@ cdef extern from "astra/ProjectionGeometry3D.h" namespace "astra":
         int getDetectorColCount()
         int getDetectorRowCount()
 
+cdef extern from "astra/Float32VolumeData3D.h" namespace "astra":
+    cdef cppclass CFloat32VolumeData3D(CFloat32Data3D):
+        pass
+
+cdef extern from "astra/Float32ProjectionData3D.h" namespace "astra":
+    cdef cppclass CFloat32ProjectionData3D(CFloat32Data3D):
+        pass
+
 
 cdef extern from "astra/Float32VolumeData3DMemory.h" namespace "astra":
-    cdef cppclass CFloat32VolumeData3DMemory:
+    cdef cppclass CFloat32VolumeData3DMemory(CFloat32VolumeData3D):
         CFloat32VolumeData3DMemory(CVolumeGeometry3D*)
         CFloat32VolumeData3DMemory(CVolumeGeometry3D*, CFloat32CustomMemory*)
         CVolumeGeometry3D* getGeometry()
@@ -266,7 +274,7 @@ cdef extern from "astra/ConeVecProjectionGeometry3D.h" namespace "astra":
         CConeVecProjectionGeometry3D()
 
 cdef extern from "astra/Float32ProjectionData3DMemory.h" namespace "astra":
-    cdef cppclass CFloat32ProjectionData3DMemory:
+    cdef cppclass CFloat32ProjectionData3DMemory(CFloat32ProjectionData3D):
         CFloat32ProjectionData3DMemory(CProjectionGeometry3D*)
         CFloat32ProjectionData3DMemory(CConeProjectionGeometry3D*)
         CFloat32ProjectionData3DMemory(CProjectionGeometry3D*, CFloat32CustomMemory*)
@@ -280,7 +288,7 @@ cdef extern from "astra/Float32ProjectionData3DMemory.h" namespace "astra":
 
 IF HAVE_CUDA==True:
     cdef extern from "astra/Float32VolumeData3DGPU.h" namespace "astra":
-        cdef cppclass CFloat32VolumeData3DGPU:
+        cdef cppclass CFloat32VolumeData3DGPU(CFloat32VolumeData3D):
             CFloat32VolumeData3DGPU(CVolumeGeometry3D*, MemHandle3D)
             CVolumeGeometry3D* getGeometry()
             void changeGeometry(CVolumeGeometry3D*)
@@ -290,7 +298,7 @@ IF HAVE_CUDA==True:
             bool isInitialized()
 
     cdef extern from "astra/Float32ProjectionData3DGPU.h" namespace "astra":
-        cdef cppclass CFloat32ProjectionData3DGPU:
+        cdef cppclass CFloat32ProjectionData3DGPU(CFloat32ProjectionData3D):
             CFloat32ProjectionData3DGPU(CProjectionGeometry3D*, MemHandle3D)
             CProjectionGeometry3D* getGeometry()
             void changeGeometry(CProjectionGeometry3D*)
diff --git a/python/astra/data2d.py b/python/astra/data2d.py
index 188ff69..6ab458f 100644
--- a/python/astra/data2d.py
+++ b/python/astra/data2d.py
@@ -65,12 +65,7 @@ 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")
+    checkArrayForLink(data)
     return d.create(datatype,geometry,data,True)
 
 def store(i, data):
diff --git a/python/astra/data3d.py b/python/astra/data3d.py
index b0d54b2..3eea0e3 100644
--- a/python/astra/data3d.py
+++ b/python/astra/data3d.py
@@ -26,7 +26,7 @@
 from . import data3d_c as d
 import numpy as np
 
-from .pythonutils import GPULink
+from .pythonutils import GPULink, checkArrayForLink
 
 def create(datatype,geometry,data=None):
     """Create a 3D object.
@@ -57,10 +57,7 @@ def link(datatype, geometry, data):
     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")
+        checkArrayForLink(data)
     return d.create(datatype,geometry,data,True)
 
 
diff --git a/python/astra/data3d_c.pyx b/python/astra/data3d_c.pyx
index 4c8aa62..a1b9138 100644
--- a/python/astra/data3d_c.pyx
+++ b/python/astra/data3d_c.pyx
@@ -44,6 +44,7 @@ from .PyXMLDocument cimport XMLDocument
 
 cimport utils
 from .utils import wrap_from_bytes
+from .utils cimport linkVolFromGeometry, linkProjFromGeometry
 
 from .pythonutils import geom_size, GPULink
 
@@ -53,9 +54,6 @@ from six.moves import reduce
 
 include "config.pxi"
 
-cdef extern from "Python.h":
-    void* PyLong_AsVoidPtr(object)
-
 
 cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr()
 
@@ -68,16 +66,12 @@ cdef CFloat32Data3DMemory * dynamic_cast_mem_safe(CFloat32Data3D *obj) except NU
         raise RuntimeError("Not a memory 3D data object")
     return ret
 
-cdef extern from "CFloat32CustomPython.h":
-    cdef cppclass CFloat32CustomPython:
-        CFloat32CustomPython(arrIn)
 
 def create(datatype,geometry,data=None, link=False):
     cdef Config *cfg
     cdef CVolumeGeometry3D * pGeometry
     cdef CProjectionGeometry3D * ppGeometry
     cdef CFloat32Data3D * pDataObject3D
-    cdef CConeProjectionGeometry3D* pppGeometry
     cdef CFloat32CustomMemory * pCustom = NULL
     IF HAVE_CUDA==True:
         cdef MemHandle3D hnd
@@ -101,20 +95,9 @@ def create(datatype,geometry,data=None, link=False):
             del pGeometry
             raise RuntimeError('Geometry class not initialized.')
         if link:
-            if isinstance(data, np.ndarray):
-                pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
-                pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DMemory(pGeometry, pCustom)
-            elif isinstance(data, GPULink):
-                IF HAVE_CUDA==True:
-                    s = geom_size(geometry)
-                    hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
-                    pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DGPU(pGeometry, hnd)
-                ELSE:
-                    raise NotImplementedError("CUDA support is not enabled in ASTRA")
-            else:
-                raise TypeError("data should be a numpy.ndarray or a GPULink object")
+            pDataObject3D = linkVolFromGeometry(pGeometry, data)
         else:
-            pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DMemory(pGeometry)
+            pDataObject3D = new CFloat32VolumeData3DMemory(pGeometry)
         del cfg
         del pGeometry
     elif datatype == '-sino' or datatype == '-proj3d' or datatype == '-sinocone':
@@ -136,20 +119,9 @@ def create(datatype,geometry,data=None, link=False):
             del ppGeometry
             raise RuntimeError('Geometry class not initialized.')
         if link:
-            if isinstance(data, np.ndarray):
-                pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
-                pDataObject3D = <CFloat32Data3D * > new CFloat32ProjectionData3DMemory(ppGeometry, pCustom)
-            elif isinstance(data, GPULink):
-                IF HAVE_CUDA==True:
-                    s = geom_size(geometry)
-                    hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
-                    pDataObject3D = <CFloat32Data3D * > new CFloat32ProjectionData3DGPU(ppGeometry, hnd)
-                ELSE:
-                    raise NotImplementedError("CUDA support is not enabled in ASTRA")
-            else:
-                raise TypeError("data should be a numpy.ndarray or a GPULink object")
+            pDataObject3D = linkProjFromGeometry(ppGeometry, data)
         else:
-            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry)
+            pDataObject3D = new CFloat32ProjectionData3DMemory(ppGeometry)
         del ppGeometry
         del cfg
     else:
diff --git a/python/astra/experimental.pyx b/python/astra/experimental.pyx
index 8f8e47d..08d907d 100644
--- a/python/astra/experimental.pyx
+++ b/python/astra/experimental.pyx
@@ -65,6 +65,8 @@ IF HAVE_CUDA==True:
     cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr()
 
     def do_composite(projector_id, vol_ids, proj_ids, mode, t):
+        if mode != MODE_ADD and mode != MODE_SET:
+            raise RuntimeError("internal error: wrong composite mode")
         cdef vector[CFloat32VolumeData3D *] vol
         cdef CFloat32VolumeData3D * pVolObject
         cdef CFloat32ProjectionData3D * pProjObject
@@ -121,3 +123,58 @@ IF HAVE_CUDA==True:
         cdef CProjector3D * projector = manProj.get(projector_id) # may be NULL
         if not m.doFDK(projector, pVolObject, pProjObject, False, NULL, MODE_ADD):
             raise Exception("Failed to perform FDK")
+
+    cimport utils
+    from .utils cimport linkVolFromGeometry, linkProjFromGeometry
+
+    def direct_FPBP3D(projector_id, vol, proj, mode, t):
+        if mode != MODE_ADD and mode != MODE_SET:
+            raise RuntimeError("internal error: wrong composite mode")
+        cdef CProjector3D * projector = manProj.get(projector_id)
+        if projector == NULL:
+            raise Exception("Projector not found")
+        cdef CVolumeGeometry3D *pGeometry = projector.getVolumeGeometry()
+        cdef CProjectionGeometry3D *ppGeometry = projector.getProjectionGeometry()
+        cdef CFloat32VolumeData3D * pVol = linkVolFromGeometry(pGeometry, vol)
+        cdef CFloat32ProjectionData3D * pProj = linkProjFromGeometry(ppGeometry, proj)
+        cdef vector[CFloat32VolumeData3D *] vols
+        cdef vector[CFloat32ProjectionData3D *] projs
+        vols.push_back(pVol)
+        projs.push_back(pProj)
+        cdef CCompositeGeometryManager m
+        try:
+            if t == "FP":
+                if not m.doFP(projector, vols, projs, mode):
+                    raise Exception("Failed to perform FP")
+            elif t == "BP":
+                if not m.doBP(projector, vols, projs, mode):
+                    raise Exception("Failed to perform BP")
+            else:
+                raise RuntimeError("internal error: wrong op type")
+        finally:
+            del pVol
+            del pProj
+
+    def direct_FP3D(projector_id, vol, proj):
+        """Perform a 3D forward projection with pre-allocated input/output.
+
+        :param projector_id: A 3D projector object handle
+        :type datatype: :class:`int`
+        :param vol: The input data, as either a numpy array, or a GPULink object
+        :type datatype: :class:`numpy.ndarray` or :class:`astra.data3d.GPULink`
+        :param proj: The pre-allocated output data, either numpy array or GPULink
+        :type datatype: :class:`numpy.ndarray` or :class:`astra.data3d.GPULink`
+        """
+        direct_FPBP3D(projector_id, vol, proj, MODE_SET, "FP")
+
+    def direct_BP3D(projector_id, vol, proj):
+        """Perform a 3D back projection with pre-allocated input/output.
+
+        :param projector_id: A 3D projector object handle
+        :type datatype: :class:`int`
+        :param vol: The pre-allocated output data, as either a numpy array, or a GPULink object
+        :type datatype: :class:`numpy.ndarray` or :class:`astra.data3d.GPULink`
+        :param proj: The input data, either numpy array or GPULink
+        :type datatype: :class:`numpy.ndarray` or :class:`astra.data3d.GPULink`
+        """
+        direct_FPBP3D(projector_id, vol, proj, MODE_SET, "BP")
diff --git a/python/astra/pythonutils.py b/python/astra/pythonutils.py
index 715df30..ef49f97 100644
--- a/python/astra/pythonutils.py
+++ b/python/astra/pythonutils.py
@@ -29,6 +29,8 @@
 
 
 """
+import numpy as np
+
 
 def geom_size(geom, dim=None):
     """Returns the size of a volume or sinogram, based on the projection or volume geometry.
@@ -62,6 +64,19 @@ def geom_size(geom, dim=None):
 
     return s
 
+def checkArrayForLink(data):
+    """Check if a numpy array is suitable for direct usage (contiguous, etc.)
+
+    This function raises an exception if not.
+    """
+
+    if not isinstance(data, np.ndarray):
+        raise ValueError("Numpy array should be numpy.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")
+
 class GPULink(object):
     """Utility class for astra.data3d.link with a CUDA pointer
 
diff --git a/python/astra/utils.pxd b/python/astra/utils.pxd
index ea3da86..69f4e96 100644
--- a/python/astra/utils.pxd
+++ b/python/astra/utils.pxd
@@ -33,3 +33,6 @@ from .PyIncludes cimport *
 
 cdef configToDict(Config *)
 cdef Config * dictToConfig(string rootname, dc) except NULL
+cdef CFloat32VolumeData3D* linkVolFromGeometry(CVolumeGeometry3D *pGeometry, data) except NULL
+cdef CFloat32ProjectionData3D* linkProjFromGeometry(CProjectionGeometry3D *pGeometry, data) except NULL
+
diff --git a/python/astra/utils.pyx b/python/astra/utils.pyx
index b534d72..12fc38c 100644
--- a/python/astra/utils.pyx
+++ b/python/astra/utils.pyx
@@ -45,6 +45,18 @@ from .PyXMLDocument cimport XMLDocument
 from .PyXMLDocument cimport XMLNode
 from .PyIncludes cimport *
 
+from .pythonutils import GPULink, checkArrayForLink
+
+cdef extern from "CFloat32CustomPython.h":
+    cdef cppclass CFloat32CustomPython:
+        CFloat32CustomPython(arrIn)
+
+cdef extern from "Python.h":
+    void* PyLong_AsVoidPtr(object)
+
+
+include "config.pxi"
+
 
 cdef Config * dictToConfig(string rootname, dc) except NULL:
     cdef Config * cfg = new Config()
@@ -230,3 +242,55 @@ cdef XMLNode2dict(XMLNode node):
         inc(it)
     if len(opts)>0: dct['options'] = opts
     return dct
+
+cdef CFloat32VolumeData3D* linkVolFromGeometry(CVolumeGeometry3D *pGeometry, data) except NULL:
+    cdef CFloat32VolumeData3D * pDataObject3D = NULL
+    geom_shape = (pGeometry.getGridSliceCount(), pGeometry.getGridRowCount(), pGeometry.getGridColCount())
+    if isinstance(data, np.ndarray):
+        data_shape = data.shape
+    elif isinstance(data, GPULink):
+        data_shape = (data.z, data.y, data.x)
+    if geom_shape != data_shape:
+        raise ValueError(
+            "The dimensions of the data do not match those specified in the geometry: {} != {}".format(data_shape, geom_shape))
+
+    if isinstance(data, np.ndarray):
+        checkArrayForLink(data)
+        pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+        pDataObject3D = new CFloat32VolumeData3DMemory(pGeometry, pCustom)
+    elif isinstance(data, GPULink):
+        IF HAVE_CUDA==True:
+            hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
+            pDataObject3D = new CFloat32VolumeData3DGPU(pGeometry, hnd)
+        ELSE:
+            raise NotImplementedError("CUDA support is not enabled in ASTRA")
+    else:
+        raise TypeError("data should be a numpy.ndarray or a GPULink object")
+    return pDataObject3D
+
+cdef CFloat32ProjectionData3D* linkProjFromGeometry(CProjectionGeometry3D *pGeometry, data) except NULL:
+    cdef CFloat32ProjectionData3D * pDataObject3D = NULL
+    geom_shape = (pGeometry.getDetectorRowCount(), pGeometry.getProjectionCount(), pGeometry.getDetectorColCount())
+    if isinstance(data, np.ndarray):
+        data_shape = data.shape
+    elif isinstance(data, GPULink):
+        data_shape = (data.z, data.y, data.x)
+    if geom_shape != data_shape:
+        raise ValueError(
+            "The dimensions of the data do not match those specified in the geometry: {} != {}".format(data_shape, geom_shape))
+
+    if isinstance(data, np.ndarray):
+        checkArrayForLink(data)
+        pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+        pDataObject3D = new CFloat32ProjectionData3DMemory(pGeometry, pCustom)
+    elif isinstance(data, GPULink):
+        IF HAVE_CUDA==True:
+            hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
+            pDataObject3D = new CFloat32ProjectionData3DGPU(pGeometry, hnd)
+        ELSE:
+            raise NotImplementedError("CUDA support is not enabled in ASTRA")
+    else:
+        raise TypeError("data should be a numpy.ndarray or a GPULink object")
+    return pDataObject3D
+
+