Speed up RangesMatrix (#241)

* skip shape check on .full
* faster len and getitem<int>
* extract_ranges

Author: mhasself

include/Ranges.h

@@ -68,8 +68,9 @@ class Ranges {
 // Support for working with RangesMatrix, which is basically just a list of Ranges
 template <typename T>
 vector<Ranges<T>> extract_ranges(const bp::object & ival_list) {
-    vector<Ranges<T>> v(bp::len(ival_list));
-    for (int i=0; i<bp::len(ival_list); i++)
+    const int N = bp::len(ival_list);
+    vector<Ranges<T>> v(N);
+    for (int i=0; i<N ; i++)
         v[i] = bp::extract<Ranges<T>>(ival_list[i])();
     return v;
 }

python/proj/ranges.py

@@ -35,7 +35,7 @@ def __repr__(self):
         return 'RangesMatrix(' + ','.join(map(str, self.shape)) + ')'
 
     def __len__(self):
-        return self.shape[0]
+        return len(self.ranges)
 
     def copy(self):
         return RangesMatrix([x.copy() for x in self.ranges],
@@ -44,7 +44,7 @@ def copy(self):
     def zeros_like(self):
         return RangesMatrix([x.zeros_like() for x in self.ranges],
                             child_shape=self.shape[1:])
-    
+
     def ones_like(self):
         return RangesMatrix([x.ones_like() for x in self.ranges],
                             child_shape=self.shape[1:])
@@ -53,7 +53,7 @@ def buffer(self, buff):
         [x.buffer(buff) for x in self.ranges]
         ## just to make this work like Ranges.buffer()
         return self
-    
+
     def buffered(self, buff):
         out = self.copy()
         [x.buffer(buff) for x in out.ranges]
@@ -93,6 +93,10 @@ def __getitem__(self, index):
           new_rm = rm[..., :]
 
         """
+        # Short-circuit return if this is a simple integer index.
+        if isinstance(index, (int, np.int32, np.int64)):
+            return self.ranges[index]
+
         if not isinstance(index, tuple):
             index = (index,)
 
@@ -128,7 +132,7 @@ def __add__(self, x):
                 elif self.shape[0] == x.shape[0]:
                     return self.__class__([r + d for r, d in zip(self.ranges, x)], skip_shape_check=True)
             return self.__class__([r + x for r in self.ranges], skip_shape_check=True)
-        
+
     def __mul__(self, x):
         if isinstance(x, Ranges):
             return self.__class__([d * x for d in self.ranges], skip_shape_check=True)
@@ -247,7 +251,7 @@ def full(shape, fill_value):
             return r
         return RangesMatrix([RangesMatrix.full(shape[1:], fill_value)
                              for i in range(shape[0])],
-                            child_shape=shape[1:])
+                            child_shape=shape[1:], skip_shape_check=True)
 
     @classmethod
     def zeros(cls, shape):

src/Projection.cxx

@@ -1889,16 +1889,20 @@ vector<vector<vector<RangesInt32>>> derive_ranges(
         ivals.push_back(vector<vector<RangesInt32>>(1, extract_ranges<int32_t>(thread_intervals)));
     } else if(bp::extract<RangesInt32>(thread_intervals[0][0]).check()) {
         // It's a per-thread RangesMatrix (nthread,ndet,nranges). Promote to single bunch
-        vector<vector<RangesInt32>> bunch;
-        for (int i=0; i<bp::len(thread_intervals); i++)
-            bunch.push_back(extract_ranges<int32_t>(thread_intervals[i]));
+        int N = bp::len(thread_intervals);
+        vector<vector<RangesInt32>> bunch(N);
+        for (int i=0; i<N; i++)
+            bunch[i] = extract_ranges<int32_t>(thread_intervals[i]);
         ivals.push_back(bunch);
     } else if(bp::extract<RangesInt32>(thread_intervals[0][0][0]).check()) {
         // It's a full multi-bunch (nbunch,nthread,ndet,nranges) thing.
-        for (int i=0; i<bp::len(thread_intervals); i++) {
-            vector<vector<RangesInt32>> bunch;
-            for (int j=0; j<bp::len(thread_intervals[i]); j++)
-                bunch.push_back(extract_ranges<int32_t>(thread_intervals[i][j]));
+        const int N = bp::len(thread_intervals);
+        for (int i=0; i<N; i++) {
+            auto ti_i = thread_intervals[i];
+            int M = bp::len(ti_i);
+            vector<vector<RangesInt32>> bunch(M);
+            for (int j=0; j<M; j++)
+                bunch[j] = extract_ranges<int32_t>(ti_i[j]);
             ivals.push_back(bunch);
         }
     } else {