feat(cli): add support for recording events of specific elements only

Atraxus · ichinii · commit d9170fa0c749 · 2025-07-16T09:25:44.000+02:00
diff --git a/Intern/rayx-core/src/Shader/DynamicElements.cpp b/Intern/rayx-core/src/Shader/DynamicElements.cpp
@@ -60,7 +60,7 @@ void dynamicElements(const int gid, const InvState& inv, OutputEvents& outputEve
         }
 
         // write ray in local element coordinates to global memory
-        if (inv.recordElementIndex < 0 || col.elementIndex == inv.recordElementIndex) {
+        if (!inv.recordMask || inv.recordMask[col.elementIndex]) {
             outputEvents.events[gid * inv.maxEvents + numRecorded] = ray;
             ++numRecorded;
         }
diff --git a/Intern/rayx-core/src/Shader/InvocationState.h b/Intern/rayx-core/src/Shader/InvocationState.h
@@ -17,14 +17,14 @@ struct RAYX_API InvState {
     int batchSize;
     int batchStartRayIndex;
     int maxEvents;
-    int recordElementIndex;  //< Index of element, for which to record events. Others are discarded. -1 to record events of all elements.
     double randomSeed;
     Sequential sequential = Sequential::No;
 
     OpticalElement* elements;
     int numElements;
     int* materialIndices;
     double* materialTables;
+    bool* recordMask;  //< Mask that decides which elements to record events for (array length is numElements)
     Ray* rays;
 };
 
diff --git a/Intern/rayx-core/src/Tracer/DeviceTracer.h b/Intern/rayx-core/src/Tracer/DeviceTracer.h
@@ -17,7 +17,7 @@ class RAYX_API DeviceTracer {
   public:
     virtual ~DeviceTracer() = default;
 
-    virtual RaySoA trace(const Group&, Sequential sequential, const int maxBatchSize, const int maxEvents, const int recordElementIndex,
+    virtual RaySoA trace(const Group&, Sequential sequential, const int maxBatchSize, const int maxEvents, std::shared_ptr<bool[]> recordMask,
                          const RayAttrFlag attr) = 0;
 };
 
diff --git a/Intern/rayx-core/src/Tracer/MegaKernelTracer.h b/Intern/rayx-core/src/Tracer/MegaKernelTracer.h
@@ -150,7 +150,7 @@ class MegaKernelTracer : public DeviceTracer {
 
   public:
     virtual RaySoA trace(const Group& beamline, const Sequential sequential, const int maxBatchSize, const int maxEvents,
-                         const int recordElementIndex, const RayAttrFlag attr) override {
+                         std::shared_ptr<bool[]> recordMask, const RayAttrFlag attr) override {
         RAYX_PROFILE_FUNCTION_STDOUT();
 
         const auto platformHost = alpaka::PlatformCpu{};
@@ -199,8 +199,7 @@ class MegaKernelTracer : public DeviceTracer {
             alpaka::memcpy(q, *m_resources.d_rays, raysViewBatch);
 
             // trace current batch
-            traceBatch(devAcc, q, conf.numElements, conf.numRaysTotal, batchSize, batchStartRayIndex, maxEvents, recordElementIndex, randomSeed,
-                       sequential);
+            traceBatch(devAcc, q, conf.numElements, conf.numRaysTotal, batchSize, batchStartRayIndex, maxEvents, recordMask, randomSeed, sequential);
 
             // prefix sum on compactEventCounts to get compactEventOffsets
             alpaka::memcpy(q, alpaka::createView(devHost, compactEventCounts, batchSize), *m_resources.d_compactEventCounts, batchSize);
@@ -260,7 +259,7 @@ class MegaKernelTracer : public DeviceTracer {
   private:
     template <typename DevAcc, typename Queue>
     void traceBatch(DevAcc devAcc, Queue q, int numElements, int numRaysTotal, int batchSize, int batchStartRayIndex, int maxEvents,
-                    int recordElementIndex, double randomSeed, Sequential sequential) {
+                    std::shared_ptr<bool[]> recordMask, double randomSeed, Sequential sequential) {
         RAYX_PROFILE_FUNCTION_STDOUT();
 
         // inputs
@@ -270,7 +269,6 @@ class MegaKernelTracer : public DeviceTracer {
             .batchSize = batchSize,
             .batchStartRayIndex = batchStartRayIndex,
             .maxEvents = maxEvents,
-            .recordElementIndex = recordElementIndex,
             .randomSeed = randomSeed,
             .sequential = sequential,
 
@@ -279,6 +277,7 @@ class MegaKernelTracer : public DeviceTracer {
             .numElements = numElements,
             .materialIndices = alpaka::getPtrNative(*m_resources.d_materialIndices),
             .materialTables = alpaka::getPtrNative(*m_resources.d_materialTable),
+            .recordMask = recordMask.get(),
             .rays = alpaka::getPtrNative(*m_resources.d_rays),
         };
 
diff --git a/Intern/rayx-core/src/Tracer/Tracer.cpp b/Intern/rayx-core/src/Tracer/Tracer.cpp
@@ -51,12 +51,12 @@ Tracer::Tracer(const DeviceConfig& deviceConfig) {
     }
 }
 
-RaySoA Tracer::trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, int32_t recordElementIndex,
+RaySoA Tracer::trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, std::shared_ptr<bool[]> recordMask,
                      RayAttrFlag attr) {
     // in sequential tracing, maxEvents should be equal to the number of elements
     if (sequential == Sequential::Yes) maxEvents = group.numElements();
 
-    return m_deviceTracer->trace(group, sequential, static_cast<int>(maxBatchSize), static_cast<int>(maxEvents), recordElementIndex, attr);
+    return m_deviceTracer->trace(group, sequential, static_cast<int>(maxBatchSize), static_cast<int>(maxEvents), recordMask, attr);
 }
 
 int Tracer::defaultMaxEvents(const Group* group) {
diff --git a/Intern/rayx-core/src/Tracer/Tracer.h b/Intern/rayx-core/src/Tracer/Tracer.h
@@ -29,7 +29,7 @@ class RAYX_API Tracer {
     // This will call the trace implementation of a subclass
     // See `BundleHistory` for information about the return value.
     // `max_batch_size` corresponds to the maximal number of rays that will be put into `traceRaw` in one batch.
-    RaySoA trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, int32_t recordElementIndex,
+    RaySoA trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, std::shared_ptr<bool[]> recordMask = nullptr,
                  RayAttrFlag attr = RayAttrFlag::All);
 
     static int defaultMaxEvents(const Group* group);
diff --git a/Intern/rayx-core/tests/setupTests.cpp b/Intern/rayx-core/tests/setupTests.cpp
@@ -77,7 +77,7 @@ void writeToOutputCSV(const RAYX::BundleHistory& hist, std::string filename) {
 
 RAYX::BundleHistory traceRML(std::string filename) {
     const auto beamline = loadBeamline(filename);
-    const auto rays = tracer->trace(beamline, Sequential::No, DEFAULT_BATCH_SIZE, beamline.numElements() + 2, -1);
+    const auto rays = tracer->trace(beamline, Sequential::No, DEFAULT_BATCH_SIZE, beamline.numElements() + 2);
     return raySoAToBundleHistory(rays);
 }
 
@@ -177,7 +177,7 @@ std::optional<RAYX::Ray> lastSequentialHit(RayHistory ray_hist, uint32_t beamlin
 // returns the rayx rays converted to be ray-UI compatible.
 std::vector<RAYX::Ray> rayUiCompat(std::string filename, Sequential seq) {
     const auto beamline = loadBeamline(filename);
-    const auto rays = tracer->trace(beamline, seq, DEFAULT_BATCH_SIZE, beamline.numElements() + 2, -1);
+    const auto rays = tracer->trace(beamline, seq, DEFAULT_BATCH_SIZE, beamline.numElements() + 2);
     const auto hist = raySoAToBundleHistory(rays);
 
     std::vector<RAYX::Ray> out;
diff --git a/Intern/rayx-ui/src/Simulator.cpp b/Intern/rayx-ui/src/Simulator.cpp
@@ -17,9 +17,7 @@ void Simulator::runSimulation() {
         m_maxEvents = RAYX::Tracer::defaultMaxEvents(&m_Beamline);
     }
 
-    constexpr int RECORD_ALL_ELEMENTS = -1;
-    const auto rays =
-        m_Tracer->trace(m_Beamline, m_seq, m_max_batch_size, m_maxEvents, RECORD_ALL_ELEMENTS);  // TODO: implement recordElementIndex for GUI?
+    const auto rays = m_Tracer->trace(m_Beamline, m_seq, m_max_batch_size, m_maxEvents);
     const auto bundleHist = RAYX::raySoAToBundleHistory(rays);
 
     bool notEnoughEvents = false;
diff --git a/Intern/rayx/src/CommandParser.cpp b/Intern/rayx/src/CommandParser.cpp
@@ -21,10 +21,11 @@ CommandParser::CommandParser(int _argc, char* const* _argv) : m_cli11{std::make_
             m_cli11->add_flag(_name, *static_cast<bool*>(option.second.option_flag), _description);
         } else if (_type == OptionType::STRING) {
             m_cli11->add_option(_name, *static_cast<std::string*>(option.second.option_flag), _description);
-        } else if (_type == OptionType::BOOL_STRING) {  // Discarded
-            m_cli11->add_flag(_name, *static_cast<bool*>(option.second.option_flag), _description);
         } else if (_type == OptionType::INT) {
             m_cli11->add_option(_name, *static_cast<int*>(option.second.option_flag), _description);
+        } else if (_type == OptionType::INT_VEC) {
+            m_cli11->add_option(_name, *static_cast<std::vector<int>*>(option.second.option_flag), _description)
+                ->expected(-1);  // allow any number of ints
         }
     }
 
diff --git a/Intern/rayx/src/CommandParser.h b/Intern/rayx/src/CommandParser.h
@@ -41,8 +41,8 @@ class CommandParser {
         bool m_verbose = false;           // --verbose (Verbose)
         std::string m_format;             // --format
         int m_maxEvents = -1;             // -m (max events)
-        int m_recordElementIndex = -1;    // -R --record-element (element index)
-        std::string m_dump = "";          // -D (dump)
+        std::vector<int> m_recordIndices = {};
+        std::string m_dump = "";  // -D (dump)
     } m_args;
 
     static inline void getVersion() {
@@ -61,7 +61,7 @@ class CommandParser {
 
   private:
     int m_cli11_return;
-    enum OptionType { BOOL, INT, STRING, BOOL_STRING };
+    enum OptionType { BOOL, INT, STRING, INT_VEC };
     struct Options {
         // CLI::Option cli11_option;
         const OptionType type;
@@ -97,7 +97,7 @@ class CommandParser {
           "\"",
           &(m_args.m_format)}},
         {'m', {OptionType::INT, "maxEvents", "Maximum number of events per ray", &(m_args.m_maxEvents)}},
-        {'R', {OptionType::INT, "record-element", "Record events only for a specifc element", &(m_args.m_recordElementIndex)}},
+        {'R', {OptionType::INT_VEC, "record-indices", "Record events only for specifc elements", &(m_args.m_recordIndices)}},
         {'D', {OptionType::STRING, "dump", "Dump the meta data of a file (h5 or rml)", &(m_args.m_dump)}},
     };
 };
diff --git a/Intern/rayx/src/TerminalApp.cpp b/Intern/rayx/src/TerminalApp.cpp
@@ -8,6 +8,7 @@
 #include <filesystem>
 #include <memory>
 #include <stdexcept>
+#include <vector>
 
 #include "Debug/Debug.h"
 #include "Random.h"
@@ -144,17 +145,30 @@ void TerminalApp::tracePath(const std::filesystem::path& path) {
         if (m_CommandParser->m_args.m_BatchSize != 0) {
             max_batch_size = m_CommandParser->m_args.m_BatchSize;
         }
+        if (m_CommandParser->m_args.m_recordIndices.empty()) {
+            RAYX_VERB << "Passed no or empty array for which elements to record.";
+        } else {
+            RAYX_VERB << "First element of passed array: " << m_CommandParser->m_args.m_recordIndices[0];
+        }
 
         // Run rayx core
         RAYX::Sequential seq = m_CommandParser->m_args.m_sequential ? RAYX::Sequential::Yes : RAYX::Sequential::No;
         const int maxEvents =
             (m_CommandParser->m_args.m_maxEvents < 1) ? RAYX::Tracer::defaultMaxEvents(m_Beamline.get()) : m_CommandParser->m_args.m_maxEvents;
-        int recordElementIndex = m_CommandParser->m_args.m_recordElementIndex;
+        auto recordIndices = m_CommandParser->m_args.m_recordIndices;
+        size_t numElements = m_Beamline->numElements();
+        auto recordMask = std::make_shared<bool[]>(numElements);
+        std::fill_n(recordMask.get(), numElements, false);
+        for (auto idx : recordIndices) {
+            if (idx < 0) RAYX_EXIT << "Only positive indices are possible for CLI option: -R/--record-indices!";
+            if (idx > numElements - 1) RAYX_EXIT << "Index {" << idx << "} provided with -R/--record-indices does not exist in the provided file.";
+            recordMask[idx] = true;
+        }
 
         const auto attr =
             m_CommandParser->m_args.m_format.empty() ? RAYX::RayAttrFlag::All : RAYX::formatStringToRayAttrFlag(m_CommandParser->m_args.m_format);
 
-        const auto rays = m_Tracer->trace(*m_Beamline, seq, max_batch_size, maxEvents, recordElementIndex, attr);
+        const auto rays = m_Tracer->trace(*m_Beamline, seq, max_batch_size, maxEvents, recordMask, attr);
 
         if (!rays.num_events)
             std::cout << "No events were recorded!" << std::endl;

Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ void dynamicElements(const int gid, const InvState& inv, OutputEvents& outputEve`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`// write ray in local element coordinates to global memory`
`63`		`- if (inv.recordElementIndex < 0 \|\| col.elementIndex == inv.recordElementIndex) {`
	`63`	`+ if (!inv.recordMask \|\| inv.recordMask[col.elementIndex]) {`
`64`	`64`	`outputEvents.events[gid * inv.maxEvents + numRecorded] = ray;`
`65`	`65`	`++numRecorded;`
`66`	`66`	`}`
Original file line number	Diff line number	Diff line change
`@@ -51,12 +51,12 @@ Tracer::Tracer(const DeviceConfig& deviceConfig) {`
`51`	`51`	`}`
`52`	`52`	`}`
`53`	`53`
`54`		`-RaySoA Tracer::trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, int32_t recordElementIndex,`
	`54`	`+RaySoA Tracer::trace(const Group& group, Sequential sequential, uint64_t maxBatchSize, uint32_t maxEvents, std::shared_ptr<bool[]> recordMask,`
`55`	`55`	`RayAttrFlag attr) {`
`56`	`56`	`// in sequential tracing, maxEvents should be equal to the number of elements`
`57`	`57`	`if (sequential == Sequential::Yes) maxEvents = group.numElements();`
`58`	`58`
`59`		`- return m_deviceTracer->trace(group, sequential, static_cast<int>(maxBatchSize), static_cast<int>(maxEvents), recordElementIndex, attr);`
	`59`	`+ return m_deviceTracer->trace(group, sequential, static_cast<int>(maxBatchSize), static_cast<int>(maxEvents), recordMask, attr);`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`int Tracer::defaultMaxEvents(const Group* group) {`
Original file line number	Diff line number	Diff line change
`@@ -17,9 +17,7 @@ void Simulator::runSimulation() {`
`17`	`17`	`m_maxEvents = RAYX::Tracer::defaultMaxEvents(&m_Beamline);`
`18`	`18`	`}`
`19`	`19`
`20`		`- constexpr int RECORD_ALL_ELEMENTS = -1;`
`21`		`- const auto rays =`
`22`		`- m_Tracer->trace(m_Beamline, m_seq, m_max_batch_size, m_maxEvents, RECORD_ALL_ELEMENTS); // TODO: implement recordElementIndex for GUI?`
	`20`	`+ const auto rays = m_Tracer->trace(m_Beamline, m_seq, m_max_batch_size, m_maxEvents);`
`23`	`21`	`const auto bundleHist = RAYX::raySoAToBundleHistory(rays);`
`24`	`22`
`25`	`23`	`bool notEnoughEvents = false;`
Original file line number	Diff line number	Diff line change
`@@ -21,10 +21,11 @@ CommandParser::CommandParser(int _argc, char* const* _argv) : m_cli11{std::make_`
`21`	`21`	`m_cli11->add_flag(_name, static_cast<bool>(option.second.option_flag), _description);`
`22`	`22`	`} else if (_type == OptionType::STRING) {`
`23`	`23`	`m_cli11->add_option(_name, static_cast<std::string>(option.second.option_flag), _description);`
`24`		`- } else if (_type == OptionType::BOOL_STRING) { // Discarded`
`25`		`- m_cli11->add_flag(_name, static_cast<bool>(option.second.option_flag), _description);`
`26`	`24`	`} else if (_type == OptionType::INT) {`
`27`	`25`	`m_cli11->add_option(_name, static_cast<int>(option.second.option_flag), _description);`
	`26`	`+ } else if (_type == OptionType::INT_VEC) {`
	`27`	`+ m_cli11->add_option(_name, static_cast<std::vector<int>>(option.second.option_flag), _description)`
	`28`	`+ ->expected(-1); // allow any number of ints`
`28`	`29`	`}`
`29`	`30`	`}`
`30`	`31`