add spline fit to smoothen the curves

Examples/Io/Obj/include/ActsExamples/Io/Obj/ObjSimHitWriter.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Acts/Definitions/Units.hpp"
 #include "Acts/Utilities/Logger.hpp"
 #include "ActsExamples/EventData/SimHit.hpp"
 #include "ActsExamples/Framework/ProcessCode.hpp"
@@ -17,6 +18,8 @@
 #include <mutex>
 #include <string>
 
+using namespace Acts::UnitLiterals;
+
 namespace ActsExamples {
 struct AlgorithmContext;
 
@@ -44,6 +47,12 @@ class ObjSimHitWriter : public WriterT<SimHitContainer> {
     std::size_t outputPrecision = std::numeric_limits<float>::max_digits10;
     /// Draw line connections between hits
     bool drawConnections = true;
+    /// Momentum threshold for hits
+    Acts::ActsScalar momentumThreshold = 0.05_GeV;
+    /// Momentum threshold for trajectories
+    Acts::ActsScalar momentumThresholdTraj = 0.05_GeV;
+    /// Number of points to interpolate between hits
+    unsigned int nInterpolatedPoints = 10;
   };
 
   /// Construct the particle writer.

Examples/Io/Obj/src/ObjSimHitWriter.cpp

@@ -10,7 +10,6 @@
 
 #include "Acts/Definitions/Algebra.hpp"
 #include "Acts/Definitions/Common.hpp"
-#include "Acts/Definitions/Units.hpp"
 #include "Acts/Geometry/GeometryIdentifier.hpp"
 #include "ActsExamples/EventData/SimHit.hpp"
 #include "ActsExamples/Framework/AlgorithmContext.hpp"
@@ -23,6 +22,33 @@
 #include <unordered_map>
 #include <vector>
 
+#include <unsupported/Eigen/Splines>
+
+namespace {
+
+template <typename input_vector_type>
+std::vector<Acts::Vector3> interpolatedPoints(
+    const std::vector<input_vector_type>& inputs, unsigned int nPoints) {
+  Eigen::MatrixXd points(3, inputs.size());
+  for (unsigned int i = 0; i < inputs.size(); ++i) {
+    points.col(i) = inputs[i].template head<3>().transpose();
+  }
+  Eigen::Spline<Acts::ActsScalar, 3> spline3D =
+      Eigen::SplineFitting<Eigen::Spline<Acts::ActsScalar, 3>>::Interpolate(
+          points, 2);
+
+  std::vector<Acts::Vector3> output;
+  Acts::ActsScalar step = 1. / (nPoints - 1);
+  for (unsigned int i = 0; i < nPoints; ++i) {
+    Acts::ActsScalar t = i * step;
+    Eigen::Vector3d point = spline3D(t);
+    output.push_back(Acts::Vector3(point[0], point[1], point[2]));
+  }
+  return output;
+}
+
+}  // namespace
+
 ActsExamples::ObjSimHitWriter::ObjSimHitWriter(
     const ActsExamples::ObjSimHitWriter::Config& config,
     Acts::Logging::Level level)
@@ -41,11 +67,16 @@ ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::writeT(
   // open per-event file for all simhit components
   std::string pathSimHit = perEventFilepath(
       m_cfg.outputDir, m_cfg.outputStem + ".obj", ctx.eventNumber);
+  std::string pathSimTrajectory = perEventFilepath(
+      m_cfg.outputDir, m_cfg.outputStem + "_trajectory.obj", ctx.eventNumber);
 
-  std::ofstream os(pathSimHit, std::ofstream::out | std::ofstream::trunc);
-  if (!os) {
-    throw std::ios_base::failure("Could not open '" + pathSimHit +
-                                 "' to write");
+  std::ofstream osHits(pathSimHit, std::ofstream::out | std::ofstream::trunc);
+  std::ofstream osTrajectory(pathSimTrajectory,
+                             std::ofstream::out | std::ofstream::trunc);
+
+  if (!osHits || !osTrajectory) {
+    throw std::ios_base::failure("Could not open '" + pathSimHit + "' or '" +
+                                 pathSimTrajectory + "' to write");
   }
 
   // Only hit plotting
@@ -55,16 +86,34 @@ ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::writeT(
       // local simhit information in global coord.
       const Acts::Vector4& globalPos4 = simHit.fourPosition();
 
-      os << "v " << globalPos4[Acts::ePos0] / Acts::UnitConstants::mm << " "
-         << globalPos4[Acts::ePos1] / Acts::UnitConstants::mm << " "
-         << globalPos4[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
+      osHits << "v " << globalPos4[Acts::ePos0] / Acts::UnitConstants::mm << " "
+             << globalPos4[Acts::ePos1] / Acts::UnitConstants::mm << " "
+             << globalPos4[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
 
     }  // end simHit loop
   } else {
     // We need to associate first
     std::unordered_map<uint64_t, std::vector<Acts::Vector4>> particleHits;
-    // pre-loop over hits
+    // Pre-loop over hits ... write those below threshold
     for (const auto& simHit : simHits) {
+      Acts::ActsScalar momentum = simHit.momentum4Before().head<3>().norm();
+      if (momentum < m_cfg.momentumThreshold) {
+        ACTS_VERBOSE("Skipping : Hit below threshold: " << momentum);
+        continue;
+      } else if (momentum < m_cfg.momentumThresholdTraj) {
+        ACTS_VERBOSE(
+            "Skipping (trajectory): Hit below threshold: " << momentum);
+        osHits << "v "
+               << simHit.fourPosition()[Acts::ePos0] / Acts::UnitConstants::mm
+               << " "
+               << simHit.fourPosition()[Acts::ePos1] / Acts::UnitConstants::mm
+               << " "
+               << simHit.fourPosition()[Acts::ePos2] / Acts::UnitConstants::mm
+               << std::endl;
+        continue;
+      }
+      ACTS_VERBOSE("Accepting : Hit above threshold: " << momentum);
+
       if (particleHits.find(simHit.particleId().value()) ==
           particleHits.end()) {
         particleHits[simHit.particleId().value()] = {};
@@ -74,28 +123,50 @@ ActsExamples::ProcessCode ActsExamples::ObjSimHitWriter::writeT(
     }
     // Draw loop
     unsigned int lOffset = 1;
-    for (auto [pId, pHits] : particleHits) {
+    for (auto& [pId, pHits] : particleHits) {
       // Draw the particle hits
       std::sort(pHits.begin(), pHits.end(),
                 [](const Acts::Vector4& a, const Acts::Vector4& b) {
                   return a[Acts::eTime] < b[Acts::eTime];
                 });
-      os << "o particle_" << pId << std::endl;
+
+      osHits << "o particle_" << pId << std::endl;
       for (const auto& hit : pHits) {
-        os << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm << " "
-           << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
-           << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
+        osHits << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm << " "
+               << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
+               << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
+      }
+      osHits << '\n';
+
+      // Interpolate the points
+      std::vector<Acts::Vector3> trajectory;
+      if (pHits.size() < 3) {
+        for (const auto& hit : pHits) {
+          trajectory.push_back(hit.template head<3>());
+        }
+      } else {
+        trajectory =
+            interpolatedPoints(pHits, pHits.size() * m_cfg.nInterpolatedPoints);
+      }
+
+      osTrajectory << "o particle_trajectory_" << pId << std::endl;
+      for (const auto& hit : trajectory) {
+        osTrajectory << "v " << hit[Acts::ePos0] / Acts::UnitConstants::mm
+                     << " " << hit[Acts::ePos1] / Acts::UnitConstants::mm << " "
+                     << hit[Acts::ePos2] / Acts::UnitConstants::mm << std::endl;
       }
       // Draw the line
-      for (unsigned int iv = lOffset + 1; iv < lOffset + pHits.size(); ++iv) {
-        os << "l " << iv - 1 << " " << iv << '\n';
+      for (unsigned int iv = lOffset + 1; iv < lOffset + trajectory.size();
+           ++iv) {
+        osTrajectory << "l " << iv - 1 << " " << iv << '\n';
       }
-      os << '\n';
+      osTrajectory << '\n';
       // Increase the offset count
-      lOffset += pHits.size();
+      lOffset += trajectory.size();
     }
   }
-  os.close();
+  osHits.close();
+  osTrajectory.close();
 
   return ActsExamples::ProcessCode::SUCCESS;
 }