[gui][shift] Find shift for all values of filtered dimension

darfix/gui/shiftCorrectionWidget.py

@@ -87,7 +87,8 @@ class ShiftCorrectionWidget(qt.QMainWindow):
         qt.QMainWindow.__init__(self, parent)
 
         self.setWindowFlags(qt.Qt.Widget)
-        self._shift = [0, 0]
+        self._shift = numpy.array([0, 0])
+        self._filtered_shift = None
         self._dimension = None
         self._update_dataset = None
         self.indices = None
@@ -172,9 +173,13 @@ class ShiftCorrectionWidget(qt.QMainWindow):
         self.sigProgressChanged.emit(progress)
 
     def _findShift(self):
-        self.thread_detection = OperationThread(self, self._update_dataset.find_shift)
+        if self._filtered_shift is not None:
+            self.thread_detection = OperationThread(self, self._update_dataset.find_shift_along_dimension)
+            self.thread_detection.setArgs(self._dimension[0], indices=self.indices)
+        else:
+            self.thread_detection = OperationThread(self, self._update_dataset.find_shift)
+            self.thread_detection.setArgs(self._dimension, indices=self.indices)
         self._inputDock.widget._findShiftB.setEnabled(False)
-        self.thread_detection.setArgs(self._dimension, indices=self.indices)
         self.thread_detection.finished.connect(self._updateShift)
         self.thread_detection.start()
         self.computingSignal.emit(True)
@@ -182,7 +187,11 @@ class ShiftCorrectionWidget(qt.QMainWindow):
     def _updateShift(self):
         self._inputDock.widget._findShiftB.setEnabled(True)
         self.thread_detection.finished.disconnect(self._updateShift)
-        self.shift = numpy.round(self.thread_detection.data[:, 1], 5)
+        if self._filtered_shift is None:
+            self.shift = numpy.round(self.thread_detection.data[:, 1], 5)
+        else:
+            self._filtered_shift = numpy.round(self.thread_detection.data[:, :, 1], 5)
+            self.shift = self._filtered_shift[self._dimension[1][0]]
         self.computingSignal.emit(False)
 
     def _updateData(self):
@@ -226,7 +235,13 @@ class ShiftCorrectionWidget(qt.QMainWindow):
     def _filterStack(self, dim=0, val=0):
         self._inputDock.widget.checkbox.show()
         self._dimension = [dim, val]
+
         data = self._update_dataset.get_data(self.indices, self._dimension)
+        if self.dataset.dims.ndim == 2:
+            stack_size = self.dataset.dims.get(dim[0]).size
+            reset_shift = self._filtered_shift is None or self._filtered_shift.shape[0] != stack_size
+            self._filtered_shift = numpy.zeros((stack_size, 2)) if reset_shift else self._filtered_shift
+            self.shift = self._filtered_shift[val[0]]
         if data.shape[0]:
             self._sv.setStack(data)
         else:
@@ -234,6 +249,8 @@ class ShiftCorrectionWidget(qt.QMainWindow):
 
     def _wholeStack(self):
         self._dimension = None
+        self._filtered_shift = None
+        self.shift = numpy.array([0, 0])
         self._inputDock.widget.checkbox.hide()
         self.setStack(self._update_dataset)