ENH: add vector geometries, closes odlgroup#1023

examples/tomo/ray_trafo_vec_geom_3d.py

@@ -0,0 +1,78 @@
+"""Example using the ray transform a custom vector geometry.
+
+We manually build a "circle plus line trajectory" (CLT) geometry by
+extracting the vectors from a circular geometry and extending it by
+vertical shifts, starting at the initial position.
+"""
+
+import numpy as np
+import odl
+
+# Reconstruction space: discretized functions on the cube [-20, 20]^3
+# with 300 samples per dimension.
+reco_space = odl.uniform_discr(
+    min_pt=[-20, -20, -20], max_pt=[20, 20, 20], shape=[300, 300, 300],
+    dtype='float32')
+
+# First part of the geometry: a 3D single-axis parallel beam geometry with
+# flat detector
+# Angles: uniformly spaced, n = 180, min = 0, max = 2 * pi
+angle_partition = odl.uniform_partition(0, 2 * np.pi, 180)
+# Detector: uniformly sampled, n = (512, 512), min = (-30, -30), max = (30, 30)
+detector_partition = odl.uniform_partition([-30, -30], [30, 30], [512, 512])
+circle_geometry = odl.tomo.CircularConeFlatGeometry(
+    angle_partition, detector_partition, src_radius=1000, det_radius=100,
+    axis=[1, 0, 0])
+
+circle_vecs = odl.tomo.astra_conebeam_3d_geom_to_vec(circle_geometry)
+
+# Cover the whole volume vertically, somewhat undersampled though
+vert_shift_min = -22
+vert_shift_max = 22
+num_shifts = 180
+vert_shifts = np.linspace(vert_shift_min, vert_shift_max, num=num_shifts)
+inital_vecs = circle_vecs[0]
+
+# Start from the initial position of the circle vectors and add the vertical
+# shifts to the columns 2 and 5 (source and detector z positions)
+line_vecs = np.repeat(circle_vecs[0][None, :], num_shifts, axis=0)
+line_vecs[:, 2] += vert_shifts
+line_vecs[:, 5] += vert_shifts
+
+# Build the composed geometry and the corresponding ray transform
+# (= forward projection)
+composed_vecs = np.vstack([circle_vecs, line_vecs])
+composed_geom = odl.tomo.ConeVecGeometry(detector_partition.shape,
+                                         composed_vecs)
+
+ray_trafo = odl.tomo.RayTransform(reco_space, composed_geom)
+
+# Create a Shepp-Logan phantom (modified version) and projection data
+phantom = odl.phantom.shepp_logan(reco_space, True)
+proj_data = ray_trafo(phantom)
+
+# Back-projection can be done by simply calling the adjoint operator on the
+# projection data (or any element in the projection space).
+backproj = ray_trafo.adjoint(proj_data)
+
+# Show the slice z=0 of phantom and backprojection, as well as a projection
+# image at theta=0 and a sinogram at v=0 (middle detector row)
+phantom.show(coords=[None, None, 0], title='Phantom, middle z slice')
+backproj.show(coords=[None, None, 0], title='Back-projection, middle z slice')
+proj_data.show(indices=[0, slice(None), slice(None)],
+               title='Projection 0 (circle start)')
+proj_data.show(indices=[45, slice(None), slice(None)],
+               title='Projection 45 (circle 1/4)')
+proj_data.show(indices=[90, slice(None), slice(None)],
+               title='Projection 90 (circle 1/2)')
+proj_data.show(indices=[135, slice(None), slice(None)],
+               title='Projection 135 (circle 3/4)')
+proj_data.show(indices=[179, slice(None), slice(None)],
+               title='Projection 179 (circle end)')
+proj_data.show(indices=[180, slice(None), slice(None)],
+               title='Projection 180 (line start)')
+proj_data.show(indices=[270, slice(None), slice(None)],
+               title='Projection 270 (line middle)')
+proj_data.show(indices=[359, slice(None), slice(None)],
+               title='Projection 359 (line end)')
+proj_data.show(coords=[None, None, 0], title='Sinogram, middle slice')

odl/tomo/backends/astra_setup.py

@@ -52,7 +52,9 @@
 
 from odl.discr import DiscreteLp, DiscreteLpElement
 from odl.tomo.geometry import (
-    Geometry, DivergentBeamGeometry, ParallelBeamGeometry, FlatDetector)
+    Geometry, DivergentBeamGeometry, ParallelBeamGeometry,
+    ParallelVecGeometry, ConeVecGeometry,
+    FlatDetector)
 from odl.tomo.util.utility import euler_matrix
 from odl.util.utility import pkg_supports
 
@@ -211,6 +213,75 @@ def astra_volume_geometry(discr_reco):
     return vol_geom
 
 
+def vecs_odl_axes_to_astra_axes(vecs):
+    """Convert geometry vectors from ODL axis convention to ASTRA.
+
+    Parameters
+    ----------
+    vecs : array-like, shape ``(N, 6)`` or ``(N, 12)``
+        Vectors defining the geometric configuration in each
+        projection. The number ``N`` of rows determines the number
+        of projections, and the number of columns the spatial
+        dimension (6 for 2D, 12 for 3D).
+
+    Returns
+    -------
+    astra_vecs : `numpy.ndarray`, same shape as ``vecs``
+        The converted geometry vectors.
+    """
+    vecs = np.asarray(vecs, dtype=float)
+
+    if vecs.shape[1] == 6:
+        # 2D geometry, nothing to do since the axes are the same
+        return vecs
+    elif vecs.shape[1] == 12:
+        # 3D geometry
+        # ASTRA has (z, y, x) axis convention, in contrast to (x, y, z) in ODL,
+        # so we need to adapt to this by changing the order.
+        newind = []
+        for i in range(4):
+            newind.extend([2 + 3 * i, 1 + 3 * i, 0 + 3 * i])
+        return vecs[:, newind]
+    else:
+        raise ValueError('`vecs` must have shape (N, 6) or (N, 12), got '
+                         'array with shape {}'.format(vecs.shape))
+
+
+def vecs_astra_axes_to_odl_axes(vecs):
+    """Convert geometry vectors from ASTRA axis convention to ODL.
+
+    Parameters
+    ----------
+    vecs : array-like, shape ``(N, 6)`` or ``(N, 12)``
+        Vectors defining the geometric configuration in each
+        projection. The number ``N`` of rows determines the number
+        of projections, and the number of columns the spatial
+        dimension (6 for 2D, 12 for 3D).
+
+    Returns
+    -------
+    odl_vecs : `numpy.ndarray`, same shape as ``vecs``
+        The converted geometry vectors.
+    """
+    vecs = np.asarray(vecs, dtype=float)
+
+    if vecs.shape[1] == 6:
+        # 2D geometry, nothing to do since the axes are the same
+        return vecs
+    elif vecs.shape[1] == 12:
+        # 3D geometry
+        # ASTRA has (z, y, x) axis convention, in contrast to (x, y, z) in ODL,
+        # so we need to adapt to this by changing the order.
+        newind = []
+        for i in range(4):
+            newind.extend([2 + 3 * i, 1 + 3 * i, 0 + 3 * i])
+        newind = np.argsort(newind).tolist()
+        return vecs[:, newind]
+    else:
+        raise ValueError('`vecs` must have shape (N, 6) or (N, 12), got '
+                         'array with shape {}'.format(vecs.shape))
+
+
 def astra_conebeam_3d_geom_to_vec(geometry):
     """Create vectors for ASTRA projection geometries from ODL geometry.
 
@@ -261,14 +332,7 @@ def astra_conebeam_3d_geom_to_vec(geometry):
     vectors[:, 6:9] = det_axes[0] * px_sizes[0]
     vectors[:, 9:12] = det_axes[1] * px_sizes[1]
 
-    # ASTRA has (z, y, x) axis convention, in contrast to (x, y, z) in ODL,
-    # so we need to adapt to this by changing the order.
-    newind = []
-    for i in range(4):
-        newind += [2 + 3 * i, 1 + 3 * i, 0 + 3 * i]
-    vectors = vectors[:, newind]
-
-    return vectors
+    return vecs_odl_axes_to_astra_axes(vectors)
 
 
 def astra_conebeam_2d_geom_to_vec(geometry):
@@ -325,7 +389,7 @@ def astra_conebeam_2d_geom_to_vec(geometry):
     px_size = geometry.det_partition.cell_sides[0]
     vectors[:, 4:6] = det_axis * px_size
 
-    return vectors
+    return vecs_odl_axes_to_astra_axes(vectors)
 
 
 def astra_parallel_3d_geom_to_vec(geometry):
@@ -379,13 +443,7 @@ def astra_parallel_3d_geom_to_vec(geometry):
     vectors[:, 6:9] = det_axes[0] * px_sizes[0]
     vectors[:, 9:12] = det_axes[1] * px_sizes[1]
 
-    # ASTRA has (z, y, x) axis convention, in contrast to (x, y, z) in ODL,
-    # so we need to adapt to this by changing the order.
-    new_ind = []
-    for i in range(4):
-        new_ind += [2 + 3 * i, 1 + 3 * i, 0 + 3 * i]
-    vectors = vectors[:, new_ind]
-    return vectors
+    return vecs_odl_axes_to_astra_axes(vectors)
 
 
 def astra_projection_geometry(geometry):
@@ -435,6 +493,22 @@ def astra_projection_geometry(geometry):
         vec = astra_conebeam_2d_geom_to_vec(geometry)
         proj_geom = astra.create_proj_geom('fanflat_vec', det_count, vec)
 
+    elif isinstance(geometry, ParallelVecGeometry) and geometry.ndim == 2:
+        det_count = geometry.detector.size
+        vec = geometry.vectors
+        # TODO: flip axes?
+        if not astra_supports('parallel2d_vec_geometry'):
+            raise NotImplementedError(
+                "'parallel_vec' geometry not supported by ASTRA "
+                'v{}'.format(ASTRA_VERSION))
+        proj_geom = astra.create_proj_geom('parallel_vec', det_count, vec)
+
+    elif isinstance(geometry, ConeVecGeometry) and geometry.ndim == 2:
+        det_count = geometry.detector.size
+        vec = geometry.vectors
+        # TODO: flip axes?
+        proj_geom = astra.create_proj_geom('fanflat_vec', det_count, vec)
+
     elif (isinstance(geometry, ParallelBeamGeometry) and
           isinstance(geometry.detector, FlatDetector) and
           geometry.ndim == 3):
@@ -452,6 +526,23 @@ def astra_projection_geometry(geometry):
         vec = astra_conebeam_3d_geom_to_vec(geometry)
         proj_geom = astra.create_proj_geom('cone_vec', det_row_count,
                                            det_col_count, vec)
+
+    elif isinstance(geometry, ParallelVecGeometry) and geometry.ndim == 3:
+        det_row_count = geometry.det_partition.shape[1]
+        det_col_count = geometry.det_partition.shape[0]
+        vec = geometry.vectors
+        # TODO: flip axes?
+        proj_geom = astra.create_proj_geom('parallel3d_vec', det_row_count,
+                                           det_col_count, vec)
+
+    elif isinstance(geometry, ConeVecGeometry) and geometry.ndim == 3:
+        det_row_count = geometry.det_partition.shape[1]
+        det_col_count = geometry.det_partition.shape[0]
+        vec = geometry.vectors
+        # TODO: flip axes?
+        proj_geom = astra.create_proj_geom('cone_vec', det_row_count,
+                                           det_col_count, vec)
+
     else:
         raise NotImplementedError('unknown ASTRA geometry type {!r}'
                                   ''.format(geometry))