Skip to content

Commit 369125d

Browse files
lobislobis
authored
feat: new awkward array layout (#10)
* working new layout * save event fields to application class * use unique pointers * add more event fields * setting event fields * adding all fields * more fields * all fields * test loading all fields * able to not set track or step fields * save strings as categoricals * save strings as categoricals * add pyarrow as required * add pandas as dependency * use unsigned int vs uint (windows problem) * remove explicit pandas from dependencies * increase awkward version * add string builder * material properties * compute energy in volume * do not convert to categorical * compute hits * add some more analysis * tqdm not working in python 3.12
1 parent 4c31759 commit 369125d

16 files changed

+1033
-213
lines changed

‎CMakeLists.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ set(CMAKE_CXX_STANDARD 20)
4040
set(CMAKE_CXX_STANDARD_REQUIRED ON)
4141
set(CMAKE_CXX_EXTENSIONS OFF)
4242

43-
set(AWKWARD_VERSION "v2.5.0")
43+
set(AWKWARD_VERSION "v2.5.1")
4444
set(PYBIND11_VERSION "v2.11.1")
4545

4646
find_package(

‎pyproject.toml

+2-3
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ classifiers = [
3131
]
3232
dependencies = [
3333
"awkward",
34+
"pyarrow",
35+
"awkward-pandas",
3436
"numpy",
3537
"requests",
3638
"tqdm",
@@ -45,9 +47,6 @@ version = "0.0.2.dev1"
4547
dev = [
4648
"boost_histogram",
4749
"hist",
48-
"pandas",
49-
"awkward-pandas",
50-
"pyarrow",
5150
"uproot",
5251
"fsspec",
5352
]

‎src/geant4_application/include/geant4_application/Application.h

+6
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ class Application {
3131
static Application* pInstance;
3232
void SetupRandomEngine();
3333

34+
static std::unordered_set<std::string> eventFieldsComplete;
35+
static std::unordered_set<std::string> eventFields;
36+
3437
public:
3538
Application();
3639
~Application();
@@ -50,6 +53,9 @@ class Application {
5053
bool IsSetup() const;
5154
bool IsInitialized() const;
5255

56+
inline static std::unordered_set<std::string> GetEventFieldsComplete() { return eventFieldsComplete; }
57+
inline static std::unordered_set<std::string> GetEventFields() { return eventFields; }
58+
static void SetEventFields(const std::unordered_set<std::string>& fields);
5359
inline long GetRandomSeed() const { return randomSeed; }
5460

5561
static void Command(const std::string& command);

‎src/geant4_application/include/geant4_application/DataModel.h

+71-103
Original file line numberDiff line numberDiff line change
@@ -33,116 +33,84 @@ using IndexedBuilder = awkward::LayoutBuilder::IndexedOption<PRIMITIVE, BUILDER>
3333
template<class PRIMITIVE>
3434
using NumpyBuilder = awkward::LayoutBuilder::Numpy<PRIMITIVE>;
3535

36-
enum Field : std::size_t {
37-
runId,
38-
eventId,
36+
template<class T>
37+
using TrackFieldBuilder = ListOffsetBuilder<unsigned int, NumpyBuilder<T>>;
38+
39+
template<class PRIMITIVE>
40+
class StringBuilder : public ListOffsetBuilder<PRIMITIVE, NumpyBuilder<uint8_t>> {
41+
public:
42+
StringBuilder() : ListOffsetBuilder<PRIMITIVE, NumpyBuilder<uint8_t>>() {
43+
this->set_parameters(R"""("__array__": "string")""");
44+
this->content().set_parameters(R"""("__array__": "char")""");
45+
}
46+
47+
void append_string(const std::string& value) {
48+
this->begin_list();
49+
for (const auto c: value) {
50+
this->content().append(c);
51+
}
52+
this->end_list();
53+
}
54+
};
55+
56+
template<class T>
57+
using StepFieldBuilder = ListOffsetBuilder<unsigned int, ListOffsetBuilder<unsigned int, NumpyBuilder<T>>>;
58+
59+
using TrackStringBuilder = ListOffsetBuilder<unsigned int, StringBuilder<unsigned int>>;
60+
using StepStringBuilder = ListOffsetBuilder<unsigned int, ListOffsetBuilder<unsigned int, StringBuilder<unsigned int>>>;
61+
62+
struct Builders {
63+
std::unordered_set<std::string> fields;
64+
NumpyBuilder<unsigned int> run;
65+
NumpyBuilder<unsigned int> id;
3966
//
40-
trackId,
41-
trackParentId,
42-
trackParticle,
43-
// trackParticleType,
44-
trackInitialEnergy,
45-
trackInitialTime,
46-
// trackCreatorProcess,
47-
// trackCreatorProcessType,
48-
trackInitialPositionX,
49-
trackInitialPositionY,
50-
trackInitialPositionZ,
51-
trackInitialMomentumX,
52-
trackInitialMomentumY,
53-
trackInitialMomentumZ,
54-
trackWeight,
67+
TrackFieldBuilder<unsigned int> track_id;
68+
TrackFieldBuilder<unsigned int> track_parent_id;
69+
TrackFieldBuilder<double> track_initial_energy;
70+
TrackFieldBuilder<double> track_initial_time;
71+
TrackFieldBuilder<double> track_initial_position_x;
72+
TrackFieldBuilder<double> track_initial_position_y;
73+
TrackFieldBuilder<double> track_initial_position_z;
74+
TrackFieldBuilder<double> track_initial_momentum_x;
75+
TrackFieldBuilder<double> track_initial_momentum_y;
76+
TrackFieldBuilder<double> track_initial_momentum_z;
77+
StepFieldBuilder<unsigned int> track_children_ids;// this is a track field, but it has the same structure of a step field
78+
TrackStringBuilder track_particle;
79+
TrackStringBuilder track_particle_type;
80+
TrackStringBuilder track_creator_process;
81+
TrackStringBuilder track_creator_process_type;
82+
TrackFieldBuilder<double> track_weight;
5583
//
56-
stepEnergy,
57-
stepTime,
58-
// stepProcess,
59-
// stepProcessType,
60-
// stepVolume,
61-
stepPositionX,
62-
stepPositionY,
63-
stepPositionZ,
64-
stepTrackKineticEnergy,
84+
StepFieldBuilder<double> step_energy;
85+
StepFieldBuilder<double> step_time;
86+
StepFieldBuilder<double> step_track_kinetic_energy;
87+
StepStringBuilder step_process;
88+
StepStringBuilder step_process_type;
89+
StepStringBuilder step_volume;
90+
StepStringBuilder step_volume_post;
91+
StepStringBuilder step_nucleus;
92+
StepFieldBuilder<double> step_position_x;
93+
StepFieldBuilder<double> step_position_y;
94+
StepFieldBuilder<double> step_position_z;
95+
StepFieldBuilder<double> step_momentum_x;
96+
StepFieldBuilder<double> step_momentum_y;
97+
StepFieldBuilder<double> step_momentum_z;
98+
99+
Builders(const std::unordered_set<std::string>& fields) : fields(fields){};
65100
};
66101

67-
inline static const UserDefinedMap fieldToNameEvent = {
68-
{Field::runId, "run_id"},
69-
{Field::eventId, "event_id"},
70-
};
71102

72-
inline static const UserDefinedMap fieldToNameTrack = {
73-
{Field::trackId, "track.id"},
74-
{Field::trackParentId, "track.parent_id"},
75-
{Field::trackParticle, "track.particle"},
76-
// {Field::trackParticleType, "track.particle_type"},
77-
{Field::trackInitialEnergy, "track.energy"},
78-
{Field::trackInitialTime, "track.time"},
79-
{Field::trackInitialPositionX, "track.position.x"},
80-
{Field::trackInitialPositionY, "track.position.y"},
81-
{Field::trackInitialPositionZ, "track.position.z"},
82-
{Field::trackInitialMomentumX, "track.momentum.x"},
83-
{Field::trackInitialMomentumY, "track.momentum.y"},
84-
{Field::trackInitialMomentumZ, "track.momentum.z"},
85-
{Field::trackWeight, "track.weight"},
86-
};
103+
void InsertEventBegin(const G4Event* event, Builders& builder);
104+
void InsertEventEnd(const G4Event* event, Builders& builder);
87105

88-
inline static const UserDefinedMap fieldToNameStep = {
89-
{Field::stepEnergy, "track.step.energy"},
90-
{Field::stepTime, "track.step.time"},
91-
{Field::stepPositionX, "track.step.position.x"},
92-
{Field::stepPositionY, "track.step.position.y"},
93-
{Field::stepPositionZ, "track.step.position.z"},
94-
{Field::stepTrackKineticEnergy, "track.step.track_kinetic_energy"},
95-
96-
// {Field::stepProcess, "track.step.process"}},
97-
// {Field::stepProcessType, "track.step.process_type"}},
98-
// {Field::stepVolume, "track.step.volume"}},
99-
};
106+
void InsertTrackBegin(const G4Track* track, Builders& builder);
107+
void InsertTrackEnd(const G4Track* track, Builders& builder);
108+
109+
void InsertEvent(const G4Event* event, Builders& builder);
110+
void InsertTrack(const G4Track* track, Builders& builder);
111+
void InsertStep(const G4Step* step, Builders& builder);
100112

101-
typedef unsigned int id;
102-
using Builder = RecordBuilder<
103-
RecordField<Field::runId, NumpyBuilder<id>>,
104-
RecordField<Field::eventId, NumpyBuilder<id>>,
105-
//
106-
RecordField<Field::trackId, ListOffsetBuilder<id, NumpyBuilder<id>>>,
107-
RecordField<Field::trackParentId, ListOffsetBuilder<id, NumpyBuilder<id>>>,
108-
RecordField<Field::trackParticle, ListOffsetBuilder<id, NumpyBuilder<bool>>>,// this should be a string
109-
// RecordField<Field::trackParticleType,ListOffsetBuilder<id, NumpyBuilder<std::string>>>,
110-
RecordField<Field::trackInitialEnergy, ListOffsetBuilder<id, NumpyBuilder<float>>>,
111-
RecordField<Field::trackInitialTime, ListOffsetBuilder<id, NumpyBuilder<float>>>,
112-
// RecordField<Field::trackCreatorProcess,ListOffsetBuilder<id, NumpyBuilder<std::string>>>,
113-
RecordField<Field::trackInitialPositionX, ListOffsetBuilder<id, NumpyBuilder<float>>>,
114-
RecordField<Field::trackInitialPositionY, ListOffsetBuilder<id, NumpyBuilder<float>>>,
115-
RecordField<Field::trackInitialPositionZ, ListOffsetBuilder<id, NumpyBuilder<float>>>,
116-
RecordField<Field::trackInitialMomentumX, ListOffsetBuilder<id, NumpyBuilder<float>>>,
117-
RecordField<Field::trackInitialMomentumY, ListOffsetBuilder<id, NumpyBuilder<float>>>,
118-
RecordField<Field::trackInitialMomentumZ, ListOffsetBuilder<id, NumpyBuilder<float>>>,
119-
RecordField<Field::trackWeight, ListOffsetBuilder<id, NumpyBuilder<float>>>,
120-
//
121-
RecordField<Field::stepEnergy, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>,
122-
RecordField<Field::stepTime, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>,
123-
// RecordField<Field::stepVolume,ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<std::string>>>>
124-
// RecordField<Field::stepProcess,ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<std::string>>>>
125-
// RecordField<Field::stepProcessType,ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<std::string>>>>
126-
RecordField<Field::stepPositionX, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>,
127-
RecordField<Field::stepPositionY, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>,
128-
RecordField<Field::stepPositionZ, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>,
129-
RecordField<Field::stepTrackKineticEnergy, ListOffsetBuilder<id, ListOffsetBuilder<id, NumpyBuilder<float>>>>
130-
//
131-
>;
132-
133-
Builder MakeBuilder();
134-
135-
void InsertEventBegin(const G4Event* event, Builder& builder);
136-
void InsertEventEnd(const G4Event* event, Builder& builder);
137-
138-
void InsertTrackBegin(const G4Track* track, Builder& builder);
139-
void InsertTrackEnd(const G4Track* track, Builder& builder);
140-
141-
void InsertEvent(const G4Event* event, Builder& builder);
142-
void InsertTrack(const G4Track* track, Builder& builder);
143-
void InsertStep(const G4Step* step, Builder& builder);
144-
145-
py::object SnapshotBuilder(Builder& builder);
113+
py::object SnapshotBuilder(Builders& builder);
146114

147115
namespace units {
148116
static constexpr auto energy = CLHEP::keV;

‎src/geant4_application/include/geant4_application/DetectorConstruction.h

+4
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@ class DetectorConstruction : public G4VUserDetectorConstruction {
3232
static std::set<std::string> GetLogicalVolumeNames();
3333
static std::set<std::string> GetPhysicalVolumeNames();
3434

35+
static std::set<std::string> GetPhysicalVolumesFromLogicalVolume(const std::string& logicalVolumeName);
36+
static std::string GetLogicalVolumeFromPhysicalVolume(const std::string& physicalVolumeName);
37+
static std::string GetMaterialFromVolume(const std::string& volumeName);
38+
3539
bool IsConstructed() const { return world != nullptr; }
3640

3741
private:

‎src/geant4_application/include/geant4_application/RunAction.h

+4-3
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include <G4RunManager.hh>
55
#include <G4UserRunAction.hh>
66

7+
#include "geant4_application/Application.h"
78
#include "geant4_application/DataModel.h"
89

910
namespace py = pybind11;
@@ -18,14 +19,14 @@ class RunAction : public G4UserRunAction {
1819
void EndOfRunAction(const G4Run*) override;
1920

2021
/// Only one instance of RunAction is created for each thread.
21-
static data::Builder& GetBuilder();
22+
static data::Builders& GetBuilder();
2223
static std::vector<py::object> GetEvents();
2324

2425
private:
25-
data::Builder builder = data::MakeBuilder();
26+
std::unique_ptr<data::Builders> builder = nullptr;
2627
std::mutex mutex;
2728
static std::unique_ptr<std::vector<py::object>> events;
28-
static std::vector<data::Builder*> buildersToSnapshot;
29+
static std::vector<std::unique_ptr<data::Builders>> buildersToSnapshot;
2930
};
3031

3132
}// namespace geant4_app

‎src/geant4_application/src/Application.cpp

+16
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@
99
#include <G4RunManager.hh>
1010
#include <G4RunManagerFactory.hh>
1111

12+
#include <algorithm>
1213
#include <random>
14+
#include <unordered_set>
1315

1416
using namespace std;
1517
using namespace geant4_app;
@@ -112,6 +114,9 @@ vector<py::object> Application::Run(int nEvents) {
112114
if (!IsInitialized()) {
113115
Initialize();
114116
}
117+
if (eventFields.empty()) {
118+
throw runtime_error("Event fields cannot be empty");
119+
}
115120
runManager->BeamOn(nEvents);
116121
py::gil_scoped_acquire acquire;
117122
return RunAction::GetEvents();
@@ -191,3 +196,14 @@ filesystem::path Application::GetTemporaryApplicationDirectory() {
191196
}
192197
return dir;
193198
}
199+
200+
void Application::SetEventFields(const unordered_set<string>& fields) {
201+
for (const auto& field: fields) {
202+
if (!eventFieldsComplete.contains(field)) {
203+
throw runtime_error("Invalid event field: " + field);
204+
}
205+
}
206+
eventFields = fields;
207+
}
208+
209+
unordered_set<string> Application::eventFields = {};

0 commit comments

Comments
 (0)