Record the fitting outcome

Author: beomki-yeo

core/include/traccc/edm/track_state.hpp

@@ -19,11 +19,22 @@
 
 namespace traccc {
 
+enum class fitter_outcome : uint32_t {
+    UNKNOWN,
+    SUCCESS,
+    FAILURE_NON_POSITIVE_NDF,
+    FAILURE_NOT_ALL_SMOOTHED,
+    MAX_OUTCOME
+};
+
 /// Fitting result per track
 template <typename algebra_t>
 struct fitting_result {
     using scalar_type = detray::dscalar<algebra_t>;
 
+    /// Fitting outcome
+    fitter_outcome fit_outcome = fitter_outcome::UNKNOWN;
+
     /// Fitted track parameter
     detray::bound_track_parameters<algebra_t> fit_params;
 
@@ -205,6 +216,7 @@ struct track_state {
 
     public:
     bool is_hole{true};
+    bool is_smoothed{false};
 
     private:
     detray::geometry::barcode m_surface_link;
@@ -227,4 +239,19 @@ using track_state_container_types =
     container_types<fitting_result<default_algebra>,
                     track_state<default_algebra>>;
 
+inline std::size_t count_fitted_tracks(
+    const track_state_container_types::host& track_states) {
+
+    const std::size_t n_tracks = track_states.size();
+    std::size_t n_fitted_tracks = 0u;
+
+    for (std::size_t i = 0; i < n_tracks; i++) {
+        if (track_states.at(i).header.fit_outcome == fitter_outcome::SUCCESS) {
+            n_fitted_tracks++;
+        }
+    }
+
+    return n_fitted_tracks;
+}
+
 }  // namespace traccc

core/include/traccc/fitting/kalman_filter/gain_matrix_smoother.hpp

@@ -148,6 +148,7 @@ struct gain_matrix_smoother {
             matrix::transpose(residual) * matrix::inverse(R) * residual;
 
         cur_state.smoothed_chi2() = getter::element(chi2, 0, 0);
+        cur_state.is_smoothed = true;
 
         return kalman_fitter_status::SUCCESS;
     }

core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp

@@ -148,6 +148,11 @@ class kalman_fitter {
             // Reset the iterator of kalman actor
             fitter_state.m_fit_actor_state.reset();
 
+            // TODO: For multiple iterations, seed parameter should be set to
+            // the first track state which has either filtered or smoothed
+            // state. If the first track state is a hole, we need to back
+            // extrapolate from the filtered or smoothed state of next valid
+            // track state.
             auto seed_params_cpy =
                 (i == 0) ? seed_params
                          : fitter_state.m_fit_actor_state.m_track_states[0]
@@ -161,6 +166,8 @@ class kalman_fitter {
                 res != kalman_fitter_status::SUCCESS) {
                 return res;
             }
+
+            check_fitting_result(fitter_state);
         }
 
         return kalman_fitter_status::SUCCESS;
@@ -303,8 +310,13 @@ class kalman_fitter {
         auto& fit_res = fitter_state.m_fit_res;
         auto& track_states = fitter_state.m_fit_actor_state.m_track_states;
 
-        // Fit parameter = smoothed track parameter at the first surface
-        fit_res.fit_params = track_states[0].smoothed();
+        // Fit parameter = smoothed track parameter of the first smoothed track
+        // state
+        for (const auto& st : track_states) {
+            if (st.is_smoothed) {
+                fit_res.fit_params = st.smoothed();
+            }
+        }
 
         for (const auto& trk_state : track_states) {
 
@@ -321,6 +333,31 @@ class kalman_fitter {
         fit_res.n_holes = fitter_state.m_fit_actor_state.n_holes;
     }
 
+    TRACCC_HOST_DEVICE
+    void check_fitting_result(state& fitter_state) {
+        auto& fit_res = fitter_state.m_fit_res;
+        const auto& track_states =
+            fitter_state.m_fit_actor_state.m_track_states;
+
+        // NDF should always be positive for fitting
+        if (fit_res.ndf > 0) {
+            for (const auto& trk_state : track_states) {
+                // Fitting fails if any of non-hole track states is not smoothed
+                if (!trk_state.is_hole && !trk_state.is_smoothed) {
+                    fit_res.fit_outcome =
+                        fitter_outcome::FAILURE_NOT_ALL_SMOOTHED;
+                    return;
+                }
+            }
+
+            // Fitting succeeds if any of non-hole track states is not smoothed
+            fit_res.fit_outcome = fitter_outcome::SUCCESS;
+        }
+
+        fit_res.fit_outcome = fitter_outcome::FAILURE_NON_POSITIVE_NDF;
+        return;
+    }
+
     private:
     // Detector object
     const detector_type& m_detector;

core/include/traccc/fitting/kalman_filter/statistics_updater.hpp

@@ -38,13 +38,15 @@ struct statistics_updater {
             // Measurement dimension
             const unsigned int D = trk_state.get_measurement().meas_dim;
 
-            // NDoF = NDoF + number of coordinates per measurement
-            fit_res.ndf += static_cast<scalar_type>(D);
-
-            // total_chi2 = total_chi2 + chi2
             if (use_backward_filter) {
-                fit_res.chi2 += trk_state.backward_chi2();
+                if (trk_state.is_smoothed) {
+                    // NDoF = NDoF + number of coordinates per measurement
+                    fit_res.ndf += static_cast<scalar_type>(D);
+                    fit_res.chi2 += trk_state.backward_chi2();
+                }
             } else {
+                // NDoF = NDoF + number of coordinates per measurement
+                fit_res.ndf += static_cast<scalar_type>(D);
                 fit_res.chi2 += trk_state.filtered_chi2();
             }
         }

core/include/traccc/fitting/kalman_filter/two_filters_smoother.hpp

@@ -169,6 +169,7 @@ struct two_filters_smoother {
         // Wrap the phi in the range of [-pi, pi]
         wrap_phi(bound_params);
 
+        trk_state.is_smoothed = true;
         return kalman_fitter_status::SUCCESS;
     }
 };

examples/run/cpu/truth_finding_example.cpp

@@ -150,8 +150,9 @@ int seq_run(const traccc::opts::track_finding& finding_opts,
         auto track_states =
             host_fitting(detector, field, traccc::get_data(track_candidates));
 
-        std::cout << "Number of fitted tracks: " << track_states.size()
-                  << std::endl;
+        std::cout << "Number of fitted tracks: ( "
+                  << count_fitted_tracks(track_states) << " / "
+                  << track_states.size() << " ) " << std::endl;
 
         const std::size_t n_fitted_tracks = track_states.size();
 

examples/run/cpu/truth_fitting_example.cpp

@@ -130,8 +130,9 @@ int main(int argc, char* argv[]) {
         auto track_states = host_fitting(
             host_det, field, traccc::get_data(truth_track_candidates));
 
-        std::cout << "Number of fitted tracks: " << track_states.size()
-                  << std::endl;
+        std::cout << "Number of fitted tracks: ( "
+                  << count_fitted_tracks(track_states) << " / "
+                  << track_states.size() << " ) " << std::endl;
 
         const decltype(track_states)::size_type n_fitted_tracks =
             track_states.size();

performance/include/traccc/resolution/fitting_performance_writer.hpp

@@ -58,6 +58,17 @@ class fitting_performance_writer {
                const fitting_result<traccc::default_algebra>& fit_res,
                const detector_t& det, event_data& evt_data) {
 
+        if (fit_res.fit_outcome != fitter_outcome::SUCCESS) {
+            return;
+        }
+
+        // Get the first smoothed track state
+        const auto& trk_state = *std::find_if(
+            track_states_per_track.begin(), track_states_per_track.end(),
+            [](const auto& st) { return st.is_smoothed; });
+        assert(!trk_state.is_hole);
+        assert(trk_state.is_smoothed);
+
         std::map<measurement, std::map<particle, std::size_t>> meas_to_ptc_map;
         std::map<measurement, std::pair<point3, point3>> meas_to_param_map;
 
@@ -69,8 +80,6 @@ class fitting_performance_writer {
             meas_to_param_map = evt_data.m_meas_to_param_map;
         }
 
-        // Get the track state at the first surface
-        const auto& trk_state = track_states_per_track[0];
         const measurement meas = trk_state.get_measurement();
 
         // Find the contributing particle
@@ -103,7 +112,9 @@ class fitting_performance_writer {
         if (fit_res.ndf > 0 && !trk_state.is_hole) {
             write_res(truth_param, trk_state.smoothed(), ptc);
         }
-        write_stat(fit_res, track_states_per_track);
+        if (fit_res.fit_outcome == fitter_outcome::SUCCESS) {
+            write_stat(fit_res, track_states_per_track);
+        }
     }
 
     /// Writing caches into the file