from _Ax import _Ax_ext
import numpy as np
import copy

def Ax(img, geo, angles, projection_type="Siddon"):

    if img.dtype != np.float32:
        raise TypeError("Input data should be float32, not "+ str(img.dtype))
    if not np.isreal(img).all():
        raise ValueError("Complex types not compatible for projection.")
    if any(img.shape != geo.nVoxel):
        raise ValueError("Input data should be of shape geo.nVoxel: "+ str(geo.nVoxel) +
                         " not:" + str(img.shape))
    geox = copy.deepcopy(geo)
    geox.check_geo(angles)
    """
    Here we cast all values in geo to single point precision float. This way we 
    know what behaviour to expect from pytigre to Cuda and can change 
    single parameters accordingly. 
    """
    geox.cast_to_single()
    #geox.checknans()


    return _Ax_ext(img, geox, geox.angles, projection_type, geox.mode)

from _Ax import _Ax_ext
import numpy as np
import copy

def Ax(img, geo, angles, projection_type="Siddon"):

    if img.dtype != np.float32:
        raise TypeError("Input data should be float32, not "+ str(img.dtype))
    if not np.isreal(img).all():
        raise ValueError("Complex types not compatible for projection.")
    geox = copy.deepcopy(geo)
    geox.check_geo(angles)
    """
    Here we cast all values in geo to single point precision float. This way we 
    know what behaviour to expect from pytigre to Cuda and can change 
    single parameters accordingly. 
    """
    geox.cast_to_single()
    #geox.checknans()

    if abs(img.shape - geox.nVoxel).max()>1e-8:
        raise ValueError("Input data should be of shape geo.nVoxel: "+ str(geox.nVoxel) +
                         " not:" + str(img.shape))

    return _Ax_ext(img, geox, geox.angles, projection_type, geox.mode)

--- 
+++ 
@@ -8,6 +8,9 @@
         raise TypeError("Input data should be float32, not "+ str(img.dtype))
     if not np.isreal(img).all():
         raise ValueError("Complex types not compatible for projection.")
+    if any(img.shape != geo.nVoxel):
+        raise ValueError("Input data should be of shape geo.nVoxel: "+ str(geo.nVoxel) +
+                         " not:" + str(img.shape))
     geox = copy.deepcopy(geo)
     geox.check_geo(angles)
     """
@@ -18,8 +21,5 @@
     geox.cast_to_single()
     #geox.checknans()
 
-    if abs(img.shape - geox.nVoxel).max()>1e-8:
-        raise ValueError("Input data should be of shape geo.nVoxel: "+ str(geox.nVoxel) +
-                         " not:" + str(img.shape))
 
     return _Ax_ext(img, geox, geox.angles, projection_type, geox.mode)
Check the shape of input data earlier

Using geo.nVoxel to check the input img shape earlier, before geo is casted to float32 (geox). We should use any() instead of all(), since  "!=" is used?