Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updating EvtCalc interface to work with any N-body decay #581

Merged
merged 6 commits into from
Dec 18, 2024
Merged

Updating EvtCalc interface to work with any N-body decay #581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fa25dc7
Updating EvtCalc interface to Nbody decays
mferril Dec 17, 2024
4359da2
Adding EvtCalc file converter
mferril Dec 17, 2024
f586857
include @olantwin's comments
mferril Dec 18, 2024
5d95813
Fix typo
mferril Dec 18, 2024
eb5eff5
Fix comment formatting in convertEvtCalc
olantwin Dec 18, 2024
2a72ac2
Update changelog
mferril Dec 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ it in future.
* Set Decay Volume Medium as helium (previously vacuums),can be explicitly switched to vacuum with --vacuums.
* Medium of SST boxes will be the same as DecayVolumeMedium (previously, always vacuum)
* Don't prune tracks (before we were using the CFL option to Track::prune, see https://github.com/GenFit/GenFit/blob/e81adeb07c8643301a1d9f7ae25048557cc72dff/core/include/Track.h#L298)
* **EventCalc LLP event generator**
This modification to the EventCalc interface accommodates for generic N-body LLP decays.

### Removed

Expand Down
74 changes: 30 additions & 44 deletions macro/convertEvtCalc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ def parse_file(infile):
try:
with open(infile, "r") as file:
lines = file.readlines()
lines = lines[2:]
olantwin marked this conversation as resolved.
Show resolved Hide resolved

parsed_data = []
current_block = []
Expand Down Expand Up @@ -91,37 +92,14 @@ def convert_file(infile, outdir):
"vy",
"vz",
]
vars_names += [
"px_prod1",
"py_prod1",
"pz_prod1",
"e_prod1",
"mass_prod1",
"pdg_prod1",
"charge_prod1",
"stability_prod1",
daughter_vars = [
"px_prod",
"py_prod",
"pz_prod",
"e_prod",
"mass_prod",
"pdg_prod",
]
vars_names += [
"px_prod2",
"py_prod2",
"pz_prod2",
"e_prod2",
"mass_prod2",
"pdg_prod2",
"charge_prod2",
"stability_prod2",
]
vars_names += [
"px_prod3",
"py_prod3",
"pz_prod3",
"e_prod3",
"mass_prod3",
"pdg_prod3",
"charge_prod3",
"stability_prod3",
]

fname = infile.split("/")[-1]
command = f"cp {infile} {outdir}/{fname}"

Expand All @@ -133,36 +111,43 @@ def convert_file(infile, outdir):
infile = f"{outdir}/{fname}"
parsed_data = parse_file(infile)
outfile = infile.split(".dat")[0] + ".root"
ncols = len(parsed_data[0][2])
nvardau = 6 # qualifiers for each daughter
remaining_vars = ncols - len(vars_names)

if (remaining_vars % nvardau)!=0:
raise ValueError(f"- convertEvtCalc - Error: number of daughters is not exact.")

ndau = remaining_vars // nvardau
print(f"- convertEvtCalc - Max multiplicity of daughters: {ndau}")

vars_names.extend(
f"{var}{i}" for i in range(1, ndau + 1) for var in daughter_vars
)

try:
check_consistency_infile(nvars=len(vars_names), ncols=len(parsed_data[0][2]))
check_consistency_infile(nvars=len(vars_names), ncols=ncols)
except ValueError as e:
print(f"- convertEvtCalc - Error: {e}")
sys.exit(1)

vars_names += ["ndau"]

if parsed_data:
root_file = r.TFile.Open(outfile, "RECREATE")

tree = r.TTree("LLP_tree", "LLP_tree")

branch_i = {}
branch_f = {}
for var in vars_names:
if "pdg_" in var:
branch_i[var] = np.zeros(1, dtype=int)
tree.Branch(var, branch_i[var], f"{var}/I")
else:
branch_f[var] = np.zeros(1, dtype=float)
tree.Branch(var, branch_f[var], f"{var}/D")
branch_f[var] = np.zeros(1, dtype=float)
tree.Branch(var, branch_f[var], f"{var}/D")

for pt, sp, vars in parsed_data:
for row in zip(*vars):
for i, value in enumerate(row):
if i < len(vars_names):
if "pdg_" in vars_names[i]:
branch_i[vars_names[i]][0] = value
else:
branch_f[vars_names[i]][0] = value
if i < len(vars_names)-1:
branch_f[vars_names[i]][0] = value
branch_f["ndau"][0] = ndau/1.
tree.Fill()

tree.Write()
Expand Down Expand Up @@ -200,3 +185,4 @@ def main():

if __name__ == "__main__":
main()

201 changes: 80 additions & 121 deletions shipgen/EvtCalcGenerator.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,6 @@

using namespace ShipUnit;

// read events from ntuples produced

// ----- Default constructor -------------------------------------------
EvtCalcGenerator::EvtCalcGenerator() {}
// -------------------------------------------------------------------------
Expand All @@ -32,91 +30,64 @@ Bool_t EvtCalcGenerator::Init(const char* fileName, const int firstEvent)
fNevents = fTree->GetEntries();
fn = firstEvent;

if (fTree->FindBranch("px_llp")) {
fTree->SetBranchAddress("px_llp", &px_llp);
}
if (fTree->FindBranch("py_llp")) {
fTree->SetBranchAddress("py_llp", &py_llp);
}
if (fTree->FindBranch("pz_llp")) {
fTree->SetBranchAddress("pz_llp", &pz_llp);
}
if (fTree->FindBranch("e_llp")) {
fTree->SetBranchAddress("e_llp", &e_llp);
}
if (fTree->FindBranch("pdg_llp")) {
fTree->SetBranchAddress("pdg_llp", &pdg_llp);
}
if (fTree->FindBranch("decay_prob")) {
fTree->SetBranchAddress("decay_prob", &decay_prob);
}
if (fTree->FindBranch("vx")) {
fTree->SetBranchAddress("vx", &vx);
}
if (fTree->FindBranch("vy")) {
fTree->SetBranchAddress("vy", &vy);
}
if (fTree->FindBranch("vz")) {
fTree->SetBranchAddress("vz", &vz);
}

if (fTree->FindBranch("px_prod1")) {
fTree->SetBranchAddress("px_prod1", &px_prod1);
}
if (fTree->FindBranch("py_prod1")) {
fTree->SetBranchAddress("py_prod1", &py_prod1);
}
if (fTree->FindBranch("pz_prod1")) {
fTree->SetBranchAddress("pz_prod1", &pz_prod1);
}
if (fTree->FindBranch("e_prod1")) {
fTree->SetBranchAddress("e_prod1", &e_prod1);
}
if (fTree->FindBranch("pdg_prod1")) {
fTree->SetBranchAddress("pdg_prod1", &pdg_prod1);
}
auto *branches = fTree->GetListOfBranches();
nBranches = branches->GetEntries();
branchVars.resize(nBranches);

if (fTree->FindBranch("px_prod2")) {
fTree->SetBranchAddress("px_prod2", &px_prod2);
}
if (fTree->FindBranch("py_prod2")) {
fTree->SetBranchAddress("py_prod2", &py_prod2);
}
if (fTree->FindBranch("pz_prod2")) {
fTree->SetBranchAddress("pz_prod2", &pz_prod2);
}
if (fTree->FindBranch("e_prod2")) {
fTree->SetBranchAddress("e_prod2", &e_prod2);
}
if (fTree->FindBranch("pdg_prod2")) {
fTree->SetBranchAddress("pdg_prod2", &pdg_prod2);
}

if (fTree->FindBranch("px_prod3")) {
fTree->SetBranchAddress("px_prod3", &px_prod3);
}
if (fTree->FindBranch("py_prod3")) {
fTree->SetBranchAddress("py_prod3", &py_prod3);
}
if (fTree->FindBranch("pz_prod3")) {
fTree->SetBranchAddress("pz_prod3", &pz_prod3);
}
if (fTree->FindBranch("e_prod3")) {
fTree->SetBranchAddress("e_prod3", &e_prod3);
}
if (fTree->FindBranch("pdg_prod3")) {
fTree->SetBranchAddress("pdg_prod3", &pdg_prod3);
for (int i = 0; i < nBranches; ++i) {
auto *branch = dynamic_cast<TBranch *>(branches->At(i));
if (fTree->FindBranch(branch->GetName())) {
fTree->SetBranchAddress(branch->GetName(), &branchVars[i]);
}
}

return kTRUE;
}
// -------------------------------------------------------------------------

// ----- Destructor ----------------------------------------------------
EvtCalcGenerator::~EvtCalcGenerator() {}

// -- Generalized branch access --------------------------------------------
Double_t EvtCalcGenerator::GetBranchValue(const std::unique_ptr<TTree>& tree, int index) const {
if (index < branchVars.size()) {
return branchVars[index];
} else {
throw std::out_of_range("Branch index out of range");
}
}
// -- Generalized daughter variable access ---------------------------------
Double_t EvtCalcGenerator::GetDaughterValue(const std::unique_ptr<TTree>& tree, int dauID, int offset) const {
int baseIndex = 10 + (dauID * 6);
return GetBranchValue(tree, baseIndex + offset);
}

// -- Wrapper functions ----------------------------------------------------
// -------------------------------------------------------------------------
Double_t EvtCalcGenerator::GetNdaughters(const std::unique_ptr<TTree>& tree) const {
return GetBranchValue(tree, nBranches-1);
}

// -- LLP properties ------------------------------------------------------
Double_t EvtCalcGenerator::GetMotherPx(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::MotherPx)); }
Double_t EvtCalcGenerator::GetMotherPy(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::MotherPy)); }
Double_t EvtCalcGenerator::GetMotherPz(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::MotherPz)); }
Double_t EvtCalcGenerator::GetMotherE(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::MotherE)); }

// -- Vertex properties ---------------------------------------------------
Double_t EvtCalcGenerator::GetVx(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::Vx)); }
Double_t EvtCalcGenerator::GetVy(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::Vy)); }
Double_t EvtCalcGenerator::GetVz(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::Vz)); }

// -- Decay probability ---------------------------------------------------
Double_t EvtCalcGenerator::GetDecayProb(const std::unique_ptr<TTree>& tree) const { return GetBranchValue(tree, static_cast<int>(BranchIndices::DecayProb)); }

// -- Daughter properties ------------------------------------------------
Double_t EvtCalcGenerator::GetDauPx(const std::unique_ptr<TTree>& tree, int dauID) const { return GetDaughterValue(tree, dauID, 0); }
Double_t EvtCalcGenerator::GetDauPy(const std::unique_ptr<TTree>& tree, int dauID) const { return GetDaughterValue(tree, dauID, 1); }
Double_t EvtCalcGenerator::GetDauPz(const std::unique_ptr<TTree>& tree, int dauID) const { return GetDaughterValue(tree, dauID, 2); }
Double_t EvtCalcGenerator::GetDauE(const std::unique_ptr<TTree>& tree, int dauID) const { return GetDaughterValue(tree, dauID, 3); }
Double_t EvtCalcGenerator::GetDauPDG(const std::unique_ptr<TTree>& tree, int dauID) const { return GetDaughterValue(tree, dauID, 5); }

// ----- Passing the event ---------------------------------------------
// ----- Passing the event -------------------------------------------
Bool_t EvtCalcGenerator::ReadEvent(FairPrimaryGenerator* cpg)
{
if (fn == fNevents) {
Expand All @@ -129,61 +100,49 @@ Bool_t EvtCalcGenerator::ReadEvent(FairPrimaryGenerator* cpg)
if (fn % 100 == 0)
LOGF(info, "Info EvtCalcGenerator: event nr %s", fn);

Ndau = GetNdaughters(fTree);
// Vertex coordinates in the SHiP reference frame, expressed in [cm]
Double_t space_unit_conv = 100.; // m to cm
Double_t coord_shift = (zDecayVolume - ztarget) / space_unit_conv; // units m
Double_t vx_transf = vx * space_unit_conv; // units cm
Double_t vy_transf = vy * space_unit_conv; // units cm
Double_t vz_transf = (vz - coord_shift) * space_unit_conv; // units cm
Double_t vx_transf = GetVx(fTree) * space_unit_conv; // units cm
Double_t vy_transf = GetVy(fTree) * space_unit_conv; // units cm
Double_t vz_transf = (GetVz(fTree) - coord_shift) * space_unit_conv; // units cm

Double_t c = 2.99792458e+10; // speed of light [cm/s]
Double_t tof = TMath::Sqrt(vx_transf * vx_transf + vy_transf * vy_transf + vz_transf * vz_transf) / c;
Double_t decay_prob = GetDecayProb(fTree);
Double_t pdg_dau = 0;

// Mother LLP
Bool_t wanttracking = false;
pdg_llp = 999; // Geantino, placeholder
Double_t pdg_llp = 999; // Geantino, placeholder

cpg->AddTrack(
pdg_llp, px_llp, py_llp, pz_llp, vx_transf, vy_transf, vz_transf, -1., wanttracking, e_llp, tof, decay_prob);
pdg_llp, GetMotherPx(fTree), GetMotherPy(fTree), GetMotherPz(fTree),
vx_transf, vy_transf, vz_transf, -1., wanttracking,
GetMotherE(fTree), tof, decay_prob);

// Secondaries
wanttracking = true;
cpg->AddTrack(pdg_prod1,
px_prod1,
py_prod1,
pz_prod1,
vx_transf,
vy_transf,
vz_transf,
0.,
wanttracking,
e_prod1,
tof,
decay_prob);
cpg->AddTrack(pdg_prod2,
px_prod2,
py_prod2,
pz_prod2,
vx_transf,
vy_transf,
vz_transf,
0.,
wanttracking,
e_prod2,
tof,
decay_prob);
if (pdg_prod3 != -999) {
cpg->AddTrack(pdg_prod3,
px_prod3,
py_prod3,
pz_prod3,
vx_transf,
vy_transf,
vz_transf,
0.,
wanttracking,
e_prod3,
tof,
decay_prob);

// Secondaries
for (int iPart = 0; iPart < Ndau; ++iPart) {
pdg_dau = GetDauPDG(fTree, iPart);
if ( pdg_dau != -999) {
cpg->AddTrack(pdg_dau,
GetDauPx(fTree, iPart),
GetDauPy(fTree, iPart),
GetDauPz(fTree, iPart),
vx_transf,
vy_transf,
vz_transf,
0.,
wanttracking,
GetDauE(fTree, iPart),
tof,
decay_prob);
}
}

return kTRUE;
}

Loading
Loading