STM diffusion model resampler updates, plumbing, workflows, and tools

[YongyiBWu]

This PR reflects the new developments for the STM diffusion model resampler. Additional features were added for more dynamic model training. Additional bases were added for pre-training coordinate transformation. The plumbing between the DM implementation and the STM applications were completed. A full workflow was set up and some tools were added. Scope of change is strictly limited to STMMC and DM toolkit.

[FNALbuild]

Hi @YongyiBWu,
You have proposed changes to files in these packages:

• STMMC
• MachineLearningTools

which require these tests: build.

@Mu2e/fnalbuild-users, @Mu2e/write have access to CI actions on main.

⌛ The following tests have been triggered for 5bb835f: build (Build queue - API unavailable)

About FNALbuild. Code review on Mu2e/Offline.

[FNALbuild]

☀️ The build tests passed at 5bb835f.

Test	Result	Details
test with	✅	Command did not list any other PRs to include
merge	✅	Merged 5bb835f at 0ef02a6
build (prof)	✅	Log file. Build time: 15 min 35 sec
ceSimReco	✅	Log file.
g4test_03MT	✅	Log file.
transportOnly	✅	Log file.
POT	✅	Log file.
g4study	✅	Log file.
cosmicSimReco	✅	Log file.
cosmicOffSpill	✅	Log file.
ceSteps	✅	Log file.
ceDigi	✅	Log file.
muDauSteps	✅	Log file.
ceMix	✅	Log file.
rootOverlaps	✅	Log file.
g4surfaceCheck	✅	Log file.
trigger	❓	Log file. Return Code 1.
check_cmake	✅	Log file.
FIXME, TODO	➡️	TODO (6) FIXME (3) in 19 files
clang-tidy	➡️	20 errors 138 warnings
whitespace check	➡️	found whitespace errors

N.B. These results were obtained from a build of this Pull Request at 5bb835f after being merged into the base branch at 0ef02a6.

For more information, please check the job page here.
Build artifacts are deleted after 5 days. If this is not desired, select Keep this build forever on the job page.

[rlcee]

Too much to check in any detail. I do have one high-level question, perhaps for @oksuzian or @brownd1978 . So far, I think, we've kept ML training out of the Offline library. I think the idea is generally if it doesn't run in sim or reco, we shouldn't build it here, since everyone builds it all every time. I think MLTrain repo was created along these lines, and maybe some other training is in personal repos. I think the trend to the future is "modularize" more than "gather".

[YongyiBWu]

I want to clarify on locating this repo here rather than in the MLTrain . As a resampler, both input and output of this code is in the sim workflow, and this module is important to run any STM simulations. The model implementation here is fully customized (C++) and depend on no external packages. Outputs are in customized binary files. So in a sense it can be considered as a sim production using ML method. Of course, if you want this to be moved elsewhere, please let me know and we can discuss.

[oksuzian]

PR Review Summary — #1960

Reviewed at head 5bb835f. First pass; no prior reviews on the PR.

Decision

• 🟡 comment only

Scope understood

• A score-based diffusion model resampler for the STM: MachineLearningTools/ScoreBasedDiffusionModel gains SDE/ODE samplers, EMA networks and a versioned checkpoint format, and STMMC gains the configure → train → generate → validate workflow around it, plus a new LaBrTree analyzer.
• The PR body's claim that scope is limited to STMMC and the DM toolkit holds: ScoreBasedDiffusionModel has exactly four consumers, all inside STMMC and all touched here, and nothing else in Offline includes MachineLearningTools. Nothing outside this workflow can regress at merge, which is why the findings below are a comment rather than a change request.
• Production#576 supplies the job fcl and needs this to merge first.

Findings

1. 🟠 [S1] A training run that never converges writes the model file the generate step reads, and the job exits 0

   • Evidence: VDResamplerTrainCommon.hh:2137-2147 logs an mf::LogWarning, then calls model.saveModel(outFile, basisTag) and returns. outFile is s.stage1ModelFile / s.stage2ModelFile / s.allAtOnceModelFile — the exact paths the generator resolves. runTraining returns void and is called from endJob (VDResamplerTrain_module.cc:359, VDResamplerTrainFromRoot_module.cc:360), so art reports success. The checkpoint records basisTag but nothing about convergence.
   • The message says "Stopping training", and that is not what happens. trainLoop is a lambda declared at :1837 and called three times (:2190 stage 1, :2221 stage 2, :2240 all-at-once); returning from it drops control back into runTraining, which proceeds to the if (s.trainStage2) block and trains stage 2 against a stage-1 model that never converged.
   • Impact: this is the "a warning is not error handling" case, and the file's own PositionBasis comment block records that it has already happened once — runs "were later found to have stopped with a training loss still above 0.99, i.e. undertrained" after being mistaken for a coordinate bug. One LogWarning among thousands of per-epoch LogInfo lines in a grid log is not a failure signal.
   • Suggested fix: throw, or write the non-converged state to <base>.notConverged.dat and not to outFile, so the downstream step fails on a missing input rather than succeeding on a bad one. Either way, correct the "Stopping training" wording or make it true.

2. 🟠 [S1] VDResamplerGenerateMix takes the VD geometry from its own fcl defaults, contradicting the invariant the PR states elsewhere

   • Evidence: VDResamplerConfigure_module.cc:74-76 says "VirtualDetectorID (and VDz0/VDr) are NOT module parameters: they come from the training plan's common_training_config, the single source of truth", and VDResamplerRegeneratePlots_module.cc:273-277 honours that, reading common.get<double>("VDr") and common.get<double>("VDz0"). VDResamplerGenerateMix_module.cc:197-206 instead declares them as fhicl::Atom<double> with hardcoded defaults 37700.39 and 2000.0, and cross-checks only the VD id, against the summary file name (:626).
   • Impact: GenerateMix is the one module that emits physics. VDr is baked into invertPosition and invertMomentumV2FromPtot through InverseParams (:856), so a campaign trained with a different VDr and generated from a fcl that omits the key produces every position and every momentum slope wrong by that ratio, with rc=0 and no warning.
   • Suggested fix: give GenerateMix the same common_training_config read RegeneratePlots has — it already loads trainingPlanFile_ for the peak tags at :374-378 — and drop the three atoms. Persisting VDr/VDz0 in the checkpoint next to basisTag and throwing on disagreement would be stronger still.

3. 🟠 [S1] loadPDGIdFromFileName searches the whole path, so a directory name can decide which particle is generated

   • Evidence: VDResamplerGenerateFromModel_module.cc:64 is const size_t pdgPos = fileName.find("pdg"); with no basename strip. The result feeds pdt_->particle(pdgId_).mass() (:498, :532), PDGCode::type(pdgId_) (:540) and the resampler's source filter (:411), and is never cross-checked against the checkpoint.
   • Impact: a model at .../train_pdg22/nts.mu2e.STMVDResamplerModel_VD116_pdg2112_stage2....dat yields pdgId_ = 22. The job then emits photons, with photon masses, drawn from a neutron model, and selects photons out of the mother file. Grid output areas are commonly named after the training job, so this is not a contrived path.
   • Suggested fix: strip the directory before the search — fileName.substr(fileName.find_last_of("/\\") + 1) — and, better, store pdgId in the checkpoint alongside basisTag and verify it.

4. 🟠 [S1] Nine SetBranchAddress return values are ignored over uninitialized locals

   • Evidence: VDResamplerTrainFromRoot_module.cc:314-325 declares double time, x, y, z, px, py, pz; int stepPdgId; ULong64_t vdId; and binds each with an unchecked ttree->SetBranchAddress(...). The same pattern is in VDResamplerPtotResampler.hh:75-86. CI agrees independently: clang-tidy on this head reports variable 'time' is not initialized and eight siblings at exactly VDResamplerTrainFromRoot_module.cc:314 and :316.
   • Impact: SetBranchAddress returns kMissingBranch and binds nothing when a branch is absent, so GetEntry(i) leaves the locals as stack garbage and the selection at :329 compares it. Point either module at a dump written by an older Configure, or at VirtualDetectorTree output, and the model trains on unrelated numbers. It is also formally undefined behaviour.
   • Suggested fix: check each return and throw naming the missing branch, and initialize the locals at declaration.

5. 🟠 [S1] Every generated training fcl ends with an unconditional InputRootFile : @nil, and on the art path that key does not exist

   • Evidence: VDResamplerConfigureCommon.hh:509-511 emits the resolved InputRootFile : "<path>" inside the analyzer block when the plan sets it, and then :578 unconditionally appends physics.analyzers.<moduleName>.InputRootFile : @nil as a trailing override that, by the comment's own reasoning at :575-577, wins. So the inputRootFileSet branch can never survive into the emitted file. Separately, trainModule is VDResamplerTrain when trainingFromROOTFile is false (:474), and VDResamplerTrain_module.cc contains no InputRootFile parameter at all — zero occurrences — while using art::EDAnalyzer::Table<Config> (:158), which rejects unrecognised keys.
   • Impact: the art path is a documented, required plan option (trainingFromROOTFile is on the required-key list at :408, and the plan has an "ART-source settings" section at VDResamplerTrainingPlan.fcl:63), and every fcl generated for it fails configuration validation. On the ROOT path the placeholder appears to be deliberate, but it silently discards a value the plan supplied.
   • Suggested fix: emit the trailing @nil only when the plan left it unset and the path is ROOT; emit the resolved value otherwise; and emit nothing on the art path.

6. 🟠 [S1] The axis-report loop indexes a six-element array with slot 6

   • Evidence: kRhoTransSlot = 6 (VDResamplerValidationPlots.hh:281) and using TransformedStatsBySlot = std::array<TransformedDimStats, 6> (VDResamplerGenerateCommon.hh:75), so valid indices are 0-5. rhoTrans is pushed with transformed = true and slot = kRhoTransSlot at :428, and the loop at :751-753 does const TransformedDimStats& st = (*stats)[s.slot]; with no bound check. The comment at :280 even notes the slot sits outside the array's range.
   • Impact: reads past the end of a caller-stack local on any job with doValidationPlots and model statistics — the normal path. st.valid is whatever follows, so the log prints either a fabricated statistics line or the right fallback by luck; under _GLIBCXX_ASSERTIONS it aborts. The sizing loop at :416-418 avoids this deliberately; the reporting loop added later did not get the same guard.
   • Suggested fix: if (s.slot >= static_cast<int>(stats->size())) continue; at the top of the loop body.

7. 🟠 [S1] An installed Offline fcl includes a Production fcl it does not use

   • Evidence: STMMC/fcl/VDResamplerTrainingPlan.fcl:34-35 includes Offline/fcl/standardServices.fcl and Production/JobConfig/pileup/STM/prolog.fcl. The file defines no process_name, source, services or physics — it is a pure data file read through ParameterSetFromFile — and it contains zero @local:: or @sequence:: references, so neither include contributes anything.
   • Impact: the file is installed (CMakeLists.txt:161) and is the default trainingPlanFile for VDResamplerReconfigure.fcl:36 and VDResamplerRegeneratePlots.fcl:55. In an Offline-only working area, with Production not on FHICL_FILE_PATH, those jobs and VDResamplerConfigure die at construction on an unresolvable include rather than on anything about the plan. It also inverts the repo dependency: Production depends on Offline, not the reverse.
   • Suggested fix: delete both include lines.

8. 🟠 [S1] The basis tag does not cover the constants the bases are built from, and an unknown basis falls through to V1

   • Evidence: packBasisTag (VDResamplerTransforms.hh:292-298) encodes only the layout and the two basis enums. The numeric parameters those enums carry are compile-time constants whose comments invite tuning — kUrSlopeScale/kUphiSlopeScale at :51-52 ("so each can be tuned to its own spread"), kTBulkCenter/kTTailScale at :63-64 ("kTTailScale sets the asinh width"), and kP0 at :19 ("tunable momentum scale"). Retune one and every previously trained model is decoded with a different map while checkModelLayout, the peak-tag check and the layout check all still pass. Separately, unpackPositionBasis and unpackMomentumBasis (:302-307) are unchecked static_casts, and both radialForward (:386-389) and radialInverse (:400-403) end in case V1_Atanh: default:, so an out-of-range value silently gets the V1 map. Layout is guarded (VDResamplerGenerateFromModel_module.cc:100); the two bases are not.
   • Impact: kUrSlopeScale changing from 0.05 to 0.02 makes every generated slope 2.5x too small, which propagates into pz = pTot/sqrt(1+ur²+uphi²), with no error at any level. And adding a V4 enumerator later compiles clean and behaves as V1.
   • Suggested fix: bump the enumerator whenever a scale changes, or persist the scales in the checkpoint and compare on load; and drop the default: arms so -Wswitch catches a new enumerator at compile time, adding an explicit range check where the tag is unpacked.

9. 🟠 [S1] forwardTransformSampleV2 uses two different values of pz, and the instrumented one is the harmless half

   • Evidence: VDResamplerTransforms.hh:601 calls extrapolateAndCenter(x, y, z, px, py, pz, ...) with the raw pz, which divides by it unguarded at :364. Eight lines later, :608-612 floors the same pz at kPzSafetyEpsilon for the slope division and records the fallback through PzFallbackStats.
   • Impact: both training entry points cut pz <= 0 (VDResamplerTrain_module.cc:322, VDResamplerTrainFromRoot_module.cc:329), but pz = 1e-30 passes. The extrapolation factor goes to ~1e30, the radius with it, and forwardPosition's rho = std::min(rho, 1.0 - kRhoClampEpsilon) at :413-415 then silently relocates the hit to the detector rim — that std::min is doing two jobs, the documented rho == 1 guard and an undocumented out-of-detector truncation. The endJob warning meanwhile reports only the slope-division fallback. PzFallbackStats is exactly the right pattern; it is simply attached to the wrong half.
   • Suggested fix: use pzSafe for the extrapolation too, and give the radial clamp its own counter so an out-of-radius hit produces a summary warning the way the pz fallback does.

10. 🟡 [S2] Smaller items, each with a one-line fix

• No output stream is ever checked after writing. saveModel (ScoreBasedDiffusionModel.cc:1703-1841) and saveModelCsv (:1848-2052) check only that the open succeeded and never test out.good(); VDResamplerConfigure_module.cc:257-315 does the same for the hit summary. A grid worker hitting quota mid-write produces a truncated checkpoint or summary and exits 0.
• The binary loader omits three shape checks the CSV loader performs. numLayers (ScoreBasedDiffusionModel.cc:2207) is never compared to layers (:2159), so the copy loops at :2345 and :2373 index loadedNetwork/loadedEmaNetwork out of bounds when a file disagrees. The CSV path checks all of this.
• Four modules construct mu2e::ParameterSetFromFile (VDResamplerConfigure_module.cc:138, VDResamplerReconfigure_module.cc:107, VDResamplerGenerateMix_module.cc:376, VDResamplerRegeneratePlots_module.cc:264) but only VDResamplerRegeneratePlots declares Offline::GeneralUtilities in CMakeLists.txt. The other three link through a transitive edge, which the coding standard rules out. src/SConscript:51 does declare it, so the two build systems disagree.
• VDResamplerConfigureCommon.hh:572-573 seeds the generated fcls from std::time(nullptr). All of them are written in one endJob, so every particle in a source gets the same baseSeed — the opposite of the comment's stated intent — and re-running Configure on the same summary produces a different set of fcls, so a training result cannot be reproduced from its inputs. Deriving the seed from (runNumber, sourceIndex, pdg, VD id) gives both properties.

11. ⚪ [S3] Collected

• Physics-affecting fcl parameters carry C++ defaults, which Mu2e reserves for verbosityLevel/diagLevel: pdgID defaults to 22 and VirtualDetectorID/VDz0/VDr to 116/37700.39/2000.0 in VDResamplerTrain_module.cc:52-57 and the same set in VDResamplerTrainFromRoot, VDResamplerGenerateFromModel and VDResamplerGenerateMix. A hand-written training fcl that omits pdgID trains photons and writes them under whatever model name was given.
• STMMC/fcl/MakeTree.fcl:43-48 keeps the superseded Stage2LaBr block commented out directly above its replacement. The explanatory comment at :40-42 is worth keeping; the dead block is not.
• LaBrTree_module.cc:56 and :67 carry an unused ParticleDataList.hh include and an unused VolumeId_type typedef, both inherited from HPGeTree_module.cc.
• VDResamplerConfigure_module.cc:234 uses mf::LogInfo log("Virtual Detector Resampler Training Configuration Summary"). A message-facility category is a fcl routing key, so one containing spaces cannot be addressed in destinations.*.categories.

Verified

• 🟢 The blast radius is genuinely confined. ScoreBasedDiffusionModel is included by exactly four files, all in STMMC and all in this PR; a repository-wide search for MachineLearningTools returns only this package, STMMC, and the two build files. Nothing outside this workflow changes behaviour at merge.
• 🟢 MachineLearningTools/CMakeLists.txt needs no edit: its CLHEP / MF_MessageLogger / cetlib_except list still covers every non-standard include in the file, and <limits> is the only include this PR adds.
• 🟢 All ten *_module.cc files are registered in CMakeLists.txt, including the new LaBrTree (:47), VDResamplerReconfigure (:102), VDResamplerTrainFromRoot (:118) and VDResamplerRegeneratePlots (:150); src/SConscript globs *_module.cc through helper.make_plugins, so nothing is missing there.
• 🟢 Every module uses validated FHiCL — art::ED{Analyzer,Producer}::Table<Config> at VDResamplerConfigure:84, GenerateFromModel:225, GenerateMix:223, Reconfigure:76, RegeneratePlots:216, Train:158, TrainFromRoot:153. None is left on a bare ParameterSet.
• 🟢 Random numbers are handled correctly. Every engine comes from createEngine(ServiceHandle<SeedService>()->getSeed()), and a search across all nine module files for std::rand, srand, mt19937, TRandom and gRandom returns nothing, so two grid jobs cannot produce identical samples. (The seed problem in finding 10 is in the generated fcl, not the engines.)
• 🟢 No forbidden handle caching: no art::Handle/ValidHandle members, no GeomHandle/ConditionsHandle/ProditionsHandle anywhere, products fetched per event through ProductToken. The cached GlobalConstantsHandle<ParticleDataList> members are the permitted service-handle case.
• 🟢 normalizeData does guard a zero standard deviation — ScoreBasedDiffusionModel.cc:559-561 throws a cet::exception rather than dividing — so a constant coordinate cannot poison the whole training set.
• 🟢 The mixing weights in GenerateMix are normalized correctly, and disabling a source renormalizes the survivors: referencePots_ is the maximum over enabled sources only (:527-546) and totalSourceWeight_ sums only enabled entries, so the joint probability is proportional to each species' per-POT yield. The absolute-rate consequence of disabling a source is stated in the fcl comment (:118-124) and quantified at endJob by reportPotEquivalence (:926-976).
• 🟢 LaBrTree is not a careless copy of HPGeTree on the point that matters most. HPGeTree_module.cc:180-182 keeps x <= xBeamCentre for HPGe and x >= xBeamCentre for LaBr; LaBrTree_module.cc:193 keeps x >= xBeamCentre, matching the donor's LaBr branch rather than its HPGe branch. The Detector parameter and its validation vector are correctly dropped. Its energy limitation is documented rather than hidden — a FIXME block at :11-27, a constructor LogWarning at :113-117, an inline FIXME at :204 and an endJob reminder at :251-252.
• 🟢 The HPGeTree change is a real fix, not a workaround. The old while (std::find(..., parent->parent()->id()) ...) dereferenced an unchecked art::Ptr; :121-124 and :132-133 add null and availability guards, the early return yields the same answer the next test would have, and nNoParent is reported at endJob (:240-242) rather than swallowed.
• 🟢 Every VD id used is 116 or below (VirtualDetectorTree, the plan's VirtualDetectorID : 116, and the 88/89/90/100/101/116 list in the Production job), so none of this is exposed to the VirtualDetectorId renumbering at v13_36_00 where 117-135 silently mislabel.
• 🟢 The trigger test's :question: / return code 1 is not this PR's doing — it shows the same result on PRs #1955, #1957 and #1949.

Validation check

• Build/tests run: not by me. FNALbuild is green at 5bb835f for the build and every physics job. Three informational results are worth acting on: the whitespace check reports trailing whitespace at ScoreBasedDiffusionModel.cc:22 and on twelve lines of VDResamplerTrainingPlan.fcl, plus VDResamplerReconfigure.fcl with no terminal newline; clang-tidy reports 20 errors and 138 warnings, largely the uninitialized locals in finding 4; and LaBrTree_module.cc carries 3 FIXME and 1 TODO, which are the honest documentation noted above rather than leftovers.
• Config contract check: partial. Every plan key was resolved against the fhicl::Name declarations in both train modules, and the four-source versionTag list matches dataSourceNames() in order and length. No fhicl-dump was produced and no job was run.
• Cross-repo consistency: merge-order note only. Production#576 supplies the job fcl and depends on Offline/STMMC/fcl/VDResamplerTrainingPlan.fcl, which this PR adds, so this merges first. I reviewed #576 separately.

Not checked

• Nothing was built, run or profiled; every finding above is static reading.
• Whether the model actually reproduces the STM VD116 distributions. No validation output, loss curve or comparison plot is attached to the PR, and that is the evidence this change ultimately rests on.
• Whether VDz0 = 37700.39 and VDr = 2000.0 are the current correct values for VD116 — I did not consult a geometry file. Findings 2 and 11 are about where the numbers live, not about the numbers being wrong.
• ParticleDataList::particle(int) behaviour for an unknown pdgId. If it returns a default rather than throwing, the six pdt_->particle(pdgId).mass() call sites are silent-degradation points feeding a wrong mass into E = sqrt(p²+m²) - m. One line for you to confirm.
• The empirical claims embedded in the PositionBasis comment block — the measured coreFrac, the 0/100000 clamp measurement, the per-band u-range table. The analytic entries I spot-checked are right; the measured ones I cannot reproduce from source.

Author follow-ups

1. Decide what a non-converged training run should do (finding 1); it is the one that can quietly put a bad model into a production campaign.
2. Make GenerateMix read the VD geometry from common_training_config like the other modules (finding 2).
3. Fix the two path/branch parsing bugs — the basename strip in loadPDGIdFromFileName, and the checked SetBranchAddress returns over initialized locals (findings 3 and 4).
4. Fix the generated-fcl emission so the art path does not write a key its module lacks (finding 5), and drop the two unused includes from VDResamplerTrainingPlan.fcl (finding 7).
5. Attach whatever validation output you have — a loss curve and one generated-vs-source comparison for a single species would let a reviewer judge the physics rather than only the plumbing.

[oksuzian]

This is a massive PR - 12,000 lines of code in a single PR.
It's impossible to review by a human, and even hard to absorb by AI.
I'm not sure if some portions of the PR belong to Offline at all.
For example, the following is a pure analysis/validation:
STMMC/inc/VDResamplerValidationPlots.hh
...and it's a very large file by itself.
I would consider splitting this PR into several, and consider moving non essential (validation) components somewhere else.