Add slicing of ImageSetData

Needed to make sliced image sets when the list of models held is only as long as the number of indices.

src/dxtbx/boost_python/imageset_ext.cc

@@ -610,6 +610,7 @@ namespace dxtbx { namespace boost_python {
       .def("set_params", &ImageSetData_set_params)
       .def("get_format_class", &ImageSetData_get_format)
       .def("set_format_class", &ImageSetData_set_format)
+      .def("partial_data", &ImageSetData::partial_data)
       .add_property(
         "external_lookup",
         make_function(&ImageSetData::external_lookup, return_internal_reference<>()))

src/dxtbx/format/Format.py

@@ -66,7 +66,7 @@ def __len__(self):
         return len(self._filenames)
 
     def copy(self, filenames):
-        return Reader(self.format_class, filenames)
+        return Reader(self.format_class, filenames, **self._kwargs)
 
     def is_single_file_reader(self):
         return False

src/dxtbx/format/FormatMultiImage.py

@@ -50,8 +50,8 @@ def paths(self):
     def __len__(self):
         return self._num_images
 
-    def copy(self, filenames, indices=None):
-        return Reader(self.format_class, filenames, indices)
+    def copy(self, filenames, num_images=None):
+        return Reader(self.format_class, filenames, num_images, **self.kwargs)
 
     def identifiers(self):
         return ["%s-%d" % (self._filename, index) for index in range(len(self))]

src/dxtbx/imageset.h

@@ -175,6 +175,28 @@ class ImageSetData {
         scans_(boost::python::len(reader)),
         reject_(boost::python::len(reader)) {}
 
+  /**
+   * Copy constructor for the imageset data object
+   */
+  /*
+  ImageSetData(boost::python::object reader, masker_ptr masker,
+               scitbx::af::shared<beam_ptr> beams, scitbx::af::shared<detector_ptr> detectors,
+               scitbx::af::shared<goniometer_ptr> goniometers, scitbx::af::shared<scan_ptr> scans,
+               scitbx::af::shared<bool> reject, ExternalLookup external_lookup,
+               std::string templatestr, std::string vendor,
+               std::string params, std::string format)
+      : reader_(reader),
+        masker_(masker),
+        beams_(beams),
+        detectors_(detectors),
+        goniometers_(goniometers),
+        scans_(scans),
+        reject_(reject),
+        external_lookup_(external_lookup),
+        template_(templatestr), vendor_(vendor),
+        params_(params), format_(format) {}
+    */
+
   /**
    * @returns The reader object
    */
@@ -434,6 +456,33 @@ class ImageSetData {
     format_ = x;
   }
 
+  ImageSetData partial_data(boost::python::object reader, std::size_t first, std::size_t last) const {
+    DXTBX_ASSERT(last > first);
+    //return ImageSetData(reader, masker_,
+    //                    scitbx::af::shared<beam_ptr>(&beams_[first], last - first),
+    //                    scitbx::af::shared<detector_ptr>(&detectors_[first], last - first),
+    //                    scitbx::af::shared<goniometer_ptr>(&goniometers_[first], last - first),
+    //                    scitbx::af::shared<scan_ptr>(&scans_[first], last - first),
+    //                    scitbx::af::shared<bool>(&reject_[first], last - first),
+    //                    external_lookup_, template_, vendor_, params_, format_);
+    //
+    ImageSetData partial = ImageSetData(reader, masker_);
+    for (size_t i = 0; i < last-first; i++) {
+      partial.beams_[i] = beams_[i+first];
+      partial.detectors_[i] = detectors_[i+first];
+      partial.goniometers_[i] = goniometers_[i+first];
+      partial.scans_[i] = scans_[i+first];
+      partial.reject_[i] = reject_[i+first];
+    }
+    partial.external_lookup_ = external_lookup_;
+    partial.template_ = template_;
+    partial.vendor_ = vendor_;
+    partial.params_ = params_;
+    partial.format_ = format_;
+
+    return partial;
+  }
+
 protected:
   ImageBuffer get_image_buffer_from_tuple(boost::python::tuple obj) {
     // Get the class name
@@ -1037,9 +1086,9 @@ class ImageSet {
    * @param last The last slice index
    * @returns The partial set
    */
-  virtual ImageSet partial_set(std::size_t first, std::size_t last) const {
+  virtual ImageSet partial_set(boost::python::object reader, std::size_t first, std::size_t last) const {
     DXTBX_ASSERT(last > first);
-    return ImageSet(data_,
+    return ImageSet(data_.partial_data(reader, first, last),
                     scitbx::af::const_ref<std::size_t>(&indices_[first], last - first));
   }
 

src/dxtbx/imageset.py

@@ -96,7 +96,11 @@ def __getitem__(self, item):
             stop = item.stop or len(self)
             if item.step is not None and item.step != 1:
                 raise IndexError("Step must be 1")
-            return self.partial_set(start, stop)
+            if self.data().has_single_file_reader():
+                reader = self.reader().copy(self.reader().paths(), stop-start)
+            else:
+                reader = self.reader().copy(self.reader().paths())
+            return self.partial_set(reader, start, stop)
         else:
             return self.get_corrected_data(item)
 
@@ -220,7 +224,15 @@ def _load_models(self, index):
 
     def __getitem__(self, item):
         if isinstance(item, slice):
-            return ImageSetLazy(self.data(), indices=self.indices()[item])
+            start = item.start or 0
+            stop = item.stop or len(self)
+            if item.step is not None and item.step != 1:
+                raise IndexError("Step must be 1")
+            if self.data().has_single_file_reader():
+                reader = self.reader().copy(self.reader().paths(), stop-start)
+            else:
+                reader = self.reader().copy(self.reader().paths())
+            return ImageSetLazy(self.data().partial_data(reader, start, stop), indices=self.indices()[item])
         self._load_models(item)
         return super().__getitem__(item)