From f71fd1d83df3d4f5af843796d6c30eb7929c3db2 Mon Sep 17 00:00:00 2001 From: THanae Date: Tue, 23 Jun 2026 16:35:24 +0200 Subject: [PATCH] feat: Add option to split pions and kaons in run_fixedTarget --- CHANGELOG.md | 1 + gconfig/g4Config.C | 16 +- macro/run_fixedTarget.py | 17 ++ muonShieldOptimization/exitHadronAbsorber.cxx | 271 ++++++++++++++++-- muonShieldOptimization/exitHadronAbsorber.h | 28 ++ shipdata/ShipStack.cxx | 32 ++- shipdata/ShipStack.h | 2 + 7 files changed, 333 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b319df3658..1997c6b62d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ it in future. ## Unreleased ### Added +* Add option to split kaons and pions right before they decay, to increase the number of muons ### Changed diff --git a/gconfig/g4Config.C b/gconfig/g4Config.C index 44706022dc..4305231622 100644 --- a/gconfig/g4Config.C +++ b/gconfig/g4Config.C @@ -30,10 +30,18 @@ void Config() { /// - stackPopper - stackPopper process /// When more than one options are selected, they should be separated with '+' /// character: eg. stepLimit+specialCuts. - TG4RunConfiguration* runConfiguration = new TG4RunConfiguration( - "geomRoot", "FTFP_BERT_HP_EMZ", "stepLimiter+specialCuts+specialControls", - false, // specialStacking (default) - false); // disable MT + + const char* env = std::getenv("KAON_PION_SPLITS"); + int32_t fNsplits = env ? std::atoi(env) : 0; + + std::string controls = "stepLimiter+specialCuts+specialControls"; + /// stackPopper is required when adding secondaries, e.g. when splitting the + /// pions and kaons + if (fNsplits != 0) controls += "+stackPopper"; + TG4RunConfiguration* runConfiguration = + new TG4RunConfiguration("geomRoot", "FTFP_BERT_HP_EMZ", controls.c_str(), + false, // specialStacking (default) + false); // disable MT /// Create the G4 VMC TGeant4* geant4 = diff --git a/macro/run_fixedTarget.py b/macro/run_fixedTarget.py index 7c33f23492..fd0bace9be 100644 --- a/macro/run_fixedTarget.py +++ b/macro/run_fixedTarget.py @@ -78,6 +78,16 @@ def get_work_dir(run_number, tag: str | None = None) -> str: ap.add_argument("-J", "--Jpsi-mainly", action=argparse.BooleanOptionalAction, dest="JpsiMainly", default=False) ap.add_argument("-b", "--boostDiMuon", type=float, default=1.0, help="boost Di-muon branching ratios") ap.add_argument("-X", "--boostFactor", type=float, default=1.0, help="boost Di-muon prod cross sections") +ap.add_argument( + "--kaon-pion-splits", + type=int, + default=0, + help="splitting factor for kaons and pions, in order to boost the number of muons stemming from their decays", +) +ap.add_argument( + "--multiple-kpi-splits", action="store_true", help="split kaons and pions multiple times along the track path" +) + ap.add_argument("-C", "--charm", action=argparse.BooleanOptionalAction, default=False, help="generate charm decays") ap.add_argument("-B", "--beauty", action=argparse.BooleanOptionalAction, default=False, help="generate beauty decays") ap.add_argument( @@ -272,6 +282,8 @@ def get_work_dir(run_number, tag: str | None = None) -> str: if args.boostFactor > 1: # Turn off UseGeneralProcess to access GammaToMuons directly when cross-sections need to be changed os.environ["SET_GENERAL_PROCESS_TO_FALSE"] = "1" +if args.kaon_pion_splits > 0: + os.environ["KAON_PION_SPLITS"] = str(args.kaon_pion_splits) run.SetUserConfig("g4Config.C") # user configuration file default g4Config.C rtdb = run.GetRuntimeDb() @@ -342,6 +354,9 @@ def get_work_dir(run_number, tag: str | None = None) -> str: sensPlaneHA = ROOT.exitHadronAbsorber() +sensPlaneHA.SetNSplits(args.kaon_pion_splits) # type: ignore[missing-attribute] +if args.multiple_kpi_splits: + sensPlaneHA.SetSplitMultipleTimes() # type: ignore[missing-attribute] sensPlaneHA.SetEnergyCut(args.ecut * u.GeV) sensPlaneHA.SetVetoPointName("PlaneHA") @@ -422,6 +437,8 @@ def get_work_dir(run_number, tag: str | None = None) -> str: fStack = gMC.GetStack() fStack.SetMinPoints(1) fStack.SetEnergyCut(-1.0) +if args.kaon_pion_splits > 0: + fStack.SetSplitting() # import AddDiMuonDecayChannelsToG4 diff --git a/muonShieldOptimization/exitHadronAbsorber.cxx b/muonShieldOptimization/exitHadronAbsorber.cxx index bf478b9157..5f3094877d 100644 --- a/muonShieldOptimization/exitHadronAbsorber.cxx +++ b/muonShieldOptimization/exitHadronAbsorber.cxx @@ -12,10 +12,12 @@ #include "FairGeoMedia.h" #include "FairGeoNode.h" #include "FairGeoVolume.h" +#include "FairLogger.h" #include "FairRootManager.h" #include "FairVolume.h" #include "ShipDetectorList.h" #include "ShipStack.h" +#include "TArrayI.h" #include "TDatabasePDG.h" #include "TGeoBBox.h" #include "TGeoBoolNode.h" @@ -24,9 +26,13 @@ #include "TGeoMaterial.h" #include "TH1D.h" #include "TH2D.h" +#include "TLorentzVector.h" +#include "TMCProcess.h" #include "TParticle.h" #include "TROOT.h" #include "TVirtualMC.h" +#include "vetoPoint.h" + using std::cout; using std::endl; @@ -41,7 +47,9 @@ exitHadronAbsorber::exitHadronAbsorber(const char* Name, Bool_t Active) fVetoName("veto"), fzPos(3E8), withNtuple(kFALSE), - fCylindricalPlane(kFALSE) {} + fCylindricalPlane(kFALSE), + fNsplits(0), + fCurrentSurvivalFactor(1) {} exitHadronAbsorber::exitHadronAbsorber() : Detector("exitHadronAbsorber", kTRUE, kVETO), @@ -52,33 +60,109 @@ exitHadronAbsorber::exitHadronAbsorber() fzPos(3E8), withNtuple(kFALSE), fCylindricalPlane(kFALSE), - fUseCaveCoordinates(kFALSE) {} + fUseCaveCoordinates(kFALSE), + fNsplits(0), + fCurrentSurvivalFactor(1) {} Bool_t exitHadronAbsorber::ProcessHits(FairVolume* vol) { /** This method is called from the MC stepping */ - if (gMC->IsTrackEntering()) { - fTrackID = gMC->GetStack()->GetCurrentTrackNumber(); - fEventID = gMC->CurrentEvent(); - TParticle* p = gMC->GetStack()->GetCurrentTrack(); - fUniqueID = p->GetUniqueID(); - Int_t pdgCode = p->GetPdgCode(); - gMC->TrackMomentum(fMom); - if (!fOnlyMuons || TMath::Abs(pdgCode) == 13) { - fTime = gMC->TrackTime() * 1.0e09; - fLength = gMC->TrackLength(); - gMC->TrackPosition(fPos); - if ((fMom.E() - fMom.M()) > EMax) { - AddHit(fEventID, fTrackID, 111, TVector3(fPos.X(), fPos.Y(), fPos.Z()), - TVector3(fMom.Px(), fMom.Py(), fMom.Pz()), fTime, fLength, 0, - pdgCode, TVector3(p->Vx(), p->Vy(), p->Vz()), - TVector3(p->Px(), p->Py(), p->Pz())); - ShipStack* stack = dynamic_cast(gMC->GetStack()); - stack->AddPoint(kVETO); + TString volName = gMC->CurrentVolName(); + + if (volName.Contains("exitHadronAbsorber")) { + if (gMC->IsTrackEntering()) { + fTrackID = gMC->GetStack()->GetCurrentTrackNumber(); + fEventID = gMC->CurrentEvent(); + TParticle* p = gMC->GetStack()->GetCurrentTrack(); + fUniqueID = p->GetUniqueID(); + Int_t pdgCode = p->GetPdgCode(); + Int_t motherId = p->GetFirstMother(); + gMC->TrackMomentum(fMom); + if (!fOnlyMuons || TMath::Abs(pdgCode) == 13) { + fTime = gMC->TrackTime() * 1.0e09; + fLength = gMC->TrackLength(); + gMC->TrackPosition(fPos); + if (((fMom.E() - fMom.M()) > EMax) && + (fDecayedParentIDs.count(motherId) == 0)) { + AddHit(fEventID, fTrackID, 111, + TVector3(fPos.X(), fPos.Y(), fPos.Z()), + TVector3(fMom.Px(), fMom.Py(), fMom.Pz()), fTime, fLength, 0, + pdgCode, TVector3(p->Vx(), p->Vy(), p->Vz()), + TVector3(p->Px(), p->Py(), p->Pz())); + auto* stack = dynamic_cast(gMC->GetStack()); + stack->AddPoint(kVETO); + } } } + + if ((!fCylindricalPlane) && fzPos > 1E8) { + gMC->StopTrack(); + } } - if ((!fCylindricalPlane) && fzPos > 1E8) { - gMC->StopTrack(); + + if (fNsplits > 0 && (!fSplitOnce)) { + Int_t currentTrackId = gMC->GetStack()->GetCurrentTrackNumber(); + + if (fCloneTracks.count(currentTrackId) > 0 || + fContinuationTracks.count(currentTrackId) > 0) { + return kTRUE; + } + + Int_t track_pid = gMC->TrackPid(); + bool kaon_or_pion = + (TMath::Abs(track_pid) == 211 || TMath::Abs(track_pid) == 321); + + if (kaon_or_pion) { + TParticle* part = gMC->GetStack()->GetCurrentTrack(); + if (!part) return kTRUE; + + Double_t delta_s = gMC->TrackStep(); + + if (delta_s > 0) { + TLorentzVector pos, mom; + gMC->TrackPosition(pos); + gMC->TrackMomentum(mom); + + Double_t mass = gMC->TrackMass(); + Double_t tau = gMC->ParticleLifeTime(track_pid); // in nanoseconds + Double_t c_tau = tau * 29.9792458; // in cm/ns + + Double_t instantP = mom.P(); + + Double_t lambda_decay = (instantP / mass) * c_tau; + Double_t P_decay = 1.0 - TMath::Exp(-delta_s / lambda_decay); + if ((P_decay > 0.0) && ((mom.E() - mom.M()) > EMax)) { + double polX = 0, polY = 0, polZ = 0; + TVector3 polVector; + part->GetPolarisation(polVector); + polX = polVector.X(); + polY = polVector.Y(); + polZ = polVector.Z(); + Int_t trueParentId = part->GetFirstMother(); + + Double_t decayBranchWeight = fCurrentSurvivalFactor * P_decay; + Double_t cloneWeight = decayBranchWeight / fNsplits; + for (int i = 0; i < fNsplits; ++i) { + TrackBuffer clone; + clone.pdg = track_pid; + clone.px = mom.Px(); + clone.py = mom.Py(); + clone.pz = mom.Pz(); + clone.e = mom.E(); + clone.x = pos.X(); + clone.y = pos.Y(); + clone.z = pos.Z(); + clone.t = pos.T(); + clone.polx = polX; + clone.poly = polY; + clone.polz = polZ; + clone.weight = cloneWeight; + clone.parentID = trueParentId; + fSecondaryBuffer.push_back(clone); + } + fCurrentSurvivalFactor *= (1.0 - P_decay); + } + } + } } return kTRUE; } @@ -152,14 +236,128 @@ void exitHadronAbsorber::Initialize() { } } +void exitHadronAbsorber::BeginEvent() { + fCloneTracks.clear(); + fContinuationTracks.clear(); + fDecayedParentIDs.clear(); + fSecondaryBuffer.clear(); +} + +void exitHadronAbsorber::PostTrack() { + Int_t currentTrackId = gMC->GetStack()->GetCurrentTrackNumber(); + + if (fCloneTracks.count(currentTrackId) > 0 || + fContinuationTracks.count(currentTrackId) > 0) { + return; + } + + Int_t track_pid = gMC->TrackPid(); + bool kaon_or_pion = + (TMath::Abs(track_pid) == 211 || TMath::Abs(track_pid) == 321); + + if (fNsplits > 0 && kaon_or_pion) { + bool isNaturalDecay = false; + bool isParticleDestroyed = gMC->IsTrackStop() || gMC->IsTrackDisappeared(); + TArrayI processes; + gMC->StepProcesses(processes); + for (int i = 0; i < processes.GetSize(); i++) { + if (processes[i] == kPDecay) { + isNaturalDecay = true; + break; + } + } + + TParticle* part = gMC->GetStack()->GetCurrentTrack(); + if (!part) return; + + TLorentzVector finalPos, finalMom; + gMC->TrackPosition(finalPos); + gMC->TrackMomentum(finalMom); + + double polX = 0, polY = 0, polZ = 0; + TVector3 polVector; + part->GetPolarisation(polVector); + polX = polVector.X(); + polY = polVector.Y(); + polZ = polVector.Z(); + Int_t trueParentId = part->GetFirstMother(); + + auto* stack = dynamic_cast(gMC->GetStack()); + + // All remaining weight used if cloning happens at point where original + // particle decays + if (isNaturalDecay) { + Double_t finalEndpointWeight = fCurrentSurvivalFactor / fNsplits; + for (int i = 0; i < fNsplits; ++i) { + TrackBuffer clone; + clone.pdg = track_pid; + clone.px = finalMom.Px(); + clone.py = finalMom.Py(); + clone.pz = finalMom.Pz(); + clone.e = finalMom.E(); + clone.x = finalPos.X(); + clone.y = finalPos.Y(); + clone.z = finalPos.Z(); + clone.t = finalPos.T(); + clone.polx = polX; + clone.poly = polY; + clone.polz = polZ; + clone.weight = finalEndpointWeight; + clone.parentID = trueParentId; + fSecondaryBuffer.push_back(clone); + } + fCurrentSurvivalFactor = 0.0; + fDecayedParentIDs.insert(currentTrackId); + } + + // tracks which do not decay are stopped with stoptrack and added back with + // a given weight + if (!isNaturalDecay && !isParticleDestroyed && stack) { + Int_t ntr; + stack->PushTrack(1, trueParentId, track_pid, finalMom.Px(), finalMom.Py(), + finalMom.Pz(), finalMom.E(), finalPos.X(), finalPos.Y(), + finalPos.Z(), finalPos.T(), polX, polY, polZ, + kPNoProcess, ntr, fCurrentSurvivalFactor, 999); + fContinuationTracks.insert(ntr); + } + + gMC->StopTrack(); + } +} + void exitHadronAbsorber::PreTrack() { + bool stackbufferisnotempty = !fSecondaryBuffer.empty(); + if (stackbufferisnotempty) { + auto* stack = dynamic_cast(gMC->GetStack()); + Int_t ntr; + for (const auto& trk : fSecondaryBuffer) { + stack->PushTrack(1, trk.parentID, trk.pdg, trk.px, trk.py, trk.pz, trk.e, + trk.x, trk.y, trk.z, trk.t, trk.polx, trk.poly, trk.polz, + kPNoProcess, ntr, trk.weight, 999); + fCloneTracks.insert(ntr); + } + // Clear the buffer so we don't duplicate them for the next track + fSecondaryBuffer.clear(); + } + + // Reset relative survival factor to 1.0 + fCurrentSurvivalFactor = 1.0; + gMC->TrackMomentum(fMom); if ((fMom.E() - fMom.M()) < EMax) { gMC->StopTrack(); return; } TParticle* p = gMC->GetStack()->GetCurrentTrack(); + Int_t currentID = gMC->GetStack()->GetCurrentTrackNumber(); + + if (fCloneTracks.find(currentID) != fCloneTracks.end()) { + // Force the decay time to 0 + gMC->ForceDecayTime(0); + } + Int_t pdgCode = p->GetPdgCode(); + // record statistics for neutrinos, electrons and photons // add pi0 111 eta 221 eta' 331 omega 223 Int_t idabs = TMath::Abs(pdgCode); @@ -252,6 +450,26 @@ void exitHadronAbsorber::FinishRun() { } } +void RegisterDaughtersRecursively(TGeoVolume* volume, + exitHadronAbsorber* detector) { + if (!volume) return; + Int_t nDaughters = volume->GetNdaughters(); + for (Int_t i = 0; i < nDaughters; ++i) { + TGeoNode* daughterNode = volume->GetNode(i); + if (daughterNode) { + TGeoVolume* daughterVol = daughterNode->GetVolume(); + if (daughterVol) { + // register daughter volume + detector->AddSensitiveVolume(daughterVol); + LOG(info) << "[exitHadronAbsorber] Registered subvolume: " + << daughterVol->GetName(); + // call function recursively + RegisterDaughtersRecursively(daughterVol, detector); + } + } + } +} + void exitHadronAbsorber::ConstructGeometry() { static FairGeoLoader* geoLoad = FairGeoLoader::Instance(); static FairGeoInterface* geoFace = geoLoad->getGeoInterface(); @@ -328,6 +546,15 @@ void exitHadronAbsorber::ConstructGeometry() { new TGeoTranslation(0, 0, 0)); AddSensitiveVolume(sensPlaneCyl); } + if ((fNsplits > 0) && (!fSplitOnce)) { + TString parentVolumeName = "/target_vacuum_box_1"; + nav->cd(parentVolumeName.Data()); + TGeoVolume* vol = nav->GetCurrentNode()->GetVolume(); + AddSensitiveVolume(vol); + int32_t nDaughters = vol->GetNdaughters(); + // register each unique daughter volume + RegisterDaughtersRecursively(vol, this); + } } void exitHadronAbsorber::Register() { diff --git a/muonShieldOptimization/exitHadronAbsorber.h b/muonShieldOptimization/exitHadronAbsorber.h index e63c4d1b4c..f590217e65 100644 --- a/muonShieldOptimization/exitHadronAbsorber.h +++ b/muonShieldOptimization/exitHadronAbsorber.h @@ -5,6 +5,7 @@ #ifndef MUONSHIELDOPTIMIZATION_EXITHADRONABSORBER_H_ #define MUONSHIELDOPTIMIZATION_EXITHADRONABSORBER_H_ +#include #include #include "Detector.h" @@ -14,6 +15,15 @@ class FairVolume; +struct TrackBuffer { + Int_t pdg; + Double_t px, py, pz, e; + Double_t x, y, z, t; + Double_t polx, poly, polz; + Double_t weight; + Int_t parentID; +}; + class exitHadronAbsorber : public SHiP::Detector { public: exitHadronAbsorber(const char* Name, Bool_t Active); @@ -29,6 +39,12 @@ class exitHadronAbsorber : public SHiP::Detector { void FinishRun() override; void PreTrack() override; + void PostTrack() override; + void BeginEvent() override; + // void FinishEvent(); + + void SetNSplits(Int_t n) { fNsplits = n; } + void SetSplitMultipleTimes() { fSplitOnce = kFALSE; } inline void SetEnergyCut(Float_t emax) { EMax = emax; } inline void SetOnlyMuons() { fOnlyMuons = kTRUE; } @@ -51,6 +67,18 @@ class exitHadronAbsorber : public SHiP::Detector { Bool_t fCylindricalPlane; //! cylindrical sensPlane flag Bool_t fUseCaveCoordinates; //! set position from cave + int32_t fNsplits; + Double_t fCurrentSurvivalFactor; // survival factor at every step, if we + // choose to split at every step + Bool_t fSplitOnce = + kTRUE; // determine if we want to split once (when the particle decays) + // or at every step (taking decay probabilities into account) + + std::vector fSecondaryBuffer; + std::set fCloneTracks; + std::set fContinuationTracks; + std::set fDecayedParentIDs; + TFile* fout; //! TClonesArray* fElectrons; //! Int_t index; diff --git a/shipdata/ShipStack.cxx b/shipdata/ShipStack.cxx index 47e120c99d..b4f25b76d5 100644 --- a/shipdata/ShipStack.cxx +++ b/shipdata/ShipStack.cxx @@ -47,7 +47,8 @@ ShipStack::ShipStack(Int_t size) fStoreSecondaries(kTRUE), fMinPoints(1), fEnergyCut(0.), - fStoreMothers(kTRUE) { + fStoreMothers(kTRUE), + fSplitting(kFALSE) { fTracks = new std::vector(); fTracks->reserve(size); } @@ -86,14 +87,33 @@ void ShipStack::PushTrack(Int_t toBeDone, Int_t parentId, Int_t pdgCode, // cout << "ShipStack: " << fNParticles << " " << pdgCode << " " << parentId // << " " << secondparentID<<" "<= 0) { + if (parentId < fParticles->GetEntriesFast()) { + TParticle* parentPart = + dynamic_cast(fParticles->At(parentId)); + if (parentPart) { + Double_t parentWeight = parentPart->GetWeight(); + weight = weight * parentWeight; + } + } + } + } + + Int_t trackId = fNParticles; + // --> Get TParticle array TClonesArray& partArray = *fParticles; - // --> Create new TParticle and add it to the TParticle array - Int_t trackId = fNParticles; + // --> Set argument variable + ntr = trackId; + Int_t nPoints = 0; Int_t daughter1Id = -1; Int_t daughter2Id = -1; + + // --> Create new TParticle and add it to the TParticle array TParticle* particle = new (partArray[fNParticles++]) TParticle(pdgCode, trackId, parentId, nPoints, daughter1Id, daughter2Id, px, py, pz, e, vx, vy, vz, time); @@ -118,12 +138,8 @@ void ShipStack::PushTrack(Int_t toBeDone, Int_t parentId, Int_t pdgCode, if (parentId < 0) { fNPrimaries++; } - - // --> Set argument variable - ntr = trackId; - - // --> Push particle on the stack if toBeDone is set if (toBeDone == 1) { + // --> Push particle on the stack if toBeDone is set particle->SetBit(kDoneBit); fStack.push(particle); } diff --git a/shipdata/ShipStack.h b/shipdata/ShipStack.h index 29f85fc7ab..fa85dc43ae 100644 --- a/shipdata/ShipStack.h +++ b/shipdata/ShipStack.h @@ -153,6 +153,7 @@ class ShipStack : public FairGenericStack { void SetMinPoints(Int_t min) { fMinPoints = min; } void SetEnergyCut(Double_t eMin) { fEnergyCut = eMin; } void StoreMothers(Bool_t choice = kTRUE) { fStoreMothers = choice; } + void SetSplitting() { fSplitting = kTRUE; } /** Increment number of points for the current track in a given detector *@param iDet Detector unique identifier @@ -204,6 +205,7 @@ class ShipStack : public FairGenericStack { Int_t fMinPoints; Double32_t fEnergyCut; Bool_t fStoreMothers; + Bool_t fSplitting; /** Mark tracks for output using selection criteria **/ void SelectTracks();