Refactor primary construction into aux data + generator

[sethrj]

In the optical transport loop, we want to be able to use Cerenkov/Scintillation "generators" (which have signature (RNG&) -> Primary) to produce optical photons on device. Like buffered track initializers, we'd like to be able to fill up track slots, pause partway (if there aren't enough track slots), and then fill any empty slots on the next loop. We'll also want to do this to support sampling for optical maps: a generator could be initialized with a volume ID and a region of space and only accept samples within that region. We should consider doing something similar for core tracks, since there is a memory usage/bandwidth advantage to initializing tracks on device rather than on host. (We'd also be able to use a consistent RNG.)

I propose we move the StepActionOrder "generate" before the "initialize" that fills empty track slots:

• Change "extend from primaries" to be the only class that handles primaries; it would take them directly from the users
• Refactor the "track initializer" buffer to likewise be an action that will fill empty slots
• New future "generator" actions can be added to replace phys/PrimaryGenerator.hh or define new generalized distributions (e.g., sample N tracks from a gaussian)

Perhaps we could also use this refactor to think of a way to make track IDs deterministic (currently, tracks from secondaries get their event ID-associated track count using an atomic_add). This would involve removing the track ID from the track initializer and instead constructing new track IDs in the initialization stage.

New classes

GeneratorInterface

• Has a GeneratorId created at setup
• Given an auxiliary vec (so that each generator class stores the remaining count as part of the state), set the number of tracks left to initialize to an index GeneratorId in a shared array of "number tracks left to generate".
• We could consider also someday having a "priority" return as well, maybe returning the amount of expected secondaries from each one (using the total energy)?
• For charge partitioning, we'll also need to split these into two counts: front and back

Interface instances:

• The ExtendFromPrimariesAction is now a generator, its aux state contains buffered primaries, all for a single event
• ExtendFromSecondaries at the end of a step does in-place replacement, and its initializers should store the "touchable" (unique volume ID) and parent ID
• A more tightly integrated Geant4 initializer could also include touchable information (and even other metadata) to reduce initialization costs
• CherenkovGeneratorAction, ScintGeneratorAction are also generator interfaces

GeneratorRegistry

• Generators are added to it during construction, using GeneratorId next_id

SlotPartitioner

This should take the place of most of the track initialization action (aside from the actual track initialization).

• Runs at the beginning of the stepping loop
• Calls all the generators and has them fill in the "number to initialize" in an array
• Maps vacancies to generator IDs plus "index within generator"
• Saves

So this class's data will be

• StateItems<GeneratorId>: j for "generate from distribution j", null if track is active or no initialization needed
• StateItems<size_type>: i in "generate the ith item of generator j into this slot"
• GeneratorItems<size_type>: number of total tracks each generator has left (updated by generators!)
• GeneratorItems<size_type>: number of total tracks generated this step, per generator (updated by partitioner)

---

Considerations:

• The 'warming up' step does a single step without any active tracks nor primaries. If generators are added as "actions" we'd have to change this to flag or some other special case; otherwise they'll generate tracks
• I don't like it, but this tightens the coupling between the generation method and the problem setup. Currently any thread can create primaries however it wants; but this way we'd have to create a new action manager and use the same method on all threads.
• To run a set of primaries after setting up the core params with a PrimaryGeneratorAction, we would have to maintain an awkward separate reference to that action and call primaries->insert(...).

cc @drbenmorgan @amandalund

[amandalund]

Thanks for this @sethrj! Sorry it's taking me a minute to fully grasp it; just a few clarifying questions:

I propose we add a new StepActionOrder "generate" before the "initialize" that fills empty track slots

• Change "extend from primaries" to be the only class that handles primaries

Does that mean a "generator" action would create new initializers, on device, directly in the initializer buffer (i.e. skip the step of creating primaries) when they can be generated from a distribution? And "extend from primaries" would still be used e.g. when reading from a HepMC3 event record?

• Refactor the "track initializer" buffer to likewise be an action that will fill empty slots

Not sure exactly what you mean by this.

The new PrimaryGeneratorAction class that replaces ExtendFromPrimariesAction would still be a CoreStepActionInterface (defining step(params, state)) but also be a AuxParamsInterface (for storing the primaries)

Or would the generator create and store primaries? Then we'd still need to convert them into initializers (except in the optical loop where primaries are initializers)?

Perhaps we could also use this refactor to think of a way to make track IDs deterministic (currently, tracks from secondaries get their event ID-associated track count using an atomic_add). This would involve removing the track ID from the track initializer and instead constructing new track IDs in the initialization stage.

Would this refactor change how we're doing the initialization (and in a way that would make this possible)?

[sethrj]

Does that mean a "generator" action would create new initializers, on device, directly in the initializer buffer (i.e. skip the step of creating primaries) when they can be generated from a distribution? And "extend from primaries" would still be used e.g. when reading from a HepMC3 event record?

Exactly. HepMC3, Geant4 offloading, and other methods for explicitly creating primaries will use the refactored ExtendFromPrimariesAction to make track initializers. Its generator operator might take the RNG but not use it (or we may need to change the signature to pass in both an RNG and an initialization index...). Other generator actions would directly create track initializers.

Not sure exactly what you mean by [Refactor the "track initializer" buffer].

Instead of the track initializers (which are different between the core and optical stepping loops) being part of the TrackInitStateData maybe it could be aux data owned by the InitializeTracksAction?

Or would the generator create and store primaries? Then we'd still need to convert them into initializers (except in the optical loop where primaries are initializers)?

Nope I think initializers. The Primary class would just be a higher-level abstraction for building track initializers. Its definition could, for example, be moved into the PrimaryGeneratorAction class file.

Would this refactor change how we're doing the initialization (and in a way that would make this possible)?

Well, maybe not all the way. But I think at least a step in that direction is to movie the TrackId out of the track initialization data, and removing it from the primary: this would remove the inconsistency between setting it from the next track_counters for secondaries or from some arbitrary user value for primaries. The track ID assignment would then be done only in InitTracksExecutor, consolidating the two different atomic_add(&state->init.track_counters...) that we currently have. This would make it easy for us to add a switch that says "if we have single-mode" (#1233) then the track ID is the counter plus the vacancy index (or whatever).

[amandalund]

Thanks, that makes a lot of sense. I guess one potential snag with the track ID stuff is the in-place initialization we do in the ProcessSecondariesExecutor...

[sethrj]

True that... perhaps if we add logical volume and surface IDs to the track initializers we can skip the in-place stuff.

[sethrj]

@amandalund Another iteration with some visual aid. Instead of each generator class having only a single generator ID, it can reserve a range of IDs. And we will probably want our "active" generator buffer to be sized more like (# optical track slots)/(avg # photons per distribution) rather than (# core track slots), with maybe another resizable buffer filling it in turn? (Or maybe we do just have a ton of generator distribution buffers)?