Improve postalignment function

5 files changed, 20 insertions, 30 deletions

diff --git a/matlab/tools/astra_geom_postalignment.m b/matlab/tools/astra_geom_postalignment.m
index c85a7a0..0385e74 100644
--- a/matlab/tools/astra_geom_postalignment.m
+++ b/matlab/tools/astra_geom_postalignment.m
@@ -9,7 +9,8 @@ function proj_geom = astra_geom_postalignment(proj_geom, factor)
 % distance to shift the detector (measured in detector pixels).
 %
 % For 3D geometries, factor is a pair of floats specifying the horizontal
-% resp. vertical distances to shift the detector.
+% resp. vertical distances to shift the detector. If only a single float is
+% specified, this is treated as an horizontal shift.
 %
 % proj_geom: input projection geometry (vector-based only, use astra_geom_2vec to convert conventional projection geometries)
 % factor: number of pixels to shift the detector
@@ -25,6 +26,8 @@ function proj_geom = astra_geom_postalignment(proj_geom, factor)
 %--------------------------------------------------------------------------
 % $Id$
 
+	proj_geom = astra_geom_2vec(proj_geom);
+
 	if strcmp(proj_geom.type,'fanflat_vec') || strcmp(proj_geom.type,'parallel_vec')
 		proj_geom.Vectors(:,3:4) = proj_geom.Vectors(:,3:4) + factor(1) * proj_geom.Vectors(:,5:6);
 
diff --git a/python/astra/__init__.py b/python/astra/__init__.py
index 9321495..a0d14b1 100644
--- a/python/astra/__init__.py
+++ b/python/astra/__init__.py
@@ -25,7 +25,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, geom_size, geom_2vec
+from .functions import data_op, add_noise_to_sino, clear, move_vol_geom, geom_size, geom_2vec, geom_postalignment
 from .extrautils import clipCircle
 from .astra import set_gpu_index, get_gpu_info, use_cuda
 from . import data2d
diff --git a/python/astra/functions.py b/python/astra/functions.py
index 5a3d453..0d359b3 100644
--- a/python/astra/functions.py
+++ b/python/astra/functions.py
@@ -280,35 +280,29 @@ def geom_postalignment(proj_geom, factor):
     For 2D geometries, the argument factor is a single float specifying the
     distance to shift the detector (measured in detector pixels).
 
-    For 3D geometries, factor is a pair of floats specifying the horizontal
-    resp. vertical distances to shift the detector.
+    For 3D geometries, factor can be a pair of floats specifying the horizontal
+    resp. vertical distances to shift the detector. If only a single float
+    is specified, this is treated as a horizontal shift.
 
-    :param proj_geom: input projection geometry (vector-based only, use astra.geom_2vec to convert conventional projection geometries)
+    :param proj_geom: input projection geometry
     :type proj_geom: :class:`dict`
     :param factor: number of pixels to shift the detector
     :type factor: :class:`float`
     """
 
+    proj_geom = geom_2vec(proj_geom)
+
     if proj_geom['type'] == 'parallel_vec' or proj_geom['type'] == 'fanflat_vec':
-        for i in range(proj_geom['Vectors'].shape[0]):
-            proj_geom['Vectors'][i,2] = proj_geom['Vectors'][i,2] + factor * proj_geom['Vectors'][i,4];
-            proj_geom['Vectors'][i,3] = proj_geom['Vectors'][i,3] + factor * proj_geom['Vectors'][i,5];
+        V = proj_geom['Vectors']
+        V[:,2:4] = V[:,2:4] + factor * V[:,4:6]
 
     elif proj_geom['type'] == 'parallel3d_vec' or proj_geom['type'] == 'cone_vec':
-
-        if len(factor) == 1:
-            for i in range(proj_geom['Vectors'].shape[0]):
-                proj_geom['Vectors'][i,3] = proj_geom['Vectors'][i,3] + factor * proj_geom['Vectors'][i,6];
-                proj_geom['Vectors'][i,4] = proj_geom['Vectors'][i,4] + factor * proj_geom['Vectors'][i,7];
-                proj_geom['Vectors'][i,5] = proj_geom['Vectors'][i,5] + factor * proj_geom['Vectors'][i,8];
-
-        elif len(factor) > 1:
-            for i in range(proj_geom['Vectors'].shape[0]):
-                proj_geom['Vectors'][i,3] = proj_geom['Vectors'][i,3] + factor[0] * proj_geom['Vectors'][i,6] + factor[1] * proj_geom['Vectors'][i, 9];
-                proj_geom['Vectors'][i,4] = proj_geom['Vectors'][i,4] + factor[0] * proj_geom['Vectors'][i,7] + factor[1] * proj_geom['Vectors'][i,10];
-                proj_geom['Vectors'][i,5] = proj_geom['Vectors'][i,5] + factor[0] * proj_geom['Vectors'][i,8] + factor[1] * proj_geom['Vectors'][i,11];
+        V = proj_geom['Vectors']
+        V[:,3:6] = V[:,3:6] + factor[0] * V[:,6:9]
+        if len(factor) > 1:
+            V[:,3:6] = V[:,3:6] + factor[1] * V[:,9:12]
     else:
-        raise ValueError('No suitable geometry for postalignment: ' + proj_geom['type'])
+        raise RuntimeError('No suitable geometry for postalignment: ' + proj_geom['type'])
 
     return proj_geom
 
diff --git a/samples/matlab/s022_fbp_cor.m b/samples/matlab/s022_fbp_cor.m
index b33d5ce..8521d7e 100644
--- a/samples/matlab/s022_fbp_cor.m
+++ b/samples/matlab/s022_fbp_cor.m
@@ -14,10 +14,7 @@ vol_geom = astra_create_vol_geom(256, 256);
 proj_geom = astra_create_proj_geom('parallel', 1.0, 256, linspace2(0,pi,180));
 
 % Projection geometry with shifted center of rotation
-% We create this by shifting the detector center V(:,3:4) by a multiple of the detector
-% orientation V(:,5:6).
-proj_geom_cor = astra_geom_2vec(proj_geom);
-proj_geom_cor.Vectors(:,3:4) = proj_geom_cor.Vectors(:,3:4) + cor_shift * proj_geom_cor.Vectors(:,5:6);
+proj_geom_cor = astra_geom_postalignment(proj_geom, cor_shift);
 
 % As before, create a sinogram from a phantom, using the shifted center of rotation
 P = phantom(256);
diff --git a/samples/python/s022_fbp_cor.py b/samples/python/s022_fbp_cor.py
index 2d811c0..bb110e5 100644
--- a/samples/python/s022_fbp_cor.py
+++ b/samples/python/s022_fbp_cor.py
@@ -32,11 +32,7 @@ vol_geom = astra.create_vol_geom(256, 256)
 proj_geom = astra.create_proj_geom('parallel', 1.0, 256, np.linspace(0,np.pi,180,False))
 
 # Projection geometry with shifted center of rotation
-# We create this by shifting the detector center V[:,2:4] by a multiple of the detector
-# orientation V[:,4:6].
-proj_geom_cor = astra.geom_2vec(proj_geom)
-V = proj_geom_cor['Vectors']
-V[:,2:4] = V[:,2:4] + cor_shift * V[:,4:6]
+proj_geom_cor = astra.geom_postalignment(proj_geom, cor_shift)
 
 # As before, create a sinogram from a phantom, using the shifted center of rotation
 import scipy.io