ENH: vectorize geometry methods

Author: Holger Kohr

odl/tomo/backends/astra_setup.py

@@ -248,19 +248,18 @@ def astra_conebeam_3d_geom_to_vec(geometry):
     angles = geometry.angles
     vectors = np.zeros((angles.size, 12))
 
-    for ang_idx, angle in enumerate(angles):
-        # Source position
-        vectors[ang_idx, 0:3] = geometry.src_position(angle)
+    # Source position
+    vectors[:, 0:3] = geometry.src_position(angles)
 
-        # Center of detector in 3D space
-        mid_pt = geometry.det_params.mid_pt
-        vectors[ang_idx, 3:6] = geometry.det_point_position(angle, mid_pt)
+    # Center of detector in 3D space
+    mid_pt = geometry.det_params.mid_pt
+    vectors[:, 3:6] = geometry.det_point_position(angles, mid_pt)
 
-        # Vectors from detector pixel (0, 0) to (1, 0) and (0, 0) to (0, 1)
-        det_axes = geometry.det_axes(angle)
-        px_sizes = geometry.det_partition.cell_sides
-        vectors[ang_idx, 6:9] = det_axes[0] * px_sizes[0]
-        vectors[ang_idx, 9:12] = det_axes[1] * px_sizes[1]
+    # Vectors from detector pixel (0, 0) to (1, 0) and (0, 0) to (0, 1)
+    det_axes = geometry.det_axes(angles)
+    px_sizes = geometry.det_partition.cell_sides
+    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.
@@ -312,19 +311,19 @@ def astra_conebeam_2d_geom_to_vec(geometry):
     angles = geometry.angles
     vectors = np.zeros((angles.size, 6))
 
-    for ang_idx, angle in enumerate(angles):
-        # Source position
-        vectors[ang_idx, 0:2] = rot_minus_90.dot(geometry.src_position(angle))
+    # Source position
+    src_pos = geometry.src_position(angles)
+    vectors[:, 0:2] = rot_minus_90.dot(src_pos.T).T  # dot along 2nd axis
 
-        # Center of detector
-        mid_pt = geometry.det_params.mid_pt
-        vectors[ang_idx, 2:4] = rot_minus_90.dot(
-            geometry.det_point_position(angle, mid_pt))
+    # Center of detector
+    mid_pt = geometry.det_params.mid_pt
+    centers = geometry.det_point_position(angles, mid_pt)
+    vectors[:, 2:4] = rot_minus_90.dot(centers.T).T
 
-        # Vector from detector pixel 0 to 1
-        det_axis = rot_minus_90.dot(geometry.det_axis(angle))
-        px_size = geometry.det_partition.cell_sides[0]
-        vectors[ang_idx, 4:6] = det_axis * px_size
+    # Vector from detector pixel 0 to 1
+    det_axis = rot_minus_90.dot(geometry.det_axis(angles).T).T
+    px_size = geometry.det_partition.cell_sides[0]
+    vectors[:, 4:6] = det_axis * px_size
 
     return vectors
 
@@ -366,20 +365,19 @@ def astra_parallel_3d_geom_to_vec(geometry):
     angles = geometry.angles
     vectors = np.zeros((angles.shape[0], 12))
 
-    for ang_idx, angle in enumerate(angles):
-        mid_pt = geometry.det_params.mid_pt
+    mid_pt = geometry.det_params.mid_pt
 
-        # Ray direction = -(detector-to-source normal vector)
-        vectors[ang_idx, 0:3] = -geometry.det_to_src(angle, mid_pt)
+    # Ray direction = -(detector-to-source normal vector)
+    vectors[:, 0:3] = -geometry.det_to_src(angles, mid_pt)
 
-        # Center of the detector in 3D space
-        vectors[ang_idx, 3:6] = geometry.det_point_position(angle, mid_pt)
+    # Center of the detector in 3D space
+    vectors[:, 3:6] = geometry.det_point_position(angles, mid_pt)
 
-        # Vectors from detector pixel (0, 0) to (1, 0) and (0, 0) to (0, 1)
-        det_axes = geometry.det_axes(angle)
-        px_sizes = geometry.det_partition.cell_sides
-        vectors[ang_idx, 6:9] = det_axes[0] * px_sizes[0]
-        vectors[ang_idx, 9:12] = det_axes[1] * px_sizes[1]
+    # Vectors from detector pixel (0, 0) to (1, 0) and (0, 0) to (0, 1)
+    det_axes = geometry.det_axes(angles)
+    px_sizes = geometry.det_partition.cell_sides
+    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.

odl/tomo/geometry/conebeam.py

@@ -183,7 +183,7 @@ def __init__(self, apart, dpart, src_radius, det_radius,
         detector = Flat1dDetector(dpart, det_axis_init)
         translation = kwargs.pop('translation', None)
         super().__init__(ndim=2, motion_part=apart, detector=detector,
-                         translation=translation)
+                         translation=translation, **kwargs)
 
         self.__src_radius = float(src_radius)
         if self.src_radius < 0:
@@ -202,11 +202,6 @@ def __init__(self, apart, dpart, src_radius, det_radius,
             raise ValueError('`apart` has dimension {}, expected 1'
                              ''.format(self.motion_partition.ndim))
 
-        # Make sure there are no leftover kwargs
-        if kwargs:
-            raise TypeError('got unexpected keyword arguments {}'
-                            ''.format(kwargs))
-
     @classmethod
     def frommatrix(cls, apart, dpart, src_radius, det_radius, init_matrix):
         """Create an instance of `FanFlatGeometry` using a matrix.
@@ -357,16 +352,19 @@ def src_position(self, angle):
         >>> np.allclose(geom.src_position(np.pi / 2), [2, 0])
         True
         """
-        if angle not in self.motion_params:
-            raise ValueError('`angle` {} is not in the valid range {}'
+        if self.check_bounds and not self.motion_params.contains_all(angle):
+            raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
 
+        angle = np.array(angle, dtype=float, copy=False, ndmin=1)
+
         # Initial vector from the rotation center to the source. It can be
         # computed this way since source and detector are at maximum distance,
         # i.e. the connecting line passes the origin.
         center_to_src_init = -self.src_radius * self.src_to_det_init
-        return (self.translation +
-                self.rotation_matrix(angle).dot(center_to_src_init))
+        pos_vec = (self.translation[None, :] +
+                   self.rotation_matrix(angle).dot(center_to_src_init))
+        return pos_vec.squeeze()
 
     def det_refpoint(self, angle):
         """Return the detector reference point position at ``angle``.
@@ -407,16 +405,18 @@ def det_refpoint(self, angle):
         >>> np.allclose(geom.det_refpoint(np.pi / 2), [-5, 0])
         True
         """
-        if angle not in self.motion_params:
-            raise ValueError('`angle` {} is not in the valid range {}'
+        if self.check_bounds and not self.motion_params.contains_all(angle):
+            raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
 
+        angle = np.array(angle, dtype=float, copy=False, ndmin=1)
         # Initial vector from the rotation center to the detector. It can be
         # computed this way since source and detector are at maximum distance,
         # i.e. the connecting line passes the origin.
         center_to_det_init = self.det_radius * self.src_to_det_init
-        return (self.translation +
-                self.rotation_matrix(angle).dot(center_to_det_init))
+        refpt = (self.translation[None, :] +
+                 self.rotation_matrix(angle).dot(center_to_det_init))
+        return refpt.squeeze()
 
     def rotation_matrix(self, angle):
         """Return the rotation matrix for ``angle``.
@@ -440,10 +440,10 @@ def rotation_matrix(self, angle):
             the local coordinate system of the detector reference point,
             expressed in the fixed system.
         """
-        angle = float(angle)
-        if angle not in self.motion_params:
+        if self.check_bounds and not self.motion_params.contains_all(angle):
             raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
+
         return euler_matrix(angle)
 
     def __repr__(self):
@@ -834,9 +834,9 @@ def angles(self):
         """Discrete angles given in this geometry."""
         return self.motion_grid.coord_vectors[0]
 
-    def det_axes(self, angles):
+    def det_axes(self, angle):
         """Return the detector axes tuple at ``angle``."""
-        return tuple(self.rotation_matrix(angles).dot(axis)
+        return tuple(self.rotation_matrix(angle).dot(axis)
                      for axis in self.det_axes_init)
 
     def det_refpoint(self, angle):
@@ -881,22 +881,25 @@ def det_refpoint(self, angle):
         >>> np.allclose(geom.det_refpoint(np.pi / 2), [-10, 0, 0.5])
         True
         """
-        angle = float(angle)
-        if angle not in self.motion_params:
-            raise ValueError('`angle` {} is not in the valid range {}'
+        if self.check_bounds and not self.motion_params.contains_all(angle):
+            raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
 
+        angle = np.array(angle, dtype=float, copy=False, ndmin=1)
+
         # Initial vector from center of rotation to detector.
         # It can be computed this way since source and detector are at
         # maximum distance, i.e. the connecting line passes the origin.
         center_to_det_init = self.det_radius * self.src_to_det_init
         circle_component = self.rotation_matrix(angle).dot(center_to_det_init)
 
         # Increment along the rotation axis according to pitch and pitch_offset
-        pitch_component = self.axis * (self.pitch_offset +
-                                       self.pitch * angle / (2 * np.pi))
+        shift_along_axis = (self.pitch_offset +
+                            self.pitch * angle / (2 * np.pi))
+        pitch_component = self.axis[None, :] * shift_along_axis[:, None]
 
-        return self.translation + circle_component + pitch_component
+        refpt = self.translation[None, :] + circle_component + pitch_component
+        return refpt.squeeze()
 
     def src_position(self, angle):
         """Return the source position at ``angle``.
@@ -940,22 +943,25 @@ def src_position(self, angle):
         >>> np.allclose(geom.src_position(np.pi / 2), [5, 0, 0.5])
         True
         """
-        angle = float(angle)
-        if angle not in self.motion_params:
-            raise ValueError('`angle` {} is not in the valid range {}'
+        if self.check_bounds and not self.motion_params.contains_all(angle):
+            raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
 
+        angle = np.array(angle, dtype=float, copy=False, ndmin=1)
+
         # Initial vector from 0 to the source (non-translated).
         # It can be computed this way since source and detector are at
         # maximum distance, i.e. the connecting line passes the origin.
         origin_to_src_init = -self.src_radius * self.src_to_det_init
         circle_component = self.rotation_matrix(angle).dot(origin_to_src_init)
 
-        # Increment by pitch
-        pitch_component = self.axis * (self.pitch_offset +
-                                       self.pitch * angle / (np.pi * 2))
+        # Increment along the rotation axis according to pitch and pitch_offset
+        shift_along_axis = (self.pitch_offset +
+                            self.pitch * angle / (2 * np.pi))
+        pitch_component = self.axis[None, :] * shift_along_axis[:, None]
 
-        return self.translation + circle_component + pitch_component
+        refpt = self.translation[None, :] + circle_component + pitch_component
+        return refpt.squeeze()
 
     def __repr__(self):
         """Return ``repr(self)``."""

odl/tomo/geometry/geometry.py

@@ -45,18 +45,26 @@ class Geometry(object):
     <https://odlgroup.github.io/odl/guide/geometry_guide.html>`_.
     """
 
-    def __init__(self, ndim, motion_part, detector, translation=None):
+    def __init__(self, ndim, motion_part, detector, translation=None,
+                 **kwargs):
         """Initialize a new instance.
 
         Parameters
         ----------
         ndim : positive int
             Number of dimensions of this geometry, i.e. dimensionality
-            of the physical space in which this geometry is embedded
+            of the physical space in which this geometry is embedded.
         motion_part : `RectPartition`
-           Partition for the set of "motion" parameters
+            Partition for the set of "motion" parameters.
         detector : `Detector`
-           The detector of this geometry
+            The detector of this geometry.
+        translation : `array-like`, optional
+            Global translation of the geometry. This is added last in any
+            method that computes an absolute vector, e.g., `det_refpoint`.
+        check_bounds : bool, optional
+            If ``True``, check if provided parameters for query functions
+            like `det_refpoint` are in the valid range.
+            Default: ``True``
         """
         ndim, ndim_in = int(ndim), ndim
         if ndim != ndim_in or ndim <= 0:
@@ -72,6 +80,7 @@ def __init__(self, ndim, motion_part, detector, translation=None):
         self.__ndim = ndim
         self.__motion_partition = motion_part
         self.__detector = detector
+        self.__check_bounds = bool(kwargs.pop('check_bounds', True))
 
         if translation is None:
             self.__translation = np.zeros(self.ndim)
@@ -84,6 +93,11 @@ def __init__(self, ndim, motion_part, detector, translation=None):
 
         self.__implementation_cache = {}
 
+        # Make sure there are no leftover kwargs
+        if kwargs:
+            raise TypeError('got unexpected keyword arguments {}'
+                            ''.format(kwargs))
+
     @property
     def ndim(self):
         """Number of dimensions of the geometry."""
@@ -154,6 +168,11 @@ def translation(self):
         """Shift of the origin of this geometry."""
         return self.__translation
 
+    @property
+    def check_bounds(self):
+        """Whether to check if method parameters are in the valid range."""
+        return self.__check_bounds
+
     def det_refpoint(self, mpar):
         """Detector reference point function.
 
@@ -271,7 +290,7 @@ def src_position(self, mpar):
         """
         raise NotImplementedError('abstract method')
 
-    def det_to_src(self, mpar, dpar, normalized=True):
+    def det_to_src(self, mparams, dparams, normalized=True):
         """Vector pointing from a detector location to the source.
 
         A function of the motion and detector parameters.
@@ -294,20 +313,22 @@ def det_to_src(self, mpar, dpar, normalized=True):
         vec : `numpy.ndarray`, shape (`ndim`,)
             (Unit) vector pointing from the detector to the source.
         """
-        if mpar not in self.motion_params:
-            raise ValueError('`mpar` {} not in the valid range {}'
-                             ''.format(mpar, self.motion_params))
-        if dpar not in self.det_params:
-            raise ValueError('`dpar` {} not in the valid range {}'
-                             ''.format(dpar, self.det_params))
+        if self.check_bounds:
+            if not self.motion_params.contains_all(mparams):
+                raise ValueError('`mparams` {} not in the valid range {}'
+                                 ''.format(mparams, self.motion_params))
+            if not self.det_params.contains_all(dparams):
+                raise ValueError('`dparams` {} not in the valid range {}'
+                                 ''.format(dparams, self.det_params))
 
-        vec = self.src_position(mpar) - self.det_point_position(mpar, dpar)
+        det_to_src_vec = (self.src_position(mparams) -
+                          self.det_point_position(mparams, dparams))
 
         if normalized:
             # axis = -1 allows this to be vectorized
-            vec /= np.linalg.norm(vec, axis=-1)
+            det_to_src_vec /= np.linalg.norm(det_to_src_vec, axis=-1)
 
-        return vec
+        return det_to_src_vec
 
 
 class AxisOrientedGeometry(object):
@@ -358,9 +379,8 @@ def rotation_matrix(self, angle):
             the local coordinate system of the detector reference point,
             expressed in the fixed system.
         """
-        angle = float(angle)
-        if angle not in self.motion_params:
-            raise ValueError('`angle` {} is not in the valid range {}'
+        if self.check_bounds and not self.motion_params.contains_all(angle):
+            raise ValueError('`angle` {} not in the valid range {}'
                              ''.format(angle, self.motion_params))
 
         return axis_rotation_matrix(self.axis, angle)