Simplify reshuffle_peak_nearest function. Check if matched peak was already matched.

batchfit/model/DataModel.py

@@ -94,43 +94,57 @@ class DataModel(BatchModel):
 
     def reshuffle_peaks_nearest(self, distance_r, distance_l, weight):
         max_peaks = max(self.fit_results.values(), key=lambda x: x.n_peaks).n_peaks
-        peak_idxs = np.zeros((len(self.fit_results), max_peaks, 2)).astype(float)
+        peak_idxs = np.zeros((len(self.fit_results), max_peaks, 3)).astype(float)
         peak_idxs[...] = np.nan
         peaks_counter = 0
+        # loop over patterns
         for i, data in enumerate(self.fit_results.items()):
             img, fit_result = data
-
-            for i_peak in range(fit_result.n_peaks):
+            #print("Process img: ", img)
+            # loop over all peaks for given pattern
+            peaks_to_process = list(range(fit_result.n_peaks))
+            for i_peak in peaks_to_process:
+                #print("Process peak: ", i, i_peak, peak_idxs[i - 1, :, 1], peaks_counter)
                 peak = fit_result.params[f'peak{i_peak}_center']
-                #print("I: ", i, i_peak, peak_idxs[i - 1, :, 1], peak, peaks_counter)
-                if i == 0:
-                    peak_idxs[i, i_peak, :] = [peaks_counter, peak]
-                    peaks_counter += 1
-                else:
-                    # compare with existing previous peaks
-                    d = peak_idxs[i - 1, :, 1] - peak
+
+                # loop over last 10 processed patterns to find match
+                for j in range(1, min(10, 1+i)):
+                    #print("Looking for match in ", i-j)
+                    d = peak_idxs[i - j, :, 1] - peak
                     sel = np.where(((d > 0) & (d < distance_l)) | ((d <= 0) & (-d < distance_r)))[0]
-                    # No peaks
-                    if len(sel) == 0:
-                        # Loop at previous images
-                        for j in range(1, min(10, i)):
-                            d = peak_idxs[i - j, :, 1] - peak
-                            sel = np.where(((d > 0) & (d < distance_l)) | ((d <= 0) & (-d < distance_r)))[0]
-                            if len(sel)>0:
-                                d[np.where(d > 0)[0]] *= weight
-                                # ToDo check if that peak was already assigned
-                                idx = sel[np.argmin(np.abs(d[sel]))]
-                                matched_i_peak = peak_idxs[i - j, idx, 0]
-                                peak_idxs[i][i_peak] = [matched_i_peak, peak]
-                        if np.isnan(peak_idxs[i, i_peak, 0]):
-                            peak_idxs[i, i_peak] = [peaks_counter, peak]
-                            peaks_counter += 1
-                    else:
+                    if len(sel) > 0:
                         d[np.where(d > 0)[0]] *= weight
-                        # ToDo check if that peak was already assigned on previous step
-                        idx = sel[np.argmin(np.abs(d[sel]))]
-                        matched_i_peak = peak_idxs[i - 1, idx, 0]
-                        peak_idxs[i][i_peak] = [matched_i_peak, peak]
+                        best_d_idx = np.argmin(np.abs(d[sel]))
+                        best_d = np.abs(d[best_d_idx])
+                        idx = sel[best_d_idx]
+                        matched_i_peak = peak_idxs[i - j, idx, 0]
+                        # check if that peak was already assigned on previous step
+                        check_idx = np.where(peak_idxs[i, :, 0] == int(matched_i_peak))[0]
+                        if len(check_idx) > 0:
+                            #print("Was already assigned to ", check_idx)
+                            prev_d = peak_idxs[i, check_idx[0], 2]
+                            #print("Matching comparing ", prev_d, best_d)
+                            if prev_d > best_d:
+                                #print("Current match better")
+                                peak_idxs[i, i_peak, :] = [matched_i_peak, peak, best_d]
+                                #print("match to img peak", i - j, i_peak, matched_i_peak, peak, best_d)
+                                # reprocess previous
+                                peak_idxs[i, check_idx, :] = np.nan
+                                peaks_to_process += [check_idx[0]]
+                                #print("reprocess ", check_idx[0], peaks_to_process)
+                                break
+                            else:
+                                continue
+                        else:
+                            peak_idxs[i, i_peak, :] = [matched_i_peak, peak, best_d]
+                            #print("match to img peak", i - j, i_peak, matched_i_peak, peak, best_d)
+                            break
+
+                # create new if peak is still not assign
+                if np.isnan(peak_idxs[i, i_peak, 0]):
+                    #print("Peak is not matched, take as new")
+                    peak_idxs[i, i_peak] = [peaks_counter, peak, 0]
+                    peaks_counter += 1
 
         peaks_array = np.zeros((len(self.fit_results), peaks_counter * 12 + 2))
         header = ['img', 'chi2']