Check data size when using 'link' function in Python

4 files changed, 98 insertions, 23 deletions

diff --git a/python/astra/data2d_c.pyx b/python/astra/data2d_c.pyx
index ac54898..29548b5 100644
--- a/python/astra/data2d_c.pyx
+++ b/python/astra/data2d_c.pyx
@@ -47,6 +47,12 @@ from .PyIncludes cimport *
 cimport utils
 from .utils import wrap_from_bytes
 
+from .pythonutils import geom_size
+
+import operator
+
+from six.moves import reduce
+
 cdef CData2DManager * man2d = <CData2DManager * >PyData2DManager.getSingletonPtr()
 
 cdef extern from "CFloat32CustomPython.h":
@@ -71,6 +77,16 @@ def create(datatype, geometry, data=None, link=False):
     cdef CProjectionGeometry2D * ppGeometry
     cdef CFloat32Data2D * pDataObject2D
     cdef CFloat32CustomMemory * pCustom
+
+    if link:
+        geomSize = geom_size(geometry)
+        if len(data.shape)==1:
+            if data.size!=reduce(operator.mul,geomSize):
+                raise Exception("The dimensions of the data do not match those specified in the geometry.")
+        else:
+            if data.shape!=geomSize:
+                raise Exception("The dimensions of the data do not match those specified in the geometry.")
+
     if datatype == '-vol':
         cfg = utils.dictToConfig(six.b('VolumeGeometry'), geometry)
         pGeometry = new CVolumeGeometry2D()
diff --git a/python/astra/data3d_c.pyx b/python/astra/data3d_c.pyx
index 84472c1..30745b4 100644
--- a/python/astra/data3d_c.pyx
+++ b/python/astra/data3d_c.pyx
@@ -45,6 +45,13 @@ from .PyXMLDocument cimport XMLDocument
 cimport utils
 from .utils import wrap_from_bytes
 
+from .pythonutils import geom_size
+
+import operator
+
+from six.moves import reduce
+
+
 cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr()
 
 cdef extern from *:
@@ -61,6 +68,16 @@ def create(datatype,geometry,data=None, link=False):
     cdef CFloat32Data3DMemory * pDataObject3D
     cdef CConeProjectionGeometry3D* pppGeometry
     cdef CFloat32CustomMemory * pCustom
+
+    if link:
+        geomSize = geom_size(geometry)
+        if len(data.shape)==1:
+            if data.size!=reduce(operator.mul,geomSize):
+                raise Exception("The dimensions of the data do not match those specified in the geometry.")
+        else:
+            if data.shape!=geomSize:
+                raise Exception("The dimensions of the data do not match those specified in the geometry.")
+
     if datatype == '-vol':
         cfg = utils.dictToConfig(six.b('VolumeGeometry'), geometry)
         pGeometry = new CVolumeGeometry3D()
diff --git a/python/astra/functions.py b/python/astra/functions.py
index 4025468..b826b86 100644
--- a/python/astra/functions.py
+++ b/python/astra/functions.py
@@ -38,6 +38,7 @@ from . import data2d
 from . import data3d
 from . import projector
 from . import algorithm
+from . import pythonutils
 
 
 
@@ -158,29 +159,7 @@ def geom_size(geom, dim=None):
     :param dim: Optional axis index to return
     :type dim: :class:`int`
     """
-
-    if 'GridSliceCount' in geom:
-        # 3D Volume geometry?
-        s = (geom['GridSliceCount'], geom[
-             'GridRowCount'], geom['GridColCount'])
-    elif 'GridColCount' in geom:
-        # 2D Volume geometry?
-        s = (geom['GridRowCount'], geom['GridColCount'])
-    elif geom['type'] == 'parallel' or geom['type'] == 'fanflat':
-        s = (len(geom['ProjectionAngles']), geom['DetectorCount'])
-    elif geom['type'] == 'parallel3d' or geom['type'] == 'cone':
-        s = (geom['DetectorRowCount'], len(
-            geom['ProjectionAngles']), geom['DetectorColCount'])
-    elif geom['type'] == 'fanflat_vec':
-        s = (geom['Vectors'].shape[0], geom['DetectorCount'])
-    elif geom['type'] == 'parallel3d_vec' or geom['type'] == 'cone_vec':
-        s = (geom['DetectorRowCount'], geom[
-             'Vectors'].shape[0], geom['DetectorColCount'])
-
-    if dim != None:
-        s = s[dim]
-
-    return s
+    return pythonutils.geom_size(geom,dim)
 
 
 def geom_2vec(proj_geom):
diff --git a/python/astra/pythonutils.py b/python/astra/pythonutils.py
new file mode 100644
index 0000000..8ea4af5
--- /dev/null
+++ b/python/astra/pythonutils.py
@@ -0,0 +1,63 @@
+#-----------------------------------------------------------------------
+# Copyright 2013 Centrum Wiskunde & Informatica, Amsterdam
+#
+# Author: Daniel M. Pelt
+# Contact: D.M.Pelt@cwi.nl
+# Website: http://dmpelt.github.io/pyastratoolbox/
+#
+#
+# This file is part of the Python interface to the
+# All Scale Tomographic Reconstruction Antwerp Toolbox ("ASTRA Toolbox").
+#
+# The Python interface to 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 Python interface to 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 Python interface to the ASTRA Toolbox. If not, see <http://www.gnu.org/licenses/>.
+#
+#-----------------------------------------------------------------------
+"""Additional purely Python functions for PyAstraToolbox.
+
+.. moduleauthor:: Daniel M. Pelt <D.M.Pelt@cwi.nl>
+
+
+"""
+
+def geom_size(geom, dim=None):
+    """Returns the size of a volume or sinogram, based on the projection or volume geometry.
+
+    :param geom: Geometry to calculate size from
+    :type geometry: :class:`dict`
+    :param dim: Optional axis index to return
+    :type dim: :class:`int`
+    """
+
+    if 'GridSliceCount' in geom:
+        # 3D Volume geometry?
+        s = (geom['GridSliceCount'], geom[
+             'GridRowCount'], geom['GridColCount'])
+    elif 'GridColCount' in geom:
+        # 2D Volume geometry?
+        s = (geom['GridRowCount'], geom['GridColCount'])
+    elif geom['type'] == 'parallel' or geom['type'] == 'fanflat':
+        s = (len(geom['ProjectionAngles']), geom['DetectorCount'])
+    elif geom['type'] == 'parallel3d' or geom['type'] == 'cone':
+        s = (geom['DetectorRowCount'], len(
+            geom['ProjectionAngles']), geom['DetectorColCount'])
+    elif geom['type'] == 'fanflat_vec':
+        s = (geom['Vectors'].shape[0], geom['DetectorCount'])
+    elif geom['type'] == 'parallel3d_vec' or geom['type'] == 'cone_vec':
+        s = (geom['DetectorRowCount'], geom[
+             'Vectors'].shape[0], geom['DetectorColCount'])
+
+    if dim != None:
+        s = s[dim]
+
+    return s