From b792de098cd3a28ff30ca1b9adafe89aafea159f Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2egpvm03.fnal.gov>
Date: Sat, 14 Mar 2026 02:17:24 -0500
Subject: [PATCH 01/15] Modules of ML diffusion model first commit

---
 .../inc/ScoreBasedDiffusionModel.hh           | 181 ++++++++++++++++++
 STMMC/src/VDResamplerTrain_module.cc          | 116 +++++++++++
 2 files changed, 297 insertions(+)
 create mode 100644 MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
 create mode 100644 STMMC/src/VDResamplerTrain_module.cc

diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
new file mode 100644
index 0000000000..b2416fd84d
--- /dev/null
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -0,0 +1,181 @@
+// Module for training and using score-based diffusion model
+// Added by Yongyi Wu
+// Mar. 2026
+
+#pragma once
+
+#include <vector>
+#include <random>
+#include <cmath>
+#include "CLHEP/Random/RandomEngine.h"
+#include "CLHEP/Random/RandFlat.h"
+#include "CLHEP/Random/RandGaussQ.h"
+
+namespace mu2e{
+    struct DiffusionTrainingSample {
+        std::vector<double> x; // transformed state vector (size = dim)
+    };
+
+    class ScoreBasedDiffusionModel {
+    public:
+        // Constructor: Initialize diffusion model with random engine.
+        //
+        // Parameters:
+        //   engine        - Reference to CLHEP random engine (externally managed)
+        //   dim           - Dimensionality of the state space
+        //   hidden        - Size of hidden layers in the neural network
+        //   layers        - Number of layers in the network
+        //   betaMin       - Minimum noise schedule parameter
+        //   betaMax       - Maximum noise schedule parameter
+        //   diffusionSteps - Number of steps in the diffusion process
+        ScoreBasedDiffusionModel(
+            CLHEP::HepRandomEngine& engine,
+            int dim,
+            int hidden,
+            int layers,
+            double betaMin,
+            double betaMax,
+            int diffusionSteps
+        );
+
+        // Train the score network on a batch of samples.
+        // Uses random sampling and noise injection via the external engine.
+        //
+        // Parameters:
+        //   data         - Training samples (transformed state vectors)
+        //   epochs       - Number of training epochs
+        //   batchSize    - Samples per batch (should divide evenly into data.size())
+        //   learningRate - Gradient descent step size
+        void train(
+            const std::vector<DiffusionTrainingSample>& data,
+            int epochs,
+            int batchSize,
+            double learningRate
+        );
+
+        // Generate a new sample from the diffusion model via reverse process.
+        // Uses the external random engine for noise generation during sampling.
+        //
+        // Returns: A generated sample vector of dimension dim_
+        std::vector<double> generateSample();
+
+    private:
+
+        // ----- network -----
+
+        // Represents a single fully-connected layer with weights and biases.
+        struct Layer {
+            std::vector<std::vector<double>> W; // Weight matrix [output_size][input_size]
+            std::vector<double> b;              // Bias vector [output_size]
+        };
+
+        std::vector<Layer> network_; // Network layers in forward order
+
+        // Forward pass through the network.
+        // Computes network output given input state.
+        //
+        // Parameters:
+        //   x - Input vector of dimension dim_
+        //
+        // Returns: Output vector (dimension depends on network architecture)
+        std::vector<double> forward(
+            const std::vector<double>& x
+        );
+
+        // Backward pass for gradient computation.
+        // Computes gradients w.r.t. network parameters from output gradients.
+        //
+        // Parameters:
+        //   gradOutput - Gradient of loss w.r.t. network output
+        void backward(
+            const std::vector<double>& gradOutput
+        );
+
+        // Update network weights using computed gradients (SGD).
+        // Applied after backward pass.
+        //
+        // Parameters:
+        //   lr - Learning rate for gradient descent step
+        void updateWeights(double lr);
+
+        // ----- diffusion -----
+
+        // Noise schedule parameter beta(t) over diffusion time [0,1].
+        // Linear interpolation between betaMin_ and betaMax_.
+        //
+        // Parameters:
+        //   t - Diffusion time parameter in [0,1]
+        //
+        // Returns: Beta value for the given time step
+        double beta(double t) const;
+
+        // Standard deviation of noise at diffusion time t.
+        // Related to the noise schedule via sigma(t) = sqrt(1 - exp(-integral(beta(s) ds))).
+        //
+        // Parameters:
+        //   t - Diffusion time parameter in [0,1]
+        //
+        // Returns: Noise standard deviation at time t
+        double sigma(double t) const;
+
+        // Add Gaussian noise to a state vector at diffusion time t.
+        // Uses external engine (randGaussQ_) for reproducible noise generation.
+        // Noise sample is stored in eps for later use in training.
+        //
+        // Parameters:
+        //   x   - Original state vector
+        //   t   - Diffusion time parameter in [0,1]
+        //   eps - Output: Gaussian noise vector used for perturbation (size = dim_)
+        //
+        // Returns: Noisy state = x + sigma(t) * eps
+        std::vector<double> addNoise(
+            const std::vector<double>& x,
+            double t,
+            std::vector<double>& eps
+        );
+
+        // ----- loss -----
+
+        // Compute Mean Squared Error between predicted score and target score.
+        // Used during training to measure model performance.
+        //
+        // Parameters:
+        //   score  - Model's predicted score vector
+        //   target - Ground truth target score vector
+        //
+        // Returns: MSE loss value (scalar)
+        double computeLoss(
+            const std::vector<double>& score,
+            const std::vector<double>& target
+        ) const;
+
+        // ----- internal vars -----
+        // Model hyperparameters
+        int dim_;           // Dimensionality of state space
+        int hidden_;        // Hidden layer size
+        int layers_;        // Number of network layers
+
+        // Noise schedule parameters (beta(t) = betaMin + t*(betaMax - betaMin))
+        double betaMin_;    // Beta value at t=0
+        double betaMax_;    // Beta value at t=1
+
+        // Diffusion process discretization
+        int diffusionSteps_;  // Number of time steps for reverse process
+
+        // The random engine is NOT owned by this class. It is injected externally
+        // by the framework. Engine reference must remain valid for the lifetime
+        // of this object.
+        CLHEP::HepRandomEngine& engine_;
+
+        // CLHEP distribution wrappers for actual random number generation.
+        // These wrap the engine_ and provide specific probability distributions.
+        // - RandFlat:   Uniform distribution on [0,1)
+        // - RandGaussQ: Gaussian (normal) distribution with mean=0, sigma=1 (or custom)
+        // Both are initialized in the constructor with the injected engine_.
+        CLHEP::RandFlat   randFlat_;      // Used for uniform sampling (e.g., batch selection)
+        CLHEP::RandGaussQ randGaussQ_;    // Used for Gaussian noise in diffusion process
+
+        // Training state
+        double runningLoss_;  // Accumulated loss for monitoring during training
+    };
+}
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
new file mode 100644
index 0000000000..72f5fdd94d
--- /dev/null
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -0,0 +1,116 @@
+// stdlib includes
+#include <cmath>
+#include <iostream>
+
+// art includes
+#include "art/Framework/Core/EDAnalyzer.h"
+#include "art/Framework/Principal/Event.h"
+#include "art/Framework/Principal/Handle.h"
+#include "art/Framework/Principal/Run.h"
+
+// exception handling
+#include "cetlib_except/exception.h"
+
+// fhicl includes
+#include "canvas/Utilities/InputTag.h"
+#include "fhiclcpp/types/Atom.h"
+
+// message handling
+#include "messagefacility/MessageLogger/MessageLogger.h"
+
+// Offline includes
+#include "Offline/GlobalConstantsService/inc/GlobalConstantsHandle.hh"
+#include "Offline/GlobalConstantsService/inc/ParticleDataList.hh"
+#include "Offline/MCDataProducts/inc/SimParticle.hh"
+#include "Offline/MCDataProducts/inc/StepPointMC.hh"
+
+// ROOT includes
+#include "art_root_io/TFileService.h"
+#include "TTree.h"
+
+namespace mu2e {
+  class VDResamplerTrain : public art::EDAnalyzer {
+    public:
+      using Name=fhicl::Name;
+      using Comment=fhicl::Comment;
+      struct Config {
+        fhicl::Atom<art::InputTag> StepPointMCsTag{Name("StepPointMCsTag"), Comment("Tag identifying the StepPointMCs")};
+        fhicl::Atom<art::InputTag> SimParticlemvTag{Name("SimParticlemvTag"), Comment("Tag identifying the SimParticlemv")};
+      };
+      using Parameters = art::EDAnalyzer::Table<Config>;
+      explicit VDResamplerTrain(const Parameters& conf);
+      void analyze(const art::Event& e);
+      void endJob();
+    private:
+      art::ProductToken<StepPointMCCollection> StepPointMCsToken;
+      art::ProductToken<SimParticleCollection> SimParticlemvToken;
+      GlobalConstantsHandle<ParticleDataList> pdt;
+
+      // Declare relevant variables
+
+
+      int pdgId = 0;
+      double x = 0.0, y = 0.0, z = 0.0, mass = 0.0, E = 0.0, time = 0.0;
+      VolumeId_type virtualdetectorId = 0;
+      TTree* ttree;
+      std::map<int, int> pdgIds; // <id, count>
+
+
+
+
+  };
+
+  VirtualDetectorTree::VirtualDetectorTree(const Parameters& conf) :
+    art::EDAnalyzer(conf),
+    StepPointMCsToken(consumes<StepPointMCCollection>(conf().StepPointMCsTag())),
+    SimParticlemvToken(consumes<SimParticleCollection>(conf().SimParticlemvTag())) {
+      art::ServiceHandle<art::TFileService> tfs;
+      ttree = tfs->make<TTree>( "ttree", "Virtual Detectors ttree");
+      ttree->Branch("time", &time, "time/D"); // ns
+      ttree->Branch("virtualdetectorId", &virtualdetectorId, "virtualdetectorId/l");
+      ttree->Branch("pdgId", &pdgId, "pdgId/I");
+      ttree->Branch("x", &x, "x/D"); // mm
+      ttree->Branch("y", &y, "y/D"); // mm
+      ttree->Branch("z", &z, "z/D"); // mm
+      ttree->Branch("E", &E, "E/D"); // MeV
+    };
+
+  void VirtualDetectorTree::analyze(const art::Event& event) {
+    // Get the data products from the event
+    auto const& StepPointMCs = event.getProduct(StepPointMCsToken);
+    if (StepPointMCs.empty())
+      return;
+    auto const& SimParticles = event.getProduct(SimParticlemvToken);
+    if (SimParticles.empty())
+      return;
+
+    // Loop over all VD hits
+    for (const StepPointMC& step : StepPointMCs) {
+      // Get the associated particle
+      const SimParticle& particle = SimParticles.at(step.trackId());
+
+      // Extract the parameters
+      time = step.time();
+      virtualdetectorId = step.virtualDetectorId();
+      pdgId = particle.pdgId();
+      x = step.position().x();
+      y = step.position().y();
+      z = step.position().z();
+      mass = pdt->particle(pdgId).mass();
+      E = std::sqrt(step.momentum().mag2()+mass*mass)-mass; // Subtract the rest mass
+      if (E < 0)
+        throw cet::exception("LogicError", "Energy is negative");
+      ttree->Fill();
+
+      // Generate the data summary
+      if (pdgIds.find(pdgId) != pdgIds.end())
+        pdgIds[pdgId] += 1;
+      else
+        pdgIds.emplace(std::make_pair(pdgId, 1));
+    };
+    return;
+  };
+
+}; // end namespace mu2e
+
+DEFINE_ART_MODULE(mu2e::VDResamplerTrain)

From 34ed29759cc6972a0e445cac512d13e04da03066 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Tue, 24 Mar 2026 15:05:16 -0500
Subject: [PATCH 02/15] Add score-based diffusion model and resamplers

Introduce a new ScoreBasedDiffusionModel implementation (header + source) providing training, sampling, save/load, Adam/SGD optimizers, linear/cosine noise schedules, gradient clipping, SiLU activations, and Heun/Euler samplers. Add STMMC resampler plugins (VDResamplerConfigure, VDResamplerGenerateFromModel, VDResamplerGenerateMix) and update VDResamplerTrain and STMMC CMakeLists to build the new modules. Register a new generator name STMDownStreamGenTool in GenId.hh to support the downstream generator. The changes wire up random-engine usage and include logging and error checks for robustness.
---
 MCDataProducts/inc/GenId.hh                   |   4 +-
 .../inc/ScoreBasedDiffusionModel.hh           | 216 +++-
 .../src/ScoreBasedDiffusionModel.cc           | 996 ++++++++++++++++++
 STMMC/CMakeLists.txt                          |  36 +
 STMMC/src/VDResamplerConfigure_module.cc      | 256 +++++
 .../VDResamplerGenerateFromModel_module.cc    | 334 ++++++
 STMMC/src/VDResamplerGenerateMix_module.cc    | 495 +++++++++
 STMMC/src/VDResamplerTrain_module.cc          | 342 +++++-
 8 files changed, 2599 insertions(+), 80 deletions(-)
 create mode 100644 MachineLearningTools/src/ScoreBasedDiffusionModel.cc
 create mode 100644 STMMC/src/VDResamplerConfigure_module.cc
 create mode 100644 STMMC/src/VDResamplerGenerateFromModel_module.cc
 create mode 100644 STMMC/src/VDResamplerGenerateMix_module.cc

diff --git a/MCDataProducts/inc/GenId.hh b/MCDataProducts/inc/GenId.hh
index 1c781150a6..8ede29be0a 100644
--- a/MCDataProducts/inc/GenId.hh
+++ b/MCDataProducts/inc/GenId.hh
@@ -44,7 +44,7 @@ namespace mu2e {
       cosmicCRY,  pbarFlat, fromAscii, ExternalRMC, InternalRMC, CeLeadingLog, cosmicCORSIKA, //44
       MuCapProtonGenTool, MuCapDeuteronGenTool, DIOGenTool, MuCapNeutronGenTool, // 48
       MuCapPhotonGenTool, MuCapGammaRayGenTool, CeLeadingLogGenTool, MuplusMichelGenTool,// 52
-      gammaPairProduction, antiproton, Mu2eXGenTool,//55
+      gammaPairProduction, antiproton, Mu2eXGenTool, STMDownStreamGenTool,//56
       lastEnum //56
     };
 
@@ -64,7 +64,7 @@ namespace mu2e {
       "CosmicCRY", "pbarFlat","fromAscii","ExternalRMC","InternalRMC","CeLeadingLog", "CosmicCORSIKA", \
     "MuCapProtonGenTool", "MuCapDeuteronGenTool", "DIOGenTool", "MuCapNeutronGenTool", \
       "MuCapPhotonGenTool", "MuCapGammaRayGenTool","CeLeadingLogGenTool","MuplusMichelGenTool", \
-      "gammaPairProduction", "antiproton", "Mu2eXGenTool"
+      "gammaPairProduction", "antiproton", "Mu2eXGenTool", "STMDownStreamGenTool"
 #endif
 
   public:
diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
index b2416fd84d..ae9afea62b 100644
--- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -5,59 +5,133 @@
 #pragma once
 
 #include <vector>
-#include <random>
+#include <algorithm>
 #include <cmath>
+#include <string>
+#include <iomanip>
+#include <fstream>
+#include <sstream>
+#include <iostream>
+#include <cctype>
+
 #include "CLHEP/Random/RandomEngine.h"
 #include "CLHEP/Random/RandFlat.h"
 #include "CLHEP/Random/RandGaussQ.h"
 
+#include "Offline/SeedService/inc/SeedService.hh"
+
+#include "messagefacility/MessageLogger/MessageLogger.h"
+#include "cetlib_except/exception.h"
+
 namespace mu2e{
     struct DiffusionTrainingSample {
-        std::vector<double> x; // transformed state vector (size = dim)
+        std::vector<double> x;    // transformed state vector to diffuse (size = dim)
+        std::vector<double> cond; // optional conditioning vector (size = conditionDim)
     };
 
     class ScoreBasedDiffusionModel {
     public:
+        // Enumeration for optimizer selection
+        enum class OptimizerType {
+            SGD,   // Stochastic Gradient Descent
+            ADAM   // Adam optimizer
+        };
+
+        // Enumeration for noise schedule selection
+        enum class NoiseScheduleType {
+            LINEAR,  // Linear noise schedule
+            COSINE   // Cosine noise schedule
+        };
+
         // Constructor: Initialize diffusion model with random engine.
         //
         // Parameters:
-        //   engine        - Reference to CLHEP random engine (externally managed)
-        //   dim           - Dimensionality of the state space
-        //   hidden        - Size of hidden layers in the neural network
-        //   layers        - Number of layers in the network
-        //   betaMin       - Minimum noise schedule parameter
-        //   betaMax       - Maximum noise schedule parameter
-        //   diffusionSteps - Number of steps in the diffusion process
+        //   engine                - Reference to CLHEP random engine (externally managed)
+        //   dim                   - Dimensionality of the state space
+        //   conditionDim          - Dimensionality of the optional conditioning vector (default: 0 for unconditional model)
+        //   hidden                - Size of hidden layers in the neural network
+        //   layers                - Number of layers in the network
+        //   optimizerType         - Type of optimizer to use (SGD or ADAM, default: ADAM)
+        //   adamBeta1             - Adam optimizer beta1 parameter (default: 0.9)
+        //   adamBeta2             - Adam optimizer beta2 parameter (default: 0.999)
+        //   adamEps               - Adam optimizer epsilon parameter (default: 1e-8)
+        //   scheduleType          - Type of noise schedule (LINEAR or COSINE, default: COSINE)
+        //   betaMin               - Minimum noise schedule parameter (for LINEAR schedule, default: 1e-4)
+        //   betaMax               - Maximum noise schedule parameter (for LINEAR schedule, default: 0.02)
+        //   cosineOffset          - Offset parameter (for cosine schedule, default: 0.008)
+        //   batchSize             - Batch size for training (default: 32)
+        //   gradientClipThreshold - Threshold for gradient clipping (default: 1.0)
+        //   learningRate          - Learning rate for training (default: 1e-3)
+        //   diffusionSteps        - Number of steps in the diffusion process (default: 200)
         ScoreBasedDiffusionModel(
-            CLHEP::HepRandomEngine& engine,
+            // Network architecture parameters
+            art::RandomNumberGenerator::base_engine_t& engine,
             int dim,
+            int conditionDim,
             int hidden,
             int layers,
-            double betaMin,
-            double betaMax,
-            int diffusionSteps
+            // Optimizer configuration
+            OptimizerType optimizerType = OptimizerType::ADAM,
+            double adamBeta1 = 0.9,
+            double adamBeta2 = 0.999,
+            double adamEps = 1e-8,
+            // Noise schedule configuration
+            NoiseScheduleType scheduleType = NoiseScheduleType::COSINE,
+            double betaMin = 1e-4,
+            double betaMax = 0.02,
+            double cosineOffset = 0.008,
+            // Training configuration
+            int batchSize = 32,
+            double gradientClipThreshold = 1.0,
+            double learningRate = 1e-3,
+            // Diffusion process configuration
+            int diffusionSteps = 200
         );
 
         // Train the score network on a batch of samples.
         // Uses random sampling and noise injection via the external engine.
+        // Note that training needs to occur on all data samples. Training on multiple small subsets
+        // of data and then averaging or aggregating the model parameters is not supported and may lead to 
+        // issues as neural networks are not linear. 
         //
         // Parameters:
         //   data         - Training samples (transformed state vectors)
-        //   epochs       - Number of training epochs
-        //   batchSize    - Samples per batch (should divide evenly into data.size())
-        //   learningRate - Gradient descent step size
+        //   epochs       - Number of training epochs to perform
         void train(
             const std::vector<DiffusionTrainingSample>& data,
-            int epochs,
-            int batchSize,
-            double learningRate
+            int epochs
         );
 
         // Generate a new sample from the diffusion model via reverse process.
         // Uses the external random engine for noise generation during sampling.
         //
+        // Parameters:
+        //   condition       - Optional conditioning vector (must match conditionDim_ when enabled)
+        //   useHeun         - If true, uses Heun's method (2nd order). If false, uses Euler method (2nd order, default)
+        //   diffusionSteps  - Number of diffusion steps for sampling (default: -1 uses the model's configured diffusionSteps_)
+        //
         // Returns: A generated sample vector of dimension dim_
-        std::vector<double> generateSample();
+        std::vector<double> generateSample(
+            const std::vector<double>& condition = {},
+            bool useHeun = true,
+            int diffusionSteps = -1
+        );
+
+        // Save the model parameters to a CSV file with annotations for later use.
+        // Uses a default filename of "DiffusionModel.csv" if not specified.
+        //
+        // Parameters: 
+        //   filename - Path to the CSV file where model parameters will be saved (default: "DiffusionModel.csv")
+        void saveModel(const std::string& filename = "DiffusionModel.csv");
+
+        // Load model parameters from a file to restore a previously trained model.
+        // Note that as the adam optimizer state is not saved, it is not possible to resume training from a loaded 
+        // model and pick up the training process where it left off. The loaded model can only be used for sampling.
+        //
+        // Parameters:
+        //   engine   - Reference to CLHEP random engine (externally managed, must be valid for lifetime of model)
+        //   filename - Path to the file from which model parameters will be loaded
+        static ScoreBasedDiffusionModel loadModel(CLHEP::HepRandomEngine& engine, const std::string& filename);
 
     private:
 
@@ -67,15 +141,35 @@ namespace mu2e{
         struct Layer {
             std::vector<std::vector<double>> W; // Weight matrix [output_size][input_size]
             std::vector<double> b;              // Bias vector [output_size]
+            
+            // storage for gradients during back-propagation
+            std::vector<std::vector<double>> gradW; // gradient of loss w.r.t. weights
+            std::vector<double> gradb;              // gradient of loss w.r.t. biases
+
+            // Adam optimizer state buffers
+            std::vector<std::vector<double>> mW; // First moment estimates for weights, m_t = beta1 m_{t-1} + (1-beta1) g_t
+            std::vector<std::vector<double>> vW; // Second moment estimates for weights, v_t = beta2 v_{t-1} + (1-beta2) g_t^2
+
+            std::vector<double> mb; // First moment estimates for biases
+            std::vector<double> vb; // Second moment estimates for biases
         };
 
         std::vector<Layer> network_; // Network layers in forward order
 
+        // Storage for activations and pre-activations during forward pass (used in back-propagation)
+        // During forward pass, preactivations_[l] = W[l] * activations_[l-1] + b[l], 
+        //                     and activations_[l] = activationFunction(preactivations_[l])
+        // These are needed to compute gradients during the backward pass.
+        std::vector<std::vector<double>> activations_;
+        std::vector<std::vector<double>> preactivations_;
+
         // Forward pass through the network.
-        // Computes network output given input state.
+        // Computes the score function.
+        // input actually contains the state vector, optional condition vector, and the time embedding.
+        // dimension hence is dim_ + conditionDim_ + 1.
         //
         // Parameters:
-        //   x - Input vector of dimension dim_
+        //   x - Input vector of dimension dim_ + conditionDim_ + 1
         //
         // Returns: Output vector (dimension depends on network architecture)
         std::vector<double> forward(
@@ -83,21 +177,31 @@ namespace mu2e{
         );
 
         // Backward pass for gradient computation.
-        // Computes gradients w.r.t. network parameters from output gradients.
+        // Computes gradients w.r.t. network parameters via chain rule.
         //
         // Parameters:
         //   gradOutput - Gradient of loss w.r.t. network output
+        //
+        // Returns: Gradient of loss w.r.t. network parameters (gradW and gradb for each layer)
         void backward(
             const std::vector<double>& gradOutput
         );
 
-        // Update network weights using computed gradients (SGD).
-        // Applied after backward pass.
+        // Update network weights using computed gradients (Stochastic Gradient Descent SGD).
+        // Applied after backward pass. Alternately, an Adam optimizer step can be implemented 
+        // in adamUpdate() for better convergence.
         //
         // Parameters:
         //   lr - Learning rate for gradient descent step
         void updateWeights(double lr);
 
+        // Update network weights using Adam optimization algorithm.
+        // This method would implement the Adam update rule using the stored gradients and moment estimates.
+        //
+        // Parameters:
+        //   lr - Learning rate for Adam update
+        void adamUpdate(double lr);
+
         // ----- diffusion -----
 
         // Noise schedule parameter beta(t) over diffusion time [0,1].
@@ -149,24 +253,19 @@ namespace mu2e{
             const std::vector<double>& target
         ) const;
 
-        // ----- internal vars -----
-        // Model hyperparameters
-        int dim_;           // Dimensionality of state space
-        int hidden_;        // Hidden layer size
-        int layers_;        // Number of network layers
-
-        // Noise schedule parameters (beta(t) = betaMin + t*(betaMax - betaMin))
-        double betaMin_;    // Beta value at t=0
-        double betaMax_;    // Beta value at t=1
-
-        // Diffusion process discretization
-        int diffusionSteps_;  // Number of time steps for reverse process
+        // Clip gradients to prevent exploding gradients during training.
+        //
+        // Parameters:
+        //   maxNorm - Maximum allowed norm for the gradients. If the total norm exceeds this threshold, 
+        //             gradients are scaled down to have norm equal to maxNorm.
+        void clipGradients(double maxNorm);
 
+        // ----- internal vars -----
+        
         // The random engine is NOT owned by this class. It is injected externally
         // by the framework. Engine reference must remain valid for the lifetime
         // of this object.
-        CLHEP::HepRandomEngine& engine_;
-
+        art::RandomNumberGenerator::base_engine_t& engine_;
         // CLHEP distribution wrappers for actual random number generation.
         // These wrap the engine_ and provide specific probability distributions.
         // - RandFlat:   Uniform distribution on [0,1)
@@ -175,7 +274,46 @@ namespace mu2e{
         CLHEP::RandFlat   randFlat_;      // Used for uniform sampling (e.g., batch selection)
         CLHEP::RandGaussQ randGaussQ_;    // Used for Gaussian noise in diffusion process
 
+        // Model hyperparameters
+        int dim_;           // Dimensionality of state space
+        int conditionDim_;  // Dimensionality of the optional conditioning vector
+        int hidden_;        // Hidden layer size
+        int layers_;        // Number of network layers
+
+        // Optimizer configuration
+        OptimizerType optimizerType_;  // Type of optimizer to use (SGD or ADAM)
+
+        // Adam optimizer parameters
+        double adamBeta1_;  // First moment exponential decay rate (default: 0.9)
+        double adamBeta2_;  // Second moment exponential decay rate (default: 0.999)
+        double adamEps_;    // Small constant for numerical stability (default: 1e-8)
+
+        // Noise schedule configuration
+        NoiseScheduleType noiseScheduleType_;  // Type of noise schedule (LINEAR or COSINE)
+
+        // Linear noise schedule parameters (beta(t) = betaMin + t*(betaMax - betaMin))
+        // default betaMin = 1e-4, betaMax = 0.02 are typical values used in diffusion models, but can be tuned for specific applications.
+        double betaMin_;    // Beta value at t=0
+        double betaMax_;    // Beta value at t=1
+
+        // Cosine noise schedule parameter (offset to avoid singularity at t=0 in cosine schedule, default: 0.008)
+        double cosineOffset_;
+
+        // Training configuration
+        int batchSize_;  // Batch size for vectorized training (default: 32)
+        double gradientClipThreshold_;  // Gradient clipping threshold (default: 1.0)
+        double learningRate_; // Learning rate for training (default: 1e-3) 
+
+        // Diffusion process discretization
+        int diffusionSteps_;  // Number of time steps to generate a sample (default: 200)
+
         // Training state
         double runningLoss_;  // Accumulated loss for monitoring during training
+        int adamStep_; // Step counter for Adam optimizer (used to compute bias-corrected moment estimates)
+        int trainingSampleSize_; // Total number of training samples 
+
+        // Container for tracking training loss over epochs
+        std::vector<double> epochLosses_;
+        
     };
 }
diff --git a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
new file mode 100644
index 0000000000..67829da8f2
--- /dev/null
+++ b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
@@ -0,0 +1,996 @@
+#include "Offline/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh"
+
+namespace mu2e {
+
+    // SiLU activation function (Sigmoid Linear Unit)
+    // Defined as: silu(x) = x / (1 + exp(-x))
+    static double silu(double x) {
+        return x / (1.0 + std::exp(-x));
+    }
+
+    // Derivative of SiLU activation function, needed for back-propagation.
+    // Defined as: silu'(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
+    static double siluDeriv(double x) {
+        double s = 1.0 / (1.0 + std::exp(-x));
+        return s * (1.0 + x * (1.0 - s));
+    }
+
+    // Linear interpolation of beta(t) between betaMin_ and betaMax_. 
+    // This is only used in the linear noise schedule.
+    double ScoreBasedDiffusionModel::beta(double t) const {
+        return betaMin_ + t * (betaMax_ - betaMin_);
+    }
+
+    // Standard deviation of noise at diffusion time t.
+    // For the linear schedule, use simply sqrt(beta(t)).
+    // For the cosine schedule, this is derived from the cumulative noise schedule.
+    double ScoreBasedDiffusionModel::sigma(double t) const {
+        if (noiseScheduleType_ == NoiseScheduleType::COSINE) {
+            // Cosine noise schedule
+            double f = (t + cosineOffset_) / (1.0 + cosineOffset_);
+            double alpha_bar = std::cos(f * M_PI * 0.5);
+            alpha_bar *= alpha_bar;
+            return std::sqrt(1.0 - alpha_bar);
+        } else {
+            // Linear noise schedule (default)
+            return std::sqrt(beta(t));
+        }
+    }
+
+    
+    ScoreBasedDiffusionModel::ScoreBasedDiffusionModel(
+        art::RandomNumberGenerator::base_engine_t& engine,
+        int dim,
+        int conditionDim,
+        int hidden,
+        int layers,
+        OptimizerType optimizerType,
+        double adamBeta1,
+        double adamBeta2,
+        double adamEps,
+        NoiseScheduleType scheduleType,
+        double betaMin,
+        double betaMax,
+        double cosineOffset,
+        int batchSize,
+        double gradientClipThreshold,
+        double learningRate,
+        int diffusionSteps
+    ) : engine_(engine), randFlat_(engine_), randGaussQ_(engine_),
+        dim_(dim), conditionDim_(conditionDim), hidden_(hidden), layers_(layers),
+        optimizerType_(optimizerType),
+        adamBeta1_(adamBeta1), adamBeta2_(adamBeta2), adamEps_(adamEps),
+        noiseScheduleType_(scheduleType),
+        betaMin_(betaMin), betaMax_(betaMax), cosineOffset_(cosineOffset),
+        batchSize_(batchSize), gradientClipThreshold_(gradientClipThreshold), learningRate_(learningRate),
+        diffusionSteps_(diffusionSteps),
+        runningLoss_(0.0), adamStep_(0), epochLosses_(), trainingSampleSize_(0) {
+
+        // Validate model dimensions and parameters
+        if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) {
+            throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid model dimensions";
+        }
+        if (diffusionSteps <= 0) {
+            throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid diffusionSteps";
+        }
+
+        // ------------------------------------------------------------
+        // Network architecture
+        //
+        // Input dimension = dim_ + conditionDim_ + 1
+        // (+1 because diffusion time t is appended to the input vector)
+        // ------------------------------------------------------------
+
+        int inputSize = dim_ + conditionDim_ + 1;
+        int in = inputSize;
+
+        // Weight initialization scale (local constant so it can be tuned easily)
+        const double weightInitScale = 0.02;
+        // const double weightInitScale = std::sqrt(2.0 / in); // He initialization for ReLU activations
+
+        for (int l = 0; l < layers_; ++l) {
+
+            // Last layer outputs the score vector (dimension = dim_)
+            int out = (l == layers_ - 1) ? dim_ : hidden_;
+
+            Layer layer;
+
+            // Allocate weights
+            layer.W.resize(out, std::vector<double>(in));
+
+            // Allocate biases
+            layer.b.resize(out);
+
+            // Allocate gradient buffers
+            layer.gradW.resize(out, std::vector<double>(in, 0.0));
+            layer.gradb.resize(out, 0.0);
+
+            // Allocate Adam buffers
+            layer.mW.resize(out, std::vector<double>(in, 0.0));
+            layer.vW.resize(out, std::vector<double>(in, 0.0));
+
+            layer.mb.resize(out, 0.0);
+            layer.vb.resize(out, 0.0);
+
+            // --------------------------------------------------------
+            // Weight initialization
+            //
+            // Small Gaussian initialization improves training stability
+            // --------------------------------------------------------
+
+            for (int i = 0; i < out; ++i) {
+                for (int j = 0; j < in; ++j) {
+                    layer.W[i][j] = weightInitScale * randGaussQ_.fire(); 
+                    // this generates a Gaussian random number with mean=0 and sigma=weightInitScale
+                }
+                layer.b[i] = 0.0;
+            }
+
+            assert (layer.W[0].size() == (size_t)in); // Check that weight matrix has correct input dimension
+            assert (layer.W.size() == (size_t)out);   // Check that weight matrix has correct output dimension
+            assert (layer.b.size() == (size_t)out);  // Check that bias vector has correct dimension
+
+            network_.push_back(layer);
+
+            // Output becomes next layer input
+            in = out;
+        }
+            
+        // Print layer and model configuration
+        std::ostringstream oss;
+        oss << "ScoreBasedDiffusionModel initialized\n"
+            << "Model configuration:\n"
+            << "  Network architecture:\n"
+            << "    | dim=" << dim_ << "\n"
+            << "    | conditionDim=" << conditionDim_ << "\n"
+            << "    | hidden=" << hidden_ << "\n"
+            << "    | layers=" << layers_ << "\n"
+            << "  Optimizer configuration:\n"
+            << "    | Optimizer=" << (optimizerType_ == OptimizerType::ADAM ? "Adam" : "SGD") << "\n";
+        if (optimizerType_ == OptimizerType::ADAM) {
+            oss << "      |- AdamBeta1=" << adamBeta1_ << "\n"
+                << "      |- AdamBeta2=" << adamBeta2_ << "\n"
+                << "      |- AdamEps=" << adamEps_ << "\n";
+        }
+        oss << "  Noise schedule configuration:\n"
+            << "    | NoiseSchedule=" << (noiseScheduleType_ == NoiseScheduleType::COSINE ? "Cosine" : "Linear") << "\n";
+        if (noiseScheduleType_ == NoiseScheduleType::COSINE) {
+            oss << "      |- CosineOffset=" << cosineOffset_ << "\n";
+        } else {
+            oss << "      |- BetaMin=" << betaMin_ << "\n"
+                << "      |- BetaMax=" << betaMax_ << "\n";
+        }
+        oss << "  Training configuration:\n"
+            << "    | BatchSize=" << batchSize_ << "\n"
+            << "    | GradientClipThreshold=" << gradientClipThreshold_ << "\n"
+            << "    | LearningRate=" << learningRate_ << "\n"
+            << "  Diffusion process configuration:\n"
+            << "    | DiffusionSteps=" << diffusionSteps_ << "\n";
+        mf::LogInfo("ScoreBasedDiffusionModel::initialize") << oss.str();
+
+        // Reserve space for forward-pass caches
+        activations_.reserve(layers_ + 1);
+        preactivations_.reserve(layers_);
+    }
+
+    // forward pass to compute the score function given input (state + time embedding)
+    std::vector<double> ScoreBasedDiffusionModel::forward(
+        const std::vector<double>& input
+    )
+    {
+        assert(!network_.empty());
+        assert(!network_[0].W.empty());
+        assert(!network_[0].W[0].empty());
+        assert(input.size() == network_[0].W[0].size());
+
+        activations_.clear();
+        preactivations_.clear();
+
+        // Input vector contains both the state and the time embedding (e.g., concatenated together).
+        std::vector<double> x = input;
+
+        // Cache the input as the activation of layer 0 for use in back-propagation.
+        activations_.push_back(x);
+
+        // Forward pass through the network layers
+        for (size_t l = 0; l < network_.size(); ++l) 
+        {
+            // Compute pre-activation z = W*x + b for the current layer
+            auto& layer = network_[l];
+            std::vector<double> z(layer.W.size());
+            for (size_t i = 0; i < layer.W.size(); ++i) 
+            {
+                double v = layer.b[i];
+                for (size_t j = 0; j < layer.W[i].size(); ++j) 
+                {
+                    v += layer.W[i][j]*x[j];
+                }
+                z[i] = v;
+            }
+
+            // Cache pre-activation for back-propagation
+            preactivations_.push_back(z);
+
+            // Apply activation function (SiLU) to get the output of the current layer.
+            std::vector<double> y(z.size());
+            if(l != network_.size()-1) // No activation on the last layer (output layer), as it predicts the score directly.
+            {
+                for (size_t i=0;i<z.size();++i)
+                    y[i] = silu(z[i]);
+            }
+            else
+                y = z;
+
+            // Cache activation for back-propagation
+            activations_.push_back(y);
+
+            // Output of current layer becomes input to the next layer
+            x = y;
+        }
+
+        return x;
+    }
+
+    // backward pass to compute gradients of the loss w.r.t. network parameters using chain rule
+    void ScoreBasedDiffusionModel::backward(
+        const std::vector<double>& gradOutput
+    )
+    {
+        std::vector<double> grad = gradOutput;
+
+        // Back-propagate through layers in reverse order
+        for (int l = network_.size() - 1; l >= 0; l--) 
+        {
+            auto& layer = network_[l];
+
+            auto& aPrev = activations_[l];
+            auto& z = preactivations_[l];
+
+            std::vector<double> gradZ(grad.size());
+
+            // For the output layer, the gradient is directly from the loss w.r.t. output (no activation function).
+            if(l == network_.size()-1)
+            {
+                gradZ = grad;
+            }
+            // For hidden layers, apply the derivative of the activation function (SiLU) to the gradient.
+            else
+            {
+                for (size_t i=0;i<grad.size();++i)
+                    gradZ[i] = grad[i]*siluDeriv(z[i]);
+            }
+
+            // Compute gradients w.r.t. weights and biases for the current layer
+            for (size_t i = 0; i < layer.W.size(); i++) 
+            {
+                for (size_t j=0;j<layer.W[i].size();++j)
+                {
+                    // Gradient of loss w.r.t. weight W[i][j]
+                    layer.gradW[i][j] += gradZ[i]*aPrev[j];
+                }
+
+                // Gradient of loss w.r.t. bias b[i]
+                // Note: We accumulate gradients in gradW and gradb because we will apply the optimizer step 
+                // after processing a batch of samples.
+                layer.gradb[i] += gradZ[i];
+            }
+
+            assert(!layer.W.empty() && !layer.W[0].empty());
+            // Compute gradient w.r.t. input of the current layer to propagate to the previous layer
+            std::vector<double> gradPrev(layer.W[0].size(),0.0);
+
+            for (size_t j=0;j<gradPrev.size();++j)
+                for (size_t i=0;i<layer.W.size();++i)
+                    gradPrev[j] += layer.W[i][j]*gradZ[i];
+
+            grad = gradPrev;
+        }
+    }
+
+    // Update network weights using computed gradients (Stochastic Gradient Descent).
+    void ScoreBasedDiffusionModel::updateWeights(
+        double lr
+    )
+    {
+        for (auto& layer : network_)
+        {
+            assert(!layer.W.empty() && !layer.W[0].empty());
+            assert(!layer.b.empty());
+
+            size_t outSize = layer.W.size();
+            size_t inSize  = layer.W[0].size();
+
+            // Update weights
+            for (size_t i = 0; i < outSize; ++i)
+            {
+                for (size_t j = 0; j < inSize; ++j)
+                {
+                    // Update weight using SGD: W = W - lr * gradW
+                    layer.W[i][j] -= lr * layer.gradW[i][j];
+
+                    // reset gradient
+                    layer.gradW[i][j] = 0.0;
+                }
+
+                // Update bias
+                layer.b[i] -= lr * layer.gradb[i];
+
+                // reset gradient
+                layer.gradb[i] = 0.0;
+            }
+        }
+    }
+
+    // Update network weights using Adam optimization algorithm.
+    // Adam optimizer maintains running estimates of the first and second moments of the gradients (m and v) for each parameter,
+    // and uses these to adapt the learning rate for each parameter individually.
+    // The update rule for a parameter theta with gradient g at time step t is:
+    // m_t = b1 * m_{t-1} + (1 - b1) * g
+    // v_t = b2 * v_{t-1} + (1 - b2) * g^2
+    // mhat = m_t / (1 - b1^t)  (bias-corrected first moment estimate)
+    // vhat = v_t / (1 - b2^t)  (bias-corrected second moment estimate)
+    // theta = theta - lr * mhat / (sqrt(vhat) + eps) 
+    void ScoreBasedDiffusionModel::adamUpdate(
+        double lr
+    )
+    {
+        adamStep_++; // initialize to 0 in constructor, 
+                     // increment first to avoid division by zero in the first step when computing bias-corrected estimates 
+                     // (1 - b1^t) and (1 - b2^t). Do not modify this to increment after the update.
+
+        for (auto& layer : network_)
+        {
+            assert(!layer.W.empty() && !layer.W[0].empty());
+            assert(!layer.b.empty());
+
+            for (size_t i=0;i<layer.W.size();++i)
+            {
+                // Update weights using Adam update rule
+                for (size_t j=0;j<layer.W[i].size();++j)
+                {
+                    double g = layer.gradW[i][j];
+
+                    layer.mW[i][j] = adamBeta1_*layer.mW[i][j] + (1-adamBeta1_)*g;
+                    layer.vW[i][j] = adamBeta2_*layer.vW[i][j] + (1-adamBeta2_)*g*g;
+
+                    double mhat = layer.mW[i][j]/(1-std::pow(adamBeta1_,adamStep_));
+                    double vhat = layer.vW[i][j]/(1-std::pow(adamBeta2_,adamStep_));
+
+                    layer.W[i][j] -= lr*mhat/(std::sqrt(vhat)+adamEps_);
+
+                    layer.gradW[i][j] = 0.0;
+                }
+
+                // Update bias using Adam update rule
+                double g = layer.gradb[i];
+
+                layer.mb[i] = adamBeta1_*layer.mb[i] + (1-adamBeta1_)*g;
+                layer.vb[i] = adamBeta2_*layer.vb[i] + (1-adamBeta2_)*g*g;
+
+                double mhat = layer.mb[i]/(1-std::pow(adamBeta1_,adamStep_));
+                double vhat = layer.vb[i]/(1-std::pow(adamBeta2_,adamStep_));
+
+                layer.b[i] -= lr*mhat/(std::sqrt(vhat)+adamEps_);
+
+                layer.gradb[i] = 0.0;
+            }
+        }
+    }
+
+    // Add noise to the input state x at diffusion time t to produce the noisy sample xt.
+    std::vector<double> ScoreBasedDiffusionModel::addNoise(
+        const std::vector<double>& x,
+        double t, // Diffusion time parameter in [0,1]
+        std::vector<double>& eps 
+    ){
+        assert(x.size() == (size_t)dim_);
+
+        double s = sigma(t);
+
+        // Generate Gaussian noise vector eps of dimension dim_ using the external random engine.
+        eps.resize(dim_);
+
+        std::vector<double> xt(dim_);
+
+        for (int i = 0; i < dim_; ++i) {
+            eps[i] = randGaussQ_.fire();   // Gaussian N(0,1)
+            xt[i]  = x[i] + s * eps[i];
+        }
+
+        return xt;
+    }
+
+    // Compute Mean Squared Error per dimension between predicted score and target score.
+    double ScoreBasedDiffusionModel::computeLoss(
+        const std::vector<double>& score,
+        const std::vector<double>& target
+    ) const{
+
+        double loss = 0.0;
+
+        for (int i=0;i<dim_;++i){
+            double d = score[i]-target[i];
+            loss += d*d;
+        }
+
+        return loss/dim_;
+    }
+
+    // Clip gradients to prevent exploding gradients during training. This is done by scaling down 
+    // the gradients if their L2 norm exceeds a specified threshold.
+    void ScoreBasedDiffusionModel::clipGradients(
+        double maxNorm
+    )
+    {
+        double norm = 0.0;
+
+        for (auto& layer : network_)
+        {
+            for (auto& row : layer.gradW)
+                for (double g : row)
+                    norm += g*g;
+            for (double g : layer.gradb)
+                norm += g*g;
+        }
+        norm = std::sqrt(norm);
+
+        // If the norm is below the threshold, no clipping is needed.
+        if(norm <= maxNorm)
+            return;
+
+        // Scale down the gradients to have the specified maximum norm.
+        double scale = maxNorm / norm;
+        for (auto& layer : network_)
+        {
+            for (auto& row : layer.gradW)
+                for (auto& g : row)
+                    g *= scale;
+            for (auto& g : layer.gradb)
+                g *= scale;
+        }
+    }
+
+    // Train the diffusion model using the provided training data. The training loop iterates over epochs and batches,
+    // applying noise to the input samples, computing the score predictions, calculating the loss, and performing 
+    // back-propagation to update the model parameters.
+    // For a data sample of 5M 6-dimensional vectors of double precision (8 Byte), total memory for the data is 5M*6*8 = 240 MB, 
+    // which is manageable for in-memory training. 
+    // Much larger datasets may require streaming from disk or using mini-batches that do not fit entirely in memory.
+    void ScoreBasedDiffusionModel::train(
+        const std::vector<DiffusionTrainingSample>& data,
+        int epochs
+    )
+    {
+        // Check that the network is properly initialized before training.
+        assert(!network_.empty());
+        assert(!network_[0].W.empty());
+        assert(!network_[0].W[0].empty());
+
+        const double eps_safe = 1e-12;
+        const size_t N = data.size();
+
+        trainingSampleSize_ = N; // Store the total number of training samples
+
+        if (N <= 1)
+        {
+            throw cet::exception("ScoreBasedDiffusionModel::train")
+                << "Insufficient data samples for training (samples must be > 1)";
+        }
+
+        for (size_t i = 0; i < N; ++i) {
+            if (data[i].x.size() != static_cast<size_t>(dim_)) {
+                throw cet::exception("ScoreBasedDiffusionModel::train")
+                    << "Training sample " << i << " has x dimension " << data[i].x.size()
+                    << ", expected " << dim_;
+            }
+            if (data[i].cond.size() != static_cast<size_t>(conditionDim_)) {
+                throw cet::exception("ScoreBasedDiffusionModel::train")
+                    << "Training sample " << i << " has conditioning dimension " << data[i].cond.size()
+                    << ", expected " << conditionDim_;
+            }
+        }
+
+        // Create index vector once
+        std::vector<size_t> indices(N);
+        for (size_t i = 0; i < N; ++i)
+            indices[i] = i;
+
+        // Data shuffling and batching is performed at the epoch level to ensure that each epoch sees the data 
+        // in a different order, which can improve training convergence.
+        for (int e = 0; e < epochs; ++e)
+        {
+            // Shuffle indices using Fisher–Yates and RandFlat for reproducibility
+            for (size_t i = N - 1; i > 0; --i)
+            {
+                size_t j = static_cast<size_t>(randFlat_.fire() * (i + 1));
+                std::swap(indices[i], indices[j]);
+            }
+
+            // Counter for number of samples processed in the epoch (used for averaging loss). Avoid using N directly to 
+            // allow for early stopping or partial epoch processing if needed.
+            int n = 0; 
+
+            int batchCounter = 0;
+            double epochLoss = 0.0;
+
+            // iterate over the shuffled data samples
+            for (size_t idx = 0; idx < N; ++idx)
+            {
+                const auto& sample = data[indices[idx]];
+
+                // Sample diffusion time
+                double t = randFlat_.fire();
+                // container for noise vector eps
+                std::vector<double> eps;
+
+                auto xt = addNoise(sample.x,t,eps);
+                double s = sigma(t);
+
+                // The target score is the negative of the noise scaled by the noise standard deviation, i.e., -eps/sigma(t).
+                std::vector<double> target(dim_);
+                for (int i=0;i<dim_;++i)
+                    target[i] = -eps[i]/std::max(s, eps_safe); // add eps_safe to prevent division by zero in case of very small sigma
+
+                // Prepare input vector for the network by concatenating the noisy sample xt with the diffusion time t.
+                std::vector<double> input = xt;
+                input.insert(input.end(), sample.cond.begin(), sample.cond.end());
+                input.push_back(t);
+
+                // Check that input dimension matches expected dimension (dim_ + conditionDim_ + 1)
+                assert(input.size() == network_[0].W[0].size());
+
+                // Forward pass to compute the predicted score from the network given the input (noisy sample + time).
+                auto score = forward(input);
+
+                // Compute the loss (Mean Squared Error) between the predicted score and the target score.
+                std::vector<double> grad(dim_);
+                double loss=0.0;
+                for (int i=0;i<dim_;++i)
+                {
+                    double d = score[i]-target[i];
+                    loss += d*d;
+                    // Gradient of the loss w.r.t. the predicted score is 2*(score-target)/dim_ (the division by dim_ is for averaging the loss per dimension).
+                    grad[i] = 2*d/dim_;
+                }
+
+                // Backward pass to compute gradients of the loss w.r.t. network parameters using the computed gradient of the loss w.r.t. the predicted score.
+                backward(grad);
+
+                // Increment batch counter and apply optimizer step if batch size is reached. This allows for vectorized updates after processing a batch of 
+                // samples, which can improve training efficiency and convergence.
+                batchCounter++;
+                if(batchCounter == batchSize_)
+                {
+                    clipGradients(gradientClipThreshold_); // Clip gradients to prevent exploding gradients
+
+                    // Apply optimizer based on configuration
+                    if (optimizerType_ == OptimizerType::ADAM) {
+                        adamUpdate(learningRate_);
+                    } else {
+                        updateWeights(learningRate_);
+                    }
+                    batchCounter = 0;
+                }
+
+                epochLoss += loss; // Accumulate loss for monitoring.
+                n++;
+            }
+
+            // Final flush: apply optimizer to remaining gradients
+            if(batchCounter > 0)
+            {
+                clipGradients(gradientClipThreshold_); // Clip gradients before final update
+
+                if (optimizerType_ == OptimizerType::ADAM) {
+                    adamUpdate(learningRate_);
+                } else {
+                    updateWeights(learningRate_);
+                }
+            }
+
+            epochLoss /= n;
+            mf::LogInfo("ScoreBasedDiffusionModel::train") << "Epoch " << e << "  Loss=" << epochLoss;
+            epochLosses_.push_back(epochLoss);
+        }
+    }
+    
+    void ScoreBasedDiffusionModel::saveModel(
+        const std::string& filename
+    )
+    {
+        std::ofstream out(filename);
+
+        if (!out) {
+            throw cet::exception("ScoreBasedDiffusionModel::saveModel")
+                << "Cannot open file " << filename;
+        }
+
+        // Write CSV header and model configuration parameters with annotations
+        out << "[MODEL_PARAMETERS]\n";
+        out << "Parameter,Value\n";
+        // Model architecture
+        out << "dim," << dim_ << "\n";
+        out << "conditionDim," << conditionDim_ << "\n";
+        out << "hidden," << hidden_ << "\n";
+        out << "layers," << layers_ << "\n";
+        // Optimizer configuration
+        out << "optimizerType," << (optimizerType_ == OptimizerType::ADAM ? "ADAM" : "SGD") << "\n";
+        out << "adamBeta1," << adamBeta1_ << "\n";
+        out << "adamBeta2," << adamBeta2_ << "\n";
+        out << "adamEps," << adamEps_ << "\n";
+        // Noise schedule configuration
+        out << "noiseScheduleType," << (noiseScheduleType_ == NoiseScheduleType::COSINE ? "COSINE" : "LINEAR") << "\n";
+        out << "betaMin," << betaMin_ << "\n";
+        out << "betaMax," << betaMax_ << "\n";
+        out << "cosineOffset," << cosineOffset_ << "\n";
+        // Training configuration
+        out << "batchSize," << batchSize_ << "\n";
+        out << "gradientClipThreshold," << gradientClipThreshold_ << "\n";
+        out << "learningRate," << learningRate_ << "\n";
+        // Diffusion process configuration
+        out << "diffusionSteps," << diffusionSteps_ << "\n";
+
+        // Write network architecture header
+        out << "\n[NETWORK_PARAMETERS]\n";
+        out << "numLayers," << network_.size() << "\n";
+        out << std::fixed << std::setprecision(17); // Use fixed-point notation with high precision for weights and biases
+       
+        for (size_t layerIdx = 0; layerIdx < network_.size(); ++layerIdx) {
+            // Write layer dimensions
+            auto& layer = network_[layerIdx];
+            size_t outSize = layer.W.size();
+            size_t inSize  = layer.W[0].size();
+            out << "\nLayer" << layerIdx << "_OutSize," << outSize << "\n";
+            out << "Layer" << layerIdx << "_InSize," << inSize << "\n";
+
+            // Write weights in matrix format
+            out << "Layer" << layerIdx << "_Weights\n";
+            for (const auto& row : layer.W) {
+                for (size_t j = 0; j < row.size(); ++j) {
+                    out << row[j];
+                    if (j < row.size() - 1) out << ",";
+                }
+                out << "\n";
+            }
+
+            // Write biases
+            out << "Layer" << layerIdx << "_Biases\n";
+            for (size_t j = 0; j < layer.b.size(); ++j) {
+                out << layer.b[j];
+                if (j < layer.b.size() - 1) out << ",";
+            }
+            out << "\n";
+        }
+
+        // Write training history
+        out << "\n[TRAINING_HISTORY]\n";
+        out << "numEpochs," << epochLosses_.size() << "\n";
+        out << "trainingSampleSize," << trainingSampleSize_ << "\n";
+        out << "EpochNumber,Loss\n";
+        for (size_t i = 0; i < epochLosses_.size(); ++i) {
+            out << i << "," << epochLosses_[i] << "\n";
+        }
+    }
+
+    ScoreBasedDiffusionModel ScoreBasedDiffusionModel::loadModel(
+        CLHEP::HepRandomEngine& engine, // We need the random engine to initialize the model architecture
+        const std::string& filename
+    )
+    {
+        std::ifstream in(filename);
+
+        if (!in) {
+            throw cet::exception("ScoreBasedDiffusionModel::loadModel")
+                << "Cannot open file " << filename;
+        }
+
+        std::string line;
+
+        // Temporary storage
+        int dim = 0, conditionDim = 0, hidden = 0, layers = 0;
+        OptimizerType optimizerType = OptimizerType::ADAM;
+        double adamBeta1 = 0.0, adamBeta2 = 0.0, adamEps = 0.0;
+        NoiseScheduleType scheduleType = NoiseScheduleType::COSINE;
+        double betaMin = 0.0, betaMax = 0.0, cosineOffset = 0.0;
+        int batchSize = 1, diffusionSteps = 1;
+        double gradientClipThreshold = 0.0, learningRate = 0.0;
+
+        std::vector<Layer> loadedNetwork;
+        std::vector<double> epochLosses;
+
+        // Helper lambda to split CSV
+        auto split = [](const std::string& s) {
+            std::vector<std::string> tokens;
+            std::stringstream ss(s);
+            std::string item;
+            while (std::getline(ss, item, ',')) {
+                // trim whitespace from item (beginning and end, should not be present if csv is generated by code, but just in case)
+                item.erase(item.begin(), std::find_if(item.begin(), item.end(), [](unsigned char ch) { return !std::isspace(ch); }));
+                item.erase(std::find_if(item.rbegin(), item.rend(), [](unsigned char ch) { return !std::isspace(ch); }).base(), item.end());
+                tokens.push_back(item);
+            }
+            return tokens;
+        };
+        // Helper lambda to extract layer index from strings like "Layer10_OutSize"
+        auto getLayerIdx = [](const std::string& s) {
+            size_t start = 5; // after "Layer"
+            size_t end = s.find('_', start);
+            if (end == std::string::npos)
+                throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Malformed layer string: " << s;
+            return std::stoi(s.substr(start, end - start));
+        };
+
+        // Parse file
+        std::string section;
+
+        while (std::getline(in, line)) {
+            if (line.empty()) continue;
+
+            // Detect section headers
+            if (line[0] == '[') {
+                section = line;
+                continue;
+            }
+
+            // MODEL PARAMETERS
+            if (section == "[MODEL_PARAMETERS]") {
+                if (line == "Parameter,Value") continue;
+
+                auto tokens = split(line);
+                if (tokens.size() != 2) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "ScoreBasedDiffusionModel::loadModel: invalid line in [MODEL_PARAMETERS] section: " << line;
+
+                const std::string& key = tokens[0];
+                const std::string& val = tokens[1];
+
+                if (key == "dim") dim = std::stoi(val);
+                else if (key == "conditionDim") conditionDim = std::stoi(val);
+                else if (key == "hidden") hidden = std::stoi(val);
+                else if (key == "layers") layers = std::stoi(val);
+                else if (key == "optimizerType")
+                    optimizerType = (val == "ADAM") ? OptimizerType::ADAM : OptimizerType::SGD;
+                else if (key == "adamBeta1") adamBeta1 = std::stod(val);
+                else if (key == "adamBeta2") adamBeta2 = std::stod(val);
+                else if (key == "adamEps") adamEps = std::stod(val);
+                else if (key == "noiseScheduleType")
+                    scheduleType = (val == "COSINE") ? NoiseScheduleType::COSINE : NoiseScheduleType::LINEAR;
+                else if (key == "betaMin") betaMin = std::stod(val);
+                else if (key == "betaMax") betaMax = std::stod(val);
+                else if (key == "cosineOffset") cosineOffset = std::stod(val);
+                else if (key == "batchSize") batchSize = std::stoi(val);
+                else if (key == "gradientClipThreshold") gradientClipThreshold = std::stod(val);
+                else if (key == "learningRate") learningRate = std::stod(val);
+                else if (key == "diffusionSteps") diffusionSteps = std::stoi(val);
+            }
+
+            // NETWORK PARAMETERS
+            else if (section == "[NETWORK_PARAMETERS]") {
+
+                if (line.rfind("numLayers", 0) == 0) {
+                    int numLayers = std::stoi(split(line)[1]);
+                    loadedNetwork.resize(numLayers);
+                    if (loadedNetwork.size() != (size_t)layers) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Layer count mismatch";
+                    }
+                    continue;
+                }
+
+                // Layer sizes
+                if (line.find("_OutSize") != std::string::npos) {
+                    auto tokens = split(line);
+                    int outSize = std::stoi(tokens[1]);
+                    int layerIdx = getLayerIdx(tokens[0]);
+                    loadedNetwork[layerIdx].W.resize(outSize);
+                    loadedNetwork[layerIdx].b.resize(outSize);
+                }
+                else if (line.find("_InSize") != std::string::npos) {
+                    auto tokens = split(line);
+                    int inSize = std::stoi(tokens[1]);
+                    int layerIdx = getLayerIdx(tokens[0]);
+                    if (layerIdx == 0 && inSize != dim + conditionDim + 1) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
+                            << "Input layer input size mismatch (expected " 
+                            << (dim + conditionDim + 1) << ", got " << inSize << ")";
+                    }
+                    if (loadedNetwork[layerIdx].W.empty()) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "InSize encountered before OutSize for layer " << layerIdx;
+                    }
+                    for (auto& row : loadedNetwork[layerIdx].W)
+                        row.resize(inSize);
+                }
+                // Weights
+                else if (line.find("_Weights") != std::string::npos) {
+                    int layerIdx = getLayerIdx(line);
+                    if (layerIdx < 0 || layerIdx >= (int)loadedNetwork.size()) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Invalid layer index: " << layerIdx;
+                    }
+                    if (loadedNetwork[layerIdx].W.empty() || loadedNetwork[layerIdx].W[0].empty()) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Weights encountered before layer size definition";
+                    }
+                    for (auto& row : loadedNetwork[layerIdx].W) {
+                        if (!std::getline(in, line)) {
+                            throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading weights";
+                        }
+                        auto vals = split(line);
+                        if (vals.size() != row.size()) {
+                            throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
+                                << "Weight row size mismatch for layer " << layerIdx;
+                        }
+                        for (size_t j = 0; j < vals.size(); ++j)
+                            row[j] = std::stod(vals[j]);
+                    }
+                }
+                // Biases
+                else if (line.find("_Biases") != std::string::npos) {
+                    int layerIdx = getLayerIdx(line);
+                    if (layerIdx < 0 || layerIdx >= (int)loadedNetwork.size()) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Invalid layer index: " << layerIdx;
+                    }
+                    if (!std::getline(in, line)) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading biases";
+                    }
+                    auto vals = split(line);
+                    if (vals.size() != loadedNetwork[layerIdx].b.size()) {
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
+                            << "Bias size mismatch for layer " << layerIdx;
+                    }
+                    for (size_t j = 0; j < vals.size(); ++j)
+                        loadedNetwork[layerIdx].b[j] = std::stod(vals[j]);
+                }
+            }
+
+            // TRAINING HISTORY
+            else if (section == "[TRAINING_HISTORY]") {
+                if (line == "EpochNumber,Loss") continue;
+
+                auto tokens = split(line);
+                if (tokens.size() == 2) {
+                    if (tokens[0] == "numEpochs") {
+                        int numEpochs = std::stoi(tokens[1]);
+                        epochLosses.reserve(numEpochs);
+                    } else if (tokens[0] == "trainingSampleSize")
+                    {
+                        // We can store this if needed for analysis, but it is not used in model reconstruction, so we will ignore it for now.
+                    }else {
+                        epochLosses.push_back(std::stod(tokens[1]));
+                    }
+                }
+            }
+        }
+
+        if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) {
+            throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Invalid model parameters in file";
+        }
+        for (size_t l = 0; l < loadedNetwork.size(); ++l) {
+            if (loadedNetwork[l].W.empty() || loadedNetwork[l].b.empty()) {
+                throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Incomplete layer data in file";
+            }
+            for (const auto& row : loadedNetwork[l].W) {
+                if (row.size() != loadedNetwork[l].W[0].size()) {
+                    throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Inconsistent row size in layer";
+                }
+            }
+        }
+        for (size_t l = 1; l < loadedNetwork.size(); ++l) {
+            if (loadedNetwork[l].W[0].size() != loadedNetwork[l-1].W.size()) {
+                throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Layer size mismatch between layers";
+            }
+        }
+
+        std::ostringstream logMsg;
+        logMsg << "Model parameters loaded successfully from " << filename << "\n"
+               << "Warning: optimizer state (e.g., Adam moments) is not saved/loaded.\n"
+               << "The loaded model is suitable for inference, or for retraining with a fresh optimizer state,\n"
+               << "but does NOT resume training from the original state.";
+        mf::LogInfo("ScoreBasedDiffusionModel::loadModel") << logMsg.str();
+
+        // Reconstruct model
+        ScoreBasedDiffusionModel model(
+            engine,
+            dim,
+            conditionDim,
+            hidden,
+            layers,
+            optimizerType,
+            adamBeta1,
+            adamBeta2,
+            adamEps,
+            scheduleType,
+            betaMin,
+            betaMax,
+            cosineOffset,
+            batchSize,
+            gradientClipThreshold,
+            learningRate,
+            diffusionSteps
+        );
+
+        // Overwrite initialized weights with loaded values
+        for (size_t l = 0; l < model.network_.size(); ++l) {
+            model.network_[l].W = loadedNetwork[l].W;
+            model.network_[l].b = loadedNetwork[l].b;
+        }
+        model.epochLosses_ = epochLosses;
+
+        return model;
+    }
+
+    std::vector<double> ScoreBasedDiffusionModel::generateSample(
+        const std::vector<double>& condition,
+        bool useHeun,
+        int diffusionSteps
+    )
+    {
+        if (condition.size() != static_cast<size_t>(conditionDim_)) {
+            throw cet::exception("ScoreBasedDiffusionModel::generateSample")
+                << "Conditioning dimension mismatch: got " << condition.size()
+                << ", expected " << conditionDim_;
+        }
+
+        // Use provided diffusionSteps if positive, otherwise use the model's configured value
+        int steps = (diffusionSteps > 0) ? diffusionSteps : diffusionSteps_;
+
+        // Start from pure noise
+        std::vector<double> x(dim_);
+
+        for (int i=0;i<dim_;++i)
+            x[i] = randGaussQ_.fire();
+
+        // Reverse diffusion process
+        for (int step = steps - 1; step >= 0; --step) {
+
+            double t = (double)step/steps;
+            double dt = 1.0/steps;
+            double s = sigma(t); // as long as diffusionSteps_ is not too large, s should not become too small to cause numerical issues.
+
+            if (!useHeun) {
+                // Euler method (1st order)
+                std::vector<double> input = x;
+                input.insert(input.end(), condition.begin(), condition.end());
+                input.push_back(t);
+
+                auto score = forward(input);
+ 
+                for (int i=0;i<dim_;++i){
+                    x[i] += -s*s*score[i]*dt;
+                }
+            } else {
+                // Heun's method (2nd order)
+                // k1 = f(x,t)
+                std::vector<double> input = x;
+                input.insert(input.end(), condition.begin(), condition.end());
+                input.push_back(t);
+                auto score_k1 = forward(input);
+
+                std::vector<double> k1(dim_);
+                for (int i=0;i<dim_;++i)
+                    k1[i] = -s*s*score_k1[i];
+
+                // predictor
+                std::vector<double> x_pred(dim_);
+                for (int i=0;i<dim_;++i)
+                    x_pred[i] = x[i] + k1[i]*dt;
+
+                // next time
+                double t_next = std::max(0.0, t - dt);
+                double s_next = sigma(t_next);
+
+                // k2 = f(x_pred,t_next)
+                std::vector<double> input_next = x_pred;
+                input_next.insert(input_next.end(), condition.begin(), condition.end());
+                input_next.push_back(t_next);
+                auto score_k2 = forward(input_next);
+
+                std::vector<double> k2(dim_);
+                for (int i=0;i<dim_;++i)
+                    k2[i] = -s_next*s_next*score_k2[i];
+
+                // trapezoidal update
+                for (int i=0;i<dim_;++i)
+                    x[i] += 0.5*(k1[i]+k2[i])*dt;
+            }
+        }
+
+        return x;
+    }
+
+} // namespace mu2e
diff --git a/STMMC/CMakeLists.txt b/STMMC/CMakeLists.txt
index 7e719143d7..cb3fce1b42 100644
--- a/STMMC/CMakeLists.txt
+++ b/STMMC/CMakeLists.txt
@@ -83,6 +83,42 @@ cet_build_plugin(VirtualDetectorTree art::module
       Offline::MCDataProducts
 )
 
+cet_build_plugin(VDResamplerConfigure art::module
+    REG_SOURCE src/VDResamplerConfigure_module.cc
+    LIBRARIES REG
+      art_root_io::TFileService_service
+      Offline::GlobalConstantsService
+      Offline::MCDataProducts
+)
+
+cet_build_plugin(VDResamplerTrain art::module
+    REG_SOURCE src/VDResamplerTrain_module.cc
+    LIBRARIES REG
+      art_root_io::TFileService_service
+      Offline::GlobalConstantsService
+      Offline::MachineLearningTools
+      Offline::MCDataProducts
+)
+
+cet_build_plugin(VDResamplerGenerateFromModel art::module
+    REG_SOURCE src/VDResamplerGenerateFromModel_module.cc
+    LIBRARIES REG
+      art_root_io::TFileService_service
+      Offline::DataProducts
+      Offline::GlobalConstantsService
+      Offline::MachineLearningTools
+      Offline::MCDataProducts
+)
+
+cet_build_plugin(VDResamplerGenerateMix art::module
+    REG_SOURCE src/VDResamplerGenerateMix_module.cc
+    LIBRARIES REG
+      art_root_io::TFileService_service
+      Offline::DataProducts
+      Offline::GlobalConstantsService
+      Offline::MachineLearningTools
+      Offline::MCDataProducts
+)
 
 install_source(SUBDIRS src)
 install_fhicl(SUBDIRS fcl SUBDIRNAME Offline/STMMC/fcl)
diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
new file mode 100644
index 0000000000..8adcd5d219
--- /dev/null
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -0,0 +1,256 @@
+// VDResamplerConfigure_module.cc
+// Makes one pass of the data set and determines which particles will need training for the 
+// resampler, and generates a fcl file for training. "dataSourceTag" specifies the tag of the 
+// data source, and will be appended to the generated fcl file and then csv files to 
+// distinguish between different data sources. The generated fcl files will be stored and 
+// named "VDResamplerTrain_[dataSourceTag].fcl".
+// On the other hand, a breakdown of the number of hits for each particle type will be stored 
+// in "VDResamplerConfigure_[dataSourceTag]_HitSummary.csv" for reference.
+// Later when generating new samples using the resampler, the program will look for the 
+// generated fcl files to determine which particle source to use, which particle types to 
+// generate and which model parameters to load for each particle type. 
+// Yongyi Wu, Mar. 2026
+
+// stdlib includes
+#include <cmath>
+#include <iostream>
+#include <fstream>
+#include <iomanip>
+
+// art includes
+#include "art/Framework/Core/EDAnalyzer.h"
+#include "art/Framework/Principal/Event.h"
+#include "art/Framework/Principal/Handle.h"
+#include "art/Framework/Principal/Run.h"
+
+// exception handling
+#include "cetlib_except/exception.h"
+
+// fhicl includes
+#include "canvas/Utilities/InputTag.h"
+#include "fhiclcpp/types/Atom.h"
+
+// message handling
+#include "messagefacility/MessageLogger/MessageLogger.h"
+
+// Offline includes
+#include "Offline/GlobalConstantsService/inc/GlobalConstantsHandle.hh"
+#include "Offline/GlobalConstantsService/inc/ParticleDataList.hh"
+#include "Offline/MCDataProducts/inc/SimParticle.hh"
+#include "Offline/MCDataProducts/inc/StepPointMC.hh"
+
+// ROOT includes
+#include "art_root_io/TFileService.h"
+#include "TTree.h"
+
+typedef cet::map_vector_key key_type;
+typedef unsigned long VolumeId_type;
+
+namespace mu2e {
+  class VDResamplerConfigure : public art::EDAnalyzer {
+    public:
+      using Name=fhicl::Name;
+      using Comment=fhicl::Comment;
+      struct Config {
+        fhicl::Atom<art::InputTag> StepPointMCsTag{Name("StepPointMCsTag"), Comment("Tag identifying the StepPointMCs")};
+        fhicl::Atom<art::InputTag> SimParticlemvTag{Name("SimParticlemvTag"), Comment("Tag identifying the SimParticlemv")};
+        fhicl::Atom<int> VirtualDetectorID{Name("VirtualDetectorID"), Comment("ID of the virtual detector to train on"), 116}; 
+        fhicl::Atom<std::string> VDResamplerDir{Name("VDResamplerDir"), Comment("Directory to store the generated csv files")};
+        fhicl::Atom<std::string> fclDir{Name("fclDir"), Comment("Directory to store the generated fhicl files"), ""};
+        fhicl::Atom<std::string> dataSourceTag{Name("dataSourceTag"), Comment("A tag to distinguish different data sources, will be appended to the generated fcl and csv files")};
+        fhicl::Atom<int> trainingThreshold{Name("trainingThreshold"), Comment("Minimum number of hits for a particle type to be included in the training"), 100};
+        fhicl::Atom<bool> doROOTDump{Name("doROOTDump"), Comment("Whether to dump the VD hit info into a ROOT file for debugging and analysis"), false};
+      };
+      using Parameters = art::EDAnalyzer::Table<Config>;
+      explicit VDResamplerConfigure(const Parameters& conf);
+      void analyze(const art::Event& e);
+      void endJob();
+    private:
+      art::ProductToken<StepPointMCCollection> StepPointMCsToken;
+      art::ProductToken<SimParticleCollection> SimParticlemvToken;
+      VolumeId_type VirtualDetectorID = 0;
+      std::string VDResamplerDir;
+      std::string fclDir;
+      std::string dataSourceTag;
+      int trainingThreshold = 0;
+      bool doROOTDump;
+      GlobalConstantsHandle<ParticleDataList> pdt;
+      int pdgId = 0;
+      double x = 0.0, y = 0.0, z = 0.0, px = 0.0, py = 0.0, pz = 0.0, mass = 0.0, E = 0.0, time = 0.0;
+      VolumeId_type virtualdetectorId = 0;
+      TTree* ttree;
+      std::map<int, int> pdgIds; // <id, count>
+  };
+
+  VDResamplerConfigure::VDResamplerConfigure(const Parameters& conf) :
+    art::EDAnalyzer(conf),
+    StepPointMCsToken(consumes<StepPointMCCollection>(conf().StepPointMCsTag())),
+    SimParticlemvToken(consumes<SimParticleCollection>(conf().SimParticlemvTag())),
+    VirtualDetectorID(conf().VirtualDetectorID()),
+    VDResamplerDir(conf().VDResamplerDir()),
+    fclDir(conf().fclDir()),
+    trainingThreshold(conf().trainingThreshold()),
+    dataSourceTag(conf().dataSourceTag()),
+    doROOTDump(conf().doROOTDump()) {
+    if (doROOTDump) {
+        art::ServiceHandle<art::TFileService> tfs;
+        ttree = tfs->make<TTree>( "ttree", "Virtual Detectors Hit Summary");
+        ttree->Branch("time", &time, "time/D"); // ns
+        ttree->Branch("virtualdetectorId", &virtualdetectorId, "virtualdetectorId/l");
+        ttree->Branch("pdgId", &pdgId, "pdgId/I");
+        ttree->Branch("x", &x, "x/D"); // mm
+        ttree->Branch("y", &y, "y/D"); // mm
+        ttree->Branch("z", &z, "z/D"); // mm
+        ttree->Branch("px", &px, "px/D"); // MeV
+        ttree->Branch("py", &py, "py/D"); // MeV
+        ttree->Branch("pz", &pz, "pz/D"); // MeV
+        ttree->Branch("E", &E, "E/D"); // MeV
+    }
+  };
+
+  void VDResamplerConfigure::analyze(const art::Event& event) {
+    // Get the data products from the event
+    auto const& StepPointMCs = event.getProduct(StepPointMCsToken);
+    if (StepPointMCs.empty())
+      return;
+    auto const& SimParticles = event.getProduct(SimParticlemvToken);
+    if (SimParticles.empty())
+      return;
+
+    // Loop over all VD hits
+    for (const StepPointMC& step : StepPointMCs) {
+      // Get the associated particle
+      const SimParticle& particle = SimParticles.at(step.trackId());
+      pdgId = particle.pdgId();
+
+      virtualdetectorId = step.virtualDetectorId();
+      pz = step.momentum().z();
+      if (virtualdetectorId != VirtualDetectorID || pz <= 0)
+        continue; // Filter hits based on the virtual detector ID and pz
+      
+      if (doROOTDump) {
+          x = step.position().x();
+          y = step.position().y();
+          z = step.position().z();
+          px = step.momentum().x();
+          py = step.momentum().y();
+          pz = step.momentum().z();
+          mass = pdt->particle(pdgId).mass();
+          E = std::sqrt(step.momentum().mag2()+mass*mass)-mass; // Subtract the rest mass
+          if (E < 0)
+            throw cet::exception("LogicError", "Energy is negative");
+          
+          ttree->Fill();
+      }
+
+      // Count the number of hits for each particle type for the summary
+      if (pdgIds.find(pdgId) != pdgIds.end())
+        pdgIds[pdgId] += 1;
+      else
+        pdgIds.emplace(std::make_pair(pdgId, 1));
+    };
+    return;
+  };
+
+  void VDResamplerConfigure::endJob() {
+    mf::LogInfo log("Virtual detector tree summary");
+    log << "========= Data summary =========\n";
+    for (auto part : pdgIds)
+      log << "PDGID " << part.first << ": " << part.second << "\n";
+    log << "================================\n";
+
+    // store a summary of the number of hits for each particle type in a csv file for reference, 
+    // all particles are included, but the particle types with hits below the training threshold are 
+    // marked with an asterisk to indicate that they will not be included in the training.
+    std::string summaryFile = VDResamplerDir + "/VDResamplerConfigure_" + dataSourceTag + "_HitSummary.csv";
+    std::string fclFilePath = fclDir.empty() ? VDResamplerDir : fclDir;
+    std::string fclFile = fclFilePath + "/VDResamplerTrain_" + dataSourceTag + ".fcl";
+    std::ofstream sumOutFile(summaryFile);
+    std::ofstream fclOutFile(fclFile);
+    if (!sumOutFile) {
+        throw cet::exception("VDResamplerConfigure::endJob") << "Cannot open file " << summaryFile;
+    }
+    if (!fclOutFile) {
+        throw cet::exception("VDResamplerConfigure::endJob") << "Cannot open file " << fclFile;
+    }
+
+    std::string trainingModuleNames = "";
+
+    sumOutFile << "PDGID,HitCount,NotTrained\n";
+
+    fclOutFile << "# This fcl is generated by VDResamplerConfigure_module.cc\n";
+    fclOutFile << "# Training configuration for the VD resampler is generated for each particle type\n\n";
+    fclOutFile << "#include \"Offline/fcl/standardServices.fcl\"\n\n";
+    fclOutFile << "process_name: VDResamplerTrain\n\n";
+    fclOutFile << "source : {\n  module_type : RootInput\n  fileNames: @nil\n}\n";
+    fclOutFile << "services : {\n  message : @local::default_message\n  GlobalConstantsService : {\n    inputFile : \"Offline/GlobalConstantsService/data/globalConstants_01.txt\"\n  }\n}\n";
+    fclOutFile << "physics: {\n  analyzers : {\n";
+
+    for (const auto& part : pdgIds) {
+      sumOutFile << part.first << "," << part.second;
+      bool useTwoStageTraining = false; 
+      if (part.second < 100000) { 
+        // if data is limited, use two-stage training to improve performance and reduce the required training data size for each model.
+        useTwoStageTraining = true;
+      }
+      if (part.second < trainingThreshold) {
+        sumOutFile << ",*\n";
+        mf::LogWarning("VDResamplerConfigure::endJob") << "Particle type with PDGID " << part.first 
+            << " has only " << part.second << " hits, which is below the training threshold of " 
+            << trainingThreshold << ". This particle type will NOT be included in the training.";
+      } else {
+        // mark how many models will be trained for this particle type (1 for all-at-once, 2 for two-stage)
+        sumOutFile << "," << (useTwoStageTraining ? "2" : "1") << "\n";
+
+        // Generate the fcl file for training for this particle type
+        // if pdgID is negative, we will use "m" instead of "-" in the filename to avoid issues with file naming
+        std::string pdgIdstr = (part.first < 0) ? "m" + std::to_string(-part.first) : std::to_string(part.first);
+        std::string moduleName = "VDResamplerTrainVD"+ std::to_string(VirtualDetectorID) + dataSourceTag + "pdg" + pdgIdstr;
+        if (!trainingModuleNames.empty()) {
+            trainingModuleNames += ", ";
+        }            
+        trainingModuleNames += moduleName;
+        fclOutFile << "    " << moduleName << " : {\n"
+                    << "      module_type : VDResamplerTrain\n"
+                    << "      StepPointMCsTag : @local::SimplifyStage1Data.StepPointMCsTag\n"
+                    << "      SimParticlemvTag : @local::SimplifyStage1Data.SimParticlemvTag\n"
+                    << "      SBDMuseTwoStageTraining : " << (useTwoStageTraining ? "true" : "false") << "\n";
+        if (useTwoStageTraining) {
+          fclOutFile << "      SBDMstage1ModelFile : \"" << VDResamplerDir << "/SBDMmodel_stage1_VD" << VirtualDetectorID << "_pdg" << pdgIdstr << ".csv\"\n"
+                     << "      SBDMstage2ModelFile : \"" << VDResamplerDir << "/SBDMmodel_stage2_VD" << VirtualDetectorID << "_pdg" << pdgIdstr << ".csv\"\n";
+        } else {
+          fclOutFile << "      SBDMallAtOnceModelFile : \"" << VDResamplerDir << "/SBDMmodel_allAtOnce_VD" << VirtualDetectorID << "_pdg" << pdgIdstr << ".csv\"\n";
+        }
+        fclOutFile << "      VirtualDetectorID : " << VirtualDetectorID << "\n"
+                    // << "      VDz0 : " << std::to_string(<value>) << "\n" // include if not VD116
+                    // << "      VDr : " << std::to_string(<value>) << "\n" // include if not VD116
+                    << "      pdgID : " << part.first << "\n"
+                    // << "      SBDMhidden : 128\n"
+                    // << "      SBDMlayers : 4\n"
+                    // << "      SBDMoptimizer : \"ADAM\"\n"
+                    // << "      SBDMadamBeta1 : 0.9\n"
+                    // << "      SBDMadamBeta2 : 0.999\n"
+                    // << "      SBDMadamEps : 1e-8\n"
+                    // << "      SBDMnoiseSchedule : \"COSINE\"\n"
+                    // << "      SBDMbetaMin : 1e-4\n"
+                    // << "      SBDMbetaMax : 0.02\n"
+                    // << "      SBDMcosineOffset : 0.008\n"
+                    // << "      SBDMbatchSize : 32\n"
+                    // << "      SBDMgradientClip : 1.0\n"
+                    // << "      SBDMlearningRate : 1e-3\n"
+                    // << "      SBDMdiffusionSteps : 200\n"
+                    // << "      SBDMtrainingSize : -1\n"
+                    // << "      SBDMtrainingEpochs : 10\n"
+                    << "    }\n";
+      }
+    }
+    fclOutFile << "  }\n";
+    fclOutFile << "  end_paths: [" + trainingModuleNames + "]\n";
+    fclOutFile << "}\n\n";
+    if (doROOTDump) fclOutFile << "services.TFileService.fileName : @nil\n";
+
+    return;
+  };
+}; // end namespace mu2e
+
+DEFINE_ART_MODULE(mu2e::VDResamplerConfigure)
diff --git a/STMMC/src/VDResamplerGenerateFromModel_module.cc b/STMMC/src/VDResamplerGenerateFromModel_module.cc
new file mode 100644
index 0000000000..9a4dd45eda
--- /dev/null
+++ b/STMMC/src/VDResamplerGenerateFromModel_module.cc
@@ -0,0 +1,334 @@
+// VDResamplerGenerateFromModel_module.cc
+// This module takes either a single all-at-once model file or a pair of stage-1/stage-2
+// model files and produces new virtual-detector-like GenParticles from the trained
+// score-based diffusion model(s).
+// Yongyi Wu, Mar. 2026
+
+// stdlib includes
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <fstream>
+#include <iostream>
+#include <string>
+#include <utility>
+
+#include "Offline/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh"
+
+// art includes
+#include "art/Framework/Core/EDProducer.h"
+#include "art/Framework/Principal/Event.h"
+#include "art/Framework/Services/Optional/RandomNumberGenerator.h"
+
+// CLHEP includes
+#include "CLHEP/Vector/LorentzVector.h"
+#include "CLHEP/Vector/ThreeVector.h"
+
+// exception handling
+#include "cetlib_except/exception.h"
+
+// fhicl includes
+#include "fhiclcpp/types/Atom.h"
+
+// message handling
+#include "messagefacility/MessageLogger/MessageLogger.h"
+
+// Offline includes
+#include "Offline/DataProducts/inc/PDGCode.hh"
+#include "Offline/GlobalConstantsService/inc/GlobalConstantsHandle.hh"
+#include "Offline/GlobalConstantsService/inc/ParticleDataList.hh"
+#include "Offline/MCDataProducts/inc/GenId.hh"
+#include "Offline/MCDataProducts/inc/GenParticle.hh"
+#include "Offline/SeedService/inc/SeedService.hh"
+
+// ROOT includes
+#include "art_root_io/TFileService.h"
+#include "TTree.h"
+
+namespace mu2e {
+
+  namespace {
+    // Utility function to extract PDG ID from the model filename, which is expected to contain a substring like "pdg[optional m][digits]", e.g. "pdg13" for muons, "pdgm13" for muon pluses.
+    int loadPDGIdFromFileName(const std::string& fileName) {
+      const size_t pdgPos = fileName.find("pdg");
+      if (pdgPos == std::string::npos) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "Cannot infer PDG ID from model filename: " << fileName;
+      }
+      size_t pos = pdgPos + 3;
+      bool negative = false;
+      if (pos < fileName.size() && (fileName[pos] == 'm' || fileName[pos] == '-')) {
+        negative = true;
+        ++pos;
+      }
+      const size_t startDigits = pos;
+      while (pos < fileName.size() && std::isdigit(static_cast<unsigned char>(fileName[pos]))) {
+        ++pos;
+      }
+      if (startDigits == pos) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "Cannot infer PDG ID from model filename: " << fileName;
+      }
+      const int magnitude = std::stoi(fileName.substr(startDigits, pos - startDigits));
+      return negative ? -magnitude : magnitude;
+    }
+  }
+
+  class VDResamplerGenerateFromModel : public art::EDProducer {
+    public:
+      using Name = fhicl::Name;
+      using Comment = fhicl::Comment;
+
+      struct Config {
+        fhicl::Atom<bool> useTwoStageModel{
+          Name("useTwoStageModel"),
+          Comment("If true, use stage-1/stage-2 models. If false, use the all-at-once 6D model."),
+          true
+        };
+        fhicl::Atom<std::string> stage1ModelFile{
+          Name("stage1ModelFile"),
+          Comment("CSV filename for the stage-1 model parameters"),
+          ""
+        };
+        fhicl::Atom<std::string> stage2ModelFile{
+          Name("stage2ModelFile"),
+          Comment("CSV filename for the stage-2 model parameters"),
+          ""
+        };
+        fhicl::Atom<std::string> allAtOnceModelFile{
+          Name("allAtOnceModelFile"),
+          Comment("CSV filename for the all-at-once 6D model parameters"),
+          ""
+        };
+        fhicl::Atom<bool> useHeun{
+          Name("useHeun"),
+          Comment("If true, use Heun's method for reverse diffusion. Otherwise use Euler."),
+          true
+        };
+        fhicl::Atom<int> diffusionSteps{
+          Name("diffusionSteps"),
+          Comment("Number of reverse-diffusion steps used for sampling"),
+          200
+        };
+        fhicl::Atom<double> VDz0{
+          Name("VDz0"),
+          Comment("Nominal z coordinate of the virtual detector"),
+          37700.39
+        };
+        fhicl::Atom<double> VDr{
+          Name("VDr"),
+          Comment("Virtual detector radius used in the training transform"),
+          2000.0
+        };
+        fhicl::Atom<bool> doROOTDump{
+          Name("doROOTDump"),
+          Comment("Whether to dump generated samples into a ROOT tree"),
+          false
+        };
+      };
+
+      using Parameters = art::EDProducer::Table<Config>;
+
+      explicit VDResamplerGenerateFromModel(const Parameters& conf);
+      void produce(art::Event& event) override;
+
+    private:
+      art::RandomNumberGenerator::base_engine_t& engine_;
+      bool useTwoStageModel_;
+      std::string stage1ModelFile_;
+      std::string stage2ModelFile_;
+      std::string allAtOnceModelFile_;
+      bool useHeun_;
+      int diffusionSteps_;
+      double VDz0_;
+      double VDr_;
+      bool doROOTDump_;
+      GlobalConstantsHandle<ParticleDataList> pdt_;
+
+      std::unique_ptr<ScoreBasedDiffusionModel> allAtOnceModel_;
+      std::unique_ptr<ScoreBasedDiffusionModel> stage1Model_;
+      std::unique_ptr<ScoreBasedDiffusionModel> stage2Model_;
+
+      int pdgId_ = 0;
+
+      // Detector-center parameters used in the same transform as training.
+      double x0_ = -3904.0;
+      double y0_ = 0.0;
+      double t0_ = 1700.0;   // ns
+      double tScale_ = 1.0;
+      double p0_ = 1.0;
+
+      // Variables for optional ROOT dump.
+      TTree* outTree_ = nullptr;
+      double x_gen_ = 0.0;
+      double y_gen_ = 0.0;
+      double z_gen_ = 0.0;
+      double t_gen_ = 0.0;
+      double px_gen_ = 0.0;
+      double py_gen_ = 0.0;
+      double pz_gen_ = 0.0;
+      double mass_gen_ = 0.0;
+      double E_gen_ = 0.0;
+    };
+
+  VDResamplerGenerateFromModel::VDResamplerGenerateFromModel(const Parameters& conf)
+    : art::EDProducer{conf},
+      engine_(createEngine(art::ServiceHandle<SeedService>()->getSeed())),
+      useTwoStageModel_(conf().useTwoStageModel()),
+      stage1ModelFile_(conf().stage1ModelFile()),
+      stage2ModelFile_(conf().stage2ModelFile()),
+      allAtOnceModelFile_(conf().allAtOnceModelFile()),
+      useHeun_(conf().useHeun()),
+      diffusionSteps_(conf().diffusionSteps()),
+      VDz0_(conf().VDz0()),
+      VDr_(conf().VDr()),
+      doROOTDump_(conf().doROOTDump()) {
+
+    produces<GenParticleCollection>();
+
+    if (useTwoStageModel_) {
+      if (stage1ModelFile_.empty() || stage2ModelFile_.empty()) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "Two-stage generation requires both stage1ModelFile and stage2ModelFile.";
+      }
+
+      stage1Model_ = std::make_unique<ScoreBasedDiffusionModel>(
+        ScoreBasedDiffusionModel::loadModel(engine_, stage1ModelFile_)
+      );
+      stage2Model_ = std::make_unique<ScoreBasedDiffusionModel>(
+        ScoreBasedDiffusionModel::loadModel(engine_, stage2ModelFile_)
+      );
+      pdgId_ = loadPDGIdFromFileName(stage1ModelFile_);
+    } else {
+      if (allAtOnceModelFile_.empty()) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "All-at-once generation requires allAtOnceModelFile.";
+      }
+
+      allAtOnceModel_ = std::make_unique<ScoreBasedDiffusionModel>(
+        ScoreBasedDiffusionModel::loadModel(engine_, allAtOnceModelFile_)
+      );
+      pdgId_ = loadPDGIdFromFileName(allAtOnceModelFile_);
+    }
+
+    z_gen_ = VDz0_;
+
+    if (doROOTDump_) {
+      art::ServiceHandle<art::TFileService> tfs;
+      outTree_ = tfs->make<TTree>("ttree", "Generated samples from VD resampler model");
+      outTree_->Branch("pdgId", &pdgId_, "pdgId/I");
+      outTree_->Branch("x", &x_gen_, "x/D");
+      outTree_->Branch("y", &y_gen_, "y/D");
+      outTree_->Branch("z", &z_gen_, "z/D");
+      outTree_->Branch("time", &t_gen_, "time/D");
+      outTree_->Branch("px", &px_gen_, "px/D");
+      outTree_->Branch("py", &py_gen_, "py/D");
+      outTree_->Branch("pz", &pz_gen_, "pz/D");
+      outTree_->Branch("E", &E_gen_, "E/D");
+    }
+  }
+
+  void VDResamplerGenerateFromModel::produce(art::Event& event) {
+    auto output = std::make_unique<GenParticleCollection>();
+
+    double x_trans = 0.0;
+    double y_trans = 0.0;
+    double t_trans = 0.0;
+    double pr_t = 0.0;
+    double pphi_t = 0.0;
+    double pz_t = 0.0;
+
+    // Generate a new sample using the loaded model(s)
+    // note the values here are transformed and need to be inverted back to the original coordinates after sampling.
+    if (useTwoStageModel_) {
+      const std::vector<double> stage1Sample = stage1Model_->generateSample({}, useHeun_, diffusionSteps_);
+      if (stage1Sample.size() != 3u) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "Stage-1 model returned " << stage1Sample.size() << " values, expected 3.";
+      }
+
+      t_trans = stage1Sample[0];
+      x_trans = stage1Sample[1];
+      y_trans = stage1Sample[2];
+
+      const std::vector<double> stage2Condition = {t_trans, x_trans, y_trans};
+      const std::vector<double> stage2Sample = stage2Model_->generateSample(stage2Condition, useHeun_, diffusionSteps_);
+      if (stage2Sample.size() != 3u) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "Stage-2 model returned " << stage2Sample.size() << " values, expected 3.";
+      }
+
+      pr_t = stage2Sample[0];
+      pphi_t = stage2Sample[1];
+      pz_t = stage2Sample[2];
+    } else {
+      const std::vector<double> sample = allAtOnceModel_->generateSample({}, useHeun_, diffusionSteps_);
+      if (sample.size() != 6u) {
+        throw cet::exception("VDResamplerGenerateFromModel")
+          << "All-at-once model returned " << sample.size() << " values, expected 6.";
+      }
+
+      t_trans = sample[0];
+      x_trans = sample[1];
+      y_trans = sample[2];
+      pr_t = sample[3];
+      pphi_t = sample[4];
+      pz_t = sample[5];
+    }
+
+    // Recover detector coordinates from the transformed (x', y').
+    const double u = std::sqrt(x_trans * x_trans + y_trans * y_trans);
+    const double theta = std::atan2(y_trans, x_trans);
+    const double rho = std::tanh(u);
+    const double r = rho * VDr_;
+    const double dx = r * std::cos(theta);
+    const double dy = r * std::sin(theta);
+    x_gen_ = dx + x0_;
+    y_gen_ = dy + y0_;
+    z_gen_ = VDz0_;
+
+    // Invert the time transform.
+    t_gen_ = t0_ * std::exp(t_trans * tScale_);
+
+    // Invert the momentum transform.
+    const double pr = p0_ * std::sinh(pr_t);
+    const double pphi = p0_ * std::sinh(pphi_t);
+    pz_gen_ = p0_ * std::sinh(pz_t);
+
+    if (r > 1e-6) {
+      const double rx = dx / r;
+      const double ry = dy / r;
+      const double phix = -ry;
+      const double phiy = rx;
+      px_gen_ = pr * rx + pphi * phix;
+      py_gen_ = pr * ry + pphi * phiy;
+    } else {
+      px_gen_ = pr;
+      py_gen_ = pphi;
+    }
+
+    mass_gen_ = pdt_->particle(pdgId_).mass();
+    const CLHEP::Hep3Vector momParticle(px_gen_, py_gen_, pz_gen_);
+    const CLHEP::Hep3Vector posParticle(x_gen_, y_gen_, z_gen_);
+    const double eTotal = std::sqrt(momParticle.mag2() + mass_gen_ * mass_gen_);
+    E_gen_ = eTotal - mass_gen_;
+    const CLHEP::HepLorentzVector fourMomParticle(eTotal, momParticle);
+
+    output->emplace_back(
+      PDGCode::type(pdgId_),
+      GenId::STMDownStreamGenTool,
+      posParticle,
+      fourMomParticle,
+      t_gen_
+    );
+
+    event.put(std::move(output));
+
+    if (doROOTDump_) {
+      outTree_->Fill();
+    }
+  }
+
+} // namespace mu2e
+
+DEFINE_ART_MODULE(mu2e::VDResamplerGenerateFromModel)
diff --git a/STMMC/src/VDResamplerGenerateMix_module.cc b/STMMC/src/VDResamplerGenerateMix_module.cc
new file mode 100644
index 0000000000..031ba23511
--- /dev/null
+++ b/STMMC/src/VDResamplerGenerateMix_module.cc
@@ -0,0 +1,495 @@
+// VDResamplerGenerateMix_module.cc
+// This module takes a list of VDResamplerConfigure_[dataSourceTag]_HitSummary.csv files
+// generated by VDResamplerConfigure_module.cc, plus a matching list of POT counts.
+// It chooses which source summary to use according to the trained hit yield per POT,
+// then chooses the particle type according to the per-summary particle ratios, loads
+// either a one-stage or two-stage model for that particle, and generates a GenParticle.
+// Yongyi Wu, Mar. 2026
+
+// stdlib includes
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <fstream>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "Offline/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh"
+
+// art includes
+#include "art/Framework/Core/EDProducer.h"
+#include "art/Framework/Principal/Event.h"
+#include "art/Framework/Services/Optional/RandomNumberGenerator.h"
+
+// CLHEP includes
+#include "CLHEP/Random/RandFlat.h"
+#include "CLHEP/Vector/LorentzVector.h"
+#include "CLHEP/Vector/ThreeVector.h"
+
+// exception handling
+#include "cetlib_except/exception.h"
+
+// fhicl includes
+#include "fhiclcpp/types/Atom.h"
+#include "fhiclcpp/types/Sequence.h"
+
+// message handling
+#include "messagefacility/MessageLogger/MessageLogger.h"
+
+// Offline includes
+#include "Offline/DataProducts/inc/PDGCode.hh"
+#include "Offline/GlobalConstantsService/inc/GlobalConstantsHandle.hh"
+#include "Offline/GlobalConstantsService/inc/ParticleDataList.hh"
+#include "Offline/MCDataProducts/inc/GenId.hh"
+#include "Offline/MCDataProducts/inc/GenParticle.hh"
+#include "Offline/SeedService/inc/SeedService.hh"
+
+// ROOT includes
+#include "art_root_io/TFileService.h"
+#include "TTree.h"
+
+namespace mu2e {
+
+  namespace {
+
+    // Utility function to trim whitespace from a string
+    std::string trim(const std::string& s) {
+      const auto begin = std::find_if_not(s.begin(), s.end(), [](unsigned char ch) { return std::isspace(ch); });
+      if (begin == s.end()) {
+        return "";
+      }
+      const auto end = std::find_if_not(s.rbegin(), s.rend(), [](unsigned char ch) { return std::isspace(ch); }).base();
+      return std::string(begin, end);
+    }
+
+    // Utility function to split a CSV line into fields, trimming whitespace
+    std::vector<std::string> splitCSVLine(const std::string& line) {
+      std::vector<std::string> fields;
+      std::stringstream ss(line);
+      std::string item;
+      while (std::getline(ss, item, ',')) {
+        fields.push_back(trim(item));
+      }
+      return fields;
+    }
+
+    std::string pdgIdToFileToken(const int pdgId) {
+      return (pdgId < 0) ? ("m" + std::to_string(-pdgId)) : std::to_string(pdgId);
+    }
+  }
+
+  class VDResamplerGenerateMix : public art::EDProducer {
+    public:
+      using Name = fhicl::Name;
+      using Comment = fhicl::Comment;
+
+      struct Config {
+        fhicl::Sequence<std::string> hitSummaryFiles{
+          Name("hitSummaryFiles"),
+          Comment("List of VDResamplerConfigure_[dataSourceTag]_HitSummary.csv files")
+        };
+        fhicl::Sequence<double> potsPerFile{
+          Name("potsPerFile"),
+          Comment("List of POT counts corresponding to hitSummaryFiles")
+        };
+        fhicl::Atom<std::string> ModelFileDir{
+          Name("ModelFileDir"),
+          Comment("Directory where the trained SBDM model files are stored")
+        };
+        fhicl::Atom<int> VirtualDetectorID{
+          Name("VirtualDetectorID"),
+          Comment("Virtual detector ID encoded in the model filenames"),
+          116
+        };
+        fhicl::Atom<bool> useHeun{
+          Name("useHeun"),
+          Comment("Whether to use Heun's method for reverse diffusion"),
+          true
+        };
+        fhicl::Atom<int> diffusionSteps{
+          Name("diffusionSteps"),
+          Comment("Number of diffusion steps for sampling"),
+          200
+        };
+        fhicl::Atom<double> VDz0{
+          Name("VDz0"),
+          Comment("Nominal z coordinate of the virtual detector"),
+          37700.39
+        };
+        fhicl::Atom<double> VDr{
+          Name("VDr"),
+          Comment("Virtual detector radius used in the training transform"),
+          2000.0
+        };
+        fhicl::Atom<bool> doROOTDump{
+          Name("doROOTDump"),
+          Comment("Whether to dump generated samples into a ROOT tree"),
+          false
+        };
+      };
+
+      using Parameters = art::EDProducer::Table<Config>;
+
+      explicit VDResamplerGenerateMix(const Parameters& conf);
+      void produce(art::Event& event) override;
+
+    private:
+      struct ParticleModelSet {
+        int pdgId = 0;
+        int hitCount = 0;
+        bool useTwoStageModel = false;
+        std::unique_ptr<ScoreBasedDiffusionModel> allAtOnceModel;
+        std::unique_ptr<ScoreBasedDiffusionModel> stage1Model;
+        std::unique_ptr<ScoreBasedDiffusionModel> stage2Model;
+      };
+
+      struct SourceSummary {
+        std::string hitSummaryFile;
+        double pots = 0.0;
+        double sourceWeight = 0.0;
+        int trainedHitCount = 0;
+        std::vector<ParticleModelSet> particles;
+      };
+
+      art::RandomNumberGenerator::base_engine_t& engine_;
+      CLHEP::RandFlat randFlat_;
+      std::string modelFileDir_;
+      int virtualDetectorID_ = 0;
+      bool useHeun_ = true;
+      int diffusionSteps_ = 0;
+      double VDz0_ = 0.0;
+      double VDr_ = 0.0;
+      bool doROOTDump_ = false;
+      GlobalConstantsHandle<ParticleDataList> pdt_;
+
+      std::vector<SourceSummary> sources_;
+      double totalSourceWeight_ = 0.0;
+
+      double x0_ = -3904.0;
+      double y0_ = 0.0;
+      double t0_ = 1700.0;
+      double tScale_ = 1.0;
+      double p0_ = 1.0;
+
+      TTree* outTree_ = nullptr;
+      int summaryIndex_ = -1;
+      int pdgId_ = 0;
+      int generationMode_ = 0; // 1 = all-at-once, 2 = two-stage
+      double x_gen_ = 0.0;
+      double y_gen_ = 0.0;
+      double z_gen_ = 0.0;
+      double t_gen_ = 0.0;
+      double px_gen_ = 0.0;
+      double py_gen_ = 0.0;
+      double pz_gen_ = 0.0;
+      double mass_gen_ = 0.0;
+      double E_gen_ = 0.0;
+
+      SourceSummary loadSourceSummary(const std::string& summaryFile, double pots);
+      void generateFromModels(
+        const ParticleModelSet& modelSet,
+        double& xTrans,
+        double& yTrans,
+        double& tTrans,
+        double& prTrans,
+        double& pphiTrans,
+        double& pzTrans
+      ) const;
+    };
+
+  VDResamplerGenerateMix::VDResamplerGenerateMix(const Parameters& conf)
+    : art::EDProducer{conf},
+      engine_(createEngine(art::ServiceHandle<SeedService>()->getSeed())),
+      randFlat_(engine_),
+      modelFileDir_(conf().ModelFileDir()),
+      virtualDetectorID_(conf().VirtualDetectorID()),
+      useHeun_(conf().useHeun()),
+      diffusionSteps_(conf().diffusionSteps()),
+      VDz0_(conf().VDz0()),
+      VDr_(conf().VDr()),
+      doROOTDump_(conf().doROOTDump()) {
+
+    produces<GenParticleCollection>();
+
+    const auto& hitSummaryFiles = conf().hitSummaryFiles();
+    const auto& potsPerFile = conf().potsPerFile();
+    if (hitSummaryFiles.size() != potsPerFile.size()) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "hitSummaryFiles and potsPerFile must have the same length.";
+    }
+    if (hitSummaryFiles.empty()) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "At least one hit summary file must be provided.";
+    }
+
+    double referencePots = 0.0;
+    for (size_t i = 0; i < hitSummaryFiles.size(); ++i) {
+      sources_.push_back(loadSourceSummary(hitSummaryFiles[i], potsPerFile[i]));
+      referencePots = std::max(referencePots, sources_.back().pots);
+    }
+
+    for (auto& source : sources_) {
+      // Use a common POT reference scale so we preserve the desired
+      // proportionality trainedHitCount / pots without ever storing tiny
+      // per-source weights like O(1e-11).
+      source.sourceWeight = static_cast<double>(source.trainedHitCount) * (referencePots / source.pots);
+      totalSourceWeight_ += source.sourceWeight;
+    }
+
+    if (totalSourceWeight_ <= 0.0) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "No trained particles with positive weight were found across the provided summaries.";
+    }
+
+    if (doROOTDump_) {
+      art::ServiceHandle<art::TFileService> tfs;
+      outTree_ = tfs->make<TTree>("ttree", "Generated samples from mixed VD resampler models");
+      outTree_->Branch("summaryIndex", &summaryIndex_, "summaryIndex/I");
+      outTree_->Branch("pdgId", &pdgId_, "pdgId/I");
+      outTree_->Branch("generationMode", &generationMode_, "generationMode/I");
+      outTree_->Branch("x", &x_gen_, "x/D");
+      outTree_->Branch("y", &y_gen_, "y/D");
+      outTree_->Branch("z", &z_gen_, "z/D");
+      outTree_->Branch("time", &t_gen_, "time/D");
+      outTree_->Branch("px", &px_gen_, "px/D");
+      outTree_->Branch("py", &py_gen_, "py/D");
+      outTree_->Branch("pz", &pz_gen_, "pz/D");
+      outTree_->Branch("E", &E_gen_, "E/D");
+    }
+  }
+
+  VDResamplerGenerateMix::SourceSummary VDResamplerGenerateMix::loadSourceSummary(
+    const std::string& summaryFile,
+    const double pots
+  ) {
+    if (pots <= 0.0) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "potsPerFile must be positive for " << summaryFile;
+    }
+
+    std::ifstream in(summaryFile);
+    if (!in) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "Cannot open hit summary file " << summaryFile;
+    }
+
+    SourceSummary source;
+    source.hitSummaryFile = summaryFile;
+    source.pots = pots;
+
+    std::string line;
+    bool firstLine = true;
+    while (std::getline(in, line)) {
+      line = trim(line);
+      if (line.empty()) {
+        continue;
+      }
+      if (firstLine) {
+        firstLine = false;
+        continue; // header
+      }
+
+      const std::vector<std::string> fields = splitCSVLine(line);
+      if (fields.size() < 3) {
+        throw cet::exception("VDResamplerGenerateMix")
+          << "Malformed summary line in " << summaryFile << ": " << line;
+      }
+
+      const int pdgId = std::stoi(fields[0]);
+      const int hitCount = std::stoi(fields[1]);
+      const std::string modeToken = fields[2];
+      if (modeToken == "*") {
+        continue;
+      }
+
+      ParticleModelSet particle;
+      particle.pdgId = pdgId;
+      particle.hitCount = hitCount;
+      particle.useTwoStageModel = (modeToken == "2");
+      if (modeToken != "1" && modeToken != "2") {
+        throw cet::exception("VDResamplerGenerateMix")
+          << "Unexpected model mode token '" << modeToken << "' in " << summaryFile;
+      }
+
+      const std::string pdgToken = pdgIdToFileToken(pdgId);
+      if (particle.useTwoStageModel) {
+        const std::string stage1ModelFile =
+          modelFileDir_ + "/SBDMmodel_stage1_VD" + std::to_string(virtualDetectorID_) + "_pdg" + pdgToken + ".csv";
+        const std::string stage2ModelFile =
+          modelFileDir_ + "/SBDMmodel_stage2_VD" + std::to_string(virtualDetectorID_) + "_pdg" + pdgToken + ".csv";
+
+        particle.stage1Model = std::make_unique<ScoreBasedDiffusionModel>(
+          ScoreBasedDiffusionModel::loadModel(engine_, stage1ModelFile)
+        );
+        particle.stage2Model = std::make_unique<ScoreBasedDiffusionModel>(
+          ScoreBasedDiffusionModel::loadModel(engine_, stage2ModelFile)
+        );
+      } else {
+        const std::string modelFile =
+          modelFileDir_ + "/SBDMmodel_allAtOnce_VD" + std::to_string(virtualDetectorID_) + "_pdg" + pdgToken + ".csv";
+        particle.allAtOnceModel = std::make_unique<ScoreBasedDiffusionModel>(
+          ScoreBasedDiffusionModel::loadModel(engine_, modelFile)
+        );
+      }
+
+      source.trainedHitCount += hitCount;
+      source.particles.push_back(std::move(particle));
+    }
+
+    if (source.particles.empty()) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "Summary file " << summaryFile << " contains no trained particle entries.";
+    }
+
+    source.sourceWeight = 0.0;
+    return source;
+  }
+
+  void VDResamplerGenerateMix::generateFromModels(
+    const ParticleModelSet& modelSet,
+    double& xTrans,
+    double& yTrans,
+    double& tTrans,
+    double& prTrans,
+    double& pphiTrans,
+    double& pzTrans
+  ) const {
+    if (modelSet.useTwoStageModel) {
+      const std::vector<double> stage1Sample = modelSet.stage1Model->generateSample({}, useHeun_, diffusionSteps_);
+      if (stage1Sample.size() != 3u) {
+        throw cet::exception("VDResamplerGenerateMix")
+          << "Stage-1 model returned " << stage1Sample.size() << " values, expected 3.";
+      }
+
+      tTrans = stage1Sample[0];
+      xTrans = stage1Sample[1];
+      yTrans = stage1Sample[2];
+
+      const std::vector<double> condition = {tTrans, xTrans, yTrans};
+      const std::vector<double> stage2Sample = modelSet.stage2Model->generateSample(condition, useHeun_, diffusionSteps_);
+      if (stage2Sample.size() != 3u) {
+        throw cet::exception("VDResamplerGenerateMix")
+          << "Stage-2 model returned " << stage2Sample.size() << " values, expected 3.";
+      }
+
+      prTrans = stage2Sample[0];
+      pphiTrans = stage2Sample[1];
+      pzTrans = stage2Sample[2];
+    } else {
+      const std::vector<double> sample = modelSet.allAtOnceModel->generateSample({}, useHeun_, diffusionSteps_);
+      if (sample.size() != 6u) {
+        throw cet::exception("VDResamplerGenerateMix")
+          << "All-at-once model returned " << sample.size() << " values, expected 6.";
+      }
+
+      tTrans = sample[0];
+      xTrans = sample[1];
+      yTrans = sample[2];
+      prTrans = sample[3];
+      pphiTrans = sample[4];
+      pzTrans = sample[5];
+    }
+  }
+
+  void VDResamplerGenerateMix::produce(art::Event& event) {
+    auto output = std::make_unique<GenParticleCollection>();
+
+    // Choose which source summary to use according to trained hit yield per POT.
+    double draw = randFlat_.fire() * totalSourceWeight_;
+    size_t chosenSourceIndex = 0;
+    for (; chosenSourceIndex + 1 < sources_.size(); ++chosenSourceIndex) {
+      draw -= sources_[chosenSourceIndex].sourceWeight;
+      if (draw < 0.0) {
+        break;
+      }
+    }
+    const SourceSummary& source = sources_[chosenSourceIndex];
+
+    // Choose particle type according to the hit ratios within the selected source.
+    double totalParticleWeight = 0.0;
+    for (const auto& particle : source.particles) {
+      totalParticleWeight += particle.hitCount;
+    }
+    if (totalParticleWeight <= 0.0) {
+      throw cet::exception("VDResamplerGenerateMix")
+        << "Selected source " << source.hitSummaryFile << " has non-positive particle weight.";
+    }
+
+    draw = randFlat_.fire() * totalParticleWeight;
+    size_t chosenParticleIndex = 0;
+    for (; chosenParticleIndex + 1 < source.particles.size(); ++chosenParticleIndex) {
+      draw -= source.particles[chosenParticleIndex].hitCount;
+      if (draw < 0.0) {
+        break;
+      }
+    }
+    const ParticleModelSet& modelSet = source.particles[chosenParticleIndex];
+
+    double xTrans = 0.0;
+    double yTrans = 0.0;
+    double tTrans = 0.0;
+    double prTrans = 0.0;
+    double pphiTrans = 0.0;
+    double pzTrans = 0.0;
+    generateFromModels(modelSet, xTrans, yTrans, tTrans, prTrans, pphiTrans, pzTrans);
+
+    const double u = std::sqrt(xTrans * xTrans + yTrans * yTrans);
+    const double theta = std::atan2(yTrans, xTrans);
+    const double rho = std::tanh(u);
+    const double r = rho * VDr_;
+    const double dx = r * std::cos(theta);
+    const double dy = r * std::sin(theta);
+    x_gen_ = dx + x0_;
+    y_gen_ = dy + y0_;
+    z_gen_ = VDz0_;
+
+    t_gen_ = t0_ * std::exp(tTrans * tScale_);
+
+    const double pr = p0_ * std::sinh(prTrans);
+    const double pphi = p0_ * std::sinh(pphiTrans);
+    pz_gen_ = p0_ * std::sinh(pzTrans);
+
+    if (r > 1e-6) {
+      const double rx = dx / r;
+      const double ry = dy / r;
+      const double phix = -ry;
+      const double phiy = rx;
+      px_gen_ = pr * rx + pphi * phix;
+      py_gen_ = pr * ry + pphi * phiy;
+    } else {
+      px_gen_ = pr;
+      py_gen_ = pphi;
+    }
+
+    pdgId_ = modelSet.pdgId;
+    generationMode_ = modelSet.useTwoStageModel ? 2 : 1;
+    summaryIndex_ = static_cast<int>(chosenSourceIndex);
+
+    mass_gen_ = pdt_->particle(pdgId_).mass();
+    const CLHEP::Hep3Vector momParticle(px_gen_, py_gen_, pz_gen_);
+    const CLHEP::Hep3Vector posParticle(x_gen_, y_gen_, z_gen_);
+    const double eTotal = std::sqrt(momParticle.mag2() + mass_gen_ * mass_gen_);
+    E_gen_ = eTotal - mass_gen_;
+    const CLHEP::HepLorentzVector fourMomParticle(eTotal, momParticle);
+
+    output->emplace_back(
+      PDGCode::type(pdgId_),
+      GenId::STMDownStreamGenTool,
+      posParticle,
+      fourMomParticle,
+      t_gen_
+    );
+
+    event.put(std::move(output));
+
+    if (doROOTDump_) {
+      outTree_->Fill();
+    }
+  }
+
+} // namespace mu2e
+
+DEFINE_ART_MODULE(mu2e::VDResamplerGenerateMix)
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
index 72f5fdd94d..e7a0be2463 100644
--- a/STMMC/src/VDResamplerTrain_module.cc
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -1,18 +1,34 @@
+// VDResamplerTrain_module.cc
+// For StepPointMCs on the designated virtual detectors, train either:
+//   1) a two-stage score-based diffusion model with
+//      Stage 1: unconditional model for (t', x', y')
+//      Stage 2: conditional model for (p_r', p_phi', p_z' | t', x', y')
+//   2) a single unconditional score-based diffusion model for
+//      (t', x', y', p_r', p_phi', p_z')
+// and store the trained model parameters in CSV files.
+// note that p_z are filtered and only hits with positive p_z are kept
+// Yongyi Wu, Mar. 2026 
+
 // stdlib includes
 #include <cmath>
 #include <iostream>
+#include <fstream>
+
+#include "Offline/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh"
 
 // art includes
 #include "art/Framework/Core/EDAnalyzer.h"
 #include "art/Framework/Principal/Event.h"
 #include "art/Framework/Principal/Handle.h"
 #include "art/Framework/Principal/Run.h"
+#include "art/Framework/Services/Optional/RandomNumberGenerator.h"
 
 // exception handling
 #include "cetlib_except/exception.h"
 
 // fhicl includes
 #include "canvas/Utilities/InputTag.h"
+#include "fhiclcpp/ParameterSet.h"
 #include "fhiclcpp/types/Atom.h"
 
 // message handling
@@ -23,19 +39,47 @@
 #include "Offline/GlobalConstantsService/inc/ParticleDataList.hh"
 #include "Offline/MCDataProducts/inc/SimParticle.hh"
 #include "Offline/MCDataProducts/inc/StepPointMC.hh"
+#include "Offline/SeedService/inc/SeedService.hh"
 
 // ROOT includes
 #include "art_root_io/TFileService.h"
 #include "TTree.h"
 
+typedef cet::map_vector_key key_type;
+typedef unsigned long VolumeId_type;
+
 namespace mu2e {
   class VDResamplerTrain : public art::EDAnalyzer {
     public:
       using Name=fhicl::Name;
       using Comment=fhicl::Comment;
       struct Config {
-        fhicl::Atom<art::InputTag> StepPointMCsTag{Name("StepPointMCsTag"), Comment("Tag identifying the StepPointMCs")};
-        fhicl::Atom<art::InputTag> SimParticlemvTag{Name("SimParticlemvTag"), Comment("Tag identifying the SimParticlemv")};
+        fhicl::Atom<art::InputTag> StepPointMCsTag{   Name("StepPointMCsTag"),         Comment("Tag identifying the StepPointMCs")};
+        fhicl::Atom<art::InputTag> SimParticlemvTag{  Name("SimParticlemvTag"),        Comment("Tag identifying the SimParticlemv")};
+        fhicl::Atom<bool> SBDMuseTwoStageTraining{    Name("SBDMuseTwoStageTraining"), Comment("If true, train the two-stage factorized model. If false, train all 6 dimensions at once."), true};
+        fhicl::Atom<std::string> SBDMallAtOnceModelFile{ Name("SBDMallAtOnceModelFile"), Comment("CSV filename for the all-at-once 6D SBDM model parameters"), ""};
+        fhicl::Atom<std::string> SBDMstage1ModelFile{ Name("SBDMstage1ModelFile"),     Comment("CSV filename for the trained stage-1 SBDM model parameters"), ""};
+        fhicl::Atom<std::string> SBDMstage2ModelFile{ Name("SBDMstage2ModelFile"),     Comment("CSV filename for the trained stage-2 SBDM model parameters"), ""};
+        fhicl::Atom<int> VirtualDetectorID{           Name("VirtualDetectorID"),       Comment("ID of the virtual detector to train on"),                        116};
+        fhicl::Atom<double> VDz0{                     Name("VDz0"),                    Comment("z coordinate of the virtual detector"),                          37700.39};
+        fhicl::Atom<double> VDr{                      Name("VDr"),                     Comment("VD radius"),                                                     2000.0 };
+        fhicl::Atom<int> pdgID{                       Name("pdgID"),                   Comment("pdgID of the particle to train on"),                             22};
+        fhicl::Atom<int> SBDMhidden{                  Name("SBDMhidden"),              Comment("Size of hidden layers in the SBDM neural network"),              128};
+        fhicl::Atom<int> SBDMlayers{                  Name("SBDMlayers"),              Comment("Number of layers in the SBDM neural network"),                   4};
+        fhicl::Atom<std::string> SBDMoptimizer{       Name("SBDMoptimizer"),           Comment("Optimizer for training the SBDM neural network (SGD or ADAM)"), "ADAM"};
+        fhicl::Atom<double> SBDMadamBeta1{            Name("SBDMadamBeta1"),           Comment("Adam optimizer beta1 parameter"),                                0.9};
+        fhicl::Atom<double> SBDMadamBeta2{            Name("SBDMadamBeta2"),           Comment("Adam optimizer beta2 parameter"),                                0.999};
+        fhicl::Atom<double> SBDMadamEps{              Name("SBDMadamEps"),             Comment("Adam optimizer epsilon parameter"),                              1e-8};
+        fhicl::Atom<std::string> SBDMnoiseSchedule{   Name("SBDMnoiseSchedule"),       Comment("Noise schedule for the SBDM (LINEAR or COSINE)"),                "COSINE"};
+        fhicl::Atom<double> SBDMbetaMin{              Name("SBDMbetaMin"),             Comment("Minimum noise schedule parameter (for LINEAR schedule)") ,       1e-4};
+        fhicl::Atom<double> SBDMbetaMax{              Name("SBDMbetaMax"),             Comment("Maximum noise schedule parameter (for LINEAR schedule)"),        0.02};
+        fhicl::Atom<double> SBDMcosineOffset{         Name("SBDMcosineOffset"),        Comment("Offset parameter (for cosine schedule)"),                        0.008};
+        fhicl::Atom<int> SBDMbatchSize{               Name("SBDMbatchSize"),           Comment("Batch size for training the SBDM"),                              32};
+        fhicl::Atom<double> SBDMgradientClip{         Name("SBDMgradientClip"),        Comment("Gradient clipping threshold for training the SBDM"),             1.0};
+        fhicl::Atom<double> SBDMlearningRate{         Name("SBDMlearningRate"),        Comment("Learning rate for training the SBDM"),                           1e-3};
+        fhicl::Atom<int> SBDMdiffusionSteps{          Name("SBDMdiffusionSteps"),      Comment("Number of steps in the diffusion process for the SBDM"),         200};
+        fhicl::Atom<int> SBDMtrainingSize{            Name("SBDMtrainingSize"),        Comment("Size of the training data for the SBDM"),                        -1}; // -1 means use all available data
+        fhicl::Atom<int> SBDMtrainingEpochs{          Name("SBDMtrainingEpochs"),      Comment("Number of epochs to train the selected SBDM mode"),              10};
       };
       using Parameters = art::EDAnalyzer::Table<Config>;
       explicit VDResamplerTrain(const Parameters& conf);
@@ -44,38 +88,139 @@ namespace mu2e {
     private:
       art::ProductToken<StepPointMCCollection> StepPointMCsToken;
       art::ProductToken<SimParticleCollection> SimParticlemvToken;
-      GlobalConstantsHandle<ParticleDataList> pdt;
+      art::RandomNumberGenerator::base_engine_t& engine;
+      
+      // SBDM model + data
+      bool useTwoStageTraining = true;
+      std::unique_ptr<ScoreBasedDiffusionModel> allAtOnceModel;
+      std::unique_ptr<ScoreBasedDiffusionModel> stage1Model;
+      std::unique_ptr<ScoreBasedDiffusionModel> stage2Model;
+      std::vector<DiffusionTrainingSample> allAtOnceTrainingData;
+      std::vector<DiffusionTrainingSample> stage1TrainingData;
+      std::vector<DiffusionTrainingSample> stage2TrainingData;
+      std::string SBDMallAtOnceModelFile;
+      std::string SBDMstage1ModelFile;
+      std::string SBDMstage2ModelFile;
+      VolumeId_type VirtualDetectorID = 0;
+      double VDz0 = 0.0;
+      double VDr = 0.0;
+      int pdgID = 0;
+      int trainingEpochs = 0;
+      int trainingSize = -1;
 
-      // Declare relevant variables
+      // variables to read from the art event
+      int stepPdgId = 0;
+      double x = 0.0, y = 0.0, z = 0.0, time = 0.0;
+      double px = 0.0, py = 0.0, pz = 0.0; 
+      VolumeId_type virtualdetectorId = 0;
 
+      // transform variables for training data preparation
+      double x0 = -3904.0;
+      double y0 = 0.0;
+      // time scaling
+      double t0 = 1700.0;   // ns
+      double tScale = 1.0;
+  };
 
-      int pdgId = 0;
-      double x = 0.0, y = 0.0, z = 0.0, mass = 0.0, E = 0.0, time = 0.0;
-      VolumeId_type virtualdetectorId = 0;
-      TTree* ttree;
-      std::map<int, int> pdgIds; // <id, count>
+  VDResamplerTrain::VDResamplerTrain(const Parameters& conf) :
+    art::EDAnalyzer(conf),
+    engine(createEngine( art::ServiceHandle<SeedService>()->getSeed())),
+    StepPointMCsToken(consumes<StepPointMCCollection>(conf().StepPointMCsTag())),
+    SimParticlemvToken(consumes<SimParticleCollection>(conf().SimParticlemvTag())),
+    useTwoStageTraining(conf().SBDMuseTwoStageTraining()),
+    SBDMallAtOnceModelFile(conf().SBDMallAtOnceModelFile()),
+    SBDMstage1ModelFile(conf().SBDMstage1ModelFile()),
+    SBDMstage2ModelFile(conf().SBDMstage2ModelFile()),
+    VirtualDetectorID(conf().VirtualDetectorID()),
+    VDz0(conf().VDz0()),
+    VDr(conf().VDr()),
+    pdgID(conf().pdgID()),
+    trainingEpochs(conf().SBDMtrainingEpochs()),
+    trainingSize(conf().SBDMtrainingSize())
+  {
+    // optimizer selection
+    ScoreBasedDiffusionModel::OptimizerType opt =
+        (conf().SBDMoptimizer() == "SGD") ?
+        ScoreBasedDiffusionModel::OptimizerType::SGD :
+        ScoreBasedDiffusionModel::OptimizerType::ADAM;
 
+    // noise schedule
+    ScoreBasedDiffusionModel::NoiseScheduleType sched =
+        (conf().SBDMnoiseSchedule() == "LINEAR") ?
+        ScoreBasedDiffusionModel::NoiseScheduleType::LINEAR :
+        ScoreBasedDiffusionModel::NoiseScheduleType::COSINE;
 
+    if (useTwoStageTraining) {
+      // create stage-1 model for (t', x', y')
+      stage1Model = std::make_unique<ScoreBasedDiffusionModel>(
+          engine,
+          3, // dim
+          0, // conditionDim
+          conf().SBDMhidden(),
+          conf().SBDMlayers(),
+          opt,
+          conf().SBDMadamBeta1(),
+          conf().SBDMadamBeta2(),
+          conf().SBDMadamEps(),
+          sched,
+          conf().SBDMbetaMin(),
+          conf().SBDMbetaMax(),
+          conf().SBDMcosineOffset(),
+          conf().SBDMbatchSize(),
+          conf().SBDMgradientClip(),
+          conf().SBDMlearningRate(),
+          conf().SBDMdiffusionSteps()
+      );
 
+      // create stage-2 model for (p_r', p_phi', p_z' | t', x', y')
+      stage2Model = std::make_unique<ScoreBasedDiffusionModel>(
+          engine,
+          3, // dim
+          3, // conditionDim
+          conf().SBDMhidden(),
+          conf().SBDMlayers(),
+          opt,
+          conf().SBDMadamBeta1(),
+          conf().SBDMadamBeta2(),
+          conf().SBDMadamEps(),
+          sched,
+          conf().SBDMbetaMin(),
+          conf().SBDMbetaMax(),
+          conf().SBDMcosineOffset(),
+          conf().SBDMbatchSize(),
+          conf().SBDMgradientClip(),
+          conf().SBDMlearningRate(),
+          conf().SBDMdiffusionSteps()
+      );
 
-  };
+      stage1TrainingData.reserve(1000000);
+      stage2TrainingData.reserve(1000000);
+    } else {
+      allAtOnceModel = std::make_unique<ScoreBasedDiffusionModel>(
+          engine,
+          6, // dim
+          0, // conditionDim
+          conf().SBDMhidden(),
+          conf().SBDMlayers(),
+          opt,
+          conf().SBDMadamBeta1(),
+          conf().SBDMadamBeta2(),
+          conf().SBDMadamEps(),
+          sched,
+          conf().SBDMbetaMin(),
+          conf().SBDMbetaMax(),
+          conf().SBDMcosineOffset(),
+          conf().SBDMbatchSize(),
+          conf().SBDMgradientClip(),
+          conf().SBDMlearningRate(),
+          conf().SBDMdiffusionSteps()
+      );
 
-  VirtualDetectorTree::VirtualDetectorTree(const Parameters& conf) :
-    art::EDAnalyzer(conf),
-    StepPointMCsToken(consumes<StepPointMCCollection>(conf().StepPointMCsTag())),
-    SimParticlemvToken(consumes<SimParticleCollection>(conf().SimParticlemvTag())) {
-      art::ServiceHandle<art::TFileService> tfs;
-      ttree = tfs->make<TTree>( "ttree", "Virtual Detectors ttree");
-      ttree->Branch("time", &time, "time/D"); // ns
-      ttree->Branch("virtualdetectorId", &virtualdetectorId, "virtualdetectorId/l");
-      ttree->Branch("pdgId", &pdgId, "pdgId/I");
-      ttree->Branch("x", &x, "x/D"); // mm
-      ttree->Branch("y", &y, "y/D"); // mm
-      ttree->Branch("z", &z, "z/D"); // mm
-      ttree->Branch("E", &E, "E/D"); // MeV
-    };
+      allAtOnceTrainingData.reserve(1000000);
+    }
+  };
 
-  void VirtualDetectorTree::analyze(const art::Event& event) {
+  void VDResamplerTrain::analyze(const art::Event& event) {
     // Get the data products from the event
     auto const& StepPointMCs = event.getProduct(StepPointMCsToken);
     if (StepPointMCs.empty())
@@ -90,27 +235,146 @@ namespace mu2e {
       const SimParticle& particle = SimParticles.at(step.trackId());
 
       // Extract the parameters
-      time = step.time();
+      stepPdgId = particle.pdgId();
       virtualdetectorId = step.virtualDetectorId();
-      pdgId = particle.pdgId();
-      x = step.position().x();
+      time = step.time();
+      x = step.position().x(); // Shift the x coordinate to be relative to the centre of beamline
       y = step.position().y();
       z = step.position().z();
-      mass = pdt->particle(pdgId).mass();
-      E = std::sqrt(step.momentum().mag2()+mass*mass)-mass; // Subtract the rest mass
-      if (E < 0)
-        throw cet::exception("LogicError", "Energy is negative");
-      ttree->Fill();
-
-      // Generate the data summary
-      if (pdgIds.find(pdgId) != pdgIds.end())
-        pdgIds[pdgId] += 1;
-      else
-        pdgIds.emplace(std::make_pair(pdgId, 1));
+      px = step.momentum().x();
+      py = step.momentum().y();
+      pz = step.momentum().z();
+
+      if (virtualdetectorId != VirtualDetectorID || (stepPdgId != pdgID && pdgID != 0) || pz <= 0)
+        continue; // Filter hits based on the virtual detector ID, particle type, and pz
+      
+      // as z maybe slightly different from the nominal VDz0 due to the step size, we will extrapolate the (x, y) coordinates 
+      // to the nominal VDz0 for all hits to compute the training parameters to be fed into the SBDM.
+      double extrapolationFactor = (VDz0 - z) / pz; // Assuming linear motion, this is the factor to extrapolate from current z to z0
+      double x_extrapolated = x + extrapolationFactor * px;
+      double y_extrapolated = y + extrapolationFactor * py;
+
+      // Now we have the extrapolated (x, y) at the nominal VDz0, we can compute the training parameters for the SBDM.
+      // We convert (t, x_extrapolated, y_extrapolated, px, py, pz) to (t', x_extrapolated, y_extrapolated, p_r', p_phi', p_z')
+
+      // shift to detector-centered coordinates
+      double dx = x_extrapolated - x0;
+      double dy = y_extrapolated - y0;
+      // polar coordinates
+      double r = std::sqrt(dx*dx + dy*dy);
+      double rho = r / VDr;
+      // numerical safety (avoid rho >= 1)
+      const double eps = 1e-6;
+      rho = std::min(rho, 1.0 - eps);
+      // boundary-removing transform
+      // u = atanh(r/R)
+      double u = 0.5 * std::log((1.0 + rho) / (1.0 - rho));
+      // angle
+      double theta = std::atan2(dy, dx);
+      // map back to Cartesian-like coordinates
+      double x_trans = u * std::cos(theta);
+      double y_trans = u * std::sin(theta);
+
+      // compute momentum components in the local polar coordinate system (r, phi, z)
+      double pr = 0.0;
+      double pphi = 0.0;
+      if (r > 1e-6) { // avoid division by zero, if r is very small, we can approximate pr ~ px and pphi ~ py
+        double rx = dx / r; // unit vector in the radial direction
+        double ry = dy / r;
+        double phix = -ry; // unit vector in the angular direction)
+        double phiy =  rx;
+        pr   = px*rx + py*ry; // radial momentum component
+        pphi = px*phix + py*phiy;
+      }
+      else {
+        pr = px;
+        pphi = py;
+      }
+      // momentum scaling
+      double p0 = 1.0; // tunable scale where I want best resolution
+      double pr_t   = std::asinh(pr   / p0);
+      double pphi_t = std::asinh(pphi / p0);
+      double pz_t   = std::asinh(pz   / p0);
+
+      // time transform
+      double t_safe = std::max(time, 0.1); // avoid log(0)
+      double t_trans = std::log(t_safe / t0) / tScale;
+
+      if (useTwoStageTraining) {
+        DiffusionTrainingSample stage1Sample;
+        stage1Sample.x = {t_trans, x_trans, y_trans};
+        stage1TrainingData.push_back(std::move(stage1Sample));
+
+        DiffusionTrainingSample stage2Sample;
+        stage2Sample.x = {pr_t, pphi_t, pz_t};
+        stage2Sample.cond = {t_trans, x_trans, y_trans};
+        stage2TrainingData.push_back(std::move(stage2Sample));
+      } else {
+        DiffusionTrainingSample allAtOnceSample;
+        allAtOnceSample.x = {t_trans, x_trans, y_trans, pr_t, pphi_t, pz_t};
+        allAtOnceTrainingData.push_back(std::move(allAtOnceSample));
+      }
     };
     return;
   };
 
+  void VDResamplerTrain::endJob() {
+    
+      if (useTwoStageTraining && (stage1TrainingData.empty() || stage2TrainingData.empty())) {
+        mf::LogWarning("VDResamplerTrain") << "No training data collected.";
+        return;
+      }
+      if (!useTwoStageTraining && allAtOnceTrainingData.empty()) {
+        mf::LogWarning("VDResamplerTrain") << "No training data collected.";
+        return;
+      }
+
+      // if SBDMtrainingSize is set and smaller than the collected training data, 
+      // truncate the training data to the specified size. 
+      if (useTwoStageTraining) {
+        if(trainingSize > 0 && (int)stage1TrainingData.size() > trainingSize)
+          stage1TrainingData.resize(trainingSize);
+        if(trainingSize > 0 && (int)stage2TrainingData.size() > trainingSize)
+          stage2TrainingData.resize(trainingSize);
+
+        if (SBDMstage1ModelFile.empty() || SBDMstage2ModelFile.empty()) {
+          throw cet::exception("VDResamplerTrain") << "Two-stage training requires both SBDMstage1ModelFile and SBDMstage2ModelFile.";
+        }
+
+        mf::LogInfo("VDResamplerTrain")
+            << "Training stage-1 diffusion model with " << stage1TrainingData.size()
+            << " samples and " << trainingEpochs << " epochs...";
+        stage1Model->train(stage1TrainingData, trainingEpochs);
+        stage1Model->saveModel(SBDMstage1ModelFile);
+        mf::LogInfo("VDResamplerTrain") << "Stage-1 model saved to " << SBDMstage1ModelFile;
+
+        mf::LogInfo("VDResamplerTrain")
+            << "Training stage-2 diffusion model with " << stage2TrainingData.size()
+            << " samples and " << trainingEpochs << " epochs...";
+        stage2Model->train(stage2TrainingData, trainingEpochs);
+        stage2Model->saveModel(SBDMstage2ModelFile);
+        mf::LogInfo("VDResamplerTrain") << "Stage-2 model saved to " << SBDMstage2ModelFile;
+
+      } else {
+        if(trainingSize > 0 && (int)allAtOnceTrainingData.size() > trainingSize)
+          allAtOnceTrainingData.resize(trainingSize);
+
+        if (SBDMallAtOnceModelFile.empty()) {
+          throw cet::exception("VDResamplerTrain") << "All-at-once training requires SBDMallAtOnceModelFile.";
+        }
+
+        mf::LogInfo("VDResamplerTrain")
+            << "Training all-at-once diffusion model with " << allAtOnceTrainingData.size()
+            << " samples and " << trainingEpochs << " epochs...";
+        allAtOnceModel->train(allAtOnceTrainingData, trainingEpochs);
+        allAtOnceModel->saveModel(SBDMallAtOnceModelFile);
+
+        mf::LogInfo("VDResamplerTrain") << "All-at-once model saved to " << SBDMallAtOnceModelFile;
+      }
+
+     return;
+  };
+
 }; // end namespace mu2e
 
 DEFINE_ART_MODULE(mu2e::VDResamplerTrain)

From ddfc0a382442f03ee9acbe20c05d2fc03169fcda Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2ebuild02.fnal.gov>
Date: Tue, 24 Mar 2026 17:32:12 -0500
Subject: [PATCH 03/15] Debugging update; build successful

---
 MachineLearningTools/CMakeLists.txt           | 11 ++++++++
 .../inc/ScoreBasedDiffusionModel.hh           | 26 ++++++++++---------
 MachineLearningTools/src/SConscript           | 20 ++++++++++++++
 .../src/ScoreBasedDiffusionModel.cc           | 17 +++++++-----
 STMMC/CMakeLists.txt                          |  3 +++
 STMMC/src/SConscript                          |  2 ++
 STMMC/src/VDResamplerConfigure_module.cc      |  4 +--
 .../VDResamplerGenerateFromModel_module.cc    | 13 +++++++---
 STMMC/src/VDResamplerGenerateMix_module.cc    | 10 ++++---
 STMMC/src/VDResamplerTrain_module.cc          | 19 +++++++++++---
 10 files changed, 94 insertions(+), 31 deletions(-)
 create mode 100644 MachineLearningTools/CMakeLists.txt
 create mode 100644 MachineLearningTools/src/SConscript

diff --git a/MachineLearningTools/CMakeLists.txt b/MachineLearningTools/CMakeLists.txt
new file mode 100644
index 0000000000..b837296aff
--- /dev/null
+++ b/MachineLearningTools/CMakeLists.txt
@@ -0,0 +1,11 @@
+cet_make_library(
+    SOURCE
+      src/ScoreBasedDiffusionModel.cc
+    LIBRARIES PUBLIC
+      CLHEP::CLHEP
+      messagefacility::messagefacility
+      cetlib_except::cetlib_except
+)
+
+install_source(SUBDIRS src)
+install_headers(USE_PROJECT_NAME SUBDIRS inc)
diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
index ae9afea62b..d0aed40046 100644
--- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -13,13 +13,12 @@
 #include <sstream>
 #include <iostream>
 #include <cctype>
+#include <cassert>
 
 #include "CLHEP/Random/RandomEngine.h"
 #include "CLHEP/Random/RandFlat.h"
 #include "CLHEP/Random/RandGaussQ.h"
 
-#include "Offline/SeedService/inc/SeedService.hh"
-
 #include "messagefacility/MessageLogger/MessageLogger.h"
 #include "cetlib_except/exception.h"
 
@@ -43,10 +42,11 @@ namespace mu2e{
             COSINE   // Cosine noise schedule
         };
 
-        // Constructor: Initialize diffusion model with random engine.
+        // Constructor: Initialize diffusion model with CLHEP random distributions.
         //
         // Parameters:
-        //   engine                - Reference to CLHEP random engine (externally managed)
+        //   randFlat              - Reference to CLHEP RandFlat for uniform sampling (externally managed)
+        //   randGaussQ            - Reference to CLHEP RandGaussQ for Gaussian noise (externally managed)
         //   dim                   - Dimensionality of the state space
         //   conditionDim          - Dimensionality of the optional conditioning vector (default: 0 for unconditional model)
         //   hidden                - Size of hidden layers in the neural network
@@ -65,7 +65,8 @@ namespace mu2e{
         //   diffusionSteps        - Number of steps in the diffusion process (default: 200)
         ScoreBasedDiffusionModel(
             // Network architecture parameters
-            art::RandomNumberGenerator::base_engine_t& engine,
+            CLHEP::RandFlat& randFlat,
+            CLHEP::RandGaussQ& randGaussQ,
             int dim,
             int conditionDim,
             int hidden,
@@ -129,9 +130,13 @@ namespace mu2e{
         // model and pick up the training process where it left off. The loaded model can only be used for sampling.
         //
         // Parameters:
-        //   engine   - Reference to CLHEP random engine (externally managed, must be valid for lifetime of model)
-        //   filename - Path to the file from which model parameters will be loaded
-        static ScoreBasedDiffusionModel loadModel(CLHEP::HepRandomEngine& engine, const std::string& filename);
+        //   randFlat / RandGaussQ - CLHEP random number generator wrappers being passed
+        //   filename              - Path to the file from which model parameters will be loaded
+        static ScoreBasedDiffusionModel loadModel(
+            CLHEP::RandFlat& randFlat,
+            CLHEP::RandGaussQ& randGaussQ,
+            const std::string& filename
+        );
 
     private:
 
@@ -263,10 +268,7 @@ namespace mu2e{
         // ----- internal vars -----
         
         // The random engine is NOT owned by this class. It is injected externally
-        // by the framework. Engine reference must remain valid for the lifetime
-        // of this object.
-        art::RandomNumberGenerator::base_engine_t& engine_;
-        // CLHEP distribution wrappers for actual random number generation.
+        // by the framework. Below are CLHEP distribution wrappers for actual random number generation.
         // These wrap the engine_ and provide specific probability distributions.
         // - RandFlat:   Uniform distribution on [0,1)
         // - RandGaussQ: Gaussian (normal) distribution with mean=0, sigma=1 (or custom)
diff --git a/MachineLearningTools/src/SConscript b/MachineLearningTools/src/SConscript
new file mode 100644
index 0000000000..f74c4ad579
--- /dev/null
+++ b/MachineLearningTools/src/SConscript
@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+#
+# Script to build the files found in this directory.
+#
+
+import os
+Import('env')
+Import('mu2e_helper')
+
+helper=mu2e_helper(env)
+
+mainlib = helper.make_mainlib ( [ 'CLHEP',
+                                  'cetlib',
+                                  'cetlib_except',
+                                  'MF_MessageLogger' ] )
+
+# This tells emacs to view this file in python mode.
+# Local Variables:
+# mode:python
+# End:
diff --git a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
index 67829da8f2..74d904118b 100644
--- a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
+++ b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
@@ -39,7 +39,8 @@ namespace mu2e {
 
     
     ScoreBasedDiffusionModel::ScoreBasedDiffusionModel(
-        art::RandomNumberGenerator::base_engine_t& engine,
+        CLHEP::RandFlat& randFlat,
+        CLHEP::RandGaussQ& randGaussQ,
         int dim,
         int conditionDim,
         int hidden,
@@ -56,7 +57,7 @@ namespace mu2e {
         double gradientClipThreshold,
         double learningRate,
         int diffusionSteps
-    ) : engine_(engine), randFlat_(engine_), randGaussQ_(engine_),
+    ) : randFlat_(randFlat), randGaussQ_(randGaussQ),
         dim_(dim), conditionDim_(conditionDim), hidden_(hidden), layers_(layers),
         optimizerType_(optimizerType),
         adamBeta1_(adamBeta1), adamBeta2_(adamBeta2), adamEps_(adamEps),
@@ -64,7 +65,7 @@ namespace mu2e {
         betaMin_(betaMin), betaMax_(betaMax), cosineOffset_(cosineOffset),
         batchSize_(batchSize), gradientClipThreshold_(gradientClipThreshold), learningRate_(learningRate),
         diffusionSteps_(diffusionSteps),
-        runningLoss_(0.0), adamStep_(0), epochLosses_(), trainingSampleSize_(0) {
+        runningLoss_(0.0), adamStep_(0), trainingSampleSize_(0), epochLosses_() {
 
         // Validate model dimensions and parameters
         if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) {
@@ -239,7 +240,7 @@ namespace mu2e {
         std::vector<double> grad = gradOutput;
 
         // Back-propagate through layers in reverse order
-        for (int l = network_.size() - 1; l >= 0; l--) 
+        for (int l = static_cast<int>(network_.size()) - 1; l >= 0; l--)
         {
             auto& layer = network_[l];
 
@@ -249,7 +250,7 @@ namespace mu2e {
             std::vector<double> gradZ(grad.size());
 
             // For the output layer, the gradient is directly from the loss w.r.t. output (no activation function).
-            if(l == network_.size()-1)
+            if(l == static_cast<int>(network_.size())-1)
             {
                 gradZ = grad;
             }
@@ -673,7 +674,8 @@ namespace mu2e {
     }
 
     ScoreBasedDiffusionModel ScoreBasedDiffusionModel::loadModel(
-        CLHEP::HepRandomEngine& engine, // We need the random engine to initialize the model architecture
+        CLHEP::RandFlat& randFlat,
+        CLHEP::RandGaussQ& randGaussQ,
         const std::string& filename
     )
     {
@@ -885,7 +887,8 @@ namespace mu2e {
 
         // Reconstruct model
         ScoreBasedDiffusionModel model(
-            engine,
+            randFlat,
+            randGaussQ,
             dim,
             conditionDim,
             hidden,
diff --git a/STMMC/CMakeLists.txt b/STMMC/CMakeLists.txt
index cb3fce1b42..e48d57a37e 100644
--- a/STMMC/CMakeLists.txt
+++ b/STMMC/CMakeLists.txt
@@ -98,6 +98,7 @@ cet_build_plugin(VDResamplerTrain art::module
       Offline::GlobalConstantsService
       Offline::MachineLearningTools
       Offline::MCDataProducts
+      Offline::SeedService
 )
 
 cet_build_plugin(VDResamplerGenerateFromModel art::module
@@ -108,6 +109,7 @@ cet_build_plugin(VDResamplerGenerateFromModel art::module
       Offline::GlobalConstantsService
       Offline::MachineLearningTools
       Offline::MCDataProducts
+      Offline::SeedService
 )
 
 cet_build_plugin(VDResamplerGenerateMix art::module
@@ -118,6 +120,7 @@ cet_build_plugin(VDResamplerGenerateMix art::module
       Offline::GlobalConstantsService
       Offline::MachineLearningTools
       Offline::MCDataProducts
+      Offline::SeedService
 )
 
 install_source(SUBDIRS src)
diff --git a/STMMC/src/SConscript b/STMMC/src/SConscript
index 6278047c86..812af72a9f 100644
--- a/STMMC/src/SConscript
+++ b/STMMC/src/SConscript
@@ -48,6 +48,8 @@ helper.make_plugins( [ mainlib,
                        'mu2e_MCDataProducts',
                        'mu2e_GlobalConstantsService',
                        'mu2e_GeometryService',
+                       'mu2e_MachineLearningTools',
+                       'mu2e_SeedService',
                        rootlibs
                      ] )
 
diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index 8adcd5d219..76594193ac 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -72,7 +72,7 @@ namespace mu2e {
       std::string VDResamplerDir;
       std::string fclDir;
       std::string dataSourceTag;
-      int trainingThreshold = 0;
+      int trainingThreshold;
       bool doROOTDump;
       GlobalConstantsHandle<ParticleDataList> pdt;
       int pdgId = 0;
@@ -89,8 +89,8 @@ namespace mu2e {
     VirtualDetectorID(conf().VirtualDetectorID()),
     VDResamplerDir(conf().VDResamplerDir()),
     fclDir(conf().fclDir()),
-    trainingThreshold(conf().trainingThreshold()),
     dataSourceTag(conf().dataSourceTag()),
+    trainingThreshold(conf().trainingThreshold()),
     doROOTDump(conf().doROOTDump()) {
     if (doROOTDump) {
         art::ServiceHandle<art::TFileService> tfs;
diff --git a/STMMC/src/VDResamplerGenerateFromModel_module.cc b/STMMC/src/VDResamplerGenerateFromModel_module.cc
index 9a4dd45eda..7102043022 100644
--- a/STMMC/src/VDResamplerGenerateFromModel_module.cc
+++ b/STMMC/src/VDResamplerGenerateFromModel_module.cc
@@ -19,10 +19,13 @@
 #include "art/Framework/Core/EDProducer.h"
 #include "art/Framework/Principal/Event.h"
 #include "art/Framework/Services/Optional/RandomNumberGenerator.h"
+#include "art/Framework/Services/Registry/ServiceHandle.h"
 
 // CLHEP includes
 #include "CLHEP/Vector/LorentzVector.h"
 #include "CLHEP/Vector/ThreeVector.h"
+#include "CLHEP/Random/RandFlat.h"
+#include "CLHEP/Random/RandGaussQ.h"
 
 // exception handling
 #include "cetlib_except/exception.h"
@@ -134,6 +137,8 @@ namespace mu2e {
 
     private:
       art::RandomNumberGenerator::base_engine_t& engine_;
+      CLHEP::RandFlat randFlat_;
+      CLHEP::RandGaussQ randGaussQ_;
       bool useTwoStageModel_;
       std::string stage1ModelFile_;
       std::string stage2ModelFile_;
@@ -174,6 +179,8 @@ namespace mu2e {
   VDResamplerGenerateFromModel::VDResamplerGenerateFromModel(const Parameters& conf)
     : art::EDProducer{conf},
       engine_(createEngine(art::ServiceHandle<SeedService>()->getSeed())),
+      randFlat_(engine_),
+      randGaussQ_(engine_),
       useTwoStageModel_(conf().useTwoStageModel()),
       stage1ModelFile_(conf().stage1ModelFile()),
       stage2ModelFile_(conf().stage2ModelFile()),
@@ -193,10 +200,10 @@ namespace mu2e {
       }
 
       stage1Model_ = std::make_unique<ScoreBasedDiffusionModel>(
-        ScoreBasedDiffusionModel::loadModel(engine_, stage1ModelFile_)
+        ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, stage1ModelFile_)
       );
       stage2Model_ = std::make_unique<ScoreBasedDiffusionModel>(
-        ScoreBasedDiffusionModel::loadModel(engine_, stage2ModelFile_)
+        ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, stage2ModelFile_)
       );
       pdgId_ = loadPDGIdFromFileName(stage1ModelFile_);
     } else {
@@ -206,7 +213,7 @@ namespace mu2e {
       }
 
       allAtOnceModel_ = std::make_unique<ScoreBasedDiffusionModel>(
-        ScoreBasedDiffusionModel::loadModel(engine_, allAtOnceModelFile_)
+        ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, allAtOnceModelFile_)
       );
       pdgId_ = loadPDGIdFromFileName(allAtOnceModelFile_);
     }
diff --git a/STMMC/src/VDResamplerGenerateMix_module.cc b/STMMC/src/VDResamplerGenerateMix_module.cc
index 031ba23511..c936b83840 100644
--- a/STMMC/src/VDResamplerGenerateMix_module.cc
+++ b/STMMC/src/VDResamplerGenerateMix_module.cc
@@ -23,9 +23,11 @@
 #include "art/Framework/Core/EDProducer.h"
 #include "art/Framework/Principal/Event.h"
 #include "art/Framework/Services/Optional/RandomNumberGenerator.h"
+#include "art/Framework/Services/Registry/ServiceHandle.h"
 
 // CLHEP includes
 #include "CLHEP/Random/RandFlat.h"
+#include "CLHEP/Random/RandGaussQ.h"
 #include "CLHEP/Vector/LorentzVector.h"
 #include "CLHEP/Vector/ThreeVector.h"
 
@@ -156,6 +158,7 @@ namespace mu2e {
 
       art::RandomNumberGenerator::base_engine_t& engine_;
       CLHEP::RandFlat randFlat_;
+      CLHEP::RandGaussQ randGaussQ_;
       std::string modelFileDir_;
       int virtualDetectorID_ = 0;
       bool useHeun_ = true;
@@ -204,6 +207,7 @@ namespace mu2e {
     : art::EDProducer{conf},
       engine_(createEngine(art::ServiceHandle<SeedService>()->getSeed())),
       randFlat_(engine_),
+      randGaussQ_(engine_),
       modelFileDir_(conf().ModelFileDir()),
       virtualDetectorID_(conf().VirtualDetectorID()),
       useHeun_(conf().useHeun()),
@@ -322,16 +326,16 @@ namespace mu2e {
           modelFileDir_ + "/SBDMmodel_stage2_VD" + std::to_string(virtualDetectorID_) + "_pdg" + pdgToken + ".csv";
 
         particle.stage1Model = std::make_unique<ScoreBasedDiffusionModel>(
-          ScoreBasedDiffusionModel::loadModel(engine_, stage1ModelFile)
+          ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, stage1ModelFile)
         );
         particle.stage2Model = std::make_unique<ScoreBasedDiffusionModel>(
-          ScoreBasedDiffusionModel::loadModel(engine_, stage2ModelFile)
+          ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, stage2ModelFile)
         );
       } else {
         const std::string modelFile =
           modelFileDir_ + "/SBDMmodel_allAtOnce_VD" + std::to_string(virtualDetectorID_) + "_pdg" + pdgToken + ".csv";
         particle.allAtOnceModel = std::make_unique<ScoreBasedDiffusionModel>(
-          ScoreBasedDiffusionModel::loadModel(engine_, modelFile)
+          ScoreBasedDiffusionModel::loadModel(randFlat_, randGaussQ_, modelFile)
         );
       }
 
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
index e7a0be2463..f70eb5c631 100644
--- a/STMMC/src/VDResamplerTrain_module.cc
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -16,12 +16,16 @@
 
 #include "Offline/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh"
 
+#include "CLHEP/Random/RandFlat.h"
+#include "CLHEP/Random/RandGaussQ.h"
+
 // art includes
 #include "art/Framework/Core/EDAnalyzer.h"
 #include "art/Framework/Principal/Event.h"
 #include "art/Framework/Principal/Handle.h"
 #include "art/Framework/Principal/Run.h"
 #include "art/Framework/Services/Optional/RandomNumberGenerator.h"
+#include "art/Framework/Services/Registry/ServiceHandle.h"
 
 // exception handling
 #include "cetlib_except/exception.h"
@@ -86,9 +90,11 @@ namespace mu2e {
       void analyze(const art::Event& e);
       void endJob();
     private:
+      art::RandomNumberGenerator::base_engine_t& engine;
+      CLHEP::RandFlat randFlat_;
+      CLHEP::RandGaussQ randGaussQ_;
       art::ProductToken<StepPointMCCollection> StepPointMCsToken;
       art::ProductToken<SimParticleCollection> SimParticlemvToken;
-      art::RandomNumberGenerator::base_engine_t& engine;
       
       // SBDM model + data
       bool useTwoStageTraining = true;
@@ -125,6 +131,8 @@ namespace mu2e {
   VDResamplerTrain::VDResamplerTrain(const Parameters& conf) :
     art::EDAnalyzer(conf),
     engine(createEngine( art::ServiceHandle<SeedService>()->getSeed())),
+    randFlat_(engine),
+    randGaussQ_(engine),
     StepPointMCsToken(consumes<StepPointMCCollection>(conf().StepPointMCsTag())),
     SimParticlemvToken(consumes<SimParticleCollection>(conf().SimParticlemvTag())),
     useTwoStageTraining(conf().SBDMuseTwoStageTraining()),
@@ -153,7 +161,8 @@ namespace mu2e {
     if (useTwoStageTraining) {
       // create stage-1 model for (t', x', y')
       stage1Model = std::make_unique<ScoreBasedDiffusionModel>(
-          engine,
+          randFlat_,
+          randGaussQ_,
           3, // dim
           0, // conditionDim
           conf().SBDMhidden(),
@@ -174,7 +183,8 @@ namespace mu2e {
 
       // create stage-2 model for (p_r', p_phi', p_z' | t', x', y')
       stage2Model = std::make_unique<ScoreBasedDiffusionModel>(
-          engine,
+          randFlat_,
+          randGaussQ_,
           3, // dim
           3, // conditionDim
           conf().SBDMhidden(),
@@ -197,7 +207,8 @@ namespace mu2e {
       stage2TrainingData.reserve(1000000);
     } else {
       allAtOnceModel = std::make_unique<ScoreBasedDiffusionModel>(
-          engine,
+          randFlat_,
+          randGaussQ_,
           6, // dim
           0, // conditionDim
           conf().SBDMhidden(),

From 1e4820dd1e936c4ae723ec0fe7303d0bc03c1860 Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2egpvm03.fnal.gov>
Date: Fri, 27 Mar 2026 09:54:47 -0500
Subject: [PATCH 04/15] Fixing trigger paths making in the fcl generation

---
 STMMC/src/VDResamplerConfigure_module.cc | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index 76594193ac..9a69945c7d 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -126,7 +126,11 @@ namespace mu2e {
       virtualdetectorId = step.virtualDetectorId();
       pz = step.momentum().z();
       if (virtualdetectorId != VirtualDetectorID || pz <= 0)
+      {
+        mf::LogWarning("VDResamplerConfigure") << "Thrown event\n"
+                                               << "PDG ID = " << pdgId << ", VDID = " << virtualdetectorId << ", z = " << step.position().z() << ", pz = " << pz;
         continue; // Filter hits based on the virtual detector ID and pz
+      }
       
       if (doROOTDump) {
           x = step.position().x();
@@ -174,7 +178,8 @@ namespace mu2e {
         throw cet::exception("VDResamplerConfigure::endJob") << "Cannot open file " << fclFile;
     }
 
-    std::string trainingModuleNames = "";
+    std::string trainingPathNames = "";
+    std::vector<std::pair<std::string, std::string>> trainingPaths; // pairs of (moduleName, pathName)
 
     sumOutFile << "PDGID,HitCount,NotTrained\n";
 
@@ -206,10 +211,8 @@ namespace mu2e {
         // if pdgID is negative, we will use "m" instead of "-" in the filename to avoid issues with file naming
         std::string pdgIdstr = (part.first < 0) ? "m" + std::to_string(-part.first) : std::to_string(part.first);
         std::string moduleName = "VDResamplerTrainVD"+ std::to_string(VirtualDetectorID) + dataSourceTag + "pdg" + pdgIdstr;
-        if (!trainingModuleNames.empty()) {
-            trainingModuleNames += ", ";
-        }            
-        trainingModuleNames += moduleName;
+        std::string pathName = "trainPathVD" + std::to_string(VirtualDetectorID) + "pdg" + pdgIdstr;
+        trainingPaths.push_back({moduleName, pathName});
         fclOutFile << "    " << moduleName << " : {\n"
                     << "      module_type : VDResamplerTrain\n"
                     << "      StepPointMCsTag : @local::SimplifyStage1Data.StepPointMCsTag\n"
@@ -245,7 +248,15 @@ namespace mu2e {
       }
     }
     fclOutFile << "  }\n";
-    fclOutFile << "  end_paths: [" + trainingModuleNames + "]\n";
+
+    // Create trigger paths for each training module
+    for (size_t i = 0; i < trainingPaths.size(); ++i) {
+      fclOutFile << "  " << trainingPaths[i].second << " : [" << trainingPaths[i].first << "]\n";
+      if (i > 0) trainingPathNames += ", ";
+      trainingPathNames += trainingPaths[i].second;
+    }
+
+    fclOutFile << "  end_paths: [" + trainingPathNames + "]\n";
     fclOutFile << "}\n\n";
     if (doROOTDump) fclOutFile << "services.TFileService.fileName : @nil\n";
 

From 3877610933c158eb4a54d6133805e753afae18a6 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Thu, 2 Apr 2026 18:33:28 -0400
Subject: [PATCH 05/15] Using common source and services in the prolog file

---
 STMMC/src/VDResamplerConfigure_module.cc | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index 9a69945c7d..d5b33f2cf0 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -185,10 +185,11 @@ namespace mu2e {
 
     fclOutFile << "# This fcl is generated by VDResamplerConfigure_module.cc\n";
     fclOutFile << "# Training configuration for the VD resampler is generated for each particle type\n\n";
-    fclOutFile << "#include \"Offline/fcl/standardServices.fcl\"\n\n";
+    fclOutFile << "#include \"Offline/fcl/standardServices.fcl\"\n";
+    fclOutFile << "#include \"Production/JobConfig/pileup/STM/prolog.fcl\"\n\n";
     fclOutFile << "process_name: VDResamplerTrain\n\n";
-    fclOutFile << "source : {\n  module_type : RootInput\n  fileNames: @nil\n}\n";
-    fclOutFile << "services : {\n  message : @local::default_message\n  GlobalConstantsService : {\n    inputFile : \"Offline/GlobalConstantsService/data/globalConstants_01.txt\"\n  }\n}\n";
+    fclOutFile << "source : @local::STMPileup.VDResamplerSource\n\n";
+    fclOutFile << "services : @local::Services.Sim\n\n";
     fclOutFile << "physics: {\n  analyzers : {\n";
 
     for (const auto& part : pdgIds) {

From ae4617fe736725f0eb97f0a4f713cfe9e35c0751 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Thu, 9 Apr 2026 01:39:18 -0400
Subject: [PATCH 06/15] Add Stickman production target geometry

Introduce the Stickman_v_1_0 production target: add a new geometry file (ProductionTarget_Stickman_v1_0.txt) and enable the model in geom_run1_a.txt. Extend ProductionTargetMaker (header + maker) to recognize and create the Stickman model. Implement full Stickman construction in constructTargetPS.cc (uses G4ExtrudedSolid, fillets, cached plate solids, unions for fins/lugs, rod/spacer placement and wheel support reuse) and add minimal support in constructPS.cc to select the model. Update ProductionTarget interface to expose new Stickman-specific parameters and wiring. This integrates the new target variant while reusing existing support-wheel logic.
---
 GeometryService/inc/ProductionTargetMaker.hh  |    2 +
 GeometryService/src/ProductionTargetMaker.cc  |  115 ++
 .../geom/ProductionTarget_Stickman_v1_0.txt   |  145 +++
 Mu2eG4/geom/geom_run1_a.txt                   |    3 +-
 Mu2eG4/src/constructPS.cc                     |    3 +
 Mu2eG4/src/constructTargetPS.cc               | 1001 +++++++++++++++++
 ProductionTargetGeom/inc/ProductionTarget.hh  |  134 +++
 ProductionTargetGeom/src/ProductionTarget.cc  |   82 ++
 8 files changed, 1484 insertions(+), 1 deletion(-)
 create mode 100644 Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt

diff --git a/GeometryService/inc/ProductionTargetMaker.hh b/GeometryService/inc/ProductionTargetMaker.hh
index 97eeb9d73e..136bee4bf8 100644
--- a/GeometryService/inc/ProductionTargetMaker.hh
+++ b/GeometryService/inc/ProductionTargetMaker.hh
@@ -19,10 +19,12 @@ namespace mu2e {
     static const int tier1{1};
     // version 2 is the low density hayman
     static const int hayman_v_2_0{3};
+    static const int stickman_v_1_0{4};
 
 
     static std::unique_ptr<ProductionTarget> makeTier1(const SimpleConfig& config, double solenoidOffset);
     static std::unique_ptr<ProductionTarget> makeHayman_v_2_0(const SimpleConfig& config, double solenoidOffset);
+    static std::unique_ptr<ProductionTarget> makeStickman_v_1_0(const SimpleConfig& config, double solenoidOffset);
 
 
 
diff --git a/GeometryService/src/ProductionTargetMaker.cc b/GeometryService/src/ProductionTargetMaker.cc
index 81d1f5f701..8eacd51f11 100644
--- a/GeometryService/src/ProductionTargetMaker.cc
+++ b/GeometryService/src/ProductionTargetMaker.cc
@@ -19,6 +19,9 @@ namespace mu2e {
         //        std::cout << " making Hayman in Maker" << std::endl;
         return makeHayman_v_2_0(c, solenoidOffset);
           } else
+      if (c.getString("targetPS_model") == "Stickman_v_1_0"){
+        return makeStickman_v_1_0(c, solenoidOffset);
+          } else
         {throw cet::exception("GEOM") << " illegal production target version specified = " << c.getInt("targetPS_version")  << std::endl;}
     return 0;
 
@@ -390,5 +393,117 @@ namespace mu2e {
     return tgtPS;
   }
 
+  std::unique_ptr<ProductionTarget> ProductionTargetMaker::makeStickman_v_1_0(const SimpleConfig& c, double solenoidOffset){
+
+    // Read the plate and fin parameters
+    std::vector<std::string> plateMaterial;
+    std::vector<double> plateROut;
+    std::vector<double> plateFinAngles;
+    std::vector<double> plateThickness;
+    std::vector<double> plateLugThickness;
+
+    c.getVectorString("targetPS_plateMaterial", plateMaterial);
+    c.getVectorDouble("targetPS_rOut", plateROut);
+    c.getVectorDouble("targetPS_plateFinAngles", plateFinAngles);
+    for_each(plateFinAngles.begin(), plateFinAngles.end(), [](double& elem){elem *= CLHEP::degree;});
+    c.getVectorDouble("targetPS_plateThickness", plateThickness);
+    c.getVectorDouble("targetPS_plateLugThickness", plateLugThickness);
+
+    std::unique_ptr<ProductionTarget> tgtPS
+      (new ProductionTarget(
+                            c.getString("targetPS_model","NULL"),
+                            c.getInt("targetPS_version"),
+                            c.getDouble("targetPS_productionTargetMotherOuterRadius"),
+                            c.getDouble("targetPS_productionTargetMotherHalfLength"),
+                            c.getDouble("targetPS_rotX") * CLHEP::degree,
+                            c.getDouble("targetPS_rotY") * CLHEP::degree,
+                            c.getDouble("targetPS_rotZ") * CLHEP::degree,
+                            c.getDouble("targetPS_halfStickmanLength"),
+                            CLHEP::Hep3Vector(solenoidOffset,
+                                              0,
+                                              c.getDouble("productionTarget.zNominal")
+                                              )
+                            + c.getHep3Vector("productionTarget.offset"),
+                            c.getString("targetPS_targetVacuumMaterial"),
+                            c.getInt("targetPS_numberOfPlates"),
+                            plateMaterial,
+                            plateROut,
+                            c.getInt("targetPS_nStickmanFins"),
+                            plateFinAngles,
+                            c.getDouble("targetPS_plateFinOuterRadius"),
+                            c.getDouble("targetPS_plateFinWidth"),
+                            c.getDouble("targetPS_plateCenterToLugCenter"),
+                            c.getDouble("targetPS_plateLugInnerRadius"),
+                            c.getDouble("targetPS_plateLugOuterRadius"),
+                            plateThickness,
+                            plateLugThickness,
+                            c.getString("targetPS_rodMaterial"),
+                            c.getDouble("targetPS_rodRadius"),
+                            c.getString("targetPS_spacerMaterial"),
+                            c.getDouble("targetPS_spacerHalfLength"),
+                            c.getDouble("targetPS_spacerOuterRadius"),
+                            c.getDouble("targetPS_spacerInnerRadius"),
+                            c.getString("targetPS_supportRingMaterial"),
+                            c.getDouble("targetPS_supportRingLength"),
+                            c.getDouble("targetPS_supportRingInnerRadius"),
+                            c.getDouble("targetPS_supportRingOuterRadius"),
+                            c.getDouble("targetPS_supportRingLugOuterRadius"),
+                            c.getDouble("targetPS_supportRingCutoutOffset")
+                            ));
+
+    // Configure plate fillet parameters (only if fillets will be used)
+    tgtPS->_addFilletToPlateCore = c.getBool("targetPS_addFilletToPlateCore");
+    tgtPS->_addFilletToPlateLug = c.getBool("targetPS_addFilletToPlateLug");
+    if(tgtPS->_addFilletToPlateCore || tgtPS->_addFilletToPlateLug) {
+      tgtPS->_plateFilletRadius = c.getDouble("targetPS_plateFilletRadius");
+    }
+
+    // Configure support ring lug fillet parameters (only if fillets will be used)
+    tgtPS->_addFilletToSupportRingLug = c.getBool("targetPS_addFilletToSupportRingLug");
+    if(tgtPS->_addFilletToSupportRingLug) {
+      tgtPS->_supportRingLugFilletRadius = c.getDouble("targetPS_supportRingLugFilletRadius");
+    }
+
+    // Configure support ring cutout parameters (only if cutouts will be used)
+    tgtPS->_addCutoutToSupportRing = c.getBool("targetPS_addCutoutToSupportRing");
+    if(tgtPS->_addCutoutToSupportRing) {
+      tgtPS->_nSupportRingCutouts = c.getInt("targetPS_nSupportRingCutouts");
+      c.getVectorDouble("targetPS_supportRingCutoutAngles", tgtPS->_supportRingCutoutAngles);
+      for_each(tgtPS->_supportRingCutoutAngles.begin(), tgtPS->_supportRingCutoutAngles.end(), [](double& elem){elem *= CLHEP::degree;});
+      tgtPS->_supportRingCutoutInnerRadius = c.getDouble("targetPS_supportRingCutoutInnerRadius");
+      tgtPS->_supportRingCutoutTilt = c.getDouble("targetPS_supportRingCutoutTilt") * CLHEP::degree;
+    }
+
+    // Configure support wheel/bicycle wheel parameters (reused from Hayman)
+    tgtPS->_supportsBuild = c.getBool("targetPS.supports.build", false);
+    if(tgtPS->_supportsBuild) {
+      //support wheel parameters
+      tgtPS->_supportWheelRIn      = c.getDouble("targetPS.supports.wheel.rIn");
+      tgtPS->_supportWheelROut     = c.getDouble("targetPS.supports.wheel.rOut");
+      tgtPS->_supportWheelHL       = c.getDouble("targetPS.supports.wheel.halfLength");
+      tgtPS->_supportWheelMaterial = c.getString("targetPS.supports.wheel.material");
+      //number of support rods and wires
+      tgtPS->_nSpokesPerSide = c.getInt("targetPS.supports.nSpokes");
+      //features on the wheel near each support rod
+      c.getVectorDouble("targetPS.supports.features.angles", tgtPS->_supportWheelFeatureAngles, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.features.arcs"  , tgtPS->_supportWheelFeatureArcs  , tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.features.rIns"  , tgtPS->_supportWheelFeatureRIns  , tgtPS->_nSpokesPerSide);
+      //support wheel rods parameters
+      c.getVectorDouble("targetPS.supports.rods.halfLength", tgtPS->_supportWheelRodHL, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.offset", tgtPS->_supportWheelRodOffset, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.radius", tgtPS->_supportWheelRodRadius, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.radialOffset", tgtPS->_supportWheelRodRadialOffset, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.wireOffset.downstream", tgtPS->_supportWheelRodWireOffsetD, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.wireOffset.upstream", tgtPS->_supportWheelRodWireOffsetU, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.rods.angles", tgtPS->_supportWheelRodAngles, tgtPS->_nSpokesPerSide);
+      //support wire (spokes) parameters
+      c.getVectorDouble("targetPS.supports.spokes.targetAngles.downstream", tgtPS->_spokeTargetAnglesD, tgtPS->_nSpokesPerSide);
+      c.getVectorDouble("targetPS.supports.spokes.targetAngles.upstream", tgtPS->_spokeTargetAnglesU, tgtPS->_nSpokesPerSide);
+      tgtPS->_spokeRadius = 0.5*c.getDouble("targetPS.supports.spokes.diameter");
+      tgtPS->_spokeMaterial = c.getString("targetPS.supports.spokes.material");
+    }
+    return tgtPS;
+  }
+
 
 } // namespace mu2e
diff --git a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
new file mode 100644
index 0000000000..46b042e829
--- /dev/null
+++ b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
@@ -0,0 +1,145 @@
+// A geometry file for using the version 1 Stickman Production Target.
+
+// Support wheel and wire geometry is inherited from the existing Hayman production target geometry,
+// others are overhauled to use the new Stickman geometry.
+
+// The Hayman dependence is intentionally removed. 
+
+int PSStickman.verbosityLevel = 0; //5;
+
+int targetPS_version  = 4;
+string targetPS_model = "Stickman_v_1_0";
+
+// Nominal target position in the Mu2e coordinate system: at PS center in XY, at the given Z:
+double productionTarget.zNominal = -6164.5;
+// Optional shift of production target from the nominal
+vector<double> productionTarget.offset = { 0., 0., 0.};
+
+// Mother volume
+double targetPS_productionTargetMotherOuterRadius = 200.; //mm
+double targetPS_productionTargetMotherHalfLength = 120.; //mm  full length 232.2 mm
+
+// Overall rotation
+double targetPS_rotX = 0.0;
+double targetPS_rotY = 14.0;
+double targetPS_rotZ = 0.0;
+// Dimension
+double targetPS_halfStickmanLength = 116.1; // mm. 35 * 6 mm / 2 + 8.1 mm + 3 mm. We do not need this except for debugging.
+// Configurations
+bool   targetPS.visible              = true;
+bool   targetPS.solid                = false;
+string targetPS_targetVacuumMaterial = "PSVacuum";
+
+// Start with target plates. Chose to specify certain parameters for each plate, so can easily play with plate-by-plate dimensions later on 
+int targetPS_numberOfPlates = 35;
+vector<string> targetPS_plateMaterial = { // counting along mu2e +z direction, i.e. from the proton downstream side to the upstream side
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", //  5
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", // 10
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", // 15
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", // 20
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", // 25
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", // 30
+    "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718"  // 35
+}; // does not separate fin & core material here
+vector<double> targetPS_rOut = {
+    3.15, 3.15, 3.15, 3.15, 3.15, //  5
+    3.15, 3.15, 3.15, 3.15, 3.15, // 10
+    3.15, 3.15, 3.15, 3.15, 3.15, // 15
+    3.15, 3.15, 3.15, 3.15, 3.15, // 20
+    3.15, 3.15, 3.15, 3.15, 3.15, // 25
+    3.15, 3.15, 3.15, 3.15, 3.15, // 30
+    3.15, 3.15, 3.15, 3.15, 3.15  // 35
+}; // in mm, core part radius, one value for each
+int targetPS_nStickmanFins             = 3;                   // since the plates are installed on rods, the number of fins is always the same
+vector<double> targetPS_plateFinAngles = {330.,210.,90.};  // degrees
+double targetPS_plateFinOuterRadius    = 18.075;              // mm, (39.20-3.05)/2, extends to lug-to-center radius - lug inner radius
+double targetPS_plateFinWidth          = 2.0;                 // mm, all fin has same width
+double targetPS_plateCenterToLugCenter = 19.6;                // mm, center of plate to center of lug
+double targetPS_plateLugInnerRadius    = 1.525;               // mm, inner radius of the lug
+double targetPS_plateLugOuterRadius    = 3.0;                 // mm, outer radius of the lug
+vector<double> targetPS_plateThickness = {
+    5.0, 5.0, 5.0, 5.0, 5.0, //  5
+    5.0, 5.0, 5.0, 5.0, 5.0, // 10
+    5.0, 5.0, 5.0, 5.0, 5.0, // 15
+    5.0, 5.0, 5.0, 5.0, 5.0, // 20
+    5.0, 5.0, 5.0, 5.0, 5.0, // 25
+    5.0, 5.0, 5.0, 5.0, 5.0, // 30
+    5.0, 5.0, 5.0, 5.0, 5.0  // 35
+}; // in mm, one value for each plate, core and fin thickness
+vector<double> targetPS_plateLugThickness = {
+    6.0, 6.0, 6.0, 6.0, 6.0, //  5
+    6.0, 6.0, 6.0, 6.0, 6.0, // 10
+    6.0, 6.0, 6.0, 6.0, 6.0, // 15
+    6.0, 6.0, 6.0, 6.0, 6.0, // 20
+    6.0, 6.0, 6.0, 6.0, 6.0, // 25
+    6.0, 6.0, 6.0, 6.0, 6.0, // 30
+    6.0, 6.0, 6.0, 6.0, 6.0  // 35
+}; // in mm, one value for each plate, note lugs are flush with the fins on the proton downstream side
+bool targetPS_addFilletToPlateCore = true; // whether to add fillet to the intersection between plate core and fins
+bool targetPS_addFilletToPlateLug  = true; // whether to add fillet to the intersection between fins and lugs 
+// Fillets between fins and lugs are less important than the fillets between the core and fins, since the lugs are not in the direct path of the beam.
+// Fillets due to lug thickness larger than the plate thickness is not added, due to the complexity of the geometry and the small effect on the physics.
+double targetPS_plateFilletRadius = 1.0; // mm
+
+// Rods holding plates, which are attached to the end rings.
+string targetPS_rodMaterial = "Inconel718";
+double targetPS_rodRadius = 1.5; // mm
+// Rod length will be the sum of lug thicknesses and the spacer thicknesses, to be determined in the code.
+// Actual rod is longer than this since it inserts into the end rings, but for the geometry reconstruction, that part will be treated as part of the end ring.
+
+// Spacer geometry, which is the same for all spacers. Added to both upstream and downstream end of the rods.
+string targetPS_spacerMaterial = "Inconel718";
+double targetPS_spacerHalfLength  = 1.5;  // mm, full length 3.0 mm
+double targetPS_spacerOuterRadius = 3.0;  // mm
+double targetPS_spacerInnerRadius = 1.55; // mm
+
+// End ring / support ring 
+string targetPS_supportRingMaterial        = "Inconel718";
+double targetPS_supportRingLength          = 8.1;  // full length in mm
+double targetPS_supportRingInnerRadius     = 15.0; // mm
+double targetPS_supportRingOuterRadius     = 17.0; // mm
+double targetPS_supportRingLugOuterRadius  = 3.0;  // mm
+// Lug distance to center is the same as the plate lug distance to center, not defined as a separate parameter.
+// Lug angles are constrained by the plate fin angles.
+// End rings have chamfers at the sockets where the rods inserts (aligned with the lugs). The sockets, chamfers, and fillets at the bottom of the sockets are not modeled.
+bool   targetPS_addFilletToSupportRingLug  = true;
+double targetPS_supportRingLugFilletRadius = 1.0;
+bool   targetPS_addCutoutToSupportRing     = true; // the six cylindrical cutouts at an angle
+double targetPS_supportRingCutoutOffset    = 4.78; // mm, distance from the center of the cutout to the side of the end ring closer to the target center. 
+                                                   // This value is a must-have. If no cutout is added, this determines where the support wire is attached from the inner edge of the end ring.
+int    targetPS_nSupportRingCutouts        = 6;
+vector<double> targetPS_supportRingCutoutAngles = {0., 60., 120., 180., 240., 300.}; // degrees
+// FIXME: double check rotation around the z axis
+double targetPS_supportRingCutoutInnerRadius = 2.06; // mm
+double targetPS_supportRingCutoutTilt        = 20.0; //degrees, angle of the cutout w.r.t the radial direction
+
+// Bicycle wheel parameters -- copied from the Hayman target
+bool   targetPS.supports.build            = true;
+double targetPS.supports.wheel.halfLength = 9.525;
+double targetPS.supports.wheel.rOut       = 196.85;
+double targetPS.supports.wheel.rIn        = 177.8;
+string targetPS.supports.wheel.material   =  "G4_Al";
+int    targetPS.supports.nSpokes          = 3;
+// Features on the wheel
+vector<double> targetPS.supports.features.angles = {-46.62,  79.38, 194.38}; //degrees
+vector<double> targetPS.supports.features.arcs   = { 28.,  28.,  28.}; //degrees
+vector<double> targetPS.supports.features.rIns   = {158.75, 158.75, 158.75}; //mm
+//support rods in the wheel that the wires connect to
+vector<double> targetPS.supports.rods.angles       = {-54., 71.5, 186.0}; //degrees
+vector<double> targetPS.supports.rods.halfLength   = {70., 70., 70.}; //mm
+vector<double> targetPS.supports.rods.offset       = {-40., -10., 40.}; //mm
+vector<double> targetPS.supports.rods.radius       = {9., 9., 9.}; //mm
+vector<double> targetPS.supports.rods.radialOffset = {177.8, 177.8, 177.8}; //mm
+vector<double> targetPS.supports.rods.wireOffset.downstream = {10., 10., 10.}; //mm
+vector<double> targetPS.supports.rods.wireOffset.upstream   = {10., 10., 10.}; //mm
+//spokes connecting the target to the wheel rods
+vector<double> targetPS.supports.spokes.targetAngles.downstream = {-54., 71.5, 186.0}; //degrees
+vector<double> targetPS.supports.spokes.targetAngles.upstream   = {-54., 71.5, 186.0}; //degrees
+double targetPS.supports.spokes.diameter = 2.;
+string targetPS.supports.spokes.material = "Inconel718";
+
+//
+// This tells emacs to view this file in c++ mode.
+// Local Variables:
+// mode:c++
+// End:
diff --git a/Mu2eG4/geom/geom_run1_a.txt b/Mu2eG4/geom/geom_run1_a.txt
index db5cb6dffb..9328566173 100644
--- a/Mu2eG4/geom/geom_run1_a.txt
+++ b/Mu2eG4/geom/geom_run1_a.txt
@@ -66,7 +66,8 @@ double mu2e.detectorSystemZ0 = 10171.;   // mm  G4BL: (17730-7292=9801 mm)
 #include "Offline/Mu2eG4/geom/protonAbsorber_cylindrical_v04.txt"
 #include "Offline/Mu2eG4/geom/degrader_v02.txt" // pion degrader.  Off by default
 
-#include "Offline/Mu2eG4/geom/ProductionTarget_Hayman_v2_2.txt"
+// #include "Offline/Mu2eG4/geom/ProductionTarget_Hayman_v2_2.txt"
+#include "Offline/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt"
 #include "Offline/Mu2eG4/geom/protonBeamDump_v03.txt"
 #include "Offline/Mu2eG4/geom/extmon_fnal_v02.txt"
 
diff --git a/Mu2eG4/src/constructPS.cc b/Mu2eG4/src/constructPS.cc
index 22d0bdf625..16d4cdfe3e 100644
--- a/Mu2eG4/src/constructPS.cc
+++ b/Mu2eG4/src/constructPS.cc
@@ -367,6 +367,9 @@ namespace mu2e {
     } else if (targetPS_model == "Hayman_v_2_0"){
       verbosityLevel> 0 && std::cout << __func__ << "Hayman 2.0 target" << std::endl;
       constructTargetPS(psVacuumInfo, _config );
+    } else if (targetPS_model == "Stickman_v_1_0"){
+      verbosityLevel> 0 && std::cout << __func__ << "Stickman 1.0 target" << std::endl;
+      constructTargetPS(psVacuumInfo, _config );
     } else{
       throw cet::exception("CONFIG")
         << "In constructPS.cc unrecognized production target model name: "
diff --git a/Mu2eG4/src/constructTargetPS.cc b/Mu2eG4/src/constructTargetPS.cc
index a101584893..330ff81965 100644
--- a/Mu2eG4/src/constructTargetPS.cc
+++ b/Mu2eG4/src/constructTargetPS.cc
@@ -6,6 +6,7 @@
 #include <iostream>
 #include <sstream>
 #include <cmath>
+#include <map>
 
 // art includes
 #include "art/Framework/Services/Registry/ServiceDefinitionMacros.h"
@@ -48,6 +49,7 @@
 #include "Geant4/G4Tubs.hh"
 #include "Geant4/G4VPhysicalVolume.hh"
 #include "Geant4/G4PVPlacement.hh"
+#include "Geant4/G4ExtrudedSolid.hh"
 
 #include "Geant4/G4UnionSolid.hh"
 #include "Geant4/G4IntersectionSolid.hh"
@@ -1248,6 +1250,1005 @@ namespace mu2e {
 
       } //end adding support structures
     } //end ProductionTargetMaker::hayman_v_2_0
+    else if (_config.getInt("targetPS_version") == ProductionTargetMaker::stickman_v_1_0) {
+      int verbosityLevel                  = _config.getInt("PSStickman.verbosityLevel");
+      verbosityLevel >0 &&
+        cout << __func__ << " verbosityLevel on Stickman1.0              : " << verbosityLevel  << endl;
+      G4ThreeVector _hallOriginInMu2e = parent.centerInMu2e();
+      // Create variable to avoid multiple look-up
+
+      G4GeometryOptions* geomOptions = art::ServiceHandle<GeometryService>()->geomOptions();
+
+      bool prodTargetVisible   = geomOptions->isVisible( "ProductionTarget" );
+      bool prodTargetSolid     = geomOptions->isSolid  ( "ProductionTarget" );
+      bool forceAuxEdgeVisible = geomOptions->forceAuxEdgeVisible( "ProductionTarget" );
+      bool placePV             = geomOptions->placePV( "ProductionTarget" );
+      bool doSurfaceCheck      = geomOptions->doSurfaceCheck( "ProductionTarget" );
+
+      // begin all names with ProductionTarget so when we build sensitive detectors we can make them all
+      // sensitive at once with LVname.find("ProductionTarget") !=std::string::npos
+      //
+      // Build the production target.
+
+      GeomHandle<ProductionTarget> tgt;
+      // allows a vector of plate materials so each plate can be different. Initialize to have same number of 
+      // entries as the number of plates. Fin and core materials are the same for each plate.
+      vector<G4Material*> prodTargetPlateMaterials(tgt->numberOfPlates(), nullptr);
+      for (int i = 0; i < tgt->numberOfPlates(); ++i) {
+        prodTargetPlateMaterials.at(i) = findMaterialOrThrow(tgt->plateMaterial(i));
+      }
+      G4Material* prodTargetRodMaterial = findMaterialOrThrow(tgt->rodMaterial());
+      G4Material* prodTargetSpacerMaterial = findMaterialOrThrow(tgt->spacerMaterial());
+      G4Material* prodTargetSupportRingMaterial = findMaterialOrThrow(tgt->stickmanSupportRingMaterial());
+
+      TubsParams prodTargetMotherParams( 0.
+                                         ,tgt->productionTargetMotherOuterRadius()
+                                         ,tgt->productionTargetMotherHalfLength());
+
+      G4ThreeVector _localCenter(0.0,0.0,0.0);
+      G4ThreeVector zeroTranslation(0.,0.,0.);
+      G4RotationMatrix* targetRotation = reg.add(G4RotationMatrix(tgt->productionTargetRotation().inverse()));
+      if (verbosityLevel > 2){G4cout << __PRETTY_FUNCTION__ << "target rotation  = " << *targetRotation << G4endl;}
+      VolumeInfo prodTargetMotherInfo   = nestTubs( "ProductionTargetMother",
+                                                    prodTargetMotherParams,
+                                                    parent.logical->GetMaterial(),
+                                                    0,
+                                                    tgt->stickmanProdTargetPosition() - parent.centerInMu2e(),
+                                                    parent,
+                                                    0,
+                                                    G4Colour::Blue(),
+                                                    "PS"
+                                                    );
+
+      if (verbosityLevel > 0){
+        G4cout << "target position and hall origin = "
+                  << tgt->stickmanProdTargetPosition() << "\n" <<
+          _hallOriginInMu2e << " " << parent.centerInMu2e() << G4endl;
+      }
+      if (verbosityLevel > 2){
+        G4cout << __PRETTY_FUNCTION__ << "created prodTargetMotherInfo "
+                  << tgt->productionTargetMotherOuterRadius()
+                  << " " <<tgt->productionTargetMotherHalfLength() << G4endl;
+      }
+
+      // tiny offset to avoid precision issues
+      constexpr double stickmanMagicOffset = 0.0001;
+
+      // store the fin rotations for later use
+      // since used for G4UnionSolid construction, a passive rotation is needed, hence the negative sign on the angle. 
+      // These are the rotations to go from the plate frame to the fin frame
+      std::vector<CLHEP::HepRotation*> stickmanFinRotations;
+      for (int ithFin = 0; ithFin < tgt->nStickmanFins(); ++ithFin) {
+        CLHEP::HepRotation* finRotation = reg.add(CLHEP::HepRotation::IDENTITY);
+        finRotation->rotateZ(-tgt->plateFinAngle(ithFin)); 
+        stickmanFinRotations.emplace_back(finRotation);
+      }
+
+      // cache the plate solids
+      std::map<std::string, G4VSolid*> stickmanPlateSolidCache;
+      // lambda to get the plate solid for a given plate. The solid is cached based on the plate parameters that affect the shape. 
+      // The cache key encodes these parameters; if a solid is not found in the cache, it is created and added to the cache before 
+      //being returned.
+      // Here chose to construct plates as a single union solid with different materials for each plate, 
+      // rather than separate fins, cores, and lugs, because the fins is much wider than the Hayman design.
+      // If separated, missing material between fins and cores would be large.
+      // Note that the geometry center of a G4UnionSolid is the same as the first solid in the union, 
+      // and I start with the plate core and then union on the fins and lugs. This means that when placing the 
+      // plates the position will be based on the center of the core.
+      auto getStickmanPlateSolid = [&](int ithPlate) -> G4VSolid* {
+        const double plateFilletRadius =
+          (tgt->addFilletToPlateCore() || tgt->addFilletToPlateLug()) ? tgt->plateFilletRadius() : 0.;
+        const std::string cacheKey = std::string("StickmanPlate_")
+          + std::to_string(tgt->plateROut(ithPlate)) + "_" // Core radius
+          + std::to_string(tgt->plateThickness(ithPlate)) + "_" // plate thickness
+          + std::to_string(tgt->plateLugThickness(ithPlate)) + "_" // lug thickness
+          + std::to_string(static_cast<int>(tgt->addFilletToPlateCore())) + "_" // whether to add fillet to plate core
+          + std::to_string(static_cast<int>(tgt->addFilletToPlateLug())) + "_" // whether to add fillet to plate lug
+          + std::to_string(plateFilletRadius); // fillet radius (if applicable)
+
+        // check if the solid is already in the cache
+        auto cached = stickmanPlateSolidCache.find(cacheKey);
+        if (cached != stickmanPlateSolidCache.end()) {
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << " reusing cached plate solid for plate "
+                   << ithPlate << G4endl;
+          }
+          return cached->second;
+        }
+
+        // if not in cache, create the solid and add to cache
+        const double plateHalfThickness = tgt->plateThickness(ithPlate)/2.;
+        const double lugHalfThickness = tgt->plateLugThickness(ithPlate)/2.;
+        const double finHalfWidth = tgt->plateFinWidth()/2.;
+        const double finHalfLength = tgt->plateFinOuterRadius()/2.;
+        const double lugZShift = lugHalfThickness - plateHalfThickness;
+        if (verbosityLevel > 3) {
+          G4cout << __PRETTY_FUNCTION__ << " creating new plate solid for plate "
+                 << ithPlate << " with cache key " << cacheKey << G4endl;
+        }
+
+        // solid for the plate core
+        G4VSolid* plateSolid = reg.add(new G4Tubs(cacheKey + "_Core",        //
+                                                  0.,                        // inner radius
+                                                  tgt->plateROut(ithPlate),  // outer radius
+                                                  plateHalfThickness,        // z half length
+                                                  0.,                        // starting angle
+                                                  CLHEP::twopi));            // segment angle is full circle
+        if (verbosityLevel > 3) {
+          G4cout << __PRETTY_FUNCTION__ << "  core solid parameters (rOut, halfLength) = ("
+                 << tgt->plateROut(ithPlate) << ", " << plateHalfThickness << ")" << G4endl;
+        }
+
+        // solid for the fin
+        G4Box* finBox = reg.add(new G4Box(cacheKey + "_FinBox",              //
+                                          finHalfLength,                     // x half length
+                                          finHalfWidth,                      // y half length
+                                          plateHalfThickness));              // z half length
+        if (verbosityLevel > 3) {
+          G4cout << __PRETTY_FUNCTION__ << "  fin solid parameters (halfLength, halfWidth, halfThickness) = ("
+                 << finHalfLength << ", " << finHalfWidth << ", " << plateHalfThickness << ")" << G4endl;
+        }
+
+        // solid for the lug
+        G4Tubs* lugSolid = reg.add(new G4Tubs(cacheKey + "_Lug",             //
+                                              tgt->plateLugInnerRadius(),    // inner radius
+                                              tgt->plateLugOuterRadius(),    // outer radius
+                                              lugHalfThickness,              // z half length
+                                              0.,                            // starting angle
+                                              CLHEP::twopi));                // segment angle is full circle
+        if (verbosityLevel > 3) {
+          G4cout << __PRETTY_FUNCTION__ << "  lug solid parameters (rIn, rOut, halfThickness) = ("
+                 << tgt->plateLugInnerRadius() << ", " << tgt->plateLugOuterRadius() << ", " << lugHalfThickness << ")" << G4endl;
+        }
+
+        // add fillets to the plate core 
+        // for a fin at 0 degree start with unioning a G4ExtrudedSolid with a triangular cross section that has vertices in
+        // the x-y plane at (0,0), 
+        //                  (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
+        //                  (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
+        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and 
+        // a scale of 1.0 at both ends.
+        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance) 
+        // and height plateThickness(ithPlate) that is centered at 
+        // (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius), 
+        // and the other fillet is completed by subtracting a similar cylinder at
+        // (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
+        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the 
+        // cylinders for subtraction are also rotated by the fin angle.
+        G4VSolid* plateCoreFillet = nullptr;
+        if (tgt->addFilletToPlateCore()) {
+          // get the coordinates of the fillet union vertices for the fin at 0 degree
+          const double coreFilletX = std::sqrt(std::pow(tgt->plateROut(ithPlate) + plateFilletRadius, 2)
+                                               - std::pow(finHalfWidth + plateFilletRadius, 2));
+          std::vector<G4TwoVector> coreFilletVertices = {
+            G4TwoVector(0., 0.),
+            G4TwoVector(coreFilletX, finHalfWidth + plateFilletRadius),
+            G4TwoVector(coreFilletX, -finHalfWidth - plateFilletRadius)
+          };
+          // create the extrusion solid for the fillet
+          G4ExtrudedSolid* coreFilletExtrusion = reg.add(new G4ExtrudedSolid(
+            cacheKey + "_CoreFilletExtrusion",          //
+            coreFilletVertices,                         // vertices
+            plateHalfThickness,                         // half length in z
+            G4TwoVector(0.,0.),                         // shift of polygon center at -half length
+            1.0,                                        // scale of polygon at -half length
+            G4TwoVector(0.,0.),                         // shift of polygon center at +half length
+            1.0));                                      // scale of polygon at +half length
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  core fillet extrusion vertices = ("
+                   << coreFilletVertices.at(0) << ", "
+                   << coreFilletVertices.at(1) << ", "
+                   << coreFilletVertices.at(2) << ")" << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  core fillet extrusion half length = " << plateHalfThickness << G4endl;
+          }
+          
+          // create the cutout cylinder for the fillet
+          G4Tubs* coreFilletCutout = reg.add(new G4Tubs(
+            cacheKey + "_CoreFilletCutout",             //
+            0.,                                         // inner radius
+            plateFilletRadius + stickmanMagicOffset,    // outer radius (add small offset to avoid precision issues)
+            plateHalfThickness + stickmanMagicOffset,   // half length in z
+            0.,                                         // starting angle
+            CLHEP::twopi));                             // segment angle is full circle
+          // make 2 cutouts for the fillet, one at the upper edge of the fin and one at the lower edge of the fin
+          G4VSolid* coreFilletUpper = reg.add(new G4SubtractionSolid(
+            cacheKey + "_CoreFilletUpper",              //
+            coreFilletExtrusion,                        // base solid
+            coreFilletCutout,                           // cutout solid
+            nullptr,                                    // rotation
+            CLHEP::Hep3Vector(coreFilletX,              
+                              finHalfWidth + plateFilletRadius,
+                              0.)));                    // translation of cutout solid
+          plateCoreFillet = reg.add(new G4SubtractionSolid( 
+            cacheKey + "_CoreFilletLower",              //
+            coreFilletUpper,                            // base solid
+            coreFilletCutout,                           // cutout solid
+            nullptr,                                    // rotation
+            CLHEP::Hep3Vector(coreFilletX,
+                              -finHalfWidth - plateFilletRadius,
+                              0.)));                    // translation of cutout solid
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  core fillet cutout radius = " << plateFilletRadius + stickmanMagicOffset << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  core fillet upper cutout center = (" << coreFilletX << ", " << finHalfWidth + plateFilletRadius << ", 0)" << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  core fillet lower cutout center = (" << coreFilletX << ", " << -finHalfWidth - plateFilletRadius << ", 0)" << G4endl;
+          }
+        } // end of plate core fillet construction
+
+        // add fillets to the plate lugs
+        // for a fin at 0 degree, start with unioning a G4ExtrudedSolid with a triangular cross section that has vertices in
+        // the x-y plane at (plateCenterToLugCenter, 0),
+        //                  (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
+        //                  (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius), 
+        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and 
+        // a scale of 1.0 at both ends.
+
+        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance) 
+        // and height plateThickness(ithPlate) that is centered at 
+        // (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
+        // and the other fillet is completed by subtracting a similar cylinder at
+        // (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius), 
+        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the 
+        // cylinders for subtraction are also rotated by the fin angle.
+        G4VSolid* plateLugFillet = nullptr;
+        if (tgt->addFilletToPlateLug()) {
+          // get the coordinates of the fillet union vertices for the fin at 0 degree
+          const double lugFilletX = tgt->plateCenterToLugCenter()
+            - std::sqrt(std::pow(tgt->plateLugOuterRadius() + plateFilletRadius, 2)
+                        - std::pow(finHalfWidth + plateFilletRadius, 2));
+          std::vector<G4TwoVector> lugFilletVertices = {
+            G4TwoVector(tgt->plateCenterToLugCenter(), 0.),
+            G4TwoVector(lugFilletX, finHalfWidth + plateFilletRadius),
+            G4TwoVector(lugFilletX, -finHalfWidth - plateFilletRadius)
+          };
+          // create the extrusion solid for the fillet
+          G4ExtrudedSolid* lugFilletExtrusion = reg.add(new G4ExtrudedSolid(
+            cacheKey + "_LugFilletExtrusion",     //
+            lugFilletVertices,                    // vertices
+            plateHalfThickness,                   // half length in z
+            G4TwoVector(0.,0.),                   // shift of polygon center at -half length
+            1.0,                                  // scale of polygon at -half length
+            G4TwoVector(0.,0.),                   // shift of polygon center at +half length
+            1.0));                                // scale of polygon at +half length
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  lug fillet extrusion vertices = ("
+                   << lugFilletVertices.at(0) << ", "
+                   << lugFilletVertices.at(1) << ", "
+                   << lugFilletVertices.at(2) << ")" << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  lug fillet extrusion half length = " << plateHalfThickness << G4endl;
+          }
+
+          // create the cutout cylinder for the fillet
+          G4Tubs* lugFilletCutout = reg.add(new G4Tubs(
+            cacheKey + "_LugFilletCutout",            //
+            0.,                                       // inner radius
+            plateFilletRadius + stickmanMagicOffset,  // outer radius (add small offset to avoid precision issues)
+            plateHalfThickness + stickmanMagicOffset, // half length in z
+            0.,                                       // starting angle
+            CLHEP::twopi));                           // segment angle is full circle
+          // lug center needs to be cut out too. Create a second cutout cylinder for the lug center with radius of 
+          // plateLugInnerRadius and height of plateThickness(ithPlate) centered at (plateCenterToLugCenter, 0)
+          G4Tubs* lugCenterCutout = reg.add(new G4Tubs(
+            cacheKey + "_LugCenterCutout",                    //
+            0.,                                               // inner radius
+            tgt->plateLugInnerRadius() + stickmanMagicOffset, // outer radius (add small offset to avoid precision issues)
+            plateHalfThickness + stickmanMagicOffset,         // half length in z
+            0.,                                               // starting angle
+            CLHEP::twopi));                                   // segment angle is full circle
+          // cut out the lug center from the extrusion first
+          G4VSolid* lugFilletWithCenterCutout = reg.add(new G4SubtractionSolid(
+            cacheKey + "_LugFilletWithCenterCutout",                    //
+            lugFilletExtrusion,                                         // base solid
+            lugCenterCutout,                                            // cutout solid
+            nullptr,                                                    // rotation
+            CLHEP::Hep3Vector(tgt->plateCenterToLugCenter(), 0., 0.))); // translation of cutout solid
+          // make 2 cutouts for the fillet, one at the upper edge of the fin and one at the lower edge of the fin
+          G4VSolid* lugFilletUpper = reg.add(new G4SubtractionSolid(
+            cacheKey + "_LugFilletUpper",                     //
+            lugFilletWithCenterCutout,                        // base solid
+            lugFilletCutout,                                  // cutout solid
+            nullptr,                                          // rotation
+            CLHEP::Hep3Vector(lugFilletX,
+                              finHalfWidth + plateFilletRadius,
+                              0.)));                          // translation of cutout solid
+          plateLugFillet = reg.add(new G4SubtractionSolid(
+            cacheKey + "_LugFilletLower",                     //
+            lugFilletUpper,                                   // base solid
+            lugFilletCutout,                                  // cutout solid
+            nullptr,                                          // rotation
+            CLHEP::Hep3Vector(lugFilletX,
+                              -finHalfWidth - plateFilletRadius,
+                              0.)));                          // translation of cutout solid
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  lug fillet cutout radius = " << plateFilletRadius + stickmanMagicOffset << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  lug center cutout radius = " << tgt->plateLugInnerRadius() + stickmanMagicOffset << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  lug fillet upper cutout center = (" << lugFilletX << ", " << finHalfWidth + plateFilletRadius << ", 0)" << G4endl;
+            G4cout << __PRETTY_FUNCTION__ << "  lug fillet lower cutout center = (" << lugFilletX << ", " << -finHalfWidth - plateFilletRadius << ", 0)" << G4endl;
+          }
+        } // end of plate lug fillet construction
+
+        //
+        // assemble the plate by unioning the core, fins, and lugs (with fillets if applicable)
+        for (int ithFin = 0; ithFin < tgt->nStickmanFins(); ++ithFin) {
+          const double currentFinAngle = tgt->plateFinAngle(ithFin);
+          // fin and lug shifts of the current fin
+          const CLHEP::Hep3Vector finShift(finHalfLength * std::cos(currentFinAngle),
+                                           finHalfLength * std::sin(currentFinAngle),
+                                           0.);
+          const CLHEP::Hep3Vector lugShift(tgt->plateCenterToLugCenter() * std::cos(currentFinAngle),
+                                           tgt->plateCenterToLugCenter() * std::sin(currentFinAngle),
+                                           lugZShift); 
+          // note the lug is shifted in z to be flush with the edge of the plate core
+          
+          // add fin
+          plateSolid = reg.add(new G4UnionSolid(cacheKey + "_FinUnion_" + std::to_string(ithFin),
+                                                plateSolid,                         // base solid
+                                                finBox,                             // union solid
+                                                stickmanFinRotations.at(ithFin),    // rotation of union solid
+                                                finShift));                         // translation of union solid
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  fin " << ithFin << " added with angle of " << tgt->plateFinAngle(ithFin)/CLHEP::degree << " deg" << G4endl;
+          }
+
+          // add core fillets if applicable
+          if (plateCoreFillet != nullptr) {
+            plateSolid = reg.add(new G4UnionSolid(cacheKey + "_CoreFilletUnion_" + std::to_string(ithFin),
+                                                  plateSolid,                       // base solid
+                                                  plateCoreFillet,                  // union solid
+                                                  stickmanFinRotations.at(ithFin),  // rotation of union solid
+                                                  CLHEP::Hep3Vector(0.,0.,0.)));    // translation of union solid
+          }
+          
+          // add lug
+          plateSolid = reg.add(new G4UnionSolid(cacheKey + "_LugUnion_" + std::to_string(ithFin),
+                                                plateSolid,                       // base solid
+                                                lugSolid,                         // union solid
+                                                nullptr,                          // rotation of union solid
+                                                lugShift));                       // translation of union solid
+          if (verbosityLevel > 3) {
+            G4cout << __PRETTY_FUNCTION__ << "  lug " << ithFin << " added with shift (" << lugShift.x() << ", " << lugShift.y() << ", " << lugShift.z() << ")" << G4endl;
+          }
+
+          // add lug fillets if applicable
+          if (plateLugFillet != nullptr) {
+            plateSolid = reg.add(new G4UnionSolid(cacheKey + "_LugFilletUnion_" + std::to_string(ithFin),
+                                                  plateSolid,                       // base solid
+                                                  plateLugFillet,                   // union solid
+                                                  stickmanFinRotations.at(ithFin),  // rotation of union solid
+                                                  CLHEP::Hep3Vector(0.,0.,0.)));    // translation of union solid
+          }
+        } // end of loop over fins to assemble the plate solid
+
+        stickmanPlateSolidCache.emplace(cacheKey, plateSolid);
+        if (verbosityLevel > 2) {
+          G4cout << __PRETTY_FUNCTION__ << " created cached plate solid " << cacheKey << G4endl;
+        }
+        return plateSolid;
+      }; // end of lambda to get plate solid with caching
+      
+      //
+      // The plates are constructed starting from the most negative z end and moving towards the positive z end.
+      int numberOfPlates = tgt->numberOfPlates();
+      // running z position for plate construction. Start at most negative end and move towards +z
+      double _currentZ = _localCenter.z() - tgt->halfStickmanLength() // end of most negative end
+                         + tgt->stickmanSupportRingLength() // move past support ring at end 
+                         + 2 * tgt->spacerHalfLength() ; // move past spacer. 
+                         // plate fin, core, and lug is flush on this side
+      if (verbosityLevel > 2){G4cout << __PRETTY_FUNCTION__ << " plate Z starts at " << _currentZ << G4endl;}
+      
+      // loop over plates to construct and add them to the target mother.
+      for (int ithPlate = 0; ithPlate < numberOfPlates; ++ithPlate) {
+        std::string plateName = "ProductionTargetPlate_" + std::to_string(ithPlate);
+        const double plateHalfThickness = tgt->plateThickness(ithPlate)/2.;
+        const double plateCenterZ = _currentZ + plateHalfThickness; // this is the z position of the core
+        const CLHEP::Hep3Vector localPlateCenter(0.,0.,plateCenterZ); // this is in the target frame
+        const CLHEP::Hep3Vector currentPlateCenter = tgt->productionTargetRotation().inverse()*localPlateCenter;
+        // productionTargetRotation() rotates to the target frame, so to get the position in the world frame we 
+        // apply the inverse rotation to the local plate center.
+
+        if (verbosityLevel > 0) {
+          G4cout << __PRETTY_FUNCTION__ << " plate " << ithPlate
+                 << " startZ=" << _currentZ
+                 << " centerZ=" << plateCenterZ
+                 << " endZ=" << (_currentZ + tgt->plateLugThickness(ithPlate))
+                 << " material=" << tgt->plateMaterial(ithPlate)
+                 << G4endl;
+        }
+        if (verbosityLevel > 2) {
+          G4cout << " rOut=" << tgt->plateROut(ithPlate)
+                 << " thickness=" << tgt->plateThickness(ithPlate)
+                 << " finWidth=" << tgt->plateFinWidth()
+                 << " finReach=" << tgt->plateFinOuterRadius()
+                 << " lugThickness=" << tgt->plateLugThickness(ithPlate)
+                 << " lug(inner,outer)=("
+                 << tgt->plateLugInnerRadius() << ","
+                 << tgt->plateLugOuterRadius() << ")" 
+                 << " addFilletToPlateCore=" << tgt->addFilletToPlateCore()
+                 << " addFilletToPlateLug=" << tgt->addFilletToPlateLug()
+                 << G4endl;
+        }
+        
+        // get the solid for the plate, using the lambda defined above
+        G4VSolid* plateSolid = getStickmanPlateSolid(ithPlate);
+
+        // add the plate to the target mother
+        VolumeInfo plateInfo(plateName, currentPlateCenter, prodTargetMotherInfo.centerInWorld);
+        plateInfo.solid = plateSolid;
+        finishNesting(plateInfo                                   // volume info as defined above
+                      ,prodTargetPlateMaterials.at(ithPlate)      // material for the plate, can be different for each plate
+                      ,&tgt->productionTargetRotation()           // PASSIVE rotation
+                      ,currentPlateCenter                         // translation to place the plate at the correct z position within the target mother
+                      ,prodTargetMotherInfo.logical               // parent volume logical to place the plate in
+                      ,ithPlate                                   // copy number
+                      ,G4Colour::Yellow()                         // color for visualization
+                      ,"ProductionTarget"                         // lookup token
+                      ,verbosityLevel>1);                         // pass the verbosity level for printing info in finishNesting
+
+        // update _currentZ to the END of the next plate, which is the start of the current plate + plateLugThickness(ithPlate)
+        _currentZ += tgt->plateLugThickness(ithPlate);
+      } // end of loop over plates
+
+      // add the rods and spacers
+      G4VSolid* spacerSolid = reg.add(new G4Tubs(
+        "ProductionTargetSpacerSolid", //
+        tgt->spacerInnerRadius(),      // inner radius
+        tgt->spacerOuterRadius(),      // outer radius
+        tgt->spacerHalfLength(),       // half length
+        0.,                            // start angle
+        CLHEP::twopi));                // sweep angle
+      const double rodHalfLength = tgt->rodHalfLength() - stickmanMagicOffset; // make it slightly shorter to avoid precesion issues
+      const double spacerCenterZOffset = tgt->rodHalfLength() - tgt->spacerHalfLength();
+
+      for (int ithRod = 0; ithRod < tgt->nStickmanFins(); ++ithRod) {
+        const double rodAngle = tgt->plateFinAngle(ithRod);
+        const CLHEP::Hep3Vector rodCenterInTargetFrame(
+          tgt->plateCenterToLugCenter() * std::cos(rodAngle),
+          tgt->plateCenterToLugCenter() * std::sin(rodAngle),
+          0.);
+        const CLHEP::Hep3Vector rodCenter = tgt->productionTargetRotation().inverse() * rodCenterInTargetFrame;
+
+        TubsParams rodParams(0.,                 // rIn
+                             tgt->rodRadius(),   // rOut
+                             rodHalfLength);     // half length
+        VolumeInfo rodInfo = nestTubs("ProductionTargetRod_" + std::to_string(ithRod),
+                 rodParams,
+                 prodTargetRodMaterial,
+                 &tgt->productionTargetRotation(), // passive rotation
+                 rodCenter,                        // translation
+                 prodTargetMotherInfo,             // parent volume info for the target mother
+                 ithRod,                           // copy number
+                 G4Colour::Green(),                // color for visualization
+                 "ProductionTarget");              // lookup token
+        if (verbosityLevel > 0) {
+          G4cout << __PRETTY_FUNCTION__ << " rod " << ithRod
+                 << " center=" << rodCenter
+                 << " material=" << tgt->rodMaterial()
+                 << G4endl;
+        }
+
+        // place the spacers one spacerHalfLength inward from each rod end
+        const CLHEP::Hep3Vector spacerCenterNegZInTargetFrame(
+          rodCenterInTargetFrame.x(),
+          rodCenterInTargetFrame.y(),
+          -spacerCenterZOffset);
+        const CLHEP::Hep3Vector spacerCenterPosZInTargetFrame(
+          rodCenterInTargetFrame.x(),
+          rodCenterInTargetFrame.y(),
+          spacerCenterZOffset);
+        // transform the spacer centers to the parent frame
+        const CLHEP::Hep3Vector spacerShiftNegZ =
+          tgt->productionTargetRotation().inverse() * spacerCenterNegZInTargetFrame;
+        const CLHEP::Hep3Vector spacerShiftPosZ =
+          tgt->productionTargetRotation().inverse() * spacerCenterPosZInTargetFrame;
+        if (verbosityLevel > 2) {
+          G4cout << __PRETTY_FUNCTION__ << " spacer " << ithRod
+                 << " target-frame centers neg/pos="
+                 << spacerCenterNegZInTargetFrame << " / "
+                 << spacerCenterPosZInTargetFrame
+                 << " mother-frame centers neg/pos="
+                 << spacerShiftNegZ << " / "
+                 << spacerShiftPosZ
+                 << G4endl;
+        }
+
+        // add the spacers
+        std::string spacerNameNegZ = "ProductionTargetSpacerNegZ_" + std::to_string(ithRod);
+        VolumeInfo spacerInfoNegZ(spacerNameNegZ, spacerShiftNegZ, prodTargetMotherInfo.centerInWorld);
+        spacerInfoNegZ.solid = spacerSolid;
+        finishNesting(spacerInfoNegZ,                   // volume info as defined above
+                      prodTargetSpacerMaterial,         // material for the spacer
+                      &tgt->productionTargetRotation(), // rotation
+                      spacerShiftNegZ,                  // translation
+                      prodTargetMotherInfo.logical,     // parent logical volume
+                      2*ithRod,                         // copy number
+                      G4Colour::Cyan(),                 // color for visualization
+                      "ProductionTarget",               // lookup token
+                      verbosityLevel>1);                // pass the verbosity level for printing info in finishNesting
+        std::string spacerNamePosZ = "ProductionTargetSpacerPosZ_" + std::to_string(ithRod);
+        VolumeInfo spacerInfoPosZ(spacerNamePosZ, spacerShiftPosZ, prodTargetMotherInfo.centerInWorld);
+        spacerInfoPosZ.solid = spacerSolid;
+        finishNesting(spacerInfoPosZ,                   // volume info as defined above
+                      prodTargetSpacerMaterial,         // material for the spacer
+                      &tgt->productionTargetRotation(), // rotation
+                      spacerShiftPosZ,                  // translation
+                      prodTargetMotherInfo.logical,     // parent logical volume
+                      2*ithRod + 1,                     // copy number
+                      G4Colour::Cyan(),                 // color for visualization
+                      "ProductionTarget",               // lookup token
+                      verbosityLevel>1);                // pass the verbosity level for printing info in finishNesting
+        if (verbosityLevel > 0) {
+          G4cout << __PRETTY_FUNCTION__ << " spacer for rod " << ithRod
+                 << " negZ center=" << spacerShiftNegZ
+                 << " posZ center=" << spacerShiftPosZ
+                 << " material=" << tgt->spacerMaterial()
+                 << G4endl;
+        }
+      } // end of loop over rods and spacers
+      
+      // now add the two end rings
+
+      // then two cutouts are made by subtracting cylinders with radius of the fillet radius (plus a small tolerance) and height of stickmanSupportRingLength (plus a small tolerance) 
+      // centered at (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), supportRingLugOuterRadius+filletRadius), and
+      //             (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), -supportRingLugOuterRadius-filletRadius)
+      // an additional cutout is made for the inner part of the ring by subtracting a tube with inner radius of 0, outer radius of supportRingInnerRadius, and height of stickmanSupportRingLength (plus a small tolerance) centered at (0,0).
+      // if addCutoutToSupportRing is true, then additional circular cutouts (through holes) are made. 
+      // The number of cutouts is given by nSupportRingCutouts at angles defined by supportRingCutoutAngles, and the cutouts are made by subtracting cylinders with radius of 
+      // supportRingCutoutInnerRadius. The centers of the cutouts are at supportRingCutoutOffset from the interface between the spacers and the end ring, and at the angles defined 
+      // by supportRingCutoutAngles. The angle is between the through hole and the local radius, the through hole points outward and towards the middle.  
+
+      // lambda to build the end support ring solid with lugs and fillets if applicable
+      // Note the end plate has a flat face on the inside that is flush with the end of the rods / spacers, details like chamfers at 
+      // the rod sockets and those at the cutouts (through holes) are ignored for simplicity.
+      auto buildStickmanSupportRingSolid = [&](const std::string& tag, bool positiveEnd) -> G4VSolid* {
+        const double ringHalfLength = tgt->stickmanSupportRingLength()/2.;
+        // dimensions for the lug stem
+        const double lugStemHalfLength = 0.5*(tgt->plateCenterToLugCenter() - tgt->stickmanSupportRingInnerRadius());
+        const double lugStemCenterRadius = tgt->stickmanSupportRingInnerRadius() + lugStemHalfLength;
+        if (verbosityLevel > 2) {
+          G4cout << __PRETTY_FUNCTION__ << " building support ring " << tag
+                 << " positiveEnd=" << positiveEnd
+                 << " ringHalfLength=" << ringHalfLength
+                 << " inner/outer=" << tgt->stickmanSupportRingInnerRadius()
+                 << "/" << tgt->stickmanSupportRingOuterRadius()
+                 << " lugStemHalfLength=" << lugStemHalfLength
+                 << " lugStemCenterRadius=" << lugStemCenterRadius
+                 << G4endl;
+        }
+
+        // ring body without lugs or center hole cutout
+        G4VSolid* supportRingSolid = reg.add(new G4Tubs(
+          "ProductionTargetSupportRingBase_" + tag, //
+          0.,                                       // inner radius
+          tgt->stickmanSupportRingOuterRadius(),    // outer radius
+          ringHalfLength,                           // half length
+          0.,                                       // start angle
+          CLHEP::twopi));                           // delta angle
+
+        // cutout for the center of the ring
+        // such construction avoids potential issue that the lug solid extends within the 
+        // inner radius of the ring, which can cause misrepresentation of the geometry
+        G4Tubs* ringCenterCutout = reg.add(new G4Tubs(
+          "ProductionTargetSupportRingCenterCutout_" + tag, //
+          0.,                                               // inner radius
+          tgt->stickmanSupportRingInnerRadius(),            // outer radius (avoid overlap with lug stem)
+          ringHalfLength + stickmanMagicOffset,             // half length (add small offset to avoid precision issues)
+          0.,                                               // start angle
+          CLHEP::twopi));                                   // delta angle 
+
+        // lugs at three fin angles.
+        G4Tubs* lugSolid = reg.add(new G4Tubs(
+          "ProductionTargetSupportRingLug_" + tag, //
+          0.,                                      // inner radius 
+          tgt->supportRingLugOuterRadius(),        // outer radius
+          ringHalfLength,                          // half length
+          0.,                                      // start angle
+          CLHEP::twopi));                          // delta angle
+        
+        // lug stem
+        G4Box* lugStemSolid = reg.add(new G4Box(
+          "ProductionTargetSupportRingLugStem_" + tag, // 
+          lugStemHalfLength,                           // half length in x
+          tgt->supportRingLugOuterRadius(),            // half length in y
+          ringHalfLength));                            // half length in z
+
+        for (int ithFin = 0; ithFin < tgt->nStickmanFins(); ++ithFin) {
+          // the lugs are aligned with the fins, so get the fin angle for the current fin to calculate the position of the lugs
+          const double finAngle = tgt->plateFinAngle(ithFin);
+          const CLHEP::Hep3Vector lugStemShift( // rotate the lug stem shift by the fin angle to get the correct position in x and y
+            lugStemCenterRadius * std::cos(finAngle),
+            lugStemCenterRadius * std::sin(finAngle),
+            0.);
+          const CLHEP::Hep3Vector lugShift( // rotate the lug shift by the fin angle
+            tgt->plateCenterToLugCenter() * std::cos(finAngle),
+            tgt->plateCenterToLugCenter() * std::sin(finAngle),
+            0.);
+          if (verbosityLevel > 2) {
+            G4cout << __PRETTY_FUNCTION__ << " support ring " << tag
+                   << " fin=" << ithFin
+                   << " angle(deg)=" << finAngle/CLHEP::degree
+                   << " lugStemShift=" << lugStemShift
+                   << " lugShift=" << lugShift
+                   << G4endl;
+          }
+          // first union the lug stem to the ring
+          supportRingSolid = reg.add(new G4UnionSolid(
+            "ProductionTargetSupportRingLugStemUnion_" + tag + "_" + std::to_string(ithFin),
+            supportRingSolid,                // base solid
+            lugStemSolid,                    // union solid
+            stickmanFinRotations.at(ithFin), // passive rotation
+            lugStemShift));                  // translation
+          // then union the lug to the ring+stem
+          supportRingSolid = reg.add(new G4UnionSolid(
+            "ProductionTargetSupportRingLugUnion_" + tag + "_" + std::to_string(ithFin),
+            supportRingSolid,                // base solid
+            lugSolid,                        // union solid
+            nullptr,                         // passive rotation
+            lugShift));                      // translation
+          
+          // add fillet to the lug if applicable
+          // The fillet is added by unioning with an extrusion 
+          // of a triangle with vertices at (0,0), 
+          //                                (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), supportRingLugOuterRadius+filletRadius),
+          //                                (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), -supportRingLugOuterRadius-filletRadius),
+          // extruded along z by stickmanSupportRingLength starting at the end of the target, with no twist and scale of 1.0 at both ends.
+          if (tgt->addFilletToSupportRingLug()) {
+            const double filletRadius = tgt->supportRingLugFilletRadius();
+            const double filletX = std::sqrt(
+              std::pow(tgt->stickmanSupportRingOuterRadius() + filletRadius, 2)
+              - std::pow(tgt->supportRingLugOuterRadius() + filletRadius, 2));
+            if (verbosityLevel > 2) {
+              G4cout << __PRETTY_FUNCTION__ << " support ring " << tag
+                     << " fin=" << ithFin
+                     << " filletRadius=" << filletRadius
+                     << " filletX=" << filletX
+                     << G4endl;
+            }
+            // get the vertices for the extrusion
+            std::vector<G4TwoVector> filletVertices = {
+              G4TwoVector(0., 0.),
+              G4TwoVector(filletX, tgt->supportRingLugOuterRadius() + filletRadius),
+              G4TwoVector(filletX, -tgt->supportRingLugOuterRadius() - filletRadius)
+            };
+
+            // create the extrusion solid for the fillet
+            G4ExtrudedSolid* filletExtrusion = reg.add(new G4ExtrudedSolid(
+              "ProductionTargetSupportRingLugFilletExtrusion_" + tag + "_" + std::to_string(ithFin),
+              filletVertices,      // vertices
+              ringHalfLength,      // half length in z
+              G4TwoVector(0., 0.), // shift of polygon center at -half length
+              1.0,                 // scale of polygon at -half length
+              G4TwoVector(0., 0.), // shift of polygon center at +half length
+              1.0));               // scale of polygon at +half length
+            // create the cutout cylinder for the fillet
+            G4Tubs* filletCutout = reg.add(new G4Tubs(
+              "ProductionTargetSupportRingLugFilletCutout_" + tag + "_" + std::to_string(ithFin),
+              0.,                                   // inner radius
+              filletRadius + stickmanMagicOffset,   // outer radius (add small offset)
+              ringHalfLength + stickmanMagicOffset, // half length in z (add small offset)
+              0.,                                   // start angle
+              CLHEP::twopi));                       // segment angle is full circle
+            // subtract the cutouts from the extrusion
+            G4VSolid* filletUpper = reg.add(new G4SubtractionSolid(
+              "ProductionTargetSupportRingLugFilletUpper_" + tag + "_" + std::to_string(ithFin),
+              filletExtrusion,                  // base solid
+              filletCutout,                     // cutout solid
+              nullptr,                          // rotation
+              CLHEP::Hep3Vector(filletX,
+                                tgt->supportRingLugOuterRadius() + filletRadius,
+                                0.)));          // translation of cutout solid
+            G4VSolid* filletSolid = reg.add(new G4SubtractionSolid(
+              "ProductionTargetSupportRingLugFilletLower_" + tag + "_" + std::to_string(ithFin),
+              filletUpper,                      // base solid
+              filletCutout,                     // cutout solid
+              nullptr,                          // rotation
+              CLHEP::Hep3Vector(filletX,
+                                -tgt->supportRingLugOuterRadius() - filletRadius,
+                                0.)));          // translation of cutout solid
+            
+            // add filletSolid to the union with the support ring
+            supportRingSolid = reg.add(new G4UnionSolid(
+              "ProductionTargetSupportRingFilletUnion_" + tag + "_" + std::to_string(ithFin),
+              supportRingSolid,                  // base solid
+              filletSolid,                       // union solid
+              stickmanFinRotations.at(ithFin),   // passive rotation
+              CLHEP::Hep3Vector(0., 0., 0.)));   // no translation
+          } // end of fillet construction for the lug
+        } // end of loop over fins to add lugs and fillets to the support ring
+
+        // add cutout for the center hole
+        supportRingSolid = reg.add(new G4SubtractionSolid(
+          "ProductionTargetSupportRingLugFilletRingCutout_" + tag,
+          supportRingSolid,                 // base solid
+          ringCenterCutout,                 // cutout solid
+          nullptr,                          // rotation
+          CLHEP::Hep3Vector(0., 0., 0.)));  // translation
+        
+        // add cutouts for the through holes if applicable
+        if (tgt->addCutoutToSupportRing()) {
+          const double cutoutTiltAngle = tgt->supportRingCutoutTilt();
+          const double supportRingWallThickness = tgt->stickmanSupportRingOuterRadius() - tgt->stickmanSupportRingInnerRadius();
+          // cutout center is placed at halfway between the inner and outer radius
+          const double cutoutCenterRadius = 0.5*(tgt->stickmanSupportRingInnerRadius() + tgt->stickmanSupportRingOuterRadius());
+          // this is z position in the local ring frame
+          const double localCutoutZ = (positiveEnd ? 1. : -1.) * 
+                                      (tgt->supportRingCutoutOffset() - ringHalfLength + supportRingWallThickness * 0.5 * std::tan(cutoutTiltAngle) );
+          // cutout needs to go through the whole wall thickness
+          const double cutoutHalfLength = 0.5 * supportRingWallThickness / std::cos(cutoutTiltAngle) + stickmanMagicOffset; // add small offset to avoid precision issues
+          if (verbosityLevel > 2) {
+            G4cout << __PRETTY_FUNCTION__ << " support ring " << tag
+                   << " cutoutTilt(deg)=" << cutoutTiltAngle/CLHEP::degree
+                   << " wallThickness=" << supportRingWallThickness
+                   << " cutoutCenterRadius=" << cutoutCenterRadius
+                   << " localCutoutZ=" << localCutoutZ
+                   << " cutoutHalfLength=" << cutoutHalfLength
+                   << G4endl;
+          }
+          // create the cutout solid for the through hole
+          G4Tubs* supportRingCutout = reg.add(new G4Tubs(
+            "ProductionTargetSupportRingCutout_" + tag,
+            0.,                                   // inner radius
+            tgt->supportRingCutoutInnerRadius(),  // outer radius
+            cutoutHalfLength,                     // half length in z
+            0.,                                   // start angle
+            CLHEP::twopi));                       // delta angle
+          const CLHEP::Hep3Vector localZAxis(0., 0., 1.);
+
+          // loop over the holes and subtract the cutout solid
+          for (int ithCutout = 0; ithCutout < tgt->nSupportRingCutouts(); ++ithCutout) {
+            const double cutoutPhi = tgt->supportRingCutoutAngle(ithCutout);
+            const CLHEP::Hep3Vector radialAxis(std::cos(cutoutPhi), std::sin(cutoutPhi), 0.); // pointing from ring center radially outward
+            const CLHEP::Hep3Vector cutoutAxis = std::cos(cutoutTiltAngle) * radialAxis // vector pointing outward along the cutout axis
+                                                 + (positiveEnd ? -1. : 1.) * std::sin(cutoutTiltAngle) * localZAxis;
+            G4RotationMatrix* cutoutRotation = reg.add(G4RotationMatrix(
+              CLHEP::HepRotation(cutoutAxis.cross(localZAxis), cutoutAxis.angle(localZAxis)))); // this rotates the cutout from the local z axis to the cutout axis passively
+            const CLHEP::Hep3Vector cutoutCenter(
+              cutoutCenterRadius * std::cos(cutoutPhi),
+              cutoutCenterRadius * std::sin(cutoutPhi),
+              localCutoutZ);
+            if (verbosityLevel > 2) {
+              G4cout << __PRETTY_FUNCTION__ << " support ring " << tag
+                     << " cutout=" << ithCutout
+                     << " phi(deg)=" << cutoutPhi/CLHEP::degree
+                     << " center=" << cutoutCenter
+                     << " axis=" << cutoutAxis
+                     << G4endl;
+            }
+            // make this cutout solid
+            supportRingSolid = reg.add(new G4SubtractionSolid(
+              "ProductionTargetSupportRingCutoutSubtraction_" + tag + "_" + std::to_string(ithCutout),
+              supportRingSolid,   // base solid
+              supportRingCutout,  // cutout solid
+              cutoutRotation,     // rotation to align the cutout with the cutout axis
+              cutoutCenter));     // translation
+          } // end of loop over cutouts for through holes
+        } // end of cutout construction for through holes
+
+        return supportRingSolid;
+      }; // end of lambda to build the end support ring solid
+
+      // now add the two end rings
+      const double supportRingCenterZ = tgt->halfStickmanLength() - tgt->stickmanSupportRingLength()/2.;
+      const CLHEP::Hep3Vector supportRingCenterNegZInTargetFrame(0., 0., -supportRingCenterZ);
+      const CLHEP::Hep3Vector supportRingCenterPosZInTargetFrame(0., 0., supportRingCenterZ);
+      const CLHEP::Hep3Vector supportRingCenterNegZ =
+        tgt->productionTargetRotation().inverse() * supportRingCenterNegZInTargetFrame;
+      const CLHEP::Hep3Vector supportRingCenterPosZ =
+        tgt->productionTargetRotation().inverse() * supportRingCenterPosZInTargetFrame;
+      if (verbosityLevel > 2) {
+        G4cout << __PRETTY_FUNCTION__ << " support ring centers target-frame neg/pos="
+               << supportRingCenterNegZInTargetFrame << " / "
+               << supportRingCenterPosZInTargetFrame
+               << " mother-frame neg/pos="
+               << supportRingCenterNegZ << " / "
+               << supportRingCenterPosZ
+               << G4endl;
+      }
+      // add the negative z end ring first
+      VolumeInfo supportRingNegZInfo("ProductionTargetSupportRing_Upstream",
+                                     supportRingCenterNegZ,
+                                     prodTargetMotherInfo.centerInWorld);
+      supportRingNegZInfo.solid = buildStickmanSupportRingSolid("Upstream", false);
+      finishNesting(supportRingNegZInfo,              // volume info as defined above
+                    prodTargetSupportRingMaterial,    // material for the support ring
+                    &tgt->productionTargetRotation(), // rotation
+                    supportRingCenterNegZ,            // translation
+                    prodTargetMotherInfo.logical,     // parent logical volume
+                    0,                                // copy number
+                    G4Colour::Green(),
+                    "ProductionTarget",
+                    verbosityLevel>1);
+      // add the positive z end ring
+      VolumeInfo supportRingPosZInfo("ProductionTargetSupportRing_Downstream",
+                                     supportRingCenterPosZ,
+                                     prodTargetMotherInfo.centerInWorld);
+      supportRingPosZInfo.solid = buildStickmanSupportRingSolid("Downstream", true);
+      finishNesting(supportRingPosZInfo,              // volume info as defined above
+                    prodTargetSupportRingMaterial,    // material for the support ring
+                    &tgt->productionTargetRotation(), // rotation
+                    supportRingCenterPosZ,            // translation
+                    prodTargetMotherInfo.logical,     // parent logical volume
+                    1,                                // copy number
+                    G4Colour::Green(),
+                    "ProductionTarget",
+                    verbosityLevel>1);
+      if (verbosityLevel > 0) {
+        G4cout << __PRETTY_FUNCTION__ << " end rings"
+               << " negZ center=" << supportRingCenterNegZ
+               << " posZ center=" << supportRingCenterPosZ
+               << " material=" << tgt->stickmanSupportRingMaterial()
+               << G4endl;
+      }
+
+      //Copied code from hayman_v_2_0 to add support structures. Stickman and Hayman supports are the same. Avoiding issue with variables like 
+      //prodTargetMotherInfo crossing scopes this way. They share names but are initialized differently. 
+      //In addition, this allows changing the support structure for one target without affecting the other.
+
+      //Add support structures for the production target
+      if(tgt->supportsBuild()) {
+        G4Material* suppWheelMaterial = findMaterialOrThrow(tgt->supportWheelMaterial());
+        G4ThreeVector localWheelCenter(0.0,0.0,0.0); //no offset
+        double suppWheelParams[] = {tgt->supportWheelRIn(), tgt->supportWheelROut(), tgt->supportWheelHL()};
+        //create the volume info for the support wheel+rods
+        VolumeInfo suppWheelInfo( "ProductionTargetSupportWheel", localWheelCenter, prodTargetMotherInfo.centerInMu2e());
+        suppWheelInfo.solid = new G4Tubs("ProductionTargetSupportWheel_wheel", suppWheelParams[0], suppWheelParams[1],
+                                         suppWheelParams[2], 0., CLHEP::twopi);
+                                               // suppWheelParams,
+                                               // suppWheelMaterial,
+                                               // 0,
+                                               // localWheelCenter,
+                                               // prodTargetMotherInfo,
+                                               // 0,
+                                               // G4Colour::Gray(),
+                                               // "PS"
+                                               // );
+
+        // add spokes //
+
+        //spoke info
+        const int nspokesperside = tgt->nSpokesPerSide();
+        G4Material* spokeMaterial = findMaterialOrThrow(tgt->spokeMaterial());
+        //target info
+        double rTarget = tgt->stickmanSupportRingOuterRadius(); //radius of the support ring to attach to
+        double zTarget = tgt->halfStickmanLength()
+                         - tgt->stickmanSupportRingLength()
+                         + tgt->supportRingCutoutOffset();
+                         //where along the target to attach, note the asymmetry
+        double smallGap = 0.001; //for adding small offsets to avoid overlaps due to precision
+        //initialize parameter vectors
+        //features on wheel
+        const vector<double> supportWheelFeatureAngles = tgt->supportWheelFeatureAngles();
+        const vector<double> supportWheelFeatureArcs   = tgt->supportWheelFeatureArcs  ();
+        const vector<double> supportWheelFeatureRIns   = tgt->supportWheelFeatureRIns  ();
+        //support rods in wheel
+        const vector<double> supportWheelRodHL           = tgt->supportWheelRodHL          ();
+        const vector<double> supportWheelRodOffset       = tgt->supportWheelRodOffset      ();
+        const vector<double> supportWheelRodRadius       = tgt->supportWheelRodRadius      ();
+        const vector<double> supportWheelRodRadialOffset = tgt->supportWheelRodRadialOffset();
+        const vector<double> supportWheelRodWireOffsetD  = tgt->supportWheelRodWireOffsetD ();
+        const vector<double> supportWheelRodWireOffsetU  = tgt->supportWheelRodWireOffsetU ();
+        const vector<double> supportWheelRodAngles       = tgt->supportWheelRodAngles      ();
+        //spoke (support wire) angles
+        const vector<double> spokeTargetAnglesD = tgt->spokeTargetAnglesD();
+        const vector<double> spokeTargetAnglesU = tgt->spokeTargetAnglesU();
+        if(verbosityLevel > 0)
+          std::cout << "Printing information about production target supports:\n";
+
+        const double targetAngle = tgt->rotStickmanY(); //assume target angle is only in the x-z plane for supports
+        CLHEP::HepRotation* rodRot = reg.add(CLHEP::HepRotation(CLHEP::HepRotation::IDENTITY));
+        rodRot->rotateY(-1.*targetAngle);
+
+        for(int istream = 0; istream < 2; ++istream) {
+          for(int ispoke = 0; ispoke < nspokesperside; ++ispoke) {
+            const double wheelAngle =  supportWheelRodAngles[ispoke]*CLHEP::degree;
+            //get angle of the support rod on the wheel and the angle on the target the wire connects to
+            const double targetWireAngle = (istream == 0) ? spokeTargetAnglesD[ispoke]*CLHEP::degree
+              : spokeTargetAnglesU[ispoke]*CLHEP::degree;
+            double rWheel = supportWheelRodRadialOffset[ispoke]; // radius of the wheel to attach to
+            CLHEP::Hep3Vector rodCenter(rWheel*cos(wheelAngle), rWheel*sin(wheelAngle), 0.);
+            const double rodOffset = supportWheelRodOffset[ispoke];
+            rodCenter += CLHEP::Hep3Vector(sin(targetAngle)*rodOffset, 0., cos(targetAngle)*rodOffset);
+            if(istream == 0) { //only do once
+              //add the features near the support rods in the bicycle wheel
+              const double featureAngle = supportWheelFeatureAngles[ispoke]*CLHEP::degree; //angle of feature center
+              const double featureArc   = supportWheelFeatureArcs[ispoke]*CLHEP::degree; //width in angle
+              const double featureRIn   = supportWheelFeatureRIns[ispoke]; //inner radius of feature
+              const double featureROut = tgt->supportWheelRIn() + smallGap; //ensure they overlap for union
+              // double featureR = (featureRIn + featureROut)/2.; //radius of feature center
+              // CLHEP::Hep3Vector featureCenter(featureR*cos(featureAngle), featureR*sin(featureAngle), 0.);
+              CLHEP::Hep3Vector featureCenter(localWheelCenter); //center is wheel center
+              double featureParams[] = {featureRIn, featureROut, tgt->supportWheelHL(), featureAngle - featureArc/2. /*phi0*/, featureArc /*dphi*/};
+              G4Tubs* featureTubs = new G4Tubs("ProductionTargetSupportFeature_" +std::to_string(ispoke),
+                                               featureParams[0], featureParams[1], featureParams[2], featureParams[3], featureParams[4]);
+              suppWheelInfo.solid = new G4UnionSolid("ProductionTargetSupportWheelFeature_union_"+std::to_string(ispoke),
+                                                     suppWheelInfo.solid, featureTubs, 0, featureCenter);
+              //add the support rod to the wheel
+              G4Tubs* rodTubs = new G4Tubs("ProductionTargetSupportRod_" +std::to_string(ispoke),
+                                           0., supportWheelRodRadius[ispoke], supportWheelRodHL[ispoke], 0., CLHEP::twopi);
+              suppWheelInfo.solid = new G4UnionSolid("ProductionTargetSupportWheelRod_union_"+std::to_string(ispoke),
+                                                     suppWheelInfo.solid, rodTubs, rodRot, rodCenter);
+            }
+            const int side = (1-2*istream); //+1 or -1
+            //info about wire connection
+            //get end of the rod on this side
+            CLHEP::Hep3Vector rodAxis(sin(targetAngle), 0., cos(targetAngle));
+            CLHEP::Hep3Vector wheelPos(rodCenter);
+            wheelPos += side*supportWheelRodHL[ispoke]*rodAxis;
+            //translate from rod center to edge
+            CLHEP::Hep3Vector rodCenterToWire(cos(wheelAngle)*cos(targetAngle),
+                                              sin(wheelAngle)*cos(targetAngle),
+                                              -cos(wheelAngle)*sin(targetAngle));
+            wheelPos -= supportWheelRodRadius[ispoke]*rodCenterToWire;
+            double zWireRodOffset = (istream == 0) ? supportWheelRodWireOffsetD[ispoke] : supportWheelRodWireOffsetU[ispoke];
+            wheelPos -= side*zWireRodOffset*rodAxis;
+
+            //get wire position on target
+            CLHEP::Hep3Vector targetPos(rTarget*cos(targetWireAngle), rTarget*sin(targetWireAngle), side*zTarget);
+            targetPos = tgt->productionTargetRotation().inverse()*targetPos; //rotate from target frame to mother frame
+            CLHEP::Hep3Vector spokeAxis((wheelPos-targetPos).unit());
+            CLHEP::Hep3Vector targetAxis(0.,0.,side);
+            targetAxis = tgt->productionTargetRotation().inverse()*targetAxis;
+            CLHEP::Hep3Vector zax(0.,0.,1.);
+            if(verbosityLevel > 0)
+              std::cout << "istream " << istream << " ispoke " << ispoke << std::endl
+                        << "target pos " << targetPos << "\nwheel pos " << wheelPos << std::endl
+                        << "Target axis " << targetAxis << "\nSpoke axis " << spokeAxis << std::endl
+                        << "Rod axis " << rodAxis << "\nRod center to wire axis " << rodCenterToWire << std::endl;
+            //to remove overlaps where the wire connects, need angle of wire and surface connecting to
+            //remove overlap at target
+            double wireTargetAngle = targetAxis.angle(-1.*spokeAxis);
+            double deltaLength = (abs(tan(wireTargetAngle)) > 1.e-6) ? abs(tgt->spokeRadius()/tan(wireTargetAngle)) : 0.; //give up if ~paralle
+            targetPos += (deltaLength+0.1)*spokeAxis; //subtract off the length
+            if(verbosityLevel > 0)
+              std::cout << "wire target angle " << wireTargetAngle << " delta L " << deltaLength
+                        << " target pos " << targetPos <<std::endl;
+
+            //next remove overlap at rod
+            //
+            double wireRodAngle = abs((rodCenterToWire).angle(spokeAxis));
+            // deltaLength = abs(tan(wireRodAngle)/tgt->spokeRadius());
+            // wheelPos -= (deltaLength+1.)*spokeAxis;
+            deltaLength = abs(tgt->spokeRadius()/tan(wireRodAngle)); // seems to be a typo in the hayman
+            wheelPos -= (deltaLength+0.1)*spokeAxis;
+            if(verbosityLevel > 0)
+              std::cout << "wire rod angle " << wireRodAngle << " delta L " << deltaLength
+                        << " wheel pos " << wheelPos <<std::endl;
+
+            CLHEP::Hep3Vector spokeCenter((wheelPos+targetPos)/2.);
+            double spokeLength = abs((wheelPos-targetPos).mag());
+            TubsParams spokeParams(0., tgt->spokeRadius(), 0.5*spokeLength);
+            CLHEP::HepRotation* spokeRot = reg.add(CLHEP::HepRotation(spokeAxis.cross(zax), spokeAxis.angle(zax)));
+            std::stringstream spokeName;
+            spokeName << "ProductionTargetSpokeWire_" ;
+            if(istream == 0)
+              spokeName << "Downstream_";
+            else
+              spokeName << "Upstream_";
+            spokeName << ispoke;
+
+            VolumeInfo spokeInfo   = nestTubs( spokeName.str(),
+                                               spokeParams,
+                                               spokeMaterial,
+                                               spokeRot,
+                                               spokeCenter,
+                                               prodTargetMotherInfo,
+                                               0,
+                                               G4Colour::Gray(),
+                                               "PS"
+                                               );
+
+          } //end spokes loop
+        } //end stream loop
+        finishNesting(suppWheelInfo,
+                      suppWheelMaterial,
+                      0,
+                      localWheelCenter,
+                      prodTargetMotherInfo.logical,
+                      0,
+                      G4Colour::Gray(),
+                      "PS"
+                      );
+
+      } //end adding support structures
+    } //end ProductionTargetMaker::stickman_v_1_0
   } //end constructTargetPS
 } //end namespace mu2e
 
diff --git a/ProductionTargetGeom/inc/ProductionTarget.hh b/ProductionTargetGeom/inc/ProductionTarget.hh
index 8ed4e3fac6..f7049cacf7 100644
--- a/ProductionTargetGeom/inc/ProductionTarget.hh
+++ b/ProductionTargetGeom/inc/ProductionTarget.hh
@@ -7,6 +7,7 @@
 #define PRODUCTIONTARGET_HH
 
 #include <map>
+#include <numeric>
 
 #include "canvas/Persistency/Common/Wrapper.h"
 
@@ -137,14 +138,66 @@ namespace mu2e {
     double productionTargetMotherOuterRadius() const {return _productionTargetMotherOuterRadius;}
     double productionTargetMotherHalfLength()  const {return _productionTargetMotherHalfLength;}
 
+    //accessors for Stickman_v_1_0
+    std::string stickmanTargetType()              const {return _stickmanTargetType;}
+    double      halfStickmanLength()              const {return _halfStickmanLength;}
+    int         numberOfPlates()                  const {return _numberOfPlates;}
+    std::string plateMaterial(int i)              const {return _plateMaterial[i];}
+    const std::vector<std::string>& plateMaterial() const {return _plateMaterial;}
+    double      plateROut(int i)                  const {return _plateROut[i];}
+    const std::vector<double>& plateROut()        const {return _plateROut;}
+    int         nStickmanFins()                   const {return _nStickmanFins;}
+    double      plateFinAngle(int i)              const {return _plateFinAngles[i];}
+    const std::vector<double>&      plateFinAngles()       const {return _plateFinAngles;}
+    double      plateFinOuterRadius()             const {return _plateFinOuterRadius;}
+    double      plateFinWidth()                   const {return _plateFinWidth;}
+    double      plateCenterToLugCenter()          const {return _plateCenterToLugCenter;}
+    double      plateLugInnerRadius()             const {return _plateLugInnerRadius;}
+    double      plateLugOuterRadius()             const {return _plateLugOuterRadius;}
+    double      plateThickness(int i)             const {return _plateThickness[i];}
+    const std::vector<double>&      plateThickness()       const {return _plateThickness;}
+    double      plateLugThickness(int i)          const {return _plateLugThickness[i];}
+    const std::vector<double>&      plateLugThickness()    const {return _plateLugThickness;}
+    bool        addFilletToPlateCore()            const {return _addFilletToPlateCore;}
+    bool        addFilletToPlateLug()             const {return _addFilletToPlateLug;}
+    double      plateFilletRadius()               const {return _plateFilletRadius;}
+    std::string rodMaterial()                 const {return _rodMaterial;}
+    double      rodRadius()                       const {return _rodRadius;}
+    double      rodHalfLength()                   const {return _rodHalfLength;}
+    std::string spacerMaterial()              const {return _spacerMaterial;}
+    double      spacerHalfLength()                const {return _spacerHalfLength;}
+    double      spacerOuterRadius()               const {return _spacerOuterRadius;}
+    double      spacerInnerRadius()               const {return _spacerInnerRadius;}
+    std::string stickmanSupportRingMaterial()     const {return _stickmanSupportRingMaterial;}
+    double      stickmanSupportRingLength()       const {return _stickmanSupportRingLength;}
+    double      stickmanSupportRingInnerRadius()  const {return _stickmanSupportRingInnerRadius;}
+    double      stickmanSupportRingOuterRadius()  const {return _stickmanSupportRingOuterRadius;}
+    double      supportRingLugOuterRadius()       const {return _supportRingLugOuterRadius;}
+    bool        addFilletToSupportRingLug()       const {return _addFilletToSupportRingLug;}
+    double      supportRingLugFilletRadius()      const {return _supportRingLugFilletRadius;}
+    bool        addCutoutToSupportRing()          const {return _addCutoutToSupportRing;}
+    int         nSupportRingCutouts()             const {return _nSupportRingCutouts;}
+    double      supportRingCutoutAngle(int i)     const {return _supportRingCutoutAngles[i];}
+    const std::vector<double>&      supportRingCutoutAngles()    const {return _supportRingCutoutAngles;}
+    double      supportRingCutoutInnerRadius()    const {return _supportRingCutoutInnerRadius;}
+    double      supportRingCutoutTilt()           const {return _supportRingCutoutTilt;}
+    double      supportRingCutoutOffset()         const {return _supportRingCutoutOffset;}
+    double      rotStickmanX()                    const {return _rotStickmanX;}
+    double      rotStickmanY()                    const {return _rotStickmanY;}
+    double      rotStickmanZ()                    const {return _rotStickmanZ;}
+    CLHEP::Hep3Vector stickmanProdTargetPosition() const {return _stickmanProdTargetPosition;}
+
     std::string hayman_v_2_0 = "Hayman_v_2_0";
     std::string tier1 = "MDC2018";
+    std::string stickman_v_1_0 = "Stickman_v_1_0";
 
     CLHEP::Hep3Vector targetPositionByVersion() const {
       if (_haymanTargetType == hayman_v_2_0){
         return _haymanProdTargetPosition;}
       else if  (_tier1TargetType == "MDC2018"){
         return _prodTargetPosition;}
+      else if (_stickmanTargetType == stickman_v_1_0){
+        return _stickmanProdTargetPosition;}
       else throw cet::exception("BADCONFIG")
              << "in ProductionTarget.hh, no valid target specified"<< std::endl;
     }
@@ -153,6 +206,8 @@ namespace mu2e {
         return _halfHaymanLength;}
      else if  (_tier1TargetType == "MDC2018"){
        return _halfLength;}
+     else if (_stickmanTargetType == stickman_v_1_0){
+       return _halfStickmanLength;}
      else throw cet::exception("BADCONFIG")
             << "in ProductionTarget.hh, no valid target specified"<< std::endl;
     }
@@ -205,6 +260,42 @@ namespace mu2e {
                      ,double supportRingCutoutLength
                      );
 
+    ProductionTarget(
+                     std::string stickmanTargetType, int version
+                     ,double productionTargetMotherOuterRadius
+                     ,double productionTargetMotherHalfLength
+                     ,double rotStickmanX
+                     ,double rotStickmanY
+                     ,double rotStickmanZ
+                     ,double halfStickmanLength
+                     ,const CLHEP::Hep3Vector& stickmanProdTargetPosition
+                     ,std::string targetVacuumMaterial
+                     ,int numberOfPlates
+                     ,std::vector<std::string> plateMaterial
+                     ,std::vector<double> plateROut
+                     ,int nStickmanFins
+                     ,std::vector<double> plateFinAngles
+                     ,double plateFinOuterRadius
+                     ,double plateFinWidth
+                     ,double plateCenterToLugCenter
+                     ,double plateLugInnerRadius
+                     ,double plateLugOuterRadius
+                     ,std::vector<double> plateThickness
+                     ,std::vector<double> plateLugThickness
+                     ,std::string rodMaterial
+                     ,double rodRadius
+                     ,std::string spacerMaterial
+                     ,double spacerHalfLength
+                     ,double spacerOuterRadius
+                     ,double spacerInnerRadius
+                     ,std::string stickmanSupportRingMaterial
+                     ,double stickmanSupportRingLength
+                     ,double stickmanSupportRingInnerRadius
+                     ,double stickmanSupportRingOuterRadius
+                     ,double supportRingLugOuterRadius
+                     ,double supportRingCutoutOffset
+                     );
+
     CLHEP::HepRotation _protonBeamRotation;
 
     // can't return by const ref if invert on the fly so need to store redundant data
@@ -302,6 +393,49 @@ namespace mu2e {
     std::vector<double> _spokeTargetAnglesU; //angle about the target the wire connects to (upstream)
     double              _spokeRadius       ; //radius of the wire
 
+    // Stickman_v_1_0 parameters
+    std::string _stickmanTargetType;
+    double      _halfStickmanLength;
+    double      _rotStickmanX;
+    double      _rotStickmanY;
+    double      _rotStickmanZ;
+    CLHEP::Hep3Vector _stickmanProdTargetPosition;
+    int         _numberOfPlates;
+    std::vector<std::string> _plateMaterial;
+    std::vector<double>      _plateROut;
+    int         _nStickmanFins;
+    std::vector<double>      _plateFinAngles;
+    double      _plateFinOuterRadius;
+    double      _plateFinWidth;
+    double      _plateCenterToLugCenter;
+    double      _plateLugInnerRadius;
+    double      _plateLugOuterRadius;
+    std::vector<double>      _plateThickness;
+    std::vector<double>      _plateLugThickness;
+    bool        _addFilletToPlateCore;
+    bool        _addFilletToPlateLug;
+    double      _plateFilletRadius;
+    std::string _rodMaterial;
+    double      _rodRadius;
+    double      _rodHalfLength;
+    std::string _spacerMaterial;
+    double      _spacerHalfLength;
+    double      _spacerOuterRadius;
+    double      _spacerInnerRadius;
+    std::string _stickmanSupportRingMaterial;
+    double      _stickmanSupportRingLength;
+    double      _stickmanSupportRingInnerRadius;
+    double      _stickmanSupportRingOuterRadius;
+    double      _supportRingLugOuterRadius;
+    bool        _addFilletToSupportRingLug;
+    double      _supportRingLugFilletRadius;
+    bool        _addCutoutToSupportRing;
+    int         _nSupportRingCutouts;
+    std::vector<double> _supportRingCutoutAngles;
+    double      _supportRingCutoutInnerRadius;
+    double      _supportRingCutoutTilt;
+    double      _supportRingCutoutOffset;
+
     // Needed for persistency
     template<class T> friend class art::Wrapper;
     ProductionTarget():_pHubsRgtParams(NULL), _pHubsLftParams(NULL) {}
diff --git a/ProductionTargetGeom/src/ProductionTarget.cc b/ProductionTargetGeom/src/ProductionTarget.cc
index 042fad66b4..94cf554320 100644
--- a/ProductionTargetGeom/src/ProductionTarget.cc
+++ b/ProductionTargetGeom/src/ProductionTarget.cc
@@ -99,5 +99,87 @@ namespace mu2e {
     _prodTargetPosition = _haymanProdTargetPosition;
   }
 
+  ProductionTarget::ProductionTarget(
+                                     std::string stickmanTargetType
+                                     ,int version
+                                     ,double productionTargetMotherOuterRadius
+                                     ,double productionTargetMotherHalfLength
+                                     ,double rotStickmanX
+                                     ,double rotStickmanY
+                                     ,double rotStickmanZ
+                                     ,double halfStickmanLength
+                                     ,const CLHEP::Hep3Vector& stickmanProdTargetPosition
+                                     ,std::string targetVacuumMaterial
+                                     ,int numberOfPlates
+                                     ,std::vector<std::string> plateMaterial
+                                     ,std::vector<double> plateROut
+                                     ,int nStickmanFins
+                                     ,std::vector<double> plateFinAngles
+                                     ,double plateFinOuterRadius
+                                     ,double plateFinWidth
+                                     ,double plateCenterToLugCenter
+                                     ,double plateLugInnerRadius
+                                     ,double plateLugOuterRadius
+                                     ,std::vector<double> plateThickness
+                                     ,std::vector<double> plateLugThickness
+                                     ,std::string rodMaterial
+                                     ,double rodRadius
+                                     ,std::string spacerMaterial
+                                     ,double spacerHalfLength
+                                     ,double spacerOuterRadius
+                                     ,double spacerInnerRadius
+                                     ,std::string stickmanSupportRingMaterial
+                                     ,double stickmanSupportRingLength
+                                     ,double stickmanSupportRingInnerRadius
+                                     ,double stickmanSupportRingOuterRadius
+                                     ,double supportRingLugOuterRadius
+                                     ,double supportRingCutoutOffset
+                                     )
+  : _protonBeamRotation(CLHEP::HepRotation::IDENTITY)
+    ,_stickmanTargetType(stickmanTargetType)
+    ,_version(version)
+    ,_productionTargetMotherOuterRadius(productionTargetMotherOuterRadius)
+    ,_productionTargetMotherHalfLength(productionTargetMotherHalfLength)
+    ,_rotStickmanX(rotStickmanX)
+    ,_rotStickmanY(rotStickmanY)
+    ,_rotStickmanZ(rotStickmanZ)
+    ,_halfStickmanLength(halfStickmanLength)
+    ,_stickmanProdTargetPosition(stickmanProdTargetPosition)
+    ,_targetVacuumMaterial(targetVacuumMaterial)
+    ,_numberOfPlates(numberOfPlates)
+    ,_plateMaterial(plateMaterial)
+    ,_plateROut(plateROut)
+    ,_nStickmanFins(nStickmanFins)
+    ,_plateFinAngles(plateFinAngles)
+    ,_plateFinOuterRadius(plateFinOuterRadius)
+    ,_plateFinWidth(plateFinWidth)
+    ,_plateCenterToLugCenter(plateCenterToLugCenter)
+    ,_plateLugInnerRadius(plateLugInnerRadius)
+    ,_plateLugOuterRadius(plateLugOuterRadius)
+    ,_plateThickness(plateThickness)
+    ,_plateLugThickness(plateLugThickness)
+    ,_rodMaterial(rodMaterial)
+    ,_rodRadius(rodRadius)
+    ,_spacerMaterial(spacerMaterial)
+    ,_spacerHalfLength(spacerHalfLength)
+    ,_spacerOuterRadius(spacerOuterRadius)
+    ,_spacerInnerRadius(spacerInnerRadius)
+    ,_stickmanSupportRingMaterial(stickmanSupportRingMaterial)
+    ,_stickmanSupportRingLength(stickmanSupportRingLength)
+    ,_stickmanSupportRingInnerRadius(stickmanSupportRingInnerRadius)
+    ,_stickmanSupportRingOuterRadius(stickmanSupportRingOuterRadius)
+    ,_supportRingLugOuterRadius(supportRingLugOuterRadius)
+    ,_supportRingCutoutOffset(supportRingCutoutOffset)
+  {
+    // rod half length, actual rod is longer than this since it inserts into the end rings, but for the geometry reconstruction, that part will be treated as part of the end ring.
+    _rodHalfLength = std::accumulate(_plateLugThickness.begin(), _plateLugThickness.end(), 0.0) / 2.0 + 2 * _spacerHalfLength;
+
+    // rotations
+    _protonBeamRotation.rotateX(rotStickmanX).rotateY(rotStickmanY).rotateZ(rotStickmanZ);
+    _protonBeamInverseRotation = _protonBeamRotation.inverse(); // passive rotation matrix
+    _halfLength = _productionTargetMotherHalfLength;
+    _prodTargetPosition = _stickmanProdTargetPosition;
+  }
+
 
 }

From e75dfa0f99e110622e514c83ca98d623cf85f3c7 Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2ebuild02.fnal.gov>
Date: Thu, 9 Apr 2026 12:48:04 -0500
Subject: [PATCH 07/15] bug fix

---
 Mu2eG4/src/constructTargetPS.cc              | 7 ++++++-
 ProductionTargetGeom/src/ProductionTarget.cc | 6 +++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/Mu2eG4/src/constructTargetPS.cc b/Mu2eG4/src/constructTargetPS.cc
index 330ff81965..567dd68e1d 100644
--- a/Mu2eG4/src/constructTargetPS.cc
+++ b/Mu2eG4/src/constructTargetPS.cc
@@ -1716,8 +1716,13 @@ namespace mu2e {
                  rodCenter,                        // translation
                  prodTargetMotherInfo,             // parent volume info for the target mother
                  ithRod,                           // copy number
+                 prodTargetVisible,                // visibility
                  G4Colour::Green(),                // color for visualization
-                 "ProductionTarget");              // lookup token
+                 prodTargetSolid,
+                 forceAuxEdgeVisible,
+                 placePV,
+                 doSurfaceCheck
+                 );
         if (verbosityLevel > 0) {
           G4cout << __PRETTY_FUNCTION__ << " rod " << ithRod
                  << " center=" << rodCenter
diff --git a/ProductionTargetGeom/src/ProductionTarget.cc b/ProductionTargetGeom/src/ProductionTarget.cc
index 94cf554320..8207b12527 100644
--- a/ProductionTargetGeom/src/ProductionTarget.cc
+++ b/ProductionTargetGeom/src/ProductionTarget.cc
@@ -136,16 +136,16 @@ namespace mu2e {
                                      ,double supportRingCutoutOffset
                                      )
   : _protonBeamRotation(CLHEP::HepRotation::IDENTITY)
-    ,_stickmanTargetType(stickmanTargetType)
     ,_version(version)
     ,_productionTargetMotherOuterRadius(productionTargetMotherOuterRadius)
     ,_productionTargetMotherHalfLength(productionTargetMotherHalfLength)
+    ,_targetVacuumMaterial(targetVacuumMaterial)
+    ,_stickmanTargetType(stickmanTargetType)
+    ,_halfStickmanLength(halfStickmanLength)
     ,_rotStickmanX(rotStickmanX)
     ,_rotStickmanY(rotStickmanY)
     ,_rotStickmanZ(rotStickmanZ)
-    ,_halfStickmanLength(halfStickmanLength)
     ,_stickmanProdTargetPosition(stickmanProdTargetPosition)
-    ,_targetVacuumMaterial(targetVacuumMaterial)
     ,_numberOfPlates(numberOfPlates)
     ,_plateMaterial(plateMaterial)
     ,_plateROut(plateROut)

From 1a738afa79086f1b894996852026177d1bcf98e6 Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2ebuild02.fnal.gov>
Date: Fri, 10 Apr 2026 00:48:17 -0500
Subject: [PATCH 08/15] fix support structure per docdb 30814

---
 GeometryService/src/GeometryService.cc        | 18 +++---
 .../geom/ProductionTarget_Stickman_v1_0.txt   | 24 ++++----
 Mu2eG4/src/constructTargetPS.cc               | 60 ++++++++++++-------
 3 files changed, 60 insertions(+), 42 deletions(-)

diff --git a/GeometryService/src/GeometryService.cc b/GeometryService/src/GeometryService.cc
index 5edbe1145a..cd54ff7d1c 100644
--- a/GeometryService/src/GeometryService.cc
+++ b/GeometryService/src/GeometryService.cc
@@ -240,15 +240,17 @@ namespace mu2e {
 
     //addDetector(PSShieldMaker::make(*_config, ps.psEndRefPoint(), prodTarget.position()));
 
-   if (_config->getString("targetPS_model") == "MDC2018"){
-     //      std::cout << "adding Tier1 in GeometryService" << std::endl;
+    if (_config->getString("targetPS_model") == "MDC2018"){
+      // std::cout << "adding Tier1 in GeometryService" << std::endl;
       addDetector(PSShieldMaker::make(*_config, ps.psEndRefPoint(), prodTarget.position()));
-        } else
-      if (_config->getString("targetPS_model") == "Hayman_v_2_0"){
-        //        std::cout << " adding Hayman in GeometryService" << std::endl;
-        addDetector(PSShieldMaker::make(*_config, ps.psEndRefPoint(), prodTarget.haymanProdTargetPosition()));
-          } else
-        {throw cet::exception("GEOM") << " " << static_cast<char const*>(__func__) << " illegal production target version specified in GeometryService_service = " << _config->getString("targetPS_model")  << std::endl;}
+    } else if (_config->getString("targetPS_model") == "Hayman_v_2_0"){
+      // std::cout << " adding Hayman in GeometryService" << std::endl;
+      addDetector(PSShieldMaker::make(*_config, ps.psEndRefPoint(), prodTarget.haymanProdTargetPosition()));
+    } else if (_config->getString("targetPS_model") == "Stickman_v_1_0"){
+      // std::cout << "adding Stickman in GeometryService" << std::endl;
+      addDetector(PSShieldMaker::make(*_config, ps.psEndRefPoint(), prodTarget.stickmanProdTargetPosition()));
+    } else {
+      throw cet::exception("GEOM") << " " << static_cast<char const*>(__func__) << " illegal production target version specified in GeometryService_service = " << _config->getString("targetPS_model")  << std::endl;}
 
 
 
diff --git a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
index 46b042e829..b298abc0ef 100644
--- a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
+++ b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
@@ -51,7 +51,7 @@ vector<double> targetPS_rOut = {
     3.15, 3.15, 3.15, 3.15, 3.15  // 35
 }; // in mm, core part radius, one value for each
 int targetPS_nStickmanFins             = 3;                   // since the plates are installed on rods, the number of fins is always the same
-vector<double> targetPS_plateFinAngles = {330.,210.,90.};  // degrees
+vector<double> targetPS_plateFinAngles = {285.,165.,45.};     // degrees
 double targetPS_plateFinOuterRadius    = 18.075;              // mm, (39.20-3.05)/2, extends to lug-to-center radius - lug inner radius
 double targetPS_plateFinWidth          = 2.0;                 // mm, all fin has same width
 double targetPS_plateCenterToLugCenter = 19.6;                // mm, center of plate to center of lug
@@ -107,8 +107,9 @@ double targetPS_supportRingLugFilletRadius = 1.0;
 bool   targetPS_addCutoutToSupportRing     = true; // the six cylindrical cutouts at an angle
 double targetPS_supportRingCutoutOffset    = 4.78; // mm, distance from the center of the cutout to the side of the end ring closer to the target center. 
                                                    // This value is a must-have. If no cutout is added, this determines where the support wire is attached from the inner edge of the end ring.
-int    targetPS_nSupportRingCutouts        = 6;
-vector<double> targetPS_supportRingCutoutAngles = {0., 60., 120., 180., 240., 300.}; // degrees
+int    targetPS_nSupportRingCutouts        = 3;
+vector<double> targetPS_supportRingCutoutAngles = {15., 135., 255.} // degrees
+                                                  //{-45., 75., 195.}; these have spokes inside so not cutting out
 // FIXME: double check rotation around the z axis
 double targetPS_supportRingCutoutInnerRadius = 2.06; // mm
 double targetPS_supportRingCutoutTilt        = 20.0; //degrees, angle of the cutout w.r.t the radial direction
@@ -120,21 +121,22 @@ double targetPS.supports.wheel.rOut       = 196.85;
 double targetPS.supports.wheel.rIn        = 177.8;
 string targetPS.supports.wheel.material   =  "G4_Al";
 int    targetPS.supports.nSpokes          = 3;
-// Features on the wheel
-vector<double> targetPS.supports.features.angles = {-46.62,  79.38, 194.38}; //degrees
+// Features on the wheel, changed according to docdb-30814
+vector<double> targetPS.supports.features.angles = {-37.,  83., 203.}; //degrees, center of the features
 vector<double> targetPS.supports.features.arcs   = { 28.,  28.,  28.}; //degrees
 vector<double> targetPS.supports.features.rIns   = {158.75, 158.75, 158.75}; //mm
 //support rods in the wheel that the wires connect to
-vector<double> targetPS.supports.rods.angles       = {-54., 71.5, 186.0}; //degrees
+vector<double> targetPS.supports.rods.angles       = {-45., 75., 195.0}; //degrees
 vector<double> targetPS.supports.rods.halfLength   = {70., 70., 70.}; //mm
-vector<double> targetPS.supports.rods.offset       = {-40., -10., 40.}; //mm
+vector<double> targetPS.supports.rods.offset       = {-35., -14.80, 38.62}; //mm, changed according to F10122537_A_DWG1
+                                                                            //rod center to the channel centers
 vector<double> targetPS.supports.rods.radius       = {9., 9., 9.}; //mm
 vector<double> targetPS.supports.rods.radialOffset = {177.8, 177.8, 177.8}; //mm
-vector<double> targetPS.supports.rods.wireOffset.downstream = {10., 10., 10.}; //mm
-vector<double> targetPS.supports.rods.wireOffset.upstream   = {10., 10., 10.}; //mm
+vector<double> targetPS.supports.rods.wireOffset.downstream = {13., 13., 13.}; //mm, on the surface
+vector<double> targetPS.supports.rods.wireOffset.upstream   = {13., 13., 13.}; //mm, calculated according to F10122537_A_DWG1
 //spokes connecting the target to the wheel rods
-vector<double> targetPS.supports.spokes.targetAngles.downstream = {-54., 71.5, 186.0}; //degrees
-vector<double> targetPS.supports.spokes.targetAngles.upstream   = {-54., 71.5, 186.0}; //degrees
+vector<double> targetPS.supports.spokes.targetAngles.downstream = {-45., 75., 195.0}; //degrees
+vector<double> targetPS.supports.spokes.targetAngles.upstream   = {-45., 75., 195.0}; //degrees
 double targetPS.supports.spokes.diameter = 2.;
 string targetPS.supports.spokes.material = "Inconel718";
 
diff --git a/Mu2eG4/src/constructTargetPS.cc b/Mu2eG4/src/constructTargetPS.cc
index 567dd68e1d..3bdc8d3101 100644
--- a/Mu2eG4/src/constructTargetPS.cc
+++ b/Mu2eG4/src/constructTargetPS.cc
@@ -1302,13 +1302,13 @@ namespace mu2e {
 
       if (verbosityLevel > 0){
         G4cout << "target position and hall origin = "
-                  << tgt->stickmanProdTargetPosition() << "\n" <<
-          _hallOriginInMu2e << " " << parent.centerInMu2e() << G4endl;
+               << tgt->stickmanProdTargetPosition() << "\n"
+               << _hallOriginInMu2e << " " << parent.centerInMu2e() << G4endl;
       }
       if (verbosityLevel > 2){
         G4cout << __PRETTY_FUNCTION__ << "created prodTargetMotherInfo "
-                  << tgt->productionTargetMotherOuterRadius()
-                  << " " <<tgt->productionTargetMotherHalfLength() << G4endl;
+               << tgt->productionTargetMotherOuterRadius()
+               << " " <<tgt->productionTargetMotherHalfLength() << G4endl;
       }
 
       // tiny offset to avoid precision issues
@@ -1646,7 +1646,8 @@ namespace mu2e {
         // apply the inverse rotation to the local plate center.
 
         if (verbosityLevel > 0) {
-          G4cout << __PRETTY_FUNCTION__ << " plate " << ithPlate
+          G4cout << __PRETTY_FUNCTION__
+                 << " plate " << ithPlate
                  << " startZ=" << _currentZ
                  << " centerZ=" << plateCenterZ
                  << " endZ=" << (_currentZ + tgt->plateLugThickness(ithPlate))
@@ -1654,7 +1655,8 @@ namespace mu2e {
                  << G4endl;
         }
         if (verbosityLevel > 2) {
-          G4cout << " rOut=" << tgt->plateROut(ithPlate)
+          G4cout << __PRETTY_FUNCTION__
+                 << " rOut=" << tgt->plateROut(ithPlate)
                  << " thickness=" << tgt->plateThickness(ithPlate)
                  << " finWidth=" << tgt->plateFinWidth()
                  << " finReach=" << tgt->plateFinOuterRadius()
@@ -1977,7 +1979,9 @@ namespace mu2e {
           const double localCutoutZ = (positiveEnd ? 1. : -1.) * 
                                       (tgt->supportRingCutoutOffset() - ringHalfLength + supportRingWallThickness * 0.5 * std::tan(cutoutTiltAngle) );
           // cutout needs to go through the whole wall thickness
-          const double cutoutHalfLength = 0.5 * supportRingWallThickness / std::cos(cutoutTiltAngle) + stickmanMagicOffset; // add small offset to avoid precision issues
+          const double cutoutHalfLength = 0.5 * supportRingWallThickness / std::cos(cutoutTiltAngle)
+                                          + tgt->supportRingCutoutInnerRadius() * std::tan(cutoutTiltAngle)
+                                          + stickmanMagicOffset; // add small offset to avoid precision issues
           if (verbosityLevel > 2) {
             G4cout << __PRETTY_FUNCTION__ << " support ring " << tag
                    << " cutoutTilt(deg)=" << cutoutTiltAngle/CLHEP::degree
@@ -2148,9 +2152,11 @@ namespace mu2e {
             const double targetWireAngle = (istream == 0) ? spokeTargetAnglesD[ispoke]*CLHEP::degree
               : spokeTargetAnglesU[ispoke]*CLHEP::degree;
             double rWheel = supportWheelRodRadialOffset[ispoke]; // radius of the wheel to attach to
-            CLHEP::Hep3Vector rodCenter(rWheel*cos(wheelAngle), rWheel*sin(wheelAngle), 0.);
+            CLHEP::Hep3Vector rodCenter(rWheel*std::cos(wheelAngle)*std::cos(targetAngle),
+                                        rWheel*std::sin(wheelAngle),
+                                        0.); // rod plane projected to the wheel plane
             const double rodOffset = supportWheelRodOffset[ispoke];
-            rodCenter += CLHEP::Hep3Vector(sin(targetAngle)*rodOffset, 0., cos(targetAngle)*rodOffset);
+            rodCenter += CLHEP::Hep3Vector(std::sin(targetAngle)*rodOffset, 0., std::cos(targetAngle)*rodOffset);
             if(istream == 0) { //only do once
               //add the features near the support rods in the bicycle wheel
               const double featureAngle = supportWheelFeatureAngles[ispoke]*CLHEP::degree; //angle of feature center
@@ -2174,37 +2180,44 @@ namespace mu2e {
             const int side = (1-2*istream); //+1 or -1
             //info about wire connection
             //get end of the rod on this side
-            CLHEP::Hep3Vector rodAxis(sin(targetAngle), 0., cos(targetAngle));
+            CLHEP::Hep3Vector rodAxis(std::sin(targetAngle), 0., std::cos(targetAngle));
             CLHEP::Hep3Vector wheelPos(rodCenter);
             wheelPos += side*supportWheelRodHL[ispoke]*rodAxis;
-            //translate from rod center to edge
-            CLHEP::Hep3Vector rodCenterToWire(cos(wheelAngle)*cos(targetAngle),
-                                              sin(wheelAngle)*cos(targetAngle),
-                                              -cos(wheelAngle)*sin(targetAngle));
+            //translate from rod center to edge, this is vector from center to rod center rotated by target angle
+            CLHEP::Hep3Vector rodCenterToWire(std::cos(wheelAngle)*std::cos(targetAngle),
+                                              std::sin(wheelAngle), // rorated by target angle around y so y should be unchanged
+                                              -std::cos(wheelAngle)*std::sin(targetAngle));
             wheelPos -= supportWheelRodRadius[ispoke]*rodCenterToWire;
             double zWireRodOffset = (istream == 0) ? supportWheelRodWireOffsetD[ispoke] : supportWheelRodWireOffsetU[ispoke];
             wheelPos -= side*zWireRodOffset*rodAxis;
 
             //get wire position on target
-            CLHEP::Hep3Vector targetPos(rTarget*cos(targetWireAngle), rTarget*sin(targetWireAngle), side*zTarget);
+            CLHEP::Hep3Vector targetPos(rTarget*std::cos(targetWireAngle), rTarget*std::sin(targetWireAngle), side*zTarget);
             targetPos = tgt->productionTargetRotation().inverse()*targetPos; //rotate from target frame to mother frame
+            if (verbosityLevel > 1)
+              G4cout << __PRETTY_FUNCTION__
+                     << "Sanity check for spoke " << ispoke << " on stream " << istream << ":\n"
+                     << "wheel pos " << wheelPos << "\ntarget pos " << targetPos << "\n"
+                     << "spoke length" << (wheelPos-targetPos).mag() << G4endl;
             CLHEP::Hep3Vector spokeAxis((wheelPos-targetPos).unit());
             CLHEP::Hep3Vector targetAxis(0.,0.,side);
             targetAxis = tgt->productionTargetRotation().inverse()*targetAxis;
             CLHEP::Hep3Vector zax(0.,0.,1.);
             if(verbosityLevel > 0)
-              std::cout << "istream " << istream << " ispoke " << ispoke << std::endl
-                        << "target pos " << targetPos << "\nwheel pos " << wheelPos << std::endl
-                        << "Target axis " << targetAxis << "\nSpoke axis " << spokeAxis << std::endl
-                        << "Rod axis " << rodAxis << "\nRod center to wire axis " << rodCenterToWire << std::endl;
+              G4cout << __PRETTY_FUNCTION__
+                     << "istream " << istream << " ispoke " << ispoke << G4endl
+                     << "target pos " << targetPos << "\nwheel pos " << wheelPos << G4endl
+                     << "Target axis " << targetAxis << "\nSpoke axis " << spokeAxis << G4endl
+                     << "Rod axis " << rodAxis << "\nRod center to wire axis " << rodCenterToWire << G4endl;
             //to remove overlaps where the wire connects, need angle of wire and surface connecting to
             //remove overlap at target
             double wireTargetAngle = targetAxis.angle(-1.*spokeAxis);
             double deltaLength = (abs(tan(wireTargetAngle)) > 1.e-6) ? abs(tgt->spokeRadius()/tan(wireTargetAngle)) : 0.; //give up if ~paralle
             targetPos += (deltaLength+0.1)*spokeAxis; //subtract off the length
             if(verbosityLevel > 0)
-              std::cout << "wire target angle " << wireTargetAngle << " delta L " << deltaLength
-                        << " target pos " << targetPos <<std::endl;
+              G4cout << __PRETTY_FUNCTION__
+                     << "wire target angle " << wireTargetAngle << " delta L " << deltaLength
+                     << " target pos " << targetPos << G4endl;
 
             //next remove overlap at rod
             //
@@ -2214,8 +2227,9 @@ namespace mu2e {
             deltaLength = abs(tgt->spokeRadius()/tan(wireRodAngle)); // seems to be a typo in the hayman
             wheelPos -= (deltaLength+0.1)*spokeAxis;
             if(verbosityLevel > 0)
-              std::cout << "wire rod angle " << wireRodAngle << " delta L " << deltaLength
-                        << " wheel pos " << wheelPos <<std::endl;
+              G4cout << __PRETTY_FUNCTION__
+                     << "wire rod angle " << wireRodAngle << " delta L " << deltaLength
+                     << " wheel pos " << wheelPos << G4endl;
 
             CLHEP::Hep3Vector spokeCenter((wheelPos+targetPos)/2.);
             double spokeLength = abs((wheelPos-targetPos).mag());

From e947cc760e5edde7e47f6a9dbad61a3c6d983820 Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2ebuild02.fnal.gov>
Date: Fri, 10 Apr 2026 01:01:51 -0500
Subject: [PATCH 09/15] bug fix

---
 Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
index b298abc0ef..ac6bfbe9ae 100644
--- a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
+++ b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
@@ -5,7 +5,7 @@
 
 // The Hayman dependence is intentionally removed. 
 
-int PSStickman.verbosityLevel = 0; //5;
+int PSStickman.verbosityLevel = 5;
 
 int targetPS_version  = 4;
 string targetPS_model = "Stickman_v_1_0";
@@ -108,7 +108,7 @@ bool   targetPS_addCutoutToSupportRing     = true; // the six cylindrical cutout
 double targetPS_supportRingCutoutOffset    = 4.78; // mm, distance from the center of the cutout to the side of the end ring closer to the target center. 
                                                    // This value is a must-have. If no cutout is added, this determines where the support wire is attached from the inner edge of the end ring.
 int    targetPS_nSupportRingCutouts        = 3;
-vector<double> targetPS_supportRingCutoutAngles = {15., 135., 255.} // degrees
+vector<double> targetPS_supportRingCutoutAngles = {15., 135., 255.}; // degrees
                                                   //{-45., 75., 195.}; these have spokes inside so not cutting out
 // FIXME: double check rotation around the z axis
 double targetPS_supportRingCutoutInnerRadius = 2.06; // mm

From 8cc1ead14439bfb91f90b4b0ee8c82d5835c49c4 Mon Sep 17 00:00:00 2001
From: Yongyi Wu <yongyiwu@mu2ebuild02.fnal.gov>
Date: Fri, 10 Apr 2026 01:03:25 -0500
Subject: [PATCH 10/15] make verbosity level default 0

---
 Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
index ac6bfbe9ae..0cb5ccff1d 100644
--- a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
+++ b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
@@ -5,7 +5,7 @@
 
 // The Hayman dependence is intentionally removed. 
 
-int PSStickman.verbosityLevel = 5;
+int PSStickman.verbosityLevel = 0; //5;
 
 int targetPS_version  = 4;
 string targetPS_model = "Stickman_v_1_0";

From ef193f39fe47d42b7b119e55fc648365b958e642 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Fri, 10 Apr 2026 17:49:20 -0400
Subject: [PATCH 11/15] Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 MachineLearningTools/inc/ScoreBasedDiffusionModel.hh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
index d0aed40046..ee11dab097 100644
--- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -108,7 +108,7 @@ namespace mu2e{
         //
         // Parameters:
         //   condition       - Optional conditioning vector (must match conditionDim_ when enabled)
-        //   useHeun         - If true, uses Heun's method (2nd order). If false, uses Euler method (2nd order, default)
+        //   useHeun         - If true, uses Heun's method (2nd order, default). If false, uses Euler's method (1st order)
         //   diffusionSteps  - Number of diffusion steps for sampling (default: -1 uses the model's configured diffusionSteps_)
         //
         // Returns: A generated sample vector of dimension dim_

From 3acafe6cfdcdf6a3aebaf80f6dc70f3f5ce844a9 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Fri, 10 Apr 2026 17:50:49 -0400
Subject: [PATCH 12/15] Update STMMC/src/VDResamplerConfigure_module.cc

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 STMMC/src/VDResamplerConfigure_module.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index d5b33f2cf0..c3721b994d 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -181,7 +181,7 @@ namespace mu2e {
     std::string trainingPathNames = "";
     std::vector<std::pair<std::string, std::string>> trainingPaths; // pairs of (moduleName, pathName)
 
-    sumOutFile << "PDGID,HitCount,NotTrained\n";
+    sumOutFile << "PDGID,HitCount,TrainingMode\n";
 
     fclOutFile << "# This fcl is generated by VDResamplerConfigure_module.cc\n";
     fclOutFile << "# Training configuration for the VD resampler is generated for each particle type\n\n";

From 9208a0e44239239b549d9d36f1c320c20fbd12a3 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Fri, 10 Apr 2026 18:15:26 -0400
Subject: [PATCH 13/15] Changed according to PR bot suggestions

---
 STMMC/src/VDResamplerConfigure_module.cc | 15 ++++++++------
 STMMC/src/VDResamplerTrain_module.cc     | 25 ++++++++++++++----------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index c3721b994d..9504fb9d9e 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -78,6 +78,7 @@ namespace mu2e {
       int pdgId = 0;
       double x = 0.0, y = 0.0, z = 0.0, px = 0.0, py = 0.0, pz = 0.0, mass = 0.0, E = 0.0, time = 0.0;
       VolumeId_type virtualdetectorId = 0;
+      int nThrownEvents = 0;
       TTree* ttree;
       std::map<int, int> pdgIds; // <id, count>
   };
@@ -127,8 +128,9 @@ namespace mu2e {
       pz = step.momentum().z();
       if (virtualdetectorId != VirtualDetectorID || pz <= 0)
       {
-        mf::LogWarning("VDResamplerConfigure") << "Thrown event\n"
-                                               << "PDG ID = " << pdgId << ", VDID = " << virtualdetectorId << ", z = " << step.position().z() << ", pz = " << pz;
+        // mf::LogWarning("VDResamplerConfigure") << "Thrown event\n"
+        //                                        << "PDG ID = " << pdgId << ", VDID = " << virtualdetectorId << ", z = " << step.position().z() << ", pz = " << pz;
+        nThrownEvents += 1;
         continue; // Filter hits based on the virtual detector ID and pz
       }
       
@@ -157,11 +159,12 @@ namespace mu2e {
   };
 
   void VDResamplerConfigure::endJob() {
-    mf::LogInfo log("Virtual detector tree summary");
-    log << "========= Data summary =========\n";
+    mf::LogInfo log("Virtual Detector Resampler Training Configuration Summary");
+    log << "========= Particle Summary =========\n";
     for (auto part : pdgIds)
       log << "PDGID " << part.first << ": " << part.second << "\n";
-    log << "================================\n";
+    log << "====================================\n";
+    log << "Number of thrown events: " << nThrownEvents << "\n";
 
     // store a summary of the number of hits for each particle type in a csv file for reference, 
     // all particles are included, but the particle types with hits below the training threshold are 
@@ -243,7 +246,7 @@ namespace mu2e {
                     // << "      SBDMgradientClip : 1.0\n"
                     // << "      SBDMlearningRate : 1e-3\n"
                     // << "      SBDMdiffusionSteps : 200\n"
-                    // << "      SBDMtrainingSize : -1\n"
+                    << "      SBDMtrainingSize : " << part.second << "\n"
                     // << "      SBDMtrainingEpochs : 10\n"
                     << "    }\n";
       }
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
index f70eb5c631..f71433387f 100644
--- a/STMMC/src/VDResamplerTrain_module.cc
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -45,10 +45,6 @@
 #include "Offline/MCDataProducts/inc/StepPointMC.hh"
 #include "Offline/SeedService/inc/SeedService.hh"
 
-// ROOT includes
-#include "art_root_io/TFileService.h"
-#include "TTree.h"
-
 typedef cet::map_vector_key key_type;
 typedef unsigned long VolumeId_type;
 
@@ -202,9 +198,14 @@ namespace mu2e {
           conf().SBDMlearningRate(),
           conf().SBDMdiffusionSteps()
       );
-
-      stage1TrainingData.reserve(1000000);
-      stage2TrainingData.reserve(1000000);
+      // allocate memory according to the training size (if specified, otherwise will grow dynamically)
+      if (trainingSize > 0) {
+        stage1TrainingData.reserve(trainingSize);
+        stage2TrainingData.reserve(trainingSize);
+      } else {
+        stage1TrainingData.reserve(1000);
+        stage2TrainingData.reserve(1000);
+      }
     } else {
       allAtOnceModel = std::make_unique<ScoreBasedDiffusionModel>(
           randFlat_,
@@ -226,8 +227,12 @@ namespace mu2e {
           conf().SBDMlearningRate(),
           conf().SBDMdiffusionSteps()
       );
-
-      allAtOnceTrainingData.reserve(1000000);
+      // allocate memory according to the training size (if specified, otherwise will grow dynamically)
+      if (trainingSize > 0) {
+        allAtOnceTrainingData.reserve(trainingSize);
+      } else {
+        allAtOnceTrainingData.reserve(1000);
+      }
     }
   };
 
@@ -249,7 +254,7 @@ namespace mu2e {
       stepPdgId = particle.pdgId();
       virtualdetectorId = step.virtualDetectorId();
       time = step.time();
-      x = step.position().x(); // Shift the x coordinate to be relative to the centre of beamline
+      x = step.position().x(); // This coordinate is in Mu2e frame, will be shifted to relative x w.r.t. the beamline
       y = step.position().y();
       z = step.position().z();
       px = step.momentum().x();

From 2c5c0a4cd87017443d81930c970e80373bf129a3 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Fri, 10 Apr 2026 18:59:49 -0400
Subject: [PATCH 14/15] Removing trailing whitespaces

---
 .../inc/ScoreBasedDiffusionModel.hh           |  24 ++--
 .../src/ScoreBasedDiffusionModel.cc           |  56 ++++-----
 .../geom/ProductionTarget_Stickman_v1_0.txt   |  10 +-
 Mu2eG4/src/constructTargetPS.cc               | 110 +++++++++---------
 STMMC/src/VDResamplerConfigure_module.cc      |  34 +++---
 STMMC/src/VDResamplerTrain_module.cc          |  16 +--
 6 files changed, 125 insertions(+), 125 deletions(-)

diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
index ee11dab097..ef87c2ecbc 100644
--- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -92,8 +92,8 @@ namespace mu2e{
         // Train the score network on a batch of samples.
         // Uses random sampling and noise injection via the external engine.
         // Note that training needs to occur on all data samples. Training on multiple small subsets
-        // of data and then averaging or aggregating the model parameters is not supported and may lead to 
-        // issues as neural networks are not linear. 
+        // of data and then averaging or aggregating the model parameters is not supported and may lead to
+        // issues as neural networks are not linear.
         //
         // Parameters:
         //   data         - Training samples (transformed state vectors)
@@ -121,12 +121,12 @@ namespace mu2e{
         // Save the model parameters to a CSV file with annotations for later use.
         // Uses a default filename of "DiffusionModel.csv" if not specified.
         //
-        // Parameters: 
+        // Parameters:
         //   filename - Path to the CSV file where model parameters will be saved (default: "DiffusionModel.csv")
         void saveModel(const std::string& filename = "DiffusionModel.csv");
 
         // Load model parameters from a file to restore a previously trained model.
-        // Note that as the adam optimizer state is not saved, it is not possible to resume training from a loaded 
+        // Note that as the adam optimizer state is not saved, it is not possible to resume training from a loaded
         // model and pick up the training process where it left off. The loaded model can only be used for sampling.
         //
         // Parameters:
@@ -146,7 +146,7 @@ namespace mu2e{
         struct Layer {
             std::vector<std::vector<double>> W; // Weight matrix [output_size][input_size]
             std::vector<double> b;              // Bias vector [output_size]
-            
+
             // storage for gradients during back-propagation
             std::vector<std::vector<double>> gradW; // gradient of loss w.r.t. weights
             std::vector<double> gradb;              // gradient of loss w.r.t. biases
@@ -162,7 +162,7 @@ namespace mu2e{
         std::vector<Layer> network_; // Network layers in forward order
 
         // Storage for activations and pre-activations during forward pass (used in back-propagation)
-        // During forward pass, preactivations_[l] = W[l] * activations_[l-1] + b[l], 
+        // During forward pass, preactivations_[l] = W[l] * activations_[l-1] + b[l],
         //                     and activations_[l] = activationFunction(preactivations_[l])
         // These are needed to compute gradients during the backward pass.
         std::vector<std::vector<double>> activations_;
@@ -193,7 +193,7 @@ namespace mu2e{
         );
 
         // Update network weights using computed gradients (Stochastic Gradient Descent SGD).
-        // Applied after backward pass. Alternately, an Adam optimizer step can be implemented 
+        // Applied after backward pass. Alternately, an Adam optimizer step can be implemented
         // in adamUpdate() for better convergence.
         //
         // Parameters:
@@ -261,12 +261,12 @@ namespace mu2e{
         // Clip gradients to prevent exploding gradients during training.
         //
         // Parameters:
-        //   maxNorm - Maximum allowed norm for the gradients. If the total norm exceeds this threshold, 
+        //   maxNorm - Maximum allowed norm for the gradients. If the total norm exceeds this threshold,
         //             gradients are scaled down to have norm equal to maxNorm.
         void clipGradients(double maxNorm);
 
         // ----- internal vars -----
-        
+
         // The random engine is NOT owned by this class. It is injected externally
         // by the framework. Below are CLHEP distribution wrappers for actual random number generation.
         // These wrap the engine_ and provide specific probability distributions.
@@ -304,7 +304,7 @@ namespace mu2e{
         // Training configuration
         int batchSize_;  // Batch size for vectorized training (default: 32)
         double gradientClipThreshold_;  // Gradient clipping threshold (default: 1.0)
-        double learningRate_; // Learning rate for training (default: 1e-3) 
+        double learningRate_; // Learning rate for training (default: 1e-3)
 
         // Diffusion process discretization
         int diffusionSteps_;  // Number of time steps to generate a sample (default: 200)
@@ -312,10 +312,10 @@ namespace mu2e{
         // Training state
         double runningLoss_;  // Accumulated loss for monitoring during training
         int adamStep_; // Step counter for Adam optimizer (used to compute bias-corrected moment estimates)
-        int trainingSampleSize_; // Total number of training samples 
+        int trainingSampleSize_; // Total number of training samples
 
         // Container for tracking training loss over epochs
         std::vector<double> epochLosses_;
-        
+
     };
 }
diff --git a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
index 74d904118b..b183744897 100644
--- a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
+++ b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
@@ -15,7 +15,7 @@ namespace mu2e {
         return s * (1.0 + x * (1.0 - s));
     }
 
-    // Linear interpolation of beta(t) between betaMin_ and betaMax_. 
+    // Linear interpolation of beta(t) between betaMin_ and betaMax_.
     // This is only used in the linear noise schedule.
     double ScoreBasedDiffusionModel::beta(double t) const {
         return betaMin_ + t * (betaMax_ - betaMin_);
@@ -37,7 +37,7 @@ namespace mu2e {
         }
     }
 
-    
+
     ScoreBasedDiffusionModel::ScoreBasedDiffusionModel(
         CLHEP::RandFlat& randFlat,
         CLHEP::RandGaussQ& randGaussQ,
@@ -121,7 +121,7 @@ namespace mu2e {
 
             for (int i = 0; i < out; ++i) {
                 for (int j = 0; j < in; ++j) {
-                    layer.W[i][j] = weightInitScale * randGaussQ_.fire(); 
+                    layer.W[i][j] = weightInitScale * randGaussQ_.fire();
                     // this generates a Gaussian random number with mean=0 and sigma=weightInitScale
                 }
                 layer.b[i] = 0.0;
@@ -136,7 +136,7 @@ namespace mu2e {
             // Output becomes next layer input
             in = out;
         }
-            
+
         // Print layer and model configuration
         std::ostringstream oss;
         oss << "ScoreBasedDiffusionModel initialized\n"
@@ -194,15 +194,15 @@ namespace mu2e {
         activations_.push_back(x);
 
         // Forward pass through the network layers
-        for (size_t l = 0; l < network_.size(); ++l) 
+        for (size_t l = 0; l < network_.size(); ++l)
         {
             // Compute pre-activation z = W*x + b for the current layer
             auto& layer = network_[l];
             std::vector<double> z(layer.W.size());
-            for (size_t i = 0; i < layer.W.size(); ++i) 
+            for (size_t i = 0; i < layer.W.size(); ++i)
             {
                 double v = layer.b[i];
-                for (size_t j = 0; j < layer.W[i].size(); ++j) 
+                for (size_t j = 0; j < layer.W[i].size(); ++j)
                 {
                     v += layer.W[i][j]*x[j];
                 }
@@ -262,7 +262,7 @@ namespace mu2e {
             }
 
             // Compute gradients w.r.t. weights and biases for the current layer
-            for (size_t i = 0; i < layer.W.size(); i++) 
+            for (size_t i = 0; i < layer.W.size(); i++)
             {
                 for (size_t j=0;j<layer.W[i].size();++j)
                 {
@@ -271,7 +271,7 @@ namespace mu2e {
                 }
 
                 // Gradient of loss w.r.t. bias b[i]
-                // Note: We accumulate gradients in gradW and gradb because we will apply the optimizer step 
+                // Note: We accumulate gradients in gradW and gradb because we will apply the optimizer step
                 // after processing a batch of samples.
                 layer.gradb[i] += gradZ[i];
             }
@@ -330,13 +330,13 @@ namespace mu2e {
     // v_t = b2 * v_{t-1} + (1 - b2) * g^2
     // mhat = m_t / (1 - b1^t)  (bias-corrected first moment estimate)
     // vhat = v_t / (1 - b2^t)  (bias-corrected second moment estimate)
-    // theta = theta - lr * mhat / (sqrt(vhat) + eps) 
+    // theta = theta - lr * mhat / (sqrt(vhat) + eps)
     void ScoreBasedDiffusionModel::adamUpdate(
         double lr
     )
     {
-        adamStep_++; // initialize to 0 in constructor, 
-                     // increment first to avoid division by zero in the first step when computing bias-corrected estimates 
+        adamStep_++; // initialize to 0 in constructor,
+                     // increment first to avoid division by zero in the first step when computing bias-corrected estimates
                      // (1 - b1^t) and (1 - b2^t). Do not modify this to increment after the update.
 
         for (auto& layer : network_)
@@ -382,7 +382,7 @@ namespace mu2e {
     std::vector<double> ScoreBasedDiffusionModel::addNoise(
         const std::vector<double>& x,
         double t, // Diffusion time parameter in [0,1]
-        std::vector<double>& eps 
+        std::vector<double>& eps
     ){
         assert(x.size() == (size_t)dim_);
 
@@ -417,7 +417,7 @@ namespace mu2e {
         return loss/dim_;
     }
 
-    // Clip gradients to prevent exploding gradients during training. This is done by scaling down 
+    // Clip gradients to prevent exploding gradients during training. This is done by scaling down
     // the gradients if their L2 norm exceeds a specified threshold.
     void ScoreBasedDiffusionModel::clipGradients(
         double maxNorm
@@ -452,10 +452,10 @@ namespace mu2e {
     }
 
     // Train the diffusion model using the provided training data. The training loop iterates over epochs and batches,
-    // applying noise to the input samples, computing the score predictions, calculating the loss, and performing 
+    // applying noise to the input samples, computing the score predictions, calculating the loss, and performing
     // back-propagation to update the model parameters.
-    // For a data sample of 5M 6-dimensional vectors of double precision (8 Byte), total memory for the data is 5M*6*8 = 240 MB, 
-    // which is manageable for in-memory training. 
+    // For a data sample of 5M 6-dimensional vectors of double precision (8 Byte), total memory for the data is 5M*6*8 = 240 MB,
+    // which is manageable for in-memory training.
     // Much larger datasets may require streaming from disk or using mini-batches that do not fit entirely in memory.
     void ScoreBasedDiffusionModel::train(
         const std::vector<DiffusionTrainingSample>& data,
@@ -496,7 +496,7 @@ namespace mu2e {
         for (size_t i = 0; i < N; ++i)
             indices[i] = i;
 
-        // Data shuffling and batching is performed at the epoch level to ensure that each epoch sees the data 
+        // Data shuffling and batching is performed at the epoch level to ensure that each epoch sees the data
         // in a different order, which can improve training convergence.
         for (int e = 0; e < epochs; ++e)
         {
@@ -507,9 +507,9 @@ namespace mu2e {
                 std::swap(indices[i], indices[j]);
             }
 
-            // Counter for number of samples processed in the epoch (used for averaging loss). Avoid using N directly to 
+            // Counter for number of samples processed in the epoch (used for averaging loss). Avoid using N directly to
             // allow for early stopping or partial epoch processing if needed.
-            int n = 0; 
+            int n = 0;
 
             int batchCounter = 0;
             double epochLoss = 0.0;
@@ -557,7 +557,7 @@ namespace mu2e {
                 // Backward pass to compute gradients of the loss w.r.t. network parameters using the computed gradient of the loss w.r.t. the predicted score.
                 backward(grad);
 
-                // Increment batch counter and apply optimizer step if batch size is reached. This allows for vectorized updates after processing a batch of 
+                // Increment batch counter and apply optimizer step if batch size is reached. This allows for vectorized updates after processing a batch of
                 // samples, which can improve training efficiency and convergence.
                 batchCounter++;
                 if(batchCounter == batchSize_)
@@ -594,7 +594,7 @@ namespace mu2e {
             epochLosses_.push_back(epochLoss);
         }
     }
-    
+
     void ScoreBasedDiffusionModel::saveModel(
         const std::string& filename
     )
@@ -635,7 +635,7 @@ namespace mu2e {
         out << "\n[NETWORK_PARAMETERS]\n";
         out << "numLayers," << network_.size() << "\n";
         out << std::fixed << std::setprecision(17); // Use fixed-point notation with high precision for weights and biases
-       
+
         for (size_t layerIdx = 0; layerIdx < network_.size(); ++layerIdx) {
             // Write layer dimensions
             auto& layer = network_[layerIdx];
@@ -789,8 +789,8 @@ namespace mu2e {
                     int inSize = std::stoi(tokens[1]);
                     int layerIdx = getLayerIdx(tokens[0]);
                     if (layerIdx == 0 && inSize != dim + conditionDim + 1) {
-                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
-                            << "Input layer input size mismatch (expected " 
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel")
+                            << "Input layer input size mismatch (expected "
                             << (dim + conditionDim + 1) << ", got " << inSize << ")";
                     }
                     if (loadedNetwork[layerIdx].W.empty()) {
@@ -814,7 +814,7 @@ namespace mu2e {
                         }
                         auto vals = split(line);
                         if (vals.size() != row.size()) {
-                            throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
+                            throw cet::exception("ScoreBasedDiffusionModel::loadModel")
                                 << "Weight row size mismatch for layer " << layerIdx;
                         }
                         for (size_t j = 0; j < vals.size(); ++j)
@@ -832,7 +832,7 @@ namespace mu2e {
                     }
                     auto vals = split(line);
                     if (vals.size() != loadedNetwork[layerIdx].b.size()) {
-                        throw cet::exception("ScoreBasedDiffusionModel::loadModel") 
+                        throw cet::exception("ScoreBasedDiffusionModel::loadModel")
                             << "Bias size mismatch for layer " << layerIdx;
                     }
                     for (size_t j = 0; j < vals.size(); ++j)
@@ -952,7 +952,7 @@ namespace mu2e {
                 input.push_back(t);
 
                 auto score = forward(input);
- 
+
                 for (int i=0;i<dim_;++i){
                     x[i] += -s*s*score[i]*dt;
                 }
diff --git a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
index 0cb5ccff1d..97b64dfbbe 100644
--- a/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
+++ b/Mu2eG4/geom/ProductionTarget_Stickman_v1_0.txt
@@ -3,7 +3,7 @@
 // Support wheel and wire geometry is inherited from the existing Hayman production target geometry,
 // others are overhauled to use the new Stickman geometry.
 
-// The Hayman dependence is intentionally removed. 
+// The Hayman dependence is intentionally removed.
 
 int PSStickman.verbosityLevel = 0; //5;
 
@@ -30,7 +30,7 @@ bool   targetPS.visible              = true;
 bool   targetPS.solid                = false;
 string targetPS_targetVacuumMaterial = "PSVacuum";
 
-// Start with target plates. Chose to specify certain parameters for each plate, so can easily play with plate-by-plate dimensions later on 
+// Start with target plates. Chose to specify certain parameters for each plate, so can easily play with plate-by-plate dimensions later on
 int targetPS_numberOfPlates = 35;
 vector<string> targetPS_plateMaterial = { // counting along mu2e +z direction, i.e. from the proton downstream side to the upstream side
     "Inconel718", "Inconel718", "Inconel718", "Inconel718", "Inconel718", //  5
@@ -76,7 +76,7 @@ vector<double> targetPS_plateLugThickness = {
     6.0, 6.0, 6.0, 6.0, 6.0  // 35
 }; // in mm, one value for each plate, note lugs are flush with the fins on the proton downstream side
 bool targetPS_addFilletToPlateCore = true; // whether to add fillet to the intersection between plate core and fins
-bool targetPS_addFilletToPlateLug  = true; // whether to add fillet to the intersection between fins and lugs 
+bool targetPS_addFilletToPlateLug  = true; // whether to add fillet to the intersection between fins and lugs
 // Fillets between fins and lugs are less important than the fillets between the core and fins, since the lugs are not in the direct path of the beam.
 // Fillets due to lug thickness larger than the plate thickness is not added, due to the complexity of the geometry and the small effect on the physics.
 double targetPS_plateFilletRadius = 1.0; // mm
@@ -93,7 +93,7 @@ double targetPS_spacerHalfLength  = 1.5;  // mm, full length 3.0 mm
 double targetPS_spacerOuterRadius = 3.0;  // mm
 double targetPS_spacerInnerRadius = 1.55; // mm
 
-// End ring / support ring 
+// End ring / support ring
 string targetPS_supportRingMaterial        = "Inconel718";
 double targetPS_supportRingLength          = 8.1;  // full length in mm
 double targetPS_supportRingInnerRadius     = 15.0; // mm
@@ -105,7 +105,7 @@ double targetPS_supportRingLugOuterRadius  = 3.0;  // mm
 bool   targetPS_addFilletToSupportRingLug  = true;
 double targetPS_supportRingLugFilletRadius = 1.0;
 bool   targetPS_addCutoutToSupportRing     = true; // the six cylindrical cutouts at an angle
-double targetPS_supportRingCutoutOffset    = 4.78; // mm, distance from the center of the cutout to the side of the end ring closer to the target center. 
+double targetPS_supportRingCutoutOffset    = 4.78; // mm, distance from the center of the cutout to the side of the end ring closer to the target center.
                                                    // This value is a must-have. If no cutout is added, this determines where the support wire is attached from the inner edge of the end ring.
 int    targetPS_nSupportRingCutouts        = 3;
 vector<double> targetPS_supportRingCutoutAngles = {15., 135., 255.}; // degrees
diff --git a/Mu2eG4/src/constructTargetPS.cc b/Mu2eG4/src/constructTargetPS.cc
index 3bdc8d3101..d1f73632b5 100644
--- a/Mu2eG4/src/constructTargetPS.cc
+++ b/Mu2eG4/src/constructTargetPS.cc
@@ -1271,7 +1271,7 @@ namespace mu2e {
       // Build the production target.
 
       GeomHandle<ProductionTarget> tgt;
-      // allows a vector of plate materials so each plate can be different. Initialize to have same number of 
+      // allows a vector of plate materials so each plate can be different. Initialize to have same number of
       // entries as the number of plates. Fin and core materials are the same for each plate.
       vector<G4Material*> prodTargetPlateMaterials(tgt->numberOfPlates(), nullptr);
       for (int i = 0; i < tgt->numberOfPlates(); ++i) {
@@ -1315,25 +1315,25 @@ namespace mu2e {
       constexpr double stickmanMagicOffset = 0.0001;
 
       // store the fin rotations for later use
-      // since used for G4UnionSolid construction, a passive rotation is needed, hence the negative sign on the angle. 
+      // since used for G4UnionSolid construction, a passive rotation is needed, hence the negative sign on the angle.
       // These are the rotations to go from the plate frame to the fin frame
       std::vector<CLHEP::HepRotation*> stickmanFinRotations;
       for (int ithFin = 0; ithFin < tgt->nStickmanFins(); ++ithFin) {
         CLHEP::HepRotation* finRotation = reg.add(CLHEP::HepRotation::IDENTITY);
-        finRotation->rotateZ(-tgt->plateFinAngle(ithFin)); 
+        finRotation->rotateZ(-tgt->plateFinAngle(ithFin));
         stickmanFinRotations.emplace_back(finRotation);
       }
 
       // cache the plate solids
       std::map<std::string, G4VSolid*> stickmanPlateSolidCache;
-      // lambda to get the plate solid for a given plate. The solid is cached based on the plate parameters that affect the shape. 
-      // The cache key encodes these parameters; if a solid is not found in the cache, it is created and added to the cache before 
+      // lambda to get the plate solid for a given plate. The solid is cached based on the plate parameters that affect the shape.
+      // The cache key encodes these parameters; if a solid is not found in the cache, it is created and added to the cache before
       //being returned.
-      // Here chose to construct plates as a single union solid with different materials for each plate, 
+      // Here chose to construct plates as a single union solid with different materials for each plate,
       // rather than separate fins, cores, and lugs, because the fins is much wider than the Hayman design.
       // If separated, missing material between fins and cores would be large.
-      // Note that the geometry center of a G4UnionSolid is the same as the first solid in the union, 
-      // and I start with the plate core and then union on the fins and lugs. This means that when placing the 
+      // Note that the geometry center of a G4UnionSolid is the same as the first solid in the union,
+      // and I start with the plate core and then union on the fins and lugs. This means that when placing the
       // plates the position will be based on the center of the core.
       auto getStickmanPlateSolid = [&](int ithPlate) -> G4VSolid* {
         const double plateFilletRadius =
@@ -1401,19 +1401,19 @@ namespace mu2e {
                  << tgt->plateLugInnerRadius() << ", " << tgt->plateLugOuterRadius() << ", " << lugHalfThickness << ")" << G4endl;
         }
 
-        // add fillets to the plate core 
+        // add fillets to the plate core
         // for a fin at 0 degree start with unioning a G4ExtrudedSolid with a triangular cross section that has vertices in
-        // the x-y plane at (0,0), 
+        // the x-y plane at (0,0),
         //                  (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
         //                  (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
-        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and 
+        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and
         // a scale of 1.0 at both ends.
-        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance) 
-        // and height plateThickness(ithPlate) that is centered at 
-        // (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius), 
+        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance)
+        // and height plateThickness(ithPlate) that is centered at
+        // (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
         // and the other fillet is completed by subtracting a similar cylinder at
         // (sqrt((rOut+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
-        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the 
+        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the
         // cylinders for subtraction are also rotated by the fin angle.
         G4VSolid* plateCoreFillet = nullptr;
         if (tgt->addFilletToPlateCore()) {
@@ -1441,7 +1441,7 @@ namespace mu2e {
                    << coreFilletVertices.at(2) << ")" << G4endl;
             G4cout << __PRETTY_FUNCTION__ << "  core fillet extrusion half length = " << plateHalfThickness << G4endl;
           }
-          
+
           // create the cutout cylinder for the fillet
           G4Tubs* coreFilletCutout = reg.add(new G4Tubs(
             cacheKey + "_CoreFilletCutout",             //
@@ -1456,10 +1456,10 @@ namespace mu2e {
             coreFilletExtrusion,                        // base solid
             coreFilletCutout,                           // cutout solid
             nullptr,                                    // rotation
-            CLHEP::Hep3Vector(coreFilletX,              
+            CLHEP::Hep3Vector(coreFilletX,
                               finHalfWidth + plateFilletRadius,
                               0.)));                    // translation of cutout solid
-          plateCoreFillet = reg.add(new G4SubtractionSolid( 
+          plateCoreFillet = reg.add(new G4SubtractionSolid(
             cacheKey + "_CoreFilletLower",              //
             coreFilletUpper,                            // base solid
             coreFilletCutout,                           // cutout solid
@@ -1478,16 +1478,16 @@ namespace mu2e {
         // for a fin at 0 degree, start with unioning a G4ExtrudedSolid with a triangular cross section that has vertices in
         // the x-y plane at (plateCenterToLugCenter, 0),
         //                  (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
-        //                  (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius), 
-        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and 
+        //                  (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
+        // and then extrude along z by plateThickness(ithPlate) starting at z = _currentZ. The extrusion has no twist and
         // a scale of 1.0 at both ends.
 
-        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance) 
-        // and height plateThickness(ithPlate) that is centered at 
+        // The fillet is then completed by subtracting a cylinder with radius of the fillet radius (plus a small tolerance)
+        // and height plateThickness(ithPlate) that is centered at
         // (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), plateFinWidth/2+FilletRadius),
         // and the other fillet is completed by subtracting a similar cylinder at
-        // (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius), 
-        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the 
+        // (plateCenterToLugCenter-sqrt((plateLugOuterRadius+FilletRadius)^2-(plateFinWidth/2+FilletRadius)^2), -plateFinWidth/2-FilletRadius),
+        // For fins at other angles, the same is done but the vertices are rotated by the fin angle and the centers of the
         // cylinders for subtraction are also rotated by the fin angle.
         G4VSolid* plateLugFillet = nullptr;
         if (tgt->addFilletToPlateLug()) {
@@ -1525,7 +1525,7 @@ namespace mu2e {
             plateHalfThickness + stickmanMagicOffset, // half length in z
             0.,                                       // starting angle
             CLHEP::twopi));                           // segment angle is full circle
-          // lug center needs to be cut out too. Create a second cutout cylinder for the lug center with radius of 
+          // lug center needs to be cut out too. Create a second cutout cylinder for the lug center with radius of
           // plateLugInnerRadius and height of plateThickness(ithPlate) centered at (plateCenterToLugCenter, 0)
           G4Tubs* lugCenterCutout = reg.add(new G4Tubs(
             cacheKey + "_LugCenterCutout",                    //
@@ -1576,9 +1576,9 @@ namespace mu2e {
                                            0.);
           const CLHEP::Hep3Vector lugShift(tgt->plateCenterToLugCenter() * std::cos(currentFinAngle),
                                            tgt->plateCenterToLugCenter() * std::sin(currentFinAngle),
-                                           lugZShift); 
+                                           lugZShift);
           // note the lug is shifted in z to be flush with the edge of the plate core
-          
+
           // add fin
           plateSolid = reg.add(new G4UnionSolid(cacheKey + "_FinUnion_" + std::to_string(ithFin),
                                                 plateSolid,                         // base solid
@@ -1597,7 +1597,7 @@ namespace mu2e {
                                                   stickmanFinRotations.at(ithFin),  // rotation of union solid
                                                   CLHEP::Hep3Vector(0.,0.,0.)));    // translation of union solid
           }
-          
+
           // add lug
           plateSolid = reg.add(new G4UnionSolid(cacheKey + "_LugUnion_" + std::to_string(ithFin),
                                                 plateSolid,                       // base solid
@@ -1624,17 +1624,17 @@ namespace mu2e {
         }
         return plateSolid;
       }; // end of lambda to get plate solid with caching
-      
+
       //
       // The plates are constructed starting from the most negative z end and moving towards the positive z end.
       int numberOfPlates = tgt->numberOfPlates();
       // running z position for plate construction. Start at most negative end and move towards +z
       double _currentZ = _localCenter.z() - tgt->halfStickmanLength() // end of most negative end
-                         + tgt->stickmanSupportRingLength() // move past support ring at end 
-                         + 2 * tgt->spacerHalfLength() ; // move past spacer. 
+                         + tgt->stickmanSupportRingLength() // move past support ring at end
+                         + 2 * tgt->spacerHalfLength() ; // move past spacer.
                          // plate fin, core, and lug is flush on this side
       if (verbosityLevel > 2){G4cout << __PRETTY_FUNCTION__ << " plate Z starts at " << _currentZ << G4endl;}
-      
+
       // loop over plates to construct and add them to the target mother.
       for (int ithPlate = 0; ithPlate < numberOfPlates; ++ithPlate) {
         std::string plateName = "ProductionTargetPlate_" + std::to_string(ithPlate);
@@ -1642,7 +1642,7 @@ namespace mu2e {
         const double plateCenterZ = _currentZ + plateHalfThickness; // this is the z position of the core
         const CLHEP::Hep3Vector localPlateCenter(0.,0.,plateCenterZ); // this is in the target frame
         const CLHEP::Hep3Vector currentPlateCenter = tgt->productionTargetRotation().inverse()*localPlateCenter;
-        // productionTargetRotation() rotates to the target frame, so to get the position in the world frame we 
+        // productionTargetRotation() rotates to the target frame, so to get the position in the world frame we
         // apply the inverse rotation to the local plate center.
 
         if (verbosityLevel > 0) {
@@ -1663,12 +1663,12 @@ namespace mu2e {
                  << " lugThickness=" << tgt->plateLugThickness(ithPlate)
                  << " lug(inner,outer)=("
                  << tgt->plateLugInnerRadius() << ","
-                 << tgt->plateLugOuterRadius() << ")" 
+                 << tgt->plateLugOuterRadius() << ")"
                  << " addFilletToPlateCore=" << tgt->addFilletToPlateCore()
                  << " addFilletToPlateLug=" << tgt->addFilletToPlateLug()
                  << G4endl;
         }
-        
+
         // get the solid for the plate, using the lambda defined above
         G4VSolid* plateSolid = getStickmanPlateSolid(ithPlate);
 
@@ -1790,20 +1790,20 @@ namespace mu2e {
                  << G4endl;
         }
       } // end of loop over rods and spacers
-      
+
       // now add the two end rings
 
-      // then two cutouts are made by subtracting cylinders with radius of the fillet radius (plus a small tolerance) and height of stickmanSupportRingLength (plus a small tolerance) 
+      // then two cutouts are made by subtracting cylinders with radius of the fillet radius (plus a small tolerance) and height of stickmanSupportRingLength (plus a small tolerance)
       // centered at (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), supportRingLugOuterRadius+filletRadius), and
       //             (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), -supportRingLugOuterRadius-filletRadius)
       // an additional cutout is made for the inner part of the ring by subtracting a tube with inner radius of 0, outer radius of supportRingInnerRadius, and height of stickmanSupportRingLength (plus a small tolerance) centered at (0,0).
-      // if addCutoutToSupportRing is true, then additional circular cutouts (through holes) are made. 
-      // The number of cutouts is given by nSupportRingCutouts at angles defined by supportRingCutoutAngles, and the cutouts are made by subtracting cylinders with radius of 
-      // supportRingCutoutInnerRadius. The centers of the cutouts are at supportRingCutoutOffset from the interface between the spacers and the end ring, and at the angles defined 
-      // by supportRingCutoutAngles. The angle is between the through hole and the local radius, the through hole points outward and towards the middle.  
+      // if addCutoutToSupportRing is true, then additional circular cutouts (through holes) are made.
+      // The number of cutouts is given by nSupportRingCutouts at angles defined by supportRingCutoutAngles, and the cutouts are made by subtracting cylinders with radius of
+      // supportRingCutoutInnerRadius. The centers of the cutouts are at supportRingCutoutOffset from the interface between the spacers and the end ring, and at the angles defined
+      // by supportRingCutoutAngles. The angle is between the through hole and the local radius, the through hole points outward and towards the middle.
 
       // lambda to build the end support ring solid with lugs and fillets if applicable
-      // Note the end plate has a flat face on the inside that is flush with the end of the rods / spacers, details like chamfers at 
+      // Note the end plate has a flat face on the inside that is flush with the end of the rods / spacers, details like chamfers at
       // the rod sockets and those at the cutouts (through holes) are ignored for simplicity.
       auto buildStickmanSupportRingSolid = [&](const std::string& tag, bool positiveEnd) -> G4VSolid* {
         const double ringHalfLength = tgt->stickmanSupportRingLength()/2.;
@@ -1831,7 +1831,7 @@ namespace mu2e {
           CLHEP::twopi));                           // delta angle
 
         // cutout for the center of the ring
-        // such construction avoids potential issue that the lug solid extends within the 
+        // such construction avoids potential issue that the lug solid extends within the
         // inner radius of the ring, which can cause misrepresentation of the geometry
         G4Tubs* ringCenterCutout = reg.add(new G4Tubs(
           "ProductionTargetSupportRingCenterCutout_" + tag, //
@@ -1839,20 +1839,20 @@ namespace mu2e {
           tgt->stickmanSupportRingInnerRadius(),            // outer radius (avoid overlap with lug stem)
           ringHalfLength + stickmanMagicOffset,             // half length (add small offset to avoid precision issues)
           0.,                                               // start angle
-          CLHEP::twopi));                                   // delta angle 
+          CLHEP::twopi));                                   // delta angle
 
         // lugs at three fin angles.
         G4Tubs* lugSolid = reg.add(new G4Tubs(
           "ProductionTargetSupportRingLug_" + tag, //
-          0.,                                      // inner radius 
+          0.,                                      // inner radius
           tgt->supportRingLugOuterRadius(),        // outer radius
           ringHalfLength,                          // half length
           0.,                                      // start angle
           CLHEP::twopi));                          // delta angle
-        
+
         // lug stem
         G4Box* lugStemSolid = reg.add(new G4Box(
-          "ProductionTargetSupportRingLugStem_" + tag, // 
+          "ProductionTargetSupportRingLugStem_" + tag, //
           lugStemHalfLength,                           // half length in x
           tgt->supportRingLugOuterRadius(),            // half length in y
           ringHalfLength));                            // half length in z
@@ -1890,10 +1890,10 @@ namespace mu2e {
             lugSolid,                        // union solid
             nullptr,                         // passive rotation
             lugShift));                      // translation
-          
+
           // add fillet to the lug if applicable
-          // The fillet is added by unioning with an extrusion 
-          // of a triangle with vertices at (0,0), 
+          // The fillet is added by unioning with an extrusion
+          // of a triangle with vertices at (0,0),
           //                                (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), supportRingLugOuterRadius+filletRadius),
           //                                (sqrt((supportRingOuterRadius+filletRadius)^2-(supportRingLugOuterRadius+filletRadius)^2), -supportRingLugOuterRadius-filletRadius),
           // extruded along z by stickmanSupportRingLength starting at the end of the target, with no twist and scale of 1.0 at both ends.
@@ -1950,7 +1950,7 @@ namespace mu2e {
               CLHEP::Hep3Vector(filletX,
                                 -tgt->supportRingLugOuterRadius() - filletRadius,
                                 0.)));          // translation of cutout solid
-            
+
             // add filletSolid to the union with the support ring
             supportRingSolid = reg.add(new G4UnionSolid(
               "ProductionTargetSupportRingFilletUnion_" + tag + "_" + std::to_string(ithFin),
@@ -1968,7 +1968,7 @@ namespace mu2e {
           ringCenterCutout,                 // cutout solid
           nullptr,                          // rotation
           CLHEP::Hep3Vector(0., 0., 0.)));  // translation
-        
+
         // add cutouts for the through holes if applicable
         if (tgt->addCutoutToSupportRing()) {
           const double cutoutTiltAngle = tgt->supportRingCutoutTilt();
@@ -1976,7 +1976,7 @@ namespace mu2e {
           // cutout center is placed at halfway between the inner and outer radius
           const double cutoutCenterRadius = 0.5*(tgt->stickmanSupportRingInnerRadius() + tgt->stickmanSupportRingOuterRadius());
           // this is z position in the local ring frame
-          const double localCutoutZ = (positiveEnd ? 1. : -1.) * 
+          const double localCutoutZ = (positiveEnd ? 1. : -1.) *
                                       (tgt->supportRingCutoutOffset() - ringHalfLength + supportRingWallThickness * 0.5 * std::tan(cutoutTiltAngle) );
           // cutout needs to go through the whole wall thickness
           const double cutoutHalfLength = 0.5 * supportRingWallThickness / std::cos(cutoutTiltAngle)
@@ -2087,8 +2087,8 @@ namespace mu2e {
                << G4endl;
       }
 
-      //Copied code from hayman_v_2_0 to add support structures. Stickman and Hayman supports are the same. Avoiding issue with variables like 
-      //prodTargetMotherInfo crossing scopes this way. They share names but are initialized differently. 
+      //Copied code from hayman_v_2_0 to add support structures. Stickman and Hayman supports are the same. Avoiding issue with variables like
+      //prodTargetMotherInfo crossing scopes this way. They share names but are initialized differently.
       //In addition, this allows changing the support structure for one target without affecting the other.
 
       //Add support structures for the production target
diff --git a/STMMC/src/VDResamplerConfigure_module.cc b/STMMC/src/VDResamplerConfigure_module.cc
index 9504fb9d9e..82db041969 100644
--- a/STMMC/src/VDResamplerConfigure_module.cc
+++ b/STMMC/src/VDResamplerConfigure_module.cc
@@ -1,14 +1,14 @@
 // VDResamplerConfigure_module.cc
-// Makes one pass of the data set and determines which particles will need training for the 
-// resampler, and generates a fcl file for training. "dataSourceTag" specifies the tag of the 
-// data source, and will be appended to the generated fcl file and then csv files to 
-// distinguish between different data sources. The generated fcl files will be stored and 
+// Makes one pass of the data set and determines which particles will need training for the
+// resampler, and generates a fcl file for training. "dataSourceTag" specifies the tag of the
+// data source, and will be appended to the generated fcl file and then csv files to
+// distinguish between different data sources. The generated fcl files will be stored and
 // named "VDResamplerTrain_[dataSourceTag].fcl".
-// On the other hand, a breakdown of the number of hits for each particle type will be stored 
+// On the other hand, a breakdown of the number of hits for each particle type will be stored
 // in "VDResamplerConfigure_[dataSourceTag]_HitSummary.csv" for reference.
-// Later when generating new samples using the resampler, the program will look for the 
-// generated fcl files to determine which particle source to use, which particle types to 
-// generate and which model parameters to load for each particle type. 
+// Later when generating new samples using the resampler, the program will look for the
+// generated fcl files to determine which particle source to use, which particle types to
+// generate and which model parameters to load for each particle type.
 // Yongyi Wu, Mar. 2026
 
 // stdlib includes
@@ -54,7 +54,7 @@ namespace mu2e {
       struct Config {
         fhicl::Atom<art::InputTag> StepPointMCsTag{Name("StepPointMCsTag"), Comment("Tag identifying the StepPointMCs")};
         fhicl::Atom<art::InputTag> SimParticlemvTag{Name("SimParticlemvTag"), Comment("Tag identifying the SimParticlemv")};
-        fhicl::Atom<int> VirtualDetectorID{Name("VirtualDetectorID"), Comment("ID of the virtual detector to train on"), 116}; 
+        fhicl::Atom<int> VirtualDetectorID{Name("VirtualDetectorID"), Comment("ID of the virtual detector to train on"), 116};
         fhicl::Atom<std::string> VDResamplerDir{Name("VDResamplerDir"), Comment("Directory to store the generated csv files")};
         fhicl::Atom<std::string> fclDir{Name("fclDir"), Comment("Directory to store the generated fhicl files"), ""};
         fhicl::Atom<std::string> dataSourceTag{Name("dataSourceTag"), Comment("A tag to distinguish different data sources, will be appended to the generated fcl and csv files")};
@@ -133,7 +133,7 @@ namespace mu2e {
         nThrownEvents += 1;
         continue; // Filter hits based on the virtual detector ID and pz
       }
-      
+
       if (doROOTDump) {
           x = step.position().x();
           y = step.position().y();
@@ -145,7 +145,7 @@ namespace mu2e {
           E = std::sqrt(step.momentum().mag2()+mass*mass)-mass; // Subtract the rest mass
           if (E < 0)
             throw cet::exception("LogicError", "Energy is negative");
-          
+
           ttree->Fill();
       }
 
@@ -166,8 +166,8 @@ namespace mu2e {
     log << "====================================\n";
     log << "Number of thrown events: " << nThrownEvents << "\n";
 
-    // store a summary of the number of hits for each particle type in a csv file for reference, 
-    // all particles are included, but the particle types with hits below the training threshold are 
+    // store a summary of the number of hits for each particle type in a csv file for reference,
+    // all particles are included, but the particle types with hits below the training threshold are
     // marked with an asterisk to indicate that they will not be included in the training.
     std::string summaryFile = VDResamplerDir + "/VDResamplerConfigure_" + dataSourceTag + "_HitSummary.csv";
     std::string fclFilePath = fclDir.empty() ? VDResamplerDir : fclDir;
@@ -197,15 +197,15 @@ namespace mu2e {
 
     for (const auto& part : pdgIds) {
       sumOutFile << part.first << "," << part.second;
-      bool useTwoStageTraining = false; 
-      if (part.second < 100000) { 
+      bool useTwoStageTraining = false;
+      if (part.second < 100000) {
         // if data is limited, use two-stage training to improve performance and reduce the required training data size for each model.
         useTwoStageTraining = true;
       }
       if (part.second < trainingThreshold) {
         sumOutFile << ",*\n";
-        mf::LogWarning("VDResamplerConfigure::endJob") << "Particle type with PDGID " << part.first 
-            << " has only " << part.second << " hits, which is below the training threshold of " 
+        mf::LogWarning("VDResamplerConfigure::endJob") << "Particle type with PDGID " << part.first
+            << " has only " << part.second << " hits, which is below the training threshold of "
             << trainingThreshold << ". This particle type will NOT be included in the training.";
       } else {
         // mark how many models will be trained for this particle type (1 for all-at-once, 2 for two-stage)
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
index f71433387f..9def2ecce0 100644
--- a/STMMC/src/VDResamplerTrain_module.cc
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -7,7 +7,7 @@
 //      (t', x', y', p_r', p_phi', p_z')
 // and store the trained model parameters in CSV files.
 // note that p_z are filtered and only hits with positive p_z are kept
-// Yongyi Wu, Mar. 2026 
+// Yongyi Wu, Mar. 2026
 
 // stdlib includes
 #include <cmath>
@@ -91,7 +91,7 @@ namespace mu2e {
       CLHEP::RandGaussQ randGaussQ_;
       art::ProductToken<StepPointMCCollection> StepPointMCsToken;
       art::ProductToken<SimParticleCollection> SimParticlemvToken;
-      
+
       // SBDM model + data
       bool useTwoStageTraining = true;
       std::unique_ptr<ScoreBasedDiffusionModel> allAtOnceModel;
@@ -113,7 +113,7 @@ namespace mu2e {
       // variables to read from the art event
       int stepPdgId = 0;
       double x = 0.0, y = 0.0, z = 0.0, time = 0.0;
-      double px = 0.0, py = 0.0, pz = 0.0; 
+      double px = 0.0, py = 0.0, pz = 0.0;
       VolumeId_type virtualdetectorId = 0;
 
       // transform variables for training data preparation
@@ -263,8 +263,8 @@ namespace mu2e {
 
       if (virtualdetectorId != VirtualDetectorID || (stepPdgId != pdgID && pdgID != 0) || pz <= 0)
         continue; // Filter hits based on the virtual detector ID, particle type, and pz
-      
-      // as z maybe slightly different from the nominal VDz0 due to the step size, we will extrapolate the (x, y) coordinates 
+
+      // as z maybe slightly different from the nominal VDz0 due to the step size, we will extrapolate the (x, y) coordinates
       // to the nominal VDz0 for all hits to compute the training parameters to be fed into the SBDM.
       double extrapolationFactor = (VDz0 - z) / pz; // Assuming linear motion, this is the factor to extrapolate from current z to z0
       double x_extrapolated = x + extrapolationFactor * px;
@@ -335,7 +335,7 @@ namespace mu2e {
   };
 
   void VDResamplerTrain::endJob() {
-    
+
       if (useTwoStageTraining && (stage1TrainingData.empty() || stage2TrainingData.empty())) {
         mf::LogWarning("VDResamplerTrain") << "No training data collected.";
         return;
@@ -345,8 +345,8 @@ namespace mu2e {
         return;
       }
 
-      // if SBDMtrainingSize is set and smaller than the collected training data, 
-      // truncate the training data to the specified size. 
+      // if SBDMtrainingSize is set and smaller than the collected training data,
+      // truncate the training data to the specified size.
       if (useTwoStageTraining) {
         if(trainingSize > 0 && (int)stage1TrainingData.size() > trainingSize)
           stage1TrainingData.resize(trainingSize);

From d37323615d80265c5543f6fb9069b2ba39acf1f1 Mon Sep 17 00:00:00 2001
From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com>
Date: Sat, 11 Apr 2026 17:25:57 -0400
Subject: [PATCH 15/15] Fixes addressing PR#1794 comments

See comments in PR#1794
---
 GeometryService/src/ProductionTargetMaker.cc  | 23 ++---
 MCDataProducts/inc/GenId.hh                   |  2 +-
 .../inc/ScoreBasedDiffusionModel.hh           |  6 +-
 .../src/ScoreBasedDiffusionModel.cc           | 87 +++++++++++++++----
 Mu2eG4/src/constructTargetPS.cc               |  2 +-
 ProductionTargetGeom/inc/ProductionTarget.hh  |  1 -
 ProductionTargetGeom/src/ProductionTarget.cc  |  1 +
 STMMC/inc/VDResamplerTransformDefaults.hh     | 13 +++
 .../VDResamplerGenerateFromModel_module.cc    | 11 +--
 STMMC/src/VDResamplerGenerateMix_module.cc    | 11 +--
 STMMC/src/VDResamplerTrain_module.cc          | 11 +--
 11 files changed, 114 insertions(+), 54 deletions(-)
 create mode 100644 STMMC/inc/VDResamplerTransformDefaults.hh

diff --git a/GeometryService/src/ProductionTargetMaker.cc b/GeometryService/src/ProductionTargetMaker.cc
index 8eacd51f11..eb2edcbb4b 100644
--- a/GeometryService/src/ProductionTargetMaker.cc
+++ b/GeometryService/src/ProductionTargetMaker.cc
@@ -9,22 +9,17 @@
 namespace mu2e {
 
   std::unique_ptr<ProductionTarget> ProductionTargetMaker::make(const SimpleConfig& c, double solenoidOffset) {
-
-
     if (c.getString("targetPS_model") == "MDC2018"){
-      //     std::cout << "making Tier1 in maker" << std::endl;
+      //     std::cout << " making Tier1 in maker" << std::endl;
       return makeTier1(c, solenoidOffset);
-        } else
-      if (c.getString("targetPS_model") == "Hayman_v_2_0"){
-        //        std::cout << " making Hayman in Maker" << std::endl;
-        return makeHayman_v_2_0(c, solenoidOffset);
-          } else
-      if (c.getString("targetPS_model") == "Stickman_v_1_0"){
-        return makeStickman_v_1_0(c, solenoidOffset);
-          } else
-        {throw cet::exception("GEOM") << " illegal production target version specified = " << c.getInt("targetPS_version")  << std::endl;}
-    return 0;
-
+    } else if (c.getString("targetPS_model") == "Hayman_v_2_0"){
+      //     std::cout << " making Hayman in Maker" << std::endl;
+      return makeHayman_v_2_0(c, solenoidOffset);
+    } else if (c.getString("targetPS_model") == "Stickman_v_1_0"){
+      return makeStickman_v_1_0(c, solenoidOffset);
+    } else{
+      throw cet::exception("GEOM") << " illegal production target model specified = " << c.getString("targetPS_model")  << std::endl;
+    }
   }
 
   std::unique_ptr<ProductionTarget> ProductionTargetMaker::makeTier1(const SimpleConfig& c, double solenoidOffset){
diff --git a/MCDataProducts/inc/GenId.hh b/MCDataProducts/inc/GenId.hh
index 8ede29be0a..d7feb7f715 100644
--- a/MCDataProducts/inc/GenId.hh
+++ b/MCDataProducts/inc/GenId.hh
@@ -45,7 +45,7 @@ namespace mu2e {
       MuCapProtonGenTool, MuCapDeuteronGenTool, DIOGenTool, MuCapNeutronGenTool, // 48
       MuCapPhotonGenTool, MuCapGammaRayGenTool, CeLeadingLogGenTool, MuplusMichelGenTool,// 52
       gammaPairProduction, antiproton, Mu2eXGenTool, STMDownStreamGenTool,//56
-      lastEnum //56
+      lastEnum //57
     };
 
 #ifndef SWIG
diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
index ef87c2ecbc..2bf4c7c697 100644
--- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
+++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh
@@ -272,9 +272,9 @@ namespace mu2e{
         // These wrap the engine_ and provide specific probability distributions.
         // - RandFlat:   Uniform distribution on [0,1)
         // - RandGaussQ: Gaussian (normal) distribution with mean=0, sigma=1 (or custom)
-        // Both are initialized in the constructor with the injected engine_.
-        CLHEP::RandFlat   randFlat_;      // Used for uniform sampling (e.g., batch selection)
-        CLHEP::RandGaussQ randGaussQ_;    // Used for Gaussian noise in diffusion process
+        // Both are bound in the constructor to externally managed wrappers.
+        CLHEP::RandFlat& randFlat_;       // Used for uniform sampling (e.g., batch selection)
+        CLHEP::RandGaussQ& randGaussQ_;   // Used for Gaussian noise in diffusion process
 
         // Model hyperparameters
         int dim_;           // Dimensionality of state space
diff --git a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
index b183744897..5cf8b2dc34 100644
--- a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
+++ b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc
@@ -71,6 +71,9 @@ namespace mu2e {
         if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) {
             throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid model dimensions";
         }
+        if (batchSize <= 0) {
+            throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid batchSize";
+        }
         if (diffusionSteps <= 0) {
             throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid diffusionSteps";
         }
@@ -86,8 +89,10 @@ namespace mu2e {
         int in = inputSize;
 
         // Weight initialization scale (local constant so it can be tuned easily)
-        const double weightInitScale = 0.02;
-        // const double weightInitScale = std::sqrt(2.0 / in); // He initialization for ReLU activations
+        // const double weightInitScale = 0.02; // not scalable with size, can lead to instability for larger models and too slow training for smaller models
+        // const double weightInitScale = std::sqrt(2.0 / in); // He initialization for ReLU activations, SiLU is similar to ReLU in terms of variance preservation
+        double reducedHe = 0.5 * std::sqrt(2.0 / in); // Scaled He initialization found to improve stability
+        const double weightInitScale = std::min(reducedHe, 0.3); // Cap the weight initialization scale to prevent instability for very small input sizes
 
         for (int l = 0; l < layers_; ++l) {
 
@@ -127,9 +132,12 @@ namespace mu2e {
                 layer.b[i] = 0.0;
             }
 
-            assert (layer.W[0].size() == (size_t)in); // Check that weight matrix has correct input dimension
-            assert (layer.W.size() == (size_t)out);   // Check that weight matrix has correct output dimension
-            assert (layer.b.size() == (size_t)out);  // Check that bias vector has correct dimension
+            if (layer.W.empty() || layer.W[0].size() != static_cast<size_t>(in) || //Check that weight matrix has correct input dimension
+                layer.W.size() != static_cast<size_t>(out) ||                      //Check that weight matrix has correct output dimension
+                layer.b.size() != static_cast<size_t>(out)) {                      //Check that bias vector has correct dimension
+                throw cet::exception("ScoreBasedDiffusionModel::initialization")
+                    << "Layer shape initialization mismatch";
+            }
 
             network_.push_back(layer);
 
@@ -179,10 +187,15 @@ namespace mu2e {
         const std::vector<double>& input
     )
     {
-        assert(!network_.empty());
-        assert(!network_[0].W.empty());
-        assert(!network_[0].W[0].empty());
-        assert(input.size() == network_[0].W[0].size());
+        if (network_.empty() || network_[0].W.empty() || network_[0].W[0].empty()) {
+            throw cet::exception("ScoreBasedDiffusionModel::forward")
+                << "Network is not properly initialized";
+        }
+        if (input.size() != network_[0].W[0].size()) {
+            throw cet::exception("ScoreBasedDiffusionModel::forward")
+                << "Input dimension mismatch: got " << input.size()
+                << ", expected " << network_[0].W[0].size();
+        }
 
         activations_.clear();
         preactivations_.clear();
@@ -276,7 +289,10 @@ namespace mu2e {
                 layer.gradb[i] += gradZ[i];
             }
 
-            assert(!layer.W.empty() && !layer.W[0].empty());
+            if (layer.W.empty() || layer.W[0].empty()) {
+                throw cet::exception("ScoreBasedDiffusionModel::backward")
+                    << "Encountered empty layer weights during backward pass";
+            }
             // Compute gradient w.r.t. input of the current layer to propagate to the previous layer
             std::vector<double> gradPrev(layer.W[0].size(),0.0);
 
@@ -295,8 +311,10 @@ namespace mu2e {
     {
         for (auto& layer : network_)
         {
-            assert(!layer.W.empty() && !layer.W[0].empty());
-            assert(!layer.b.empty());
+            if (layer.W.empty() || layer.W[0].empty() || layer.b.empty()) {
+                throw cet::exception("ScoreBasedDiffusionModel::updateWeights")
+                    << "Encountered malformed layer during SGD update";
+            }
 
             size_t outSize = layer.W.size();
             size_t inSize  = layer.W[0].size();
@@ -341,8 +359,10 @@ namespace mu2e {
 
         for (auto& layer : network_)
         {
-            assert(!layer.W.empty() && !layer.W[0].empty());
-            assert(!layer.b.empty());
+            if (layer.W.empty() || layer.W[0].empty() || layer.b.empty()) {
+                throw cet::exception("ScoreBasedDiffusionModel::adamUpdate")
+                    << "Encountered malformed layer during Adam update";
+            }
 
             for (size_t i=0;i<layer.W.size();++i)
             {
@@ -384,7 +404,11 @@ namespace mu2e {
         double t, // Diffusion time parameter in [0,1]
         std::vector<double>& eps
     ){
-        assert(x.size() == (size_t)dim_);
+        if (x.size() != static_cast<size_t>(dim_)) {
+            throw cet::exception("ScoreBasedDiffusionModel::addNoise")
+                << "Input x dimension mismatch: got " << x.size()
+                << ", expected " << dim_;
+        }
 
         double s = sigma(t);
 
@@ -463,9 +487,10 @@ namespace mu2e {
     )
     {
         // Check that the network is properly initialized before training.
-        assert(!network_.empty());
-        assert(!network_[0].W.empty());
-        assert(!network_[0].W[0].empty());
+        if (network_.empty() || network_[0].W.empty() || network_[0].W[0].empty()) {
+            throw cet::exception("ScoreBasedDiffusionModel::train")
+                << "Network is not properly initialized";
+        }
 
         const double eps_safe = 1e-12;
         const size_t N = data.size();
@@ -538,7 +563,11 @@ namespace mu2e {
                 input.push_back(t);
 
                 // Check that input dimension matches expected dimension (dim_ + conditionDim_ + 1)
-                assert(input.size() == network_[0].W[0].size());
+                if (input.size() != network_[0].W[0].size()) {
+                    throw cet::exception("ScoreBasedDiffusionModel::train")
+                        << "Training input dimension mismatch: got " << input.size()
+                        << ", expected " << network_[0].W[0].size();
+                }
 
                 // Forward pass to compute the predicted score from the network given the input (noisy sample + time).
                 auto score = forward(input);
@@ -562,6 +591,16 @@ namespace mu2e {
                 batchCounter++;
                 if(batchCounter == batchSize_)
                 {
+                    // Average accumulated gradients so optimizer step size is independent of batch size.
+                    const double invBatch = 1.0 / static_cast<double>(batchCounter);
+                    for (auto& layer : network_) {
+                        for (auto& row : layer.gradW)
+                            for (auto& g : row)
+                                g *= invBatch;
+                        for (auto& g : layer.gradb)
+                            g *= invBatch;
+                    }
+
                     clipGradients(gradientClipThreshold_); // Clip gradients to prevent exploding gradients
 
                     // Apply optimizer based on configuration
@@ -580,6 +619,16 @@ namespace mu2e {
             // Final flush: apply optimizer to remaining gradients
             if(batchCounter > 0)
             {
+                // Average by the true final mini-batch size (which may be < batchSize_).
+                const double invBatch = 1.0 / static_cast<double>(batchCounter);
+                for (auto& layer : network_) {
+                    for (auto& row : layer.gradW)
+                        for (auto& g : row)
+                            g *= invBatch;
+                    for (auto& g : layer.gradb)
+                        g *= invBatch;
+                }
+
                 clipGradients(gradientClipThreshold_); // Clip gradients before final update
 
                 if (optimizerType_ == OptimizerType::ADAM) {
diff --git a/Mu2eG4/src/constructTargetPS.cc b/Mu2eG4/src/constructTargetPS.cc
index d1f73632b5..fcfb8d3b90 100644
--- a/Mu2eG4/src/constructTargetPS.cc
+++ b/Mu2eG4/src/constructTargetPS.cc
@@ -972,7 +972,7 @@ namespace mu2e {
             //
             // ok here I have a little confusing fix.  The fin is built along the y-axis.  But the rotation is given wrt the x axis.  Hence I need to
             // subtract off that 90^o
-            double rotAdj = -M_PI/2 + currentFinAngle;
+            double rotAdj = -M_PI/2. + currentFinAngle;
             rotFin->rotateZ(rotAdj);
             //
             // g4 version.
diff --git a/ProductionTargetGeom/inc/ProductionTarget.hh b/ProductionTargetGeom/inc/ProductionTarget.hh
index f7049cacf7..71efd17748 100644
--- a/ProductionTargetGeom/inc/ProductionTarget.hh
+++ b/ProductionTargetGeom/inc/ProductionTarget.hh
@@ -7,7 +7,6 @@
 #define PRODUCTIONTARGET_HH
 
 #include <map>
-#include <numeric>
 
 #include "canvas/Persistency/Common/Wrapper.h"
 
diff --git a/ProductionTargetGeom/src/ProductionTarget.cc b/ProductionTargetGeom/src/ProductionTarget.cc
index 8207b12527..8d59726bbd 100644
--- a/ProductionTargetGeom/src/ProductionTarget.cc
+++ b/ProductionTargetGeom/src/ProductionTarget.cc
@@ -1,4 +1,5 @@
 #include "Offline/ProductionTargetGeom/inc/ProductionTarget.hh"
+#include <numeric>
 
 namespace mu2e {
   ProductionTarget::ProductionTarget(std::string tier1TargetType, int version, double rOut,
diff --git a/STMMC/inc/VDResamplerTransformDefaults.hh b/STMMC/inc/VDResamplerTransformDefaults.hh
new file mode 100644
index 0000000000..ada0d1d7f1
--- /dev/null
+++ b/STMMC/inc/VDResamplerTransformDefaults.hh
@@ -0,0 +1,13 @@
+#pragma once
+
+namespace mu2e {
+  namespace vdresampler {
+
+    constexpr double kX0 = -3904.0;
+    constexpr double kY0 = 0.0;
+    constexpr double kT0 = 1700.0;    // ns
+    constexpr double kTScale = 1.0;
+    constexpr double kP0 = 1.0;
+
+  }  // namespace vdresampler
+}  // namespace mu2e
diff --git a/STMMC/src/VDResamplerGenerateFromModel_module.cc b/STMMC/src/VDResamplerGenerateFromModel_module.cc
index 7102043022..9f72ccbea0 100644
--- a/STMMC/src/VDResamplerGenerateFromModel_module.cc
+++ b/STMMC/src/VDResamplerGenerateFromModel_module.cc
@@ -43,6 +43,7 @@
 #include "Offline/MCDataProducts/inc/GenId.hh"
 #include "Offline/MCDataProducts/inc/GenParticle.hh"
 #include "Offline/SeedService/inc/SeedService.hh"
+#include "Offline/STMMC/inc/VDResamplerTransformDefaults.hh"
 
 // ROOT includes
 #include "art_root_io/TFileService.h"
@@ -157,11 +158,11 @@ namespace mu2e {
       int pdgId_ = 0;
 
       // Detector-center parameters used in the same transform as training.
-      double x0_ = -3904.0;
-      double y0_ = 0.0;
-      double t0_ = 1700.0;   // ns
-      double tScale_ = 1.0;
-      double p0_ = 1.0;
+      double x0_ = vdresampler::kX0;
+      double y0_ = vdresampler::kY0;
+      double t0_ = vdresampler::kT0;
+      double tScale_ = vdresampler::kTScale;
+      double p0_ = vdresampler::kP0;
 
       // Variables for optional ROOT dump.
       TTree* outTree_ = nullptr;
diff --git a/STMMC/src/VDResamplerGenerateMix_module.cc b/STMMC/src/VDResamplerGenerateMix_module.cc
index c936b83840..a0c477b375 100644
--- a/STMMC/src/VDResamplerGenerateMix_module.cc
+++ b/STMMC/src/VDResamplerGenerateMix_module.cc
@@ -48,6 +48,7 @@
 #include "Offline/MCDataProducts/inc/GenId.hh"
 #include "Offline/MCDataProducts/inc/GenParticle.hh"
 #include "Offline/SeedService/inc/SeedService.hh"
+#include "Offline/STMMC/inc/VDResamplerTransformDefaults.hh"
 
 // ROOT includes
 #include "art_root_io/TFileService.h"
@@ -171,11 +172,11 @@ namespace mu2e {
       std::vector<SourceSummary> sources_;
       double totalSourceWeight_ = 0.0;
 
-      double x0_ = -3904.0;
-      double y0_ = 0.0;
-      double t0_ = 1700.0;
-      double tScale_ = 1.0;
-      double p0_ = 1.0;
+      double x0_ = vdresampler::kX0;
+      double y0_ = vdresampler::kY0;
+      double t0_ = vdresampler::kT0;
+      double tScale_ = vdresampler::kTScale;
+      double p0_ = vdresampler::kP0;
 
       TTree* outTree_ = nullptr;
       int summaryIndex_ = -1;
diff --git a/STMMC/src/VDResamplerTrain_module.cc b/STMMC/src/VDResamplerTrain_module.cc
index 9def2ecce0..14c0f6d843 100644
--- a/STMMC/src/VDResamplerTrain_module.cc
+++ b/STMMC/src/VDResamplerTrain_module.cc
@@ -44,6 +44,7 @@
 #include "Offline/MCDataProducts/inc/SimParticle.hh"
 #include "Offline/MCDataProducts/inc/StepPointMC.hh"
 #include "Offline/SeedService/inc/SeedService.hh"
+#include "Offline/STMMC/inc/VDResamplerTransformDefaults.hh"
 
 typedef cet::map_vector_key key_type;
 typedef unsigned long VolumeId_type;
@@ -117,11 +118,11 @@ namespace mu2e {
       VolumeId_type virtualdetectorId = 0;
 
       // transform variables for training data preparation
-      double x0 = -3904.0;
-      double y0 = 0.0;
+      double x0 = vdresampler::kX0;
+      double y0 = vdresampler::kY0;
       // time scaling
-      double t0 = 1700.0;   // ns
-      double tScale = 1.0;
+      double t0 = vdresampler::kT0;
+      double tScale = vdresampler::kTScale;
   };
 
   VDResamplerTrain::VDResamplerTrain(const Parameters& conf) :
@@ -307,7 +308,7 @@ namespace mu2e {
         pphi = py;
       }
       // momentum scaling
-      double p0 = 1.0; // tunable scale where I want best resolution
+      double p0 = vdresampler::kP0; // tunable scale where I want best resolution
       double pr_t   = std::asinh(pr   / p0);
       double pphi_t = std::asinh(pphi / p0);
       double pz_t   = std::asinh(pz   / p0);
