Merge pull request #1043 from stephenswat/refactor/more_ckf_info

Add additional information to candidate links

Author: stephenswat

core/include/traccc/finding/candidate_link.hpp

@@ -33,6 +33,12 @@ struct candidate_link {
 
     // chi2
     traccc::scalar chi2;
+
+    // chi2 sum
+    traccc::scalar chi2_sum;
+
+    // degrees of freedom
+    unsigned int ndf_sum;
 };
 
 }  // namespace traccc

core/include/traccc/finding/details/combinatorial_kalman_filter.hpp

@@ -186,6 +186,16 @@ combinatorial_kalman_filter(
                      ? 0
                      : links[step - 1][param_to_link[step - 1][in_param_id]]
                            .n_skipped);
+            const scalar prev_chi2_sum =
+                (step == 0
+                     ? 0.f
+                     : links[step - 1][param_to_link[step - 1][in_param_id]]
+                           .chi2_sum);
+            const unsigned int prev_ndf_sum =
+                (step == 0
+                     ? 0
+                     : links[step - 1][param_to_link[step - 1][in_param_id]]
+                           .ndf_sum);
 
             TRACCC_VERBOSE("Processing input parameter "
                            << in_param_id + 1 << " / " << n_in_params << ": "
@@ -273,7 +283,10 @@ combinatorial_kalman_filter(
                           .meas_idx = item_id,
                           .seed_idx = orig_param_id,
                           .n_skipped = skip_counter,
-                          .chi2 = chi2},
+                          .chi2 = chi2,
+                          .chi2_sum = prev_chi2_sum + chi2,
+                          .ndf_sum = prev_ndf_sum +
+                                     trk_state.get_measurement().meas_dim},
                          trk_state});
                 }
             }
@@ -316,7 +329,9 @@ combinatorial_kalman_filter(
                      .meas_idx = std::numeric_limits<unsigned int>::max(),
                      .seed_idx = orig_param_id,
                      .n_skipped = skip_counter + 1,
-                     .chi2 = std::numeric_limits<traccc::scalar>::max()});
+                     .chi2 = std::numeric_limits<traccc::scalar>::max(),
+                     .chi2_sum = prev_chi2_sum,
+                     .ndf_sum = prev_ndf_sum});
 
                 updated_params.push_back(in_param);
                 TRACCC_VERBOSE("updated_params["

device/common/include/traccc/finding/device/impl/find_tracks.ipp

@@ -454,14 +454,23 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                     const unsigned int seed_idx =
                         payload.step > 0 ? links.at(prev_link_idx).seed_idx
                                          : owner_global_thread_id;
+                    const scalar prev_chi2_sum =
+                        payload.step > 0 ? links.at(prev_link_idx).chi2_sum
+                                         : 0.f;
+                    const unsigned int prev_ndf_sum =
+                        payload.step > 0 ? links.at(prev_link_idx).ndf_sum : 0;
 
                     tmp_links.at(p_offset + l_pos) = {
                         .step = payload.step,
                         .previous_candidate_idx = prev_link_idx,
                         .meas_idx = meas_idx,
                         .seed_idx = seed_idx,
                         .n_skipped = n_skipped,
-                        .chi2 = chi2};
+                        .chi2 = chi2,
+                        .chi2_sum = prev_chi2_sum + chi2,
+                        .ndf_sum =
+                            prev_ndf_sum +
+                            std::get<0>(*result).get_measurement().meas_dim};
 
                     tmp_params.at(p_offset + l_pos) =
                         std::get<0>(*result).filtered();
@@ -525,6 +534,8 @@ TRACCC_HOST_DEVICE inline void find_tracks(
     unsigned int prev_link_idx = std::numeric_limits<unsigned int>::max();
     unsigned int seed_idx = std::numeric_limits<unsigned int>::max();
     unsigned int n_skipped = std::numeric_limits<unsigned int>::max();
+    unsigned int prev_ndf_sum = 0u;
+    scalar prev_chi2_sum = 0.f;
 
     unsigned int local_out_offset = 0;
     unsigned int local_num_params = 0;
@@ -542,6 +553,9 @@ TRACCC_HOST_DEVICE inline void find_tracks(
         n_skipped = payload.step == 0 ? 0 : links.at(prev_link_idx).n_skipped;
         in_param_can_create_hole =
             (n_skipped <= cfg.max_num_skipping_per_cand) && (!last_step);
+        prev_ndf_sum = payload.step == 0 ? 0 : links.at(prev_link_idx).ndf_sum;
+        prev_chi2_sum =
+            payload.step == 0 ? 0.f : links.at(prev_link_idx).chi2_sum;
     }
 
     /*
@@ -613,7 +627,9 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                     .meas_idx = std::numeric_limits<unsigned int>::max(),
                     .seed_idx = seed_idx,
                     .n_skipped = n_skipped + 1,
-                    .chi2 = std::numeric_limits<traccc::scalar>::max()};
+                    .chi2 = std::numeric_limits<traccc::scalar>::max(),
+                    .chi2_sum = prev_chi2_sum,
+                    .ndf_sum = prev_ndf_sum};
 
                 unsigned int param_pos = out_offset - payload.curr_links_idx;