From a32e69eb89f27580f7693b62bcca711b30fdef7c Mon Sep 17 00:00:00 2001 From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com> Date: Sat, 5 Sep 2026 06:40:50 -0500 Subject: [PATCH 1/2] ScoreBasedDiffusionModel: curriculum, EMA, prediction targets, format v9 The new constructor parameters are appended after initializeRandomWeights rather than placed logically, and SBDMGeneratedSample carries a temporary conversion to vector, so the VDResampler modules already in the tree keep compiling against this header. Both are undone by the later PRs that replace those callers. --- .../inc/ScoreBasedDiffusionModel.hh | 853 ++++- .../src/ScoreBasedDiffusionModel.cc | 2952 +++++++++++++++-- 2 files changed, 3451 insertions(+), 354 deletions(-) diff --git a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh index a7cfd78659..0ec4632415 100644 --- a/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh +++ b/MachineLearningTools/inc/ScoreBasedDiffusionModel.hh @@ -1,4 +1,4 @@ -// Module for training and using score-based diffusion model +// Module for training and using (variance preserving) score-based diffusion model // Added by Yongyi Wu // Mar. 2026 @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -14,6 +15,7 @@ #include #include #include +#include #include "CLHEP/Random/RandomEngine.h" #include "CLHEP/Random/RandFlat.h" @@ -28,6 +30,58 @@ namespace mu2e{ std::vector cond; // optional conditioning vector (size = conditionDim) }; + struct SBDMGeneratedSample { + std::vector zscore; // normalized data + std::vector value; // unnormalized data + + // TEMPORARY, remove together with the VDResampler update that follows this PR. + // generateSample() used to return the unnormalized vector directly; this lets the + // existing VDResampler callers keep compiling until they are updated to take the + // struct and pick zscore/value explicitly. + operator const std::vector&() const { return value; } + }; + + // One peak-importance-sampling window (see ScoreBasedDiffusionModel::train). Oversamples + // training examples whose NORMALIZED coordinate `dim` lies in [low, high) to fix the + // under-weighting of a rare but important localized feature. Multiple windows may be + // supplied; they should be disjoint (a sample is assigned to the first window it matches). + struct PeakWindow { + int dim = -1; // state-coordinate index this window is defined on + double low = 0.0; // window lower edge in TRANSFORMED (pre-z-score) units, inclusive (e.g. pz_t = log(pz/p0)) + double high = 0.0; // window upper edge in TRANSFORMED (pre-z-score) units, exclusive + double gMax = 0.0; // peak sampling-fraction ceiling at low sigma (0 < gMax < 1; sum over windows < 1) + double sigma0 = 0.0; // Gaussian sigma-taper scale (~ feature width); oversampling decays for sigma >> sigma0 + double alpha = 1.0; // 1 = unbiased (variance reduction only), <1 = deliberate up-weight + }; + + // One conditional-loss diagnostic sample for one data point at a freshly drawn diffusion time t. + // Used to profile L(sigma) inside vs outside a feature window (see + // ScoreBasedDiffusionModel::evalEpsLossSample). Reports TWO per-dimension losses: + // perDimLoss — eps-style loss (eps_hat - eps)^2, computed for ALL prediction targets + // (kept as a common lens for cross-mode comparison). + // perDimNativeLoss — loss on the model's NATIVE training target: EPS -> (out - eps)^2 (equals + // perDimLoss), V -> (out - v_target)^2, SCORE -> (score_hat - score)^2. + // The SCORE native loss scales ~ 1/sigma^2 at low sigma (score variance + // blow-up) and so is not magnitude-comparable across sigma to eps/v loss. + struct SBDMEpsLossSample { + double t = 0.0; // drawn diffusion time + double sigma = 0.0; // sigma(t) + std::vector perDimLoss; // eps-style (eps_hat - eps)^2 per dim (all modes) + std::vector perDimNativeLoss; // native-target loss per dim (mode-dependent; see above) + }; + + // First-layer input-feature-block magnitude (see ScoreBasedDiffusionModel::firstLayerBlockMagnitudes). + // Reports, for one contiguous block of the network input, the L2 norm of the first layer's + // weight columns and of the (training-gradient) gradient columns over that block. + struct SBDMFeatureBlockMagnitude { + std::string name; // human-readable block label, e.g. "fourier_state[5]" + int kind = 0; // 0 raw-state, 1 fourier-state, 2 raw-cond, 3 fourier-cond, 4 raw-time, 5 fourier-time + int coord = -1; // state/condition coordinate index this block belongs to (-1 for time) + int nCols = 0; // number of input columns in the block + double weightL2 = 0.0; + double gradL2 = 0.0; + }; + class ScoreBasedDiffusionModel { public: // Enumeration for optimizer selection @@ -39,8 +93,24 @@ namespace mu2e{ // Enumeration for noise schedule selection enum class NoiseScheduleType { LINEAR, // Linear noise schedule - COSINE // Cosine noise schedule + COSINE, // Cosine noise schedule + LOGSIG // Log-sigma (pure exponential) noise schedule: sigma(t) = sigMin*exp(k*t), k=ln(sigMax/sigMin) + }; + + // What the network regresses. Explicit values are chosen so the legacy serialized + // bool (epsPrediction: false->0, true->1) maps directly onto SCORE/EPS, keeping old + // saved models loadable. V (v-prediction) predicts v = alpha*eps - sigma*x0, which is + // well-conditioned at both small and large sigma (eps-like at high sigma, x0-like at + // low sigma) so sharp small-sigma structure carries a real gradient. + enum class PredictionTarget { + SCORE = 0, // network outputs the score -eps/sigma (legacy epsPrediction=false) + EPS = 1, // network outputs the noise eps (legacy epsPrediction=true) + V = 2 // network outputs v = alpha*eps - sigma*x0 (v-prediction) }; + // String<->enum mapping shared by CSV serialization and fhicl parsing, so both agree + // on spelling. Throws cet::exception on an unknown string. + static std::string predictionTargetName(PredictionTarget t); + static PredictionTarget predictionTargetFromName(const std::string& s); // Constructor: Initialize diffusion model with CLHEP random distributions. // @@ -49,19 +119,47 @@ namespace mu2e{ // 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) + // timeEmbeddingDim - Dimensionality of the sinusoidal time embedding applied to the diffusion time t in [0,1]. + // 0 = raw scalar t (default, backward-compatible). + // Even integer >= 2 = sinusoidal embedding: pairs [sin(2π·2^i·t), cos(2π·2^i·t)] for i=0..k/2-1. + // inputEmbeddingDims - Per-coordinate depth of the sinusoidal (Fourier) embedding applied to the dim state + // coordinates. One entry per state dimension. Accepts {} (no embedding on any dim), + // {k} (broadcast depth k to every dim), or a length-dim vector (per-dim depth). + // Each depth is 0 (raw coordinate only) or an even integer >= 2, emitting pairs + // [sin(π·2^i·x_j), cos(π·2^i·x_j)] for i=0..k/2-1 on that coordinate. + // Counteracts MLP spectral bias so the score network can represent structure much finer + // than O(1) in the normalized coordinates (e.g. a narrow peak). To resolve a feature of + // normalized width w, the top frequency π·2^(k/2-1) should be ≳ 2π/w (k=24 reaches w~1e-3). + // Use per-dim depth to give a structured coordinate (e.g. pz) a deep embedding while + // smooth coordinates get 0, avoiding injecting noise features on dims that don't need them. + // conditionEmbeddingDims - Same per-coordinate Fourier embedding, applied to the conditionDim condition + // coordinates ({} / {k} / length-conditionDim). 0 on every dim unless the conditional + // distribution changes sharply with that condition coordinate. Must be empty when conditionDim==0. // 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) + // scheduleType - Type of noise schedule (LINEAR, COSINE, or LOGSIG, 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) + // logSigMin - Minimum sigma for LOGSIG schedule (default: 1e-5) + // logSigMax - Maximum sigma for LOGSIG schedule (default: 1.0) + // predictionTarget - What the network regresses: SCORE (-eps/sigma, default), EPS (noise eps), + // or V (v = alpha*eps - sigma*x0). V is well-conditioned at small sigma + // where eps-prediction degenerates. Under V, lossWeightPower is forced to 0 + // (the v target already embeds the SNR weighting), and for the LOGSIG + // schedule logSigMax is coerced to 1.0. + // lossWeightPower - Power for weighting the loss function (default: 2.0 for quadratic weighting) // batchSize - Batch size for training (default: 32) // gradientClipThreshold - Threshold for gradient clipping (default: 1.0) // learningRate - Learning rate for training (default: 1e-3) + // useDimWeightController - If true, enables adaptive per-dimension gradient weighting (default: false) + // dimWeightEMADecay - EMA decay rate for per-dimension loss tracking (default: 0.99) + // useEMANetwork - If true, maintains a slow-moving EMA copy of network weights for inference (default: true) + // emaNetworkDecay - Decay rate for EMA network update per optimizer step (default: 0.9999) // diffusionSteps - Number of steps in the diffusion process (default: 200) // initializeRandomWeights - If true, initialize network weights from random Gaussian draws. // Set false when constructing (loading) from a saved model. @@ -71,8 +169,8 @@ namespace mu2e{ CLHEP::RandGaussQ& randGaussQ, int dim, int conditionDim, - int hidden, - int layers, + int hidden = 128, + int layers = 4, // Optimizer configuration OptimizerType optimizerType = OptimizerType::ADAM, double adamBeta1 = 0.9, @@ -80,18 +178,202 @@ namespace mu2e{ double adamEps = 1e-8, // Noise schedule configuration NoiseScheduleType scheduleType = NoiseScheduleType::COSINE, + // -- linear schedule parameters double betaMin = 1e-4, double betaMax = 0.02, + // -- cosine schedule parameters double cosineOffset = 0.008, - // Training configuration int batchSize = 32, double gradientClipThreshold = 1.0, double learningRate = 1e-3, // Diffusion process configuration int diffusionSteps = 200, - bool initializeRandomWeights = true + bool initializeRandomWeights = true, + // --------------------------------------------------------------------------- + // TEMPORARY PARAMETER ORDER. Everything above reproduces the previous signature + // exactly, and everything new is appended below, so the existing positional + // callers keep binding correctly until they are updated. The VDResampler update + // that follows this PR moves these back to their logical places: the embedding + // depths after conditionDim, the schedule bounds beside cosineOffset, and the + // training options beside learningRate. + // --------------------------------------------------------------------------- + // -- log-sigma schedule parameters + double logSigMin = 1e-5, + double logSigMax = 1.0, + // -- Training target (SCORE / EPS / V); replaces the legacy bool epsPrediction + PredictionTarget predictionTarget = PredictionTarget::SCORE, + double lossWeightPower = 2.0, + // -- Adaptive dimensional weight controller + bool useDimWeightController = false, + double dimWeightEMADecay = 0.99, + // -- EMA copy of network parameters for inference + bool useEMANetwork = true, + double emaNetworkDecay = 0.9999, + // -- Fourier embedding depths + int timeEmbeddingDim = 0, + std::vector inputEmbeddingDims = {}, + std::vector conditionEmbeddingDims = {} + ); + + // Data normalization to be applied before training + // Parameters: + // mean - Mean values for each dimension of the data (both state and condition) + // stdev - Standard deviation values for each dimension of the data (both state and condition) + // data - Training samples (transformed state vectors) + void normalizeData( + const std::vector& mean, + const std::vector& stdev, + std::vector& data ); + // functions to update certain training parameters + double updateLossWeightPower( + double value + ){ + // v-prediction's target already embeds the SNR (sigma) weighting, so a non-zero + // lossWeightPower would double-apply. Force it to 0 and warn (only when a non-zero + // value was requested, so already-zero per-phase curriculum calls don't spam). + if (predictionTarget_ == PredictionTarget::V && value != 0.0) { + mf::LogWarning("ScoreBasedDiffusionModel") + << "lossWeightPower=" << value + << " ignored under v-prediction (target already SNR-weighted); forcing 0.0"; + value = 0.0; + } + lossWeightPower_ = value; + return lossWeightPower_; + } + + double updateGradientClipThreshold( + double value + ){ + gradientClipThreshold_ = value; + return gradientClipThreshold_; + } + + double updateLearningRate( + double value + ){ + learningRate_ = value; + return learningRate_; + } + + int updateBatchSize( + int value + ){ + batchSize_ = value; + if (useEMANetwork_) { + emaNetworkDecay_ = std::pow(emaNetworkDecayBase_, + (double)batchSize_ / kEMABatchSizeRef_); + mf::LogInfo("ScoreBasedDiffusionModel") + << "Batch size updated to " << batchSize_ + << "; EMA decay rescaled to " << emaNetworkDecay_; + } + return batchSize_; + } + + // Enable/disable the adaptive per-dimension gradient weight controller mid-run + // (e.g. as a curriculum-phase change, or as an explicit override after loading a + // checkpoint that was trained with it on). Turning it OFF does NOT reset + // dimWeights_: train() keeps multiplying each dimension's gradient by the frozen + // weight, it merely stops adapting them. The frozen values are logged so the + // override is auditable. Returns the resulting flag state. + bool updateUseDimWeightController( + bool enabled + ){ + if (useDimWeightController_ && !enabled) { + std::ostringstream woss; + woss << "Dimensional weight controller turned OFF; freezing dimWeights at ["; + for (int i = 0; i < dim_; ++i) { + woss << dimWeights_[i]; + if (i < dim_ - 1) woss << ", "; + } + woss << "] (still applied to gradients, but no longer adapting)."; + mf::LogInfo("ScoreBasedDiffusionModel::updateUseDimWeightController") << woss.str(); + } + useDimWeightController_ = enabled; + return useDimWeightController_; + } + + // Return the controller to its neutral start state: all weights 1.0 and a cleared + // loss EMA. Intended to be called at a curriculum phase boundary. + // + // Two problems this addresses. First, a phase change alters the per-dimension loss + // SCALE (lossWeightPower, tLowBound/tFocus, batch size, peak sampling, an EMA + // promotion), but dimLossEMA_ carries over with a slow decay — so for many epochs + // the weights are a ratio of stale-scale to new-scale numbers, which is what + // produces the large jump seen at the start of a phase. Second, because train() + // applies dimWeights_ whether or not the controller is adapting, a skew left over + // from the previous phase would otherwise stay in force for the rest of the run + // (and in the saved model) even when the new phase disables the controller. + void resetDimWeightController() { + std::fill(dimWeights_.begin(), dimWeights_.end(), 1.0); + std::fill(dimLossEMA_.begin(), dimLossEMA_.end(), 0.0); + mf::LogInfo("ScoreBasedDiffusionModel::resetDimWeightController") + << "Dimensional weight controller reset: dimWeights=1.0, dimLossEMA cleared."; + } + + void promoteEMAToNetwork() { + if (!useEMANetwork_) { + mf::LogWarning("ScoreBasedDiffusionModel") + << "promoteEMAToNetwork called but EMA network is disabled — no-op."; + return; + } + for (size_t l = 0; l < network_.size(); ++l) { + network_[l].W = emaNetwork_[l].W; + network_[l].b = emaNetwork_[l].b; + for (auto& row : network_[l].mW) std::fill(row.begin(), row.end(), 0.0); + for (auto& row : network_[l].vW) std::fill(row.begin(), row.end(), 0.0); + std::fill(network_[l].mb.begin(), network_[l].mb.end(), 0.0); + std::fill(network_[l].vb.begin(), network_[l].vb.end(), 0.0); + } + adamStep_ = 0; + mf::LogInfo("ScoreBasedDiffusionModel::promoteEMAToNetwork") + << "EMA weights promoted to network. Adam optimizer state reset."; + } + + const std::vector& getDimWeights() const { return dimWeights_; } + + // In-memory snapshot of the trainable state. Used by the auto curriculum planner to + // capture the smoothed-best point within a phase and restore it before advancing — + // the substitute for EMA promotion when that is disabled. The snapshot captures the + // SAME mutable state that loadModel() round-trips, so an in-memory restore is + // equivalent to reloading the checkpoint: network weights + Adam moments (full + // Layer), the EMA copy, adamStep_, AND the dimension-weight controller state + // (dimWeights_ / dimLossEMA_). Omitting the controller state previously left a + // skewed dimWeights_ in force after a restore, which destabilised the next phase. + // The model cannot be reassigned in place (reference members delete operator=), so a + // value-copy snapshot is used instead of a loadModel() round-trip. + void snapshotNetwork() { + networkSnapshot_ = network_; + emaNetworkSnapshot_ = emaNetwork_; + adamStepSnapshot_ = adamStep_; + dimWeightsSnapshot_ = dimWeights_; + dimLossEMASnapshot_ = dimLossEMA_; + hasSnapshot_ = true; + } + void restoreNetwork() { + if (!hasSnapshot_) { + mf::LogWarning("ScoreBasedDiffusionModel") + << "restoreNetwork called with no snapshot — no-op."; + return; + } + network_ = networkSnapshot_; + emaNetwork_ = emaNetworkSnapshot_; + adamStep_ = adamStepSnapshot_; + dimWeights_ = dimWeightsSnapshot_; + dimLossEMA_ = dimLossEMASnapshot_; + mf::LogInfo("ScoreBasedDiffusionModel::restoreNetwork") + << "Restored network weights, optimizer, and dim-weight controller from in-memory snapshot."; + } + bool hasNetworkSnapshot() const { return hasSnapshot_; } + void clearNetworkSnapshot() { + networkSnapshot_.clear(); + emaNetworkSnapshot_.clear(); + dimWeightsSnapshot_.clear(); + dimLossEMASnapshot_.clear(); + hasSnapshot_ = false; + } + // 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 @@ -101,36 +383,321 @@ namespace mu2e{ // Parameters: // data - Training samples (transformed state vectors) // epochs - Number of training epochs to perform + // samplesDrawnPerEpoch - If > 0, the number of training samples DRAWN per epoch (default: 0, uses all data + // = one pass over the dataset). This defines the epoch as a fixed quantum of + // optimization work (samplesDrawnPerEpoch/batchSize gradient steps), which is the + // meaningful progress metric — NOT "one pass over the data". If it is <= the dataset + // size N, a random subset of that many DISTINCT samples is used (drawn without + // replacement from the per-epoch shuffle). If it EXCEEDS N, the shuffled dataset is + // cycled with reshuffle-on-wrap so the requested count is still drawn (samples are + // then reused within the epoch, with fresh noise per draw); a warning is emitted once + // since epoch no longer equals one dataset pass. (Formerly named trainSubsetDataSize.) + // biasLowSigma - If true, biases the sampling of diffusion time t towards smaller values (smaller sigma) by sampling t^2 instead of t (default: false) + // tLowBound - If > 0, enforces a lower bound on the sampled diffusion time t to focus training on larger sigma values (default: 0.0, no lower bound) + // tFocusLow/tFocusHigh/tFocusFraction + // - If tFocusFraction > 0, each training sample draws t uniformly from the window + // [tFocusLow, tFocusHigh] with probability tFocusFraction, and from the regular + // tLowBound/biasLowSigma logic otherwise. Concentrates gradient steps on a target + // sigma band (e.g. the band matching a narrow spectral feature) while keeping full + // [0,1] coverage. The t-sampling distribution only reweights the loss; it does not + // bias the learned data distribution. (defaults: 0.0, 0.0, 0.0 = disabled) + // peakWindows - Peak importance sampling (disabled when empty). A list of disjoint windows + // (see PeakWindow); each oversamples training examples whose coordinate `dim` + // lies in [low, high) — low/high given in TRANSFORMED (pre-z-score) units and + // converted to z-score internally via normalizeCoord — to fix under-weighting of rare but + // physically important features (e.g. narrow pz peaks holding a tiny data + // fraction). Per draw, t is sampled first; window k is drawn with probability + // g_eff_k(t) = max(f_k, gMax_k * exp(-sigma(t)^2 / (2 sigma0_k^2))), where f_k is + // the empirical in-window fraction — so oversampling concentrates at + // sigma <~ sigma0_k (set ~ the feature width) and decays to the natural rate at + // high sigma. A single cumulative draw selects window k or the out-of-window + // pool. The loss is reweighted by (f_k/g_eff_k)^alpha_k in window k (alpha_k = 1 + // exactly unbiased / variance reduction only; alpha_k < 1 a deliberate up-weight), + // and by (f_Q/(1-Sum g_eff)) out (unbiased continuum), with f_Q = 1 - Sum f_k. + // Requires Sum gMax_k < 1 (leaves probability for the out-of-window pool). Each + // pool is drawn without replacement via a per-epoch reshuffled cursor. void train( const std::vector& data, - int epochs + int epochs, + int samplesDrawnPerEpoch = 0, + bool biasLowSigma = false, + double tLowBound = 0.0, + double tFocusLow = 0.0, + double tFocusHigh = 0.0, + double tFocusFraction = 0.0, + const std::vector& peakWindows = {} ); + // Diagnostic: average per-sample loss at the CURRENT weights with NO optimizer step. + // Replicates train()'s per-sample loss (same t sampling / focus / bias, addNoise, + // forward, computeLoss with sigma^lossWeightPower weighting) over a subset, but does + // NOT backprop, step the optimizer, update the EMA network, touch dimLossEMA_, or + // increment adamStep_. Used to read the loss of weights as just installed by + // restoreNetwork()/promoteEMAToNetwork(), isolating weight quality from the + // destabilising effect of the first training step. Draws from the RNG (t + noise). + double evaluateAverageLoss( + const std::vector& data, + int subsetSize = 0, + bool biasLowSigma = false, + double tLowBound = 0.0, + double tFocusLow = 0.0, + double tFocusHigh = 0.0, + double tFocusFraction = 0.0 + ); + + // Public accessor for the noise level sigma(t) (pure function of the schedule), so + // diagnostics can set a log-sigma axis without duplicating the schedule. + double diffusionSigma(double t) const { return sigma(t); } + + // Convert a raw (transformed, pre-z-score) value on state dimension `dim` into the model's + // normalized (z-score) space using the stored per-dimension mean/stdev, so feature windows + // can be specified in physical transformed units (e.g. log(pz/p0)) instead of z-scores. + // Requires normalizeData() to have run (or a loaded checkpoint); identity if it has not. + double normalizeCoord(int dim, double rawValue) const { + if (dim < 0 || dim >= dim_) + throw cet::exception("ScoreBasedDiffusionModel::normalizeCoord") + << "dim " << dim << " out of range [0, " << dim_ << ")"; + return (rawValue - dataMean_[dim]) / dataStdev_[dim]; + } + + // Per-dimension training-data statistics for one STATE dimension, in RAW (transformed, + // pre-z-score) units — the units histograms of the model's coordinates are binned in. + // mean/stdev come straight from the stored normalization; min/max are the normalized + // extremes mapped back through it, so all four describe the same raw axis. + // + // These are the training set's own summary, recorded at normalizeData() and persisted in + // the checkpoint, so a consumer can size an axis to where the population actually is + // WITHOUT having generated anything yet. + struct DimStats { + double mean = 0.0, stdev = 0.0, min = 0.0, max = 0.0; + }; + DimStats dimStats(int dim) const { + if (dim < 0 || dim >= dim_) + throw cet::exception("ScoreBasedDiffusionModel::dimStats") + << "dim " << dim << " out of range [0, " << dim_ << ")"; + DimStats s; + s.mean = dataMean_[dim]; + s.stdev = dataStdev_[dim]; + // normMin_/normMax_ are in z-score space; undo the z-score to get raw units. + s.min = dataMean_[dim] + normMin_[dim] * dataStdev_[dim]; + s.max = dataMean_[dim] + normMax_[dim] * dataStdev_[dim]; + return s; + } + + // True when the stored normalization is real (normalizeData() ran, or a checkpoint + // carrying it was loaded) rather than the constructor's placeholder. The ctor seeds + // normMin_/normMax_ to -999/+999, which would otherwise be mistaken for a genuine + // (and absurdly wide) data range by anything sizing an axis from dimStats(). + bool hasDataNormalization() const { + if (dataMean_.empty() || dataStdev_.empty()) return false; + if (normMin_.empty() || normMax_.empty()) return false; + for (int d = 0; d < dim_; ++d) { + if (!(dataStdev_[d] > 0.0)) return false; // unset or degenerate + if (!(normMin_[d] < normMax_[d])) return false; // empty/never-filled range + if (normMin_[d] <= -999.0 || normMax_[d] >= 999.0) return false; // placeholder + } + return true; + } + + // Z-score a raw CONDITION coordinate. The normalization arrays store data dims + // first (0..dim_-1) then condition dims (dim_..dim_+conditionDim_-1); condIdx is + // the 0-based index within the conditioning vector. Used to feed an externally + // sourced condition (e.g. a resampled pTotal) with the exact training z-score. + // + // The categorical (class-label) dim, if any, is returned verbatim: it was never + // z-scored at training time either. Callers can therefore run every condition + // coordinate through this function uniformly without special-casing the label. + double normalizeCondition(int condIdx, double rawValue) const { + if (condIdx < 0 || condIdx >= conditionDim_) + throw cet::exception("ScoreBasedDiffusionModel::normalizeCondition") + << "condIdx " << condIdx << " out of range [0, " << conditionDim_ << ")"; + if (condIdx == categoricalConditionDim_) return rawValue; + const int idx = dim_ + condIdx; + return (rawValue - dataMean_[idx]) / dataStdev_[idx]; + } + + // Diagnostic (conditional-loss profile): draw a uniform diffusion time t and fresh noise, + // run one forward pass, and return per-dimension losses along with t and sigma(t). TWO losses + // are returned (see SBDMEpsLossSample): the eps-style (eps_hat - eps)^2 for ALL modes (eps_hat + // recovered from the output so it is comparable across sigma) AND the loss on the model's + // native training target (EPS->eps, V->v, SCORE->score). Under V the eps-style loss saturates + // near 1 at low sigma (eps is unobservable when x_t ~ a*x0), so the native v-loss is the + // honest measure of fit there. Binning these by log sigma, split by whether the sample lies + // inside a feature window, gives L_P(sigma) vs L_Q(sigma): an underfit only at low sigma points + // to sampling/under-weighting, an underfit at all sigma to capacity / Fourier resolution. + // Draws from the RNG; no optimizer step. + SBDMEpsLossSample evalEpsLossSample( + const std::vector& xNorm, + const std::vector& condition, + bool useEMANetworkIfAvailable = true); + + // Diagnostic (first-layer feature-block magnitudes): for each contiguous block of the + // network input (each raw state coord, each state coord's Fourier columns, each raw + // condition coord, each condition coord's Fourier columns, and the time block), return the + // L2 norm of the first layer's weight columns and of the gradient accumulated over nSamples + // training-like draws (uniform t, addNoise, the configured weighted eps/score loss, + // backward) with NO optimizer step. Near-zero gradL2 (or weightL2 stuck at init scale) on a + // block means that input feature is not being used/driven — e.g. dead pz Fourier columns. + // perDimLossOut is filled with the FRESH mean per-output-dimension squared residual measured + // over the same sweep (a live alternative to the checkpoint's stale dimLossEMA_). Operates + // on the base network_ and leaves the gradient buffers zeroed. + std::vector firstLayerBlockMagnitudes( + const std::vector& data, + int nSamples, + std::vector& perDimLossOut); + // 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) + // useEMANetworkIfAvailable - If true (default), uses the EMA network when available (i.e. when the model was + // configured with useEMANetwork=true). Pass false to force the base score network. // useHeun - If true, uses Heun's method (2nd order, default). If false, uses Euler's method (1st order) + // useSDE - If true, uses SDE (Stochastic Differential Equation) method. If false, uses the deterministic reverse process // 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 generateSample( + // Returns: Two generated sample vectors, zscore and value of dimensions dim_ + SBDMGeneratedSample generateSample( const std::vector& condition = {}, + bool useEMANetworkIfAvailable = true, bool useHeun = true, - int diffusionSteps = -1 + bool useSDE = true, + int diffusionSteps = -1, + double sdeToOdeSigmaThreshold = -1.0 ); - // Save the model parameters to a CSV file with annotations for later use. - // Uses a default filename of "DiffusionModel.csv" if not specified. + // One-step denoising diagnostic: perturb a normalized data sample at a fixed diffusion + // time t, run a single network forward pass, and reconstruct the denoised estimate + // x0_hat = (x_t - sigma(t) * eps_hat) / sqrt(alphabar(t)). + // Comparing the x0_hat distribution against truth separates "the score network never + // learned a feature" (absent here) from "the reverse-process sampler destroys it" + // (present here but absent in generated samples). + // + // Parameters: + // xNorm - Normalized (z-scored) state vector of size dim_, e.g. a training sample + // after normalizeData() + // condition - Normalized conditioning vector (must match conditionDim_) + // t - Fixed diffusion time in (0,1); choose sigma(t) comparable to the feature + // width under study + // useEMANetworkIfAvailable - If true (default), uses the EMA network when available + // + // Returns: zscore = x0_hat (normalized space), value = de-normalized x0_hat + // If noisedZscoreOut is non-null it receives the noised input x_t (z-score space) the network + // sees, so callers can plot the noised data alongside truth and reconstruction. + SBDMGeneratedSample denoiseOneStep( + const std::vector& xNorm, + const std::vector& condition, + double t, + bool useEMANetworkIfAvailable = true, + std::vector* noisedZscoreOut = nullptr + ); + + // Partial-reverse diagnostic: perturb a normalized data sample at diffusion time t0, + // then run the full multi-step reverse sampler from t0 down to 0 (instead of starting + // from pure noise at t=1). Scanning t0 localizes where the sampler loses a feature: + // if the feature survives a start at t0 but not a full generation, the t>t0 phase + // delivers a wrong marginal; if it is already lost starting at t0, the score or its + // integration below t0 is responsible. t0=1.0 degenerates to full generation. + // + // Parameters: + // xNorm - Normalized (z-scored) state vector of size dim_ + // condition - Normalized conditioning vector (must match conditionDim_) + // t0 - Diffusion time in (0,1] to noise the sample to; snapped to the + // sampler's discrete time grid (round(t0*steps)/steps) + // useEMANetworkIfAvailable / useHeun / useSDE / diffusionSteps / + // sdeToOdeSigmaThreshold - identical to generateSample() + // + // Returns: zscore = sampled state (normalized space), value = de-normalized state + SBDMGeneratedSample partialReverseSample( + const std::vector& xNorm, + const std::vector& condition, + double t0, + bool useEMANetworkIfAvailable = true, + bool useHeun = true, + bool useSDE = true, + int diffusionSteps = -1, + double sdeToOdeSigmaThreshold = -1.0, + std::vector* noisedZscoreOut = nullptr // receives the noised start state x_t0 (z-score) + ); + + // Save the model to a binary file (.dat) preserving full double precision. + // This is the default save format. Use saveModelCsv for human-readable output. + // + // Parameters: + // filename - Path to the binary file (default: "DiffusionModel.bin"; callers in the + // VDResampler pipeline always pass an explicit ".dat" name) + // basisTag - Opaque application-level integer (default 0) round-tripped + // verbatim through save/load. This class never interprets it; + // callers (e.g. VDResampler) use it to record which momentum + // basis / model layout the file holds. Field introduced in + // binary format version 7; v<=6 files load it as 0. + void saveModel(const std::string& filename = "DiffusionModel.bin", int basisTag = 0); + + // Opaque basis tag set on load (or 0 for pre-v7 files / fresh models). The + // class assigns no meaning to it; see saveModel. + int basisTag() const { return basisTag_; } + void setBasisTag(int tag) { basisTag_ = tag; } + + // Two more opaque application-level markers, round-tripped verbatim exactly like + // basisTag and equally uninterpreted here (binary format v9+; older files load the + // defaults, 0 and empty). What the values mean, which ids exist, and what a + // disagreement implies are entirely the caller's business. + int pdgId() const { return pdgId_; } + void setPdgId(int pdg) { pdgId_ = pdg; } + + const std::vector>& buildConstants() const { return buildConstants_; } + void setBuildConstants(const std::vector>& c) { buildConstants_ = c; } + + // CATEGORICAL condition dimension — the index of the ONE condition coordinate (if any) + // carrying a small integer CLASS LABEL rather than a physical measurement. + // + // -1 (the DEFAULT) means there is no such coordinate: every condition dim is an + // ordinary continuous measurement, z-scored, which is the historical behaviour and + // what every pre-v8 checkpoint loads as. Do not confuse that with a label VALUE of 0, + // which means the model does have a categorical dim and this event belongs to no class. + // + // A single dim suffices however many classes there are: the classes are distinguished + // by the VALUE in that dim (0 = none, 1..K = class k), not by one dim per class. That + // keeps conditionDim_ — and therefore the checkpoint layout and the first-layer width — + // fixed as classes are added. + // + // The dim is treated differently in exactly two places: + // + // 1. NOT z-scored. normalizeData() and normalizeCondition() pass it through verbatim + // and store mean=0 / stdev=1 for it. Z-scoring a label is actively harmful: + // * a training file in which every event carries the same label has stdev 0, + // which normalizeData() rejects outright; + // * otherwise the label's value at the network input depends on that label's + // POPULATION FRACTION in the particular training set, so the same physical + // class lands at a different input value in every run, and a value fed at + // generation time is only meaningful against the one checkpoint it came from. + // Raw, class k is the literal value k in every model. + // + // 2. Fourier embedding depth forced to 0 (the constructor and the loader reject a + // non-zero depth on it). The embedding exists to defeat spectral bias on a + // coordinate with structure finer than O(1); a label has no fine structure, its + // values are already O(1) apart, and sin/cos of a small integer argument are + // constants — columns perfectly collinear with the raw label, costing parameters + // and carrying no information. One raw input suffices: what a label needs from the + // first layer is a bias shift, i.e. one weight column of adequate magnitude, which + // training grows on its own. It is not competing for representational capacity. + // + // Must be -1 or a valid 0-based index < conditionDim_. Set before normalizeData(); + // persisted in binary format v8. + int categoricalConditionDim() const { return categoricalConditionDim_; } + void setCategoricalConditionDim(int condIdx); + + // Save the model parameters to a CSV file with annotations for human inspection. // // Parameters: - // filename - Path to the CSV file where model parameters will be saved (default: "DiffusionModel.csv") - void saveModel(const std::string& filename = "DiffusionModel.csv"); + // filename - Path to the CSV file (default: "DiffusionModel.csv") + void saveModelCsv(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. + // Auto-detects format by extension: ".bin" loads binary, anything else loads CSV. + // If optimizer state is present, Adam moments and step counter are restored so + // training can resume seamlessly. // // Parameters: // randFlat / RandGaussQ - CLHEP random number generator wrappers being passed @@ -141,6 +708,27 @@ namespace mu2e{ const std::string& filename ); + // Most recent per-epoch average loss (last entry appended by train()); NaN if + // epochLosses_ is empty. Each train(..., epochs=1, ...) call appends exactly one + // entry, so this returns that epoch's loss WHEN CALLED IMMEDIATELY AFTER train(). + // Note: a checkpoint loaded via loadModel() pre-fills epochLosses_ with its saved + // history, so back() is meaningful as "the current run's latest epoch" only when + // read right after a fresh train() call (which is how the curriculum planner in + // VDResamplerTrainCommon uses it). + double getLastEpochLoss() const { + return epochLosses_.empty() ? std::numeric_limits::quiet_NaN() + : epochLosses_.back(); + } + + // Mean unweighted per-event squared residual over the most recent epoch's draws that + // landed in a peak window (NaN if peak sampling was disabled or no in-window draw + // occurred). Like getLastEpochLoss(), meaningful only when read right after a fresh + // train() call; transient per-run state, not serialized. Used by the curriculum + // planner to plateau on the feature-region fit. getLastEpochPeakCount() returns the + // number of in-window draws that fed that average (0 => loss is NaN). + double getLastEpochPeakLoss() const { return lastEpochPeakLoss_; } + long getLastEpochPeakCount() const { return lastEpochPeakCount_; } + private: // ----- network ----- @@ -171,19 +759,39 @@ namespace mu2e{ std::vector> activations_; std::vector> preactivations_; - // Forward pass through the network. - // 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_ + conditionDim_ + 1 - // - // Returns: Output vector (dimension depends on network architecture) + // Construct the time input vector to be appended to the network input. + // Always includes raw t; if timeEmbeddingDim_ > 0 also appends sinusoidal features + // [sin(2π·2^i·t), cos(2π·2^i·t)] for i = 0..timeEmbeddingDim_/2-1. + std::vector timeEmbed(double t) const; + + // Assemble the full network input vector from a (noisy) state vector, the optional + // conditioning vector, and the diffusion time t. Layout: + // [x (dim_), Fourier(x) (Sigma of inputEmbeddingDims_), + // cond (conditionDim_), Fourier(cond) (Sigma of conditionEmbeddingDims_), + // t (+ time embedding)] + // Fourier features for coordinate j: [sin(π·2^i·v_j), cos(π·2^i·v_j)] for + // i = 0..dims[j]/2-1 (none when dims[j] == 0). Depths are per-coordinate. + // The embedding has no trainable parameters, so back-propagation is unaffected. + std::vector buildNetworkInput( + const std::vector& x, + const std::vector& cond, + double t + ) const; + + // Forward pass through the network (training path). + // Caches activations and pre-activations for the backward pass. + // input contains state vector, optional condition vector, and diffusion time t. std::vector forward( const std::vector& x ); + // Const forward pass through an arbitrary network (inference path). + // Does not write to activations_/preactivations_, safe to call during generateSample(). + std::vector forwardInference( + const std::vector& input, + const std::vector& net + ) const; + // Backward pass for gradient computation. // Computes gradients w.r.t. network parameters via chain rule. // @@ -195,6 +803,32 @@ namespace mu2e{ const std::vector& gradOutput ); + // Shared reverse-diffusion integrator behind generateSample() and + // partialReverseSample(): starting from state x at grid time stepStart/steps, + // iterates the Euler/Heun SDE/ODE updates down to t=0 and returns the + // de-normalized sample. generateSample() passes stepStart = steps (t=1). + // + // Parameters: + // x - Initial normalized state at time stepStart/steps (consumed) + // condition - Normalized conditioning vector + // stepStart - First reverse step index in [1, steps]; integration covers + // t = stepStart/steps, ..., 1/steps + // useEMANetworkIfAvailable / useHeun / useSDE / sdeToOdeSigmaThreshold - + // identical to generateSample() + // steps - Total number of grid steps (defines dt = 1/steps) + // + // Returns: zscore = final state (normalized space), value = de-normalized state + SBDMGeneratedSample reverseDiffuseFrom( + std::vector x, + const std::vector& condition, + int stepStart, + bool useEMANetworkIfAvailable, + bool useHeun, + bool useSDE, + int steps, + double sdeToOdeSigmaThreshold + ); + // 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. @@ -213,7 +847,9 @@ namespace mu2e{ // ----- diffusion ----- // Noise schedule parameter beta(t) over diffusion time [0,1]. - // Linear interpolation between betaMin_ and betaMax_. + // Linear scheme: interpolation between betaMin_ and betaMax_. + // Cosine scheme: derived from sigma(t) = sqrt(1 - alpha_bar(t)), + // where alpha_bar(t) = cos^2((t + cosineOffset_) / (1 + cosineOffset_) * pi/2). // // Parameters: // t - Diffusion time parameter in [0,1] @@ -221,8 +857,19 @@ namespace mu2e{ // 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))). + // Cumulative signal retention factor alpha_bar(t) over diffusion time [0,1]. + // For linear noise schedule, alpha_bar(t) = exp(-integral_0^t beta(s) ds). + // For cosine noise schedule, alpha_bar(t) = cos^2((t + cosineOffset_) / (1 + cosineOffset_) * pi/2). + // + // Parameters: + // t - Diffusion time parameter in [0,1] + // + // Returns: Cumulative signal retention factor at time t + double alphabar(double t) const; + + // Cumulative perturbation standard deviation of noise sigma(t) at diffusion time [0,1]. + // Related to the noise schedule via sigma(t) = sqrt(1 - alpha_bar(t)). + // For LOGSIG: sigma(t) = logSigMin_ * exp(k*t), k = ln(logSigMax_/logSigMin_). // // Parameters: // t - Diffusion time parameter in [0,1] @@ -230,6 +877,14 @@ namespace mu2e{ // Returns: Noise standard deviation at time t double sigma(double t) const; + // Time derivative dσ/dt. Used by beta() for the LOGSIG schedule. + // + // Parameters: + // t - Diffusion time parameter in [0,1] + // + // Returns: dσ/dt at time t + double dSigmadt(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. @@ -239,7 +894,7 @@ namespace mu2e{ // t - Diffusion time parameter in [0,1] // eps - Output: Gaussian noise vector used for perturbation (size = dim_) // - // Returns: Noisy state = x + sigma(t) * eps + // Returns: Noisy state = sqrt(alpha_bar(t)) * x + sigma(t) * eps std::vector addNoise( const std::vector& x, double t, @@ -258,9 +913,25 @@ namespace mu2e{ // Returns: MSE loss value (scalar) double computeLoss( const std::vector& score, - const std::vector& target + const std::vector& target, + double weight // use weighted loss to prevent sigma(t) at tiny t from blowing up and dominating the training. ) const; + // Target/output conversions keyed on predictionTarget_, centralizing what used to be + // inline `epsPrediction_ ? ... : ...` branches. s = sigma(t), a = sqrt(max(0,alphabar(t))). + // trainingTargetComponent: the regression target for one dimension given the drawn + // noise eps_i and clean coordinate x0_i. SCORE -> -eps_i/s; EPS -> eps_i; V -> a*eps_i - s*x0_i. + // epsHatFromOutput: recover eps_hat for one dimension from the network output out_i at + // noised coordinate xt_i. SCORE -> -out_i*s; EPS -> out_i; V -> a*out_i + s*xt_i. + // Used by the loss/denoise diagnostics that genuinely need eps_hat. + // scoreFromOutput: the score -eps_hat/sigma computed DIRECTLY per target, avoiding + // the multiply-then-divide-by-sigma round-trip the samplers would otherwise incur + // (SCORE -> out_i with no sigma touched; EPS -> -out_i/s; V -> -(a/s)*out_i - xt_i). + // Use this in the reverse sampler instead of -epsHatFromOutput(...)/s. + double trainingTargetComponent(double eps_i, double x0_i, double s, double a) const; + double epsHatFromOutput(double out_i, double xt_i, double s, double a) const; + double scoreFromOutput(double out_i, double xt_i, double s, double a) const; + // Clip gradients to prevent exploding gradients during training. // // Parameters: @@ -268,6 +939,10 @@ namespace mu2e{ // gradients are scaled down to have norm equal to maxNorm. void clipGradients(double maxNorm); + // Apply one EMA step: emaNetwork_ = decay * emaNetwork_ + (1 - decay) * network_. + // Called after each optimizer step during training. + void updateEMANetwork(); + // ----- internal vars ----- // The random engine is NOT owned by this class. It is injected externally @@ -280,10 +955,13 @@ namespace mu2e{ 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 + int dim_; // Dimensionality of state space + int conditionDim_; // Dimensionality of the optional conditioning vector + int timeEmbeddingDim_; // Sinusoidal time embedding dimensions (0 = raw scalar only) + std::vector inputEmbeddingDims_; // Per-coordinate Fourier embedding depth for each state dim (size dim_; 0 = raw only) + std::vector conditionEmbeddingDims_; // Per-coordinate Fourier embedding depth for each condition dim (size conditionDim_; 0 = raw only) + int hidden_; // Hidden layer size + int layers_; // Number of network layers // Optimizer configuration OptimizerType optimizerType_; // Type of optimizer to use (SGD or ADAM) @@ -294,7 +972,7 @@ namespace mu2e{ double adamEps_; // Small constant for numerical stability (default: 1e-8) // Noise schedule configuration - NoiseScheduleType noiseScheduleType_; // Type of noise schedule (LINEAR or COSINE) + NoiseScheduleType noiseScheduleType_; // Type of noise schedule (LINEAR, COSINE, or LOGSIG) // 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. @@ -304,14 +982,99 @@ namespace mu2e{ // Cosine noise schedule parameter (offset to avoid singularity at t=0 in cosine schedule, default: 0.008) double cosineOffset_; + // Log-sigma noise schedule parameters: sigma(t) = logSigMin_ * exp(k*t), k = ln(logSigMax_/logSigMin_) + double logSigMin_; // sigma at t=0 (default: 1e-5) + double logSigMax_; // sigma at t=1 (default: 1.0) + + // If true, network predicts noise epsilon instead of the score s = -eps/sigma. + // Prevents value explosion from 1/sigma at small t. + PredictionTarget predictionTarget_; // what the network regresses: SCORE / EPS / V + // Training configuration + double lossWeightPower_; // Power of the loss function weighting (default: 2.0) 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) + // Adaptive dimensional weight controller --- + // The raw controller output (dimLossEMA_[i]/mean) is unbounded and self-reinforcing: + // a dimension that is already well fit gets a small weight, which suppresses its + // gradient, which keeps its loss small. Left unclamped this parks well-fit + // dimensions at ~0.05 and effectively stops training them. The bounds below cap + // that dynamic range; a dimension at kDimWeightMax_ is merely trained faster, but + // one below kDimWeightMin_ is not being trained at all, so the floor is the + // important half. + static constexpr double kDimWeightMin_ = 0.5; // floor on a dimension's gradient weight + static constexpr double kDimWeightMax_ = 2.0; // ceiling on a dimension's gradient weight + + // Enforce [kDimWeightMin_, kDimWeightMax_] on dimWeights_, warning if anything + // actually moved. Called both where the controller produces new weights and where + // loadModel() restores them from a checkpoint, so the bounds are an invariant of + // the object rather than a property of one code path. The load path matters + // because checkpoints written before the bounds existed can carry arbitrarily + // skewed weights, and train() applies dimWeights_ even when the controller is + // switched off — so an unclamped restore would otherwise suppress a dimension's + // gradient for a whole run with nothing to re-clamp it. + void clampDimWeights(const char* context) { + std::ostringstream before; + bool changed = false; + for (int i = 0; i < dim_; ++i) { + const double raw = dimWeights_[i]; + const double c = std::clamp(raw, kDimWeightMin_, kDimWeightMax_); + if (c != raw) changed = true; + before << raw << (i < dim_ - 1 ? ", " : ""); + dimWeights_[i] = c; + } + if (changed) { + std::ostringstream after; + for (int i = 0; i < dim_; ++i) + after << dimWeights_[i] << (i < dim_ - 1 ? ", " : ""); + mf::LogWarning("ScoreBasedDiffusionModel::clampDimWeights") + << context << ": dimWeights outside [" + << kDimWeightMin_ << ", " << kDimWeightMax_ << "] were clamped. " + << "Before: [" << before.str() << "] After: [" << after.str() << "]. " + << "A weight far below 1 means that dimension was receiving almost no " + << "gradient; check whether the affected dimensions are the poorly fit ones."; + } + } + bool useDimWeightController_; // if true, per-dimension gradient weights are applied during training + double dimWeightEMADecay_; // EMA decay rate for per-dimension loss tracking + std::vector dimLossEMA_; // per-dim raw MSE EMA (size dim_), updated every epoch + std::vector dimWeights_; // normalized gradient weights (size dim_), init 1.0 + + // EMA copy of network parameters for inference --- + static constexpr int kEMABatchSizeRef_ = 32; // canonical reference batch size for emaNetworkDecay interpretation + bool useEMANetwork_; // if true, generateSample() uses emaNetwork_ instead of network_ + double emaNetworkDecayBase_; // user-configured decay at kEMABatchSizeRef_=32 samples/step + double emaNetworkDecay_; // effective decay per optimizer step (rescaled from emaNetworkDecayBase_) + std::vector emaNetwork_; // slow-moving EMA copy, only W and b are used + + // In-memory snapshot buffers for the curriculum planner's resume-from-best + // (see snapshotNetwork()/restoreNetwork()). Not serialized. + std::vector networkSnapshot_; + std::vector emaNetworkSnapshot_; + std::vector dimWeightsSnapshot_; + std::vector dimLossEMASnapshot_; + int adamStepSnapshot_ = 0; + bool hasSnapshot_ = false; + // Diffusion process discretization int diffusionSteps_; // Number of time steps to generate a sample (default: 200) + // Opaque application-level tag (default 0). Not interpreted by this class; + // round-tripped through save/load (binary format v7+). See saveModel/basisTag(). + int basisTag_ = 0; + + // Index of the single class-label condition dim, or -1 for none (the default, and + // what every pre-v8 checkpoint loads as). See categoricalConditionDim(). + int categoricalConditionDim_ = -1; + + // Opaque application-level markers (format v9+). Not interpreted by this class; + // see pdgId() / buildConstants(). The pair count is the vector's own size — on disk + // it is written count-prefixed, like the embedding-dim vectors. + int pdgId_ = 0; + std::vector> buildConstants_; + // Training state double runningLoss_; // Accumulated loss for monitoring during training int adamStep_; // Step counter for Adam optimizer (used to compute bias-corrected moment estimates) @@ -320,5 +1083,23 @@ namespace mu2e{ // Container for tracking training loss over epochs std::vector epochLosses_; + // Most recent epoch's peak-window loss (mean unweighted squared residual over in-window + // draws) and the in-window draw count that produced it. Transient per-run; NOT serialized. + // Read by the curriculum planner immediately after train(); see getLastEpochPeakLoss(). + double lastEpochPeakLoss_ = std::numeric_limits::quiet_NaN(); + long lastEpochPeakCount_ = 0; + + // Variables to track gradient clipping statistics for monitoring + size_t clipCount_; + size_t totalClipChecks_; + double clipScaleAccum_; + + // Data normalization containers + // Mean and variance of input data (for normalization) + std::vector dataMean_; + std::vector dataStdev_; + // Min and max of normalized data (for possible clamping or rescaling) + std::vector normMin_; + std::vector normMax_; }; } diff --git a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc index 9e0fd318a4..fdef81ac54 100644 --- a/MachineLearningTools/src/ScoreBasedDiffusionModel.cc +++ b/MachineLearningTools/src/ScoreBasedDiffusionModel.cc @@ -15,28 +15,235 @@ namespace mu2e { 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. + // Here instantaneous noise scale, cumulative signal retention factor, + // and cumulative perturbation stddev are defined for variance preserving (VP) SDEs + // for both linear and cosine noise schedules. This choice is out of the consideration + // of numerical stability, ease of implementation, and faster sampler convergence + + // Instantaneous noise scale (diffusion coefficient) at time t double ScoreBasedDiffusionModel::beta(double t) const { - return betaMin_ + t * (betaMax_ - betaMin_); + if (noiseScheduleType_ == NoiseScheduleType::COSINE) { + // Cosine noise schedule does not use betaMin and betaMax, but we can still define an effective beta if needed + double f = (t + cosineOffset_) / (1.0 + cosineOffset_); + double beta = (M_PI / (1.0 + cosineOffset_)) * std::tan(f * M_PI * 0.5); // beta(t) = pi / (1+offset) * tan(pi*f/2) + // Cap beta to prevent numerical issues at t close to 1 + return std::min(beta, 10.0); // cap beta to a large value + } else if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + // LOGSIG: beta(t) = 2*sigma*dSigmadt / (1 - sigma^2) + // Beta diverges as sigma->1 (t->1). For default sigMin=1e-5, k~11.5: + // beta=100 is first reached at sigma~0.90 (t~0.97). + // Capping at 100 keeps beta*dt <= 0.5 for 200 diffusion steps, + // analogous to the cap=10 used for COSINE (where k is ~7x smaller). + double s = sigma(t); + double sd = dSigmadt(t); + double denom = 1.0 - s * s; + if (denom < 1e-12) return 100.0; + return std::min(2.0 * s * sd / denom, 100.0); + } else { + // Linear interpolation of beta(t) between betaMin_ and betaMax_. + 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 { + // cumulative signal retention factor at time t + double ScoreBasedDiffusionModel::alphabar(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); + alpha_bar *= alpha_bar; // alpha_bar(t) = cos^2(pi*f/2) + return alpha_bar; + } else if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + // LOGSIG: alphabar(t) = 1 - sigma^2(t), clamped to [0,1] + double s = sigma(t); + return std::max(0.0, 1.0 - s * s); } else { - // Linear noise schedule (default) - return std::sqrt(beta(t)); + // Linear noise schedule + double integral = betaMin_ * t + 0.5 * (betaMax_ - betaMin_) * t * t; + return std::exp(-integral); // alpha_bar(t) = exp(-integral(beta(s) ds)) + } + } + + // Cumulative perturbation stddev (noise level) + double ScoreBasedDiffusionModel::sigma(double t) const { + if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + // Direct formula; computed here (not via alphabar) to avoid circular dependency + double k = std::log(logSigMax_ / logSigMin_); + return logSigMin_ * std::exp(k * t); + } + return std::sqrt(1.0 - alphabar(t)); // sigma(t) = sqrt(1 - alpha_bar(t)) + } + + // Time derivative dσ/dt — only defined for LOGSIG + double ScoreBasedDiffusionModel::dSigmadt(double t) const { + if (noiseScheduleType_ != NoiseScheduleType::LOGSIG) { + throw cet::exception("ScoreBasedDiffusionModel::dSigmadt") + << "dSigmadt() is only defined for the LOGSIG noise schedule"; + } + double k = std::log(logSigMax_ / logSigMin_); + return k * sigma(t); + } + + // ---- PredictionTarget helpers ------------------------------------------------------- + + std::string ScoreBasedDiffusionModel::predictionTargetName(PredictionTarget t) { + switch (t) { + case PredictionTarget::SCORE: return "SCORE"; + case PredictionTarget::EPS: return "EPS"; + case PredictionTarget::V: return "V"; } + throw cet::exception("ScoreBasedDiffusionModel::predictionTargetName") + << "Unknown PredictionTarget value " << static_cast(t); + } + + ScoreBasedDiffusionModel::PredictionTarget + ScoreBasedDiffusionModel::predictionTargetFromName(const std::string& s) { + if (s == "SCORE") return PredictionTarget::SCORE; + if (s == "EPS") return PredictionTarget::EPS; + if (s == "V") return PredictionTarget::V; + throw cet::exception("ScoreBasedDiffusionModel::predictionTargetFromName") + << "Unknown prediction target '" << s << "' (expected SCORE, EPS, or V)"; + } + + // Regression target for one dimension. s = sigma(t), a = sqrt(max(0,alphabar(t))). + double ScoreBasedDiffusionModel::trainingTargetComponent( + double eps_i, double x0_i, double s, double a) const { + switch (predictionTarget_) { + case PredictionTarget::SCORE: return -eps_i / s; + case PredictionTarget::EPS: return eps_i; + case PredictionTarget::V: return a * eps_i - s * x0_i; + } + return eps_i; // unreachable; keeps the compiler happy + } + + // Recover eps_hat for one dimension from the network output out_i at noised coordinate + // xt_i. For V: xt = a*x0 + s*eps and out = a*eps - s*x0, so (using a^2+s^2=1) + // eps_hat = a*out + s*xt. Samplers form score_i = -eps_hat/s. + double ScoreBasedDiffusionModel::epsHatFromOutput( + double out_i, double xt_i, double s, double a) const { + switch (predictionTarget_) { + case PredictionTarget::SCORE: return -out_i * s; + case PredictionTarget::EPS: return out_i; + case PredictionTarget::V: return a * out_i + s * xt_i; + } + return out_i; // unreachable + } + + // The score (-eps_hat/sigma) computed DIRECTLY per target, avoiding the redundant + // multiply-then-divide-by-sigma that -epsHatFromOutput(...)/s would do: + // SCORE: eps_hat = -out*s, so score = out (sigma cancels exactly) + // EPS: eps_hat = out, so score = -out/s + // V: eps_hat = a*out + s*xt, so score = -(a/s)*out - xt (xt term is exact, no s round-trip) + double ScoreBasedDiffusionModel::scoreFromOutput( + double out_i, double xt_i, double s, double a) const { + switch (predictionTarget_) { + case PredictionTarget::SCORE: return out_i; + case PredictionTarget::EPS: return -out_i / s; + case PredictionTarget::V: return -(a / s) * out_i - xt_i; + } + return out_i; // unreachable + } + + + std::vector ScoreBasedDiffusionModel::timeEmbed(double t) const { + std::vector emb; + emb.reserve(1 + timeEmbeddingDim_); + emb.push_back(t); + for (int i = 0; i < timeEmbeddingDim_ / 2; ++i) { + double freq = 2.0 * M_PI * std::pow(2.0, i); + emb.push_back(std::sin(freq * t)); + emb.push_back(std::cos(freq * t)); + } + return emb; + } + + // Assemble the full network input: raw state coordinates, optional per-coordinate Fourier + // features, conditioning vector, then the time input (raw t + optional time embedding). + // The Fourier features give the MLP access to high-frequency structure in the state + // coordinates (e.g. narrow spectral lines), which a plain MLP is biased against learning. + std::vector ScoreBasedDiffusionModel::buildNetworkInput( + const std::vector& x, + const std::vector& cond, + double t + ) const { + std::vector input; + int sumInputEmb = 0; for (int e : inputEmbeddingDims_) sumInputEmb += e; + int sumCondEmb = 0; for (int e : conditionEmbeddingDims_) sumCondEmb += e; + input.reserve(dim_ + sumInputEmb + conditionDim_ + sumCondEmb + 1 + timeEmbeddingDim_); + // Per-coordinate Fourier features with a per-coordinate depth: for coordinate j, + // [sin(π·2^i·v_j), cos(π·2^i·v_j)] for i = 0..dims[j]/2-1 (none when dims[j] == 0). + auto appendFourier = [&input](const std::vector& v, const std::vector& dims) { + for (size_t j = 0; j < v.size(); ++j) { + for (int i = 0; i < dims[j] / 2; ++i) { + double freq = M_PI * std::pow(2.0, i); + input.push_back(std::sin(freq * v[j])); + input.push_back(std::cos(freq * v[j])); + } + } + }; + input.insert(input.end(), x.begin(), x.end()); + if (sumInputEmb > 0) appendFourier(x, inputEmbeddingDims_); + input.insert(input.end(), cond.begin(), cond.end()); + if (sumCondEmb > 0) appendFourier(cond, conditionEmbeddingDims_); + auto tEmb = timeEmbed(t); + input.insert(input.end(), tEmb.begin(), tEmb.end()); + return input; + } + + // Format an int vector as "[a, b, c]" for human-readable LOG output only. The CSV wire + // format writes bare comma-separated values (see saveModelCsv) so the existing CSV + // splitter parses them; do not use this bracketed form there. + static std::string embDimsToString(const std::vector& v) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < v.size(); ++i) { oss << v[i]; if (i + 1 < v.size()) oss << ", "; } + oss << "]"; + return oss.str(); } + // Resolve a user-supplied per-coordinate Fourier embedding spec into a length-nDim vector + // and validate it. Accepts {} (all zeros), {k} (broadcast k to every dim), or a length-nDim + // vector. Each resolved depth must be 0 or an even integer >= 2 (an odd depth would emit an + // unpaired sin/cos and silently corrupt the input layout, so it is a hard error, not a clamp). + // Used both at construction (from fcl) and on the load path (from a checkpoint), so a + // malformed vector is rejected rather than mis-shaping the network. + static std::vector resolveEmbeddingDims( + const std::vector& in, int nDim, const char* name) + { + // A non-positive target dimension carries no embedding. nDim==0 is legitimate + // (e.g. conditionEmbeddingDims with conditionDim==0); nDim<0 only arises from a corrupt + // checkpoint and is reported by the constructor's dim>0 check right after this runs — + // return empty here to avoid a negative-size allocation in the broadcast path below. + if (nDim <= 0) { + for (int e : in) + if (e != 0) + throw cet::exception("ScoreBasedDiffusionModel::initialization") + << name << " requests depth " << e << " but the corresponding dimension is " << nDim; + return {}; + } + + std::vector out; + if (in.empty()) { + out.assign(nDim, 0); + } else if (in.size() == 1) { + out.assign(nDim, in[0]); + } else if (static_cast(in.size()) == nDim) { + out = in; + } else { + throw cet::exception("ScoreBasedDiffusionModel::initialization") + << name << " must have length 0 (none), 1 (broadcast), or " << nDim + << " (per dimension); got " << in.size(); + } + + for (int i = 0; i < nDim; ++i) { + int e = out[i]; + if (!(e == 0 || (e >= 2 && e % 2 == 0))) + throw cet::exception("ScoreBasedDiffusionModel::initialization") + << name << "[" << i << "] = " << e + << " must be 0 (raw) or an even integer >= 2"; + } + return out; + } ScoreBasedDiffusionModel::ScoreBasedDiffusionModel( CLHEP::RandFlat& randFlat, @@ -57,21 +264,50 @@ namespace mu2e { double gradientClipThreshold, double learningRate, int diffusionSteps, - bool initializeRandomWeights + bool initializeRandomWeights, + // Temporary position; see the declaration. + double logSigMin, + double logSigMax, + PredictionTarget predictionTarget, + double lossWeightPower, + bool useDimWeightController, + double dimWeightEMADecay, + bool useEMANetwork, + double emaNetworkDecay, + int timeEmbeddingDim, + std::vector inputEmbeddingDims, + std::vector conditionEmbeddingDims ) : randFlat_(randFlat), randGaussQ_(randGaussQ), - dim_(dim), conditionDim_(conditionDim), hidden_(hidden), layers_(layers), + dim_(dim), conditionDim_(conditionDim), timeEmbeddingDim_(timeEmbeddingDim), + inputEmbeddingDims_(resolveEmbeddingDims(inputEmbeddingDims, dim, "inputEmbeddingDims")), + conditionEmbeddingDims_(resolveEmbeddingDims(conditionEmbeddingDims, conditionDim, "conditionEmbeddingDims")), + 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), + logSigMin_(logSigMin), logSigMax_(logSigMax), predictionTarget_(predictionTarget), + lossWeightPower_(lossWeightPower), batchSize_(batchSize), gradientClipThreshold_(gradientClipThreshold), learningRate_(learningRate), + useDimWeightController_(useDimWeightController), dimWeightEMADecay_(dimWeightEMADecay), + dimLossEMA_(dim, 0.0), dimWeights_(dim, 1.0), + useEMANetwork_(useEMANetwork), emaNetworkDecayBase_(emaNetworkDecay), + emaNetworkDecay_(useEMANetwork ? std::pow(emaNetworkDecay, (double)batchSize / kEMABatchSizeRef_) : emaNetworkDecay), diffusionSteps_(diffusionSteps), - runningLoss_(0.0), adamStep_(0), trainingSampleSize_(0), epochLosses_() { + runningLoss_(0.0), adamStep_(0), trainingSampleSize_(0), epochLosses_(), + clipCount_(0), totalClipChecks_(0), clipScaleAccum_(0.0), + dataMean_(dim + conditionDim, 0.0), dataStdev_(dim + conditionDim, 1.0), + normMin_(dim + conditionDim, -999.0), normMax_(dim + conditionDim, 999.0) { // Validate model dimensions and parameters if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) { throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid model dimensions"; } + if (timeEmbeddingDim_ != 0 && (timeEmbeddingDim_ < 2 || timeEmbeddingDim_ % 2 != 0)) { + throw cet::exception("ScoreBasedDiffusionModel::initialization") + << "timeEmbeddingDim must be 0 (raw scalar) or an even integer >= 2"; + } + // inputEmbeddingDims_ / conditionEmbeddingDims_ were resolved and validated by + // resolveEmbeddingDims() in the member-initializer list. if (batchSize <= 0) { throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid batchSize"; } @@ -79,14 +315,48 @@ namespace mu2e { throw cet::exception("ScoreBasedDiffusionModel::initialization") << "Invalid diffusionSteps"; } + // v-prediction guardrails (must run before any schedule-derived quantity is used and + // before serialization, since the coerced values are what get saved). + if (predictionTarget_ == PredictionTarget::V) { + if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + if (logSigMax_ != 1.0) { + mf::LogWarning("ScoreBasedDiffusionModel") + << "v-prediction with LOGSIG requires logSigMax=1.0; overriding " + << logSigMax_ << " -> 1.0"; + logSigMax_ = 1.0; + } + if (logSigMin_ > 1e-3) { + mf::LogWarning("ScoreBasedDiffusionModel") + << "v-prediction with LOGSIG: logSigMin=" << logSigMin_ + << " > 1e-3; the small-sigma feature band may be under-resolved."; + } + // VP note for LOGSIG: with logSigMax coerced to 1, alpha=sqrt(1-sigma^2) + // gives alpha^2+sigma^2=1 EXACTLY (same as LINEAR/COSINE). The only deviation + // from ideal VP is at t=0: sigma(0)=logSigMin>0 so alpha(0)<1 rather than + // exactly 1. This gap is negligible for small logSigMin and only matters when + // logSigMin is large (already warned above). + mf::LogWarning("ScoreBasedDiffusionModel") + << "v-prediction on LOGSIG: alpha^2+sigma^2=1 holds exactly " + << "(logSigMax forced to 1); the only VP deviation is alpha(t=0)<1 because " + << "sigma(0)=logSigMin=" << logSigMin_ << " > 0 (matters only if logSigMin is large)."; + } + // The v target already embeds the SNR weighting; force lossWeightPower to 0 + // (updateLossWeightPower warns if a non-zero value was requested). + updateLossWeightPower(lossWeightPower_); + } + // ------------------------------------------------------------ // Network architecture // - // Input dimension = dim_ + conditionDim_ + 1 + // Input dimension = dim_ (+ per-coordinate Fourier features) + conditionDim_ + 1 // (+1 because diffusion time t is appended to the input vector) // ------------------------------------------------------------ - int inputSize = dim_ + conditionDim_ + 1; + int sumInputEmb = 0; for (int e : inputEmbeddingDims_) sumInputEmb += e; + int sumCondEmb = 0; for (int e : conditionEmbeddingDims_) sumCondEmb += e; + int inputSize = dim_ + sumInputEmb + + conditionDim_ + sumCondEmb + + 1 + timeEmbeddingDim_; int in = inputSize; // Weight initialization scale (local constant so it can be tuned easily) @@ -147,6 +417,13 @@ namespace mu2e { in = out; } + // Initialize emaNetwork_ with the same W/b as network_ so it starts from the same point + emaNetwork_.resize(network_.size()); + for (size_t l = 0; l < network_.size(); ++l) { + emaNetwork_[l].W = network_[l].W; + emaNetwork_[l].b = network_[l].b; + } + // Print layer and model configuration std::ostringstream oss; oss << "ScoreBasedDiffusionModel initialized\n" @@ -154,6 +431,9 @@ namespace mu2e { << " Network architecture:\n" << " | dim=" << dim_ << "\n" << " | conditionDim=" << conditionDim_ << "\n" + << " | timeEmbeddingDim=" << timeEmbeddingDim_ << (timeEmbeddingDim_ == 0 ? " (raw scalar)" : " (sinusoidal)") << "\n" + << " | inputEmbeddingDims=" << embDimsToString(inputEmbeddingDims_) << " (per-state-coordinate Fourier depth; 0 = raw)\n" + << " | conditionEmbeddingDims=" << embDimsToString(conditionEmbeddingDims_) << " (per-condition-coordinate Fourier depth; 0 = raw)\n" << " | hidden=" << hidden_ << "\n" << " | layers=" << layers_ << "\n" << " Optimizer configuration:\n" @@ -163,18 +443,34 @@ namespace mu2e { << " |- AdamBeta2=" << adamBeta2_ << "\n" << " |- AdamEps=" << adamEps_ << "\n"; } - oss << " Noise schedule configuration:\n" - << " | NoiseSchedule=" << (noiseScheduleType_ == NoiseScheduleType::COSINE ? "Cosine" : "Linear") << "\n"; + oss << " Noise schedule configuration:\n"; if (noiseScheduleType_ == NoiseScheduleType::COSINE) { - oss << " |- CosineOffset=" << cosineOffset_ << "\n"; + oss << " | NoiseSchedule=Cosine\n" + << " |- CosineOffset=" << cosineOffset_ << "\n"; + } else if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + oss << " | NoiseSchedule=LogSig\n" + << " |- LogSigMin=" << logSigMin_ << "\n" + << " |- LogSigMax=" << logSigMax_ << "\n"; } else { - oss << " |- BetaMin=" << betaMin_ << "\n" + oss << " | NoiseSchedule=Linear\n" + << " |- BetaMin=" << betaMin_ << "\n" << " |- BetaMax=" << betaMax_ << "\n"; } oss << " Training configuration:\n" + << " | PredictionTarget=" << predictionTargetName(predictionTarget_) << "\n" + << " | LossWeightPower=" << lossWeightPower_ << "\n" << " | BatchSize=" << batchSize_ << "\n" << " | GradientClipThreshold=" << gradientClipThreshold_ << "\n" << " | LearningRate=" << learningRate_ << "\n" + << " | DimWeightController=" << (useDimWeightController_ ? "enabled" : "disabled"); + if (useDimWeightController_) + oss << " (EMADecay=" << dimWeightEMADecay_ << ")"; + oss << "\n" + << " EMA network:\n" + << " | EMANetwork=" << (useEMANetwork_ ? "enabled" : "disabled"); + if (useEMANetwork_) + oss << " (decay=" << emaNetworkDecay_ << ")"; + oss << "\n" << " Diffusion process configuration:\n" << " | DiffusionSteps=" << diffusionSteps_ << "\n"; mf::LogInfo("ScoreBasedDiffusionModel::initialize") << oss.str(); @@ -184,6 +480,166 @@ namespace mu2e { preactivations_.reserve(layers_); } + // Declare which condition coordinate (if any) carries a class label. See the header for + // why such a dim is neither z-scored nor Fourier-embedded. Rejecting a non-zero embedding + // depth here rather than silently zeroing it keeps the network input layout exactly what + // the configuration says it is: the depth feeds the first-layer width, so quietly changing + // it would make a checkpoint's geometry disagree with the config that produced it. + void ScoreBasedDiffusionModel::setCategoricalConditionDim(int condIdx) { + if (condIdx != -1) { + if (condIdx < 0 || condIdx >= conditionDim_) + throw cet::exception("ScoreBasedDiffusionModel::setCategoricalConditionDim") + << "categorical condition dim " << condIdx << " out of range [0, " + << conditionDim_ << ") (-1 means none)"; + if (!conditionEmbeddingDims_.empty() && conditionEmbeddingDims_[condIdx] != 0) + throw cet::exception("ScoreBasedDiffusionModel::setCategoricalConditionDim") + << "condition dim " << condIdx << " is declared categorical but has Fourier " + << "embedding depth " << conditionEmbeddingDims_[condIdx] + << "; a class label must have depth 0 (its sin/cos columns would be " + << "constants, collinear with the raw label). Set that entry of " + << "conditionEmbeddingDims to 0."; + } + categoricalConditionDim_ = condIdx; + } + + // normalize input data + // note the order in the stdev and mean is always x then cond + void ScoreBasedDiffusionModel::normalizeData( + const std::vector& mean, + const std::vector& stdev, + std::vector& data + ) + { + if (data.empty()) { + throw cet::exception("ScoreBasedDiffusionModel::normalizeData") + << "No training data provided"; + } + + const size_t totalDim = dim_ + conditionDim_; + + // check mean and stdev dimensions + if (mean.size() != totalDim || stdev.size() != totalDim) { + throw cet::exception("ScoreBasedDiffusionModel::normalizeData") + << "Mean or standard deviation dimension mismatch, expected " + << totalDim << " but got " << mean.size() + << " for mean and " << stdev.size() << " for standard deviation"; + } + // store the mean and standard deviation of the original data + dataMean_ = mean; + dataStdev_ = stdev; + // The categorical (class-label) dim is not standardized, so record the IDENTITY for + // it regardless of what the caller passed. Everything downstream reads the stored + // normalization rather than re-deriving it — normalizeCondition, dimStats, the + // checkpoint — so pinning it to 0/1 here is what keeps them all consistent with the + // verbatim values written into the samples below. + if (categoricalConditionDim_ >= 0) { + const int idx = dim_ + categoricalConditionDim_; + dataMean_[idx] = 0.0; + dataStdev_[idx] = 1.0; + } + + // containers to track the mean and standard deviation of the normalized data + std::vector normMean(totalDim, 0.0); + std::vector normStdev(totalDim, 0.0); + // containers to track the sum of squares of the normalized data + std::vector M2(totalDim, 0.0); + // reset min/max tracers + normMin_.assign(totalDim, std::numeric_limits::max()); + normMax_.assign(totalDim, std::numeric_limits::lowest()); + + size_t count = 0; + + // Normalize the training data using the given mean and standard deviation + // keep track of the mean and standard deviation, min and max values + // of the normalized data + for (auto& sample : data) { + ++count; + + // deal with sample.x first + for (int i = 0; i < dim_; ++i) { + if (stdev[i] == 0.0) { + throw cet::exception("ScoreBasedDiffusionModel::normalizeData") + << "Zero standard deviation encountered at dimension " + << i; + } + + double value = (sample.x[i] - mean[i]) / stdev[i]; + sample.x[i] = value; //overwrite the original value + + // update min/max values + normMin_[i] = std::min(normMin_[i], value); + normMax_[i] = std::max(normMax_[i], value); + // Welford update + double delta = value - normMean[i]; + normMean[i] += delta / static_cast(count); + double delta2 = value - normMean[i]; + M2[i] += delta * delta2; + } + //now deal with sample.cond + for (int i = 0; i < conditionDim_; ++i) { + const int idx = dim_ + i; + double value; + if (i == categoricalConditionDim_) { + // Class label: fed to the network verbatim (see categoricalConditionDim). + // dataMean_/dataStdev_ were forced to 0/1 for this dim above, so the + // stored normalization is the identity and normalizeCondition, dimStats + // and any inverse mapping all agree with what is written here. + value = sample.cond[i]; + } else { + if (stdev[idx] == 0.0) { + throw cet::exception("ScoreBasedDiffusionModel::normalizeData") + << "Zero standard deviation encountered at dimension " + << idx; + } + value = (sample.cond[i] - mean[idx]) / stdev[idx]; + } + sample.cond[i] = value; + + // update min/max + normMin_[idx] = std::min(normMin_[idx], value); + normMax_[idx] = std::max(normMax_[idx], value); + // Welford update + double delta = value - normMean[idx]; + normMean[idx] += delta / static_cast(count); + double delta2 = value - normMean[idx]; + M2[idx] += delta * delta2; + } + } + + for (size_t i = 0; i < totalDim; ++i) { + normStdev[i] = std::sqrt(M2[i] / static_cast(count)); + // The categorical dim is deliberately left un-standardized, so its mean and stdev + // are whatever the class populations make them. Report them (they are a useful + // record of the class mix) but skip the checks below, which would otherwise fire + // on every run for a dim that is behaving exactly as intended. + if (categoricalConditionDim_ >= 0 + && static_cast(i) == dim_ + categoricalConditionDim_) { + mf::LogInfo("ScoreBasedDiffusionModel::normalizeData") + << "Dimension " << i << " is the categorical (class-label) condition dim: " + << "left un-normalized, observed mean = " << normMean[i] + << ", stdev = " << normStdev[i]; + continue; + } + mf::LogInfo("ScoreBasedDiffusionModel::normalizeData") + << "Dimension " << i << " normalized: mean = " << normMean[i] << ", stdev = " << normStdev[i]; + + // sanity checks + if (std::abs(normMean[i]) > 0.1) { + mf::LogWarning("ScoreBasedDiffusionModel::normalizeData") + << "Normalized mean deviates from 0 at dimension " + << i + << ": mean = " << normMean[i]; + } + if (std::abs(normStdev[i] - 1.0) > 0.1) { + mf::LogWarning("ScoreBasedDiffusionModel::normalizeData") + << "Normalized stdev deviates from 1 at dimension " + << i + << ": stdev = " << normStdev[i]; + } + } + return; + } + // forward pass to compute the score function given input (state + time embedding) std::vector ScoreBasedDiffusionModel::forward( const std::vector& input @@ -247,6 +703,43 @@ namespace mu2e { return x; } + // Const forward pass for inference — does not cache activations or preactivations. + // Used by generateSample() so it can use either network_ or emaNetwork_ without + // corrupting the activations needed by the training backward pass. + std::vector ScoreBasedDiffusionModel::forwardInference( + const std::vector& input, + const std::vector& net + ) const + { + if (net.empty() || net[0].W.empty() || net[0].W[0].empty()) { + throw cet::exception("ScoreBasedDiffusionModel::forwardInference") + << "Network is not properly initialized"; + } + if (input.size() != net[0].W[0].size()) { + throw cet::exception("ScoreBasedDiffusionModel::forwardInference") + << "Input dimension mismatch: got " << input.size() + << ", expected " << net[0].W[0].size(); + } + + std::vector x = input; + for (size_t l = 0; l < net.size(); ++l) { + const auto& layer = net[l]; + std::vector 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; + } + if (l != net.size() - 1) { + for (size_t i = 0; i < z.size(); ++i) + z[i] = silu(z[i]); + } + x = z; + } + return x; + } + // backward pass to compute gradients of the loss w.r.t. network parameters using chain rule void ScoreBasedDiffusionModel::backward( const std::vector& gradOutput @@ -407,7 +900,17 @@ namespace mu2e { << ", expected " << dim_; } - double s = sigma(t); + double alpha_bar, s; + if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) { + s= sigma(t); + // Clamp to >= 0 (matching alphabar()): at t=1 sigma(1)=logSigMax can round to 1+eps, + // making 1 - s*s a tiny negative, whose sqrt() below would be NaN and poison the + // whole sample (e.g. partial-reverse at t0=1.0). + alpha_bar = std::max(0.0, 1.0 - s * s); + } else { + alpha_bar = alphabar(t); + s = std::sqrt(1.0 - alpha_bar); // i.e. sigma(t), avoid recalculating alpha_bar inside the function for efficiency + } // Generate Gaussian noise vector eps of dimension dim_ using the external random engine. eps.resize(dim_); @@ -416,7 +919,7 @@ namespace mu2e { for (int i = 0; i < dim_; ++i) { eps[i] = randGaussQ_.fire(); // Gaussian N(0,1) - xt[i] = x[i] + s * eps[i]; + xt[i] = std::sqrt(alpha_bar) * x[i] + s * eps[i]; } return xt; @@ -425,7 +928,8 @@ namespace mu2e { // Compute Mean Squared Error per dimension between predicted score and target score. double ScoreBasedDiffusionModel::computeLoss( const std::vector& score, - const std::vector& target + const std::vector& target, + double weight // use weighted loss to prevent sigma(t) at tiny t from blowing up and dominating the training. ) const { if (score.size() != static_cast(dim_)) { @@ -443,7 +947,7 @@ namespace mu2e { for (int i = 0; i < dim_; ++i) { double d = score[i] - target[i]; - loss += d * d; + loss += weight * d * d; } return loss / dim_; @@ -467,13 +971,17 @@ namespace mu2e { } norm = std::sqrt(norm); + totalClipChecks_++; // If the norm is below the threshold, no clipping is needed. if (norm <= maxNorm) { + clipScaleAccum_ += 1.0; return; } + clipCount_++; // Scale down the gradients to have the specified maximum norm. double scale = maxNorm / norm; + clipScaleAccum_ += scale; for (auto& layer : network_) { for (auto& row : layer.gradW) @@ -490,9 +998,29 @@ namespace mu2e { // 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::updateEMANetwork() + { + for (size_t l = 0; l < network_.size(); ++l) { + for (size_t i = 0; i < network_[l].W.size(); ++i) { + for (size_t j = 0; j < network_[l].W[i].size(); ++j) + emaNetwork_[l].W[i][j] = emaNetworkDecay_ * emaNetwork_[l].W[i][j] + + (1.0 - emaNetworkDecay_) * network_[l].W[i][j]; + emaNetwork_[l].b[i] = emaNetworkDecay_ * emaNetwork_[l].b[i] + + (1.0 - emaNetworkDecay_) * network_[l].b[i]; + } + } + } + void ScoreBasedDiffusionModel::train( const std::vector& data, - int epochs + int epochs, + int samplesDrawnPerEpoch, + bool biasLowSigma, + double tLowBound, + double tFocusLow, + double tFocusHigh, + double tFocusFraction, + const std::vector& peakWindows ) { // Check that the network is properly initialized before training. @@ -501,6 +1029,52 @@ namespace mu2e { << "Network is not properly initialized"; } + if (tLowBound >= 1.0 || tLowBound < 0.0) { + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid tLowBound: must be in [0,1)"; + } + + if (tFocusFraction < 0.0 || tFocusFraction > 1.0) { + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid tFocusFraction: must be in [0,1]"; + } + if (tFocusFraction > 0.0 && + (tFocusLow < 0.0 || tFocusHigh > 1.0 || tFocusLow >= tFocusHigh)) { + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid t focus window: require 0 <= tFocusLow < tFocusHigh <= 1 (got [" + << tFocusLow << ", " << tFocusHigh << "])"; + } + + // Peak importance sampling is enabled by a non-empty window list; validate up front. + const bool peakRequested = !peakWindows.empty(); + if (peakRequested) { + double sumGMax = 0.0; + for (size_t k = 0; k < peakWindows.size(); ++k) { + const PeakWindow& pw = peakWindows[k]; + if (pw.dim < 0 || pw.dim >= dim_) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid peakWindows[" << k << "].dim " << pw.dim << ": must be in [0, " << dim_ << ")"; + if (pw.high <= pw.low) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid peakWindows[" << k << "] window: require low < high (got [" + << pw.low << ", " << pw.high << "])"; + if (pw.gMax <= 0.0 || pw.gMax >= 1.0) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid peakWindows[" << k << "].gMax " << pw.gMax << ": must be in (0, 1)"; + if (pw.sigma0 <= 0.0) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid peakWindows[" << k << "].sigma0 " << pw.sigma0 << ": must be > 0"; + if (pw.alpha < 0.0 || pw.alpha > 1.0) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Invalid peakWindows[" << k << "].alpha " << pw.alpha << ": must be in [0, 1]"; + sumGMax += pw.gMax; + } + if (sumGMax >= 1.0) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Sum of peak gMax (" << sumGMax << ") must be < 1 to leave sampling probability " + << "for the out-of-window pool."; + } + const double eps_safe = 1e-12; const size_t N = data.size(); @@ -525,21 +1099,134 @@ namespace mu2e { } } + // Determine how many samples to DRAW per epoch. An epoch is a fixed quantum of + // optimization work (activeN/batchSize gradient steps), which is the meaningful + // progress metric — not "one pass over the data". samplesDrawnPerEpoch == 0 means + // one full pass (activeN = N). Otherwise activeN is exactly the requested count, + // EVEN IF it exceeds N: the non-peak draw path cycles the shuffled index list with + // reshuffle-on-wrap (samples reused within the epoch, fresh noise per draw). Warn + // once when wrapping, since the epoch no longer equals a single dataset pass. + const size_t activeN = (samplesDrawnPerEpoch > 0) + ? static_cast(samplesDrawnPerEpoch) + : N; + if (samplesDrawnPerEpoch > 0 && static_cast(samplesDrawnPerEpoch) > N) { + mf::LogWarning("ScoreBasedDiffusionModel::train") + << "samplesDrawnPerEpoch=" << samplesDrawnPerEpoch << " exceeds the dataset size N=" + << N << "; the dataset will be cycled (with reshuffle-on-wrap) to draw the full count. " + << "Samples are reused within each epoch, so an epoch no longer equals one dataset pass " + << "(~" << (static_cast(samplesDrawnPerEpoch) / static_cast(N)) + << " passes/epoch). Planner budgets count these epochs."; + } + // Create index vector once std::vector indices(N); for (size_t i = 0; i < N; ++i) indices[i] = i; + // Fisher–Yates shuffle of an index vector using RandFlat (reproducible via the seed service). + auto shuffleVec = [this](std::vector& v) { + for (size_t i = v.size(); i > 1; --i) { + size_t j = static_cast(randFlat_.fire() * i); + j = std::min(j, i - 1); + std::swap(v[i - 1], v[j]); + } + }; + + // Peak importance sampling: partition the FULL data set once into per-window in-window pools + // P[k] and a single out-of-window pool Q, assigning each event to the FIRST window it matches + // (so overlapping windows stay disjoint), and record each empirical in-window fraction f[k]. + // Windows whose pool is empty are dropped. Each pool is reshuffled per epoch and drawn without + // replacement via a cursor (reshuffled again if the cursor wraps within an epoch). The number + // of in-window draws per epoch is ~ (mean g_eff_k) * activeN, which for small g and a large + // rare pool is LESS than |P[k]| — so not every in-window event is seen each epoch; the + // per-epoch reshuffle makes coverage even in expectation across epochs. Disabled if no peak + // pool survives or Q is empty (windows cover everything). + bool peakSampling = peakRequested; + std::vector pk; // surviving (non-empty) windows + std::vector> P; // in-window pools, aligned with pk + std::vector f; // in-window fractions, aligned with pk + std::vector pP; // per-pool cursors, aligned with pk + std::vector Q; // out-of-window pool + size_t pQ = 0; + double sumF = 0.0; // total in-window fraction + if (peakSampling) { + const size_t K = peakWindows.size(); + // Window edges are given in transformed (pre-z-score) units; convert to the z-scored + // space the data lives in, once, using the stored per-dimension mean/stdev. + std::vector zLow(K), zHigh(K); + for (size_t k = 0; k < K; ++k) { + zLow[k] = normalizeCoord(peakWindows[k].dim, peakWindows[k].low); + zHigh[k] = normalizeCoord(peakWindows[k].dim, peakWindows[k].high); + } + std::vector> P0(K); + for (size_t i = 0; i < N; ++i) { + int which = -1; + for (size_t k = 0; k < K; ++k) { + double v = data[i].x[peakWindows[k].dim]; + if (v >= zLow[k] && v < zHigh[k]) { which = static_cast(k); break; } + } + if (which >= 0) P0[which].push_back(i); + else Q.push_back(i); + } + for (size_t k = 0; k < K; ++k) { + if (P0[k].empty()) { + mf::LogWarning("ScoreBasedDiffusionModel::train") + << "Peak window " << k << " [" << peakWindows[k].low << ", " << peakWindows[k].high + << ") on dim " << peakWindows[k].dim << " matched 0 events; dropping it."; + continue; + } + pk.push_back(peakWindows[k]); + f.push_back(static_cast(P0[k].size()) / static_cast(N)); + P.push_back(std::move(P0[k])); + } + if (pk.empty() || Q.empty()) { + mf::LogWarning("ScoreBasedDiffusionModel::train") + << "Peak importance sampling disabled: " << pk.size() << " surviving window(s), " + << "out-of-window pool size " << Q.size() << " (need both non-empty)."; + peakSampling = false; + } else { + pP.assign(pk.size(), 0); + for (double fk : f) sumF += fk; + // geff[k] = max(f[k], gMax[k]*exp(-sigma^2/2sigma0^2)); the exp factor is in (0,1], + // so geff[k] <= gMax[k] AS LONG AS f[k] <= gMax[k], and then sum geff <= sum gMax < 1 + // (validated above). If a window's empirical fraction f[k] exceeds its gMax, that + // window's geff is pinned to f[k] > gMax[k], which can push sum geff over 1 and break + // the cumulative selection / importance weights. Warn per offending window so the + // configuration can be fixed (raise gMax or narrow the window). + for (size_t k = 0; k < pk.size(); ++k) { + if (f[k] > pk[k].gMax) + mf::LogWarning("ScoreBasedDiffusionModel::train") + << "Peak window " << k << " on dim " << pk[k].dim + << " has empirical fraction f=" << f[k] << " > gMax=" << pk[k].gMax + << "; its g_eff floors at f (not gMax), which can make sum g_eff exceed 1. " + << "Raise gMax above f or narrow the window."; + } + std::ostringstream oss; + oss << "Peak importance sampling enabled: " << pk.size() + << " window(s), total in-window fraction " << sumF + << ". Per window [dim, low, high (transformed); [zLow, zHigh] (z-score); f, gMax, sigma0, alpha]:"; + for (size_t k = 0; k < pk.size(); ++k) + oss << "\n [" << pk[k].dim << ", " << pk[k].low << ", " << pk[k].high + << "; [" << normalizeCoord(pk[k].dim, pk[k].low) << ", " << normalizeCoord(pk[k].dim, pk[k].high) + << "]; " << f[k] << ", " << pk[k].gMax << ", " << pk[k].sigma0 << ", " << pk[k].alpha << "]"; + mf::LogInfo("ScoreBasedDiffusionModel::train") << oss.str(); + } + } + std::vector geff(pk.size(), 0.0); // reusable per-draw g_eff buffer + size_t pIdx = 0; // non-peak wrapping cursor into `indices` (reset + reshuffle on wrap) + // 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(randFlat_.fire() * (i + 1)); - j = std::min(j, i); // Ensure j is within bounds - std::swap(indices[i], indices[j]); + // Reshuffle the sampling pool(s) each epoch. Peak sampling cycles each in-window pool and + // the out-of-window pool (cursors reset to 0); otherwise the single full-data index list + // is shuffled. + if (peakSampling) { + for (size_t k = 0; k < P.size(); ++k) { shuffleVec(P[k]); pP[k] = 0; } + shuffleVec(Q); pQ = 0; + } else { + shuffleVec(indices); pIdx = 0; } // Counter for number of samples processed in the epoch (used for averaging loss). Avoid using N directly to @@ -548,32 +1235,122 @@ namespace mu2e { int batchCounter = 0; double epochLoss = 0.0; - - // iterate over the shuffled data samples - for (size_t idx = 0; idx < N; ++idx) + std::vector epochDimLoss(dim_, 0.0); + // Peak-window loss accumulators: mean unweighted per-event squared residual over + // draws that landed in any peak window this epoch. Surfaced via getLastEpochPeakLoss() + // so the curriculum planner can plateau on the feature-region fit instead of the + // aggregate loss. Accumulated regardless of useDimWeightController_. + double epochPeakLoss = 0.0; + long epochPeakCount = 0; + + // iterate over the drawn samples (samplesDrawnPerEpoch if specified; cycles the data when it exceeds N) + for (size_t idx = 0; idx < activeN; ++idx) { - const auto& sample = data[indices[idx]]; + // Sample diffusion time. If tFocusFraction > 0, this sample draws t uniformly from the + // focus window [tFocusLow, tFocusHigh] with probability tFocusFraction, concentrating + // gradient steps on a target sigma band while the remaining samples keep full coverage. + // Otherwise: if tLowBound is set, we scale the uniform random number to be in + // [tLowBound, 1] instead of [0,1] to focus training on later diffusion steps with higher + // noise levels, which can improve stability and convergence in some cases. If biasLowSigma + // is true, we apply a quadratic transformation to the sampled time to bias it towards smaller + // values (less noise), which can help the model learn the score function more accurately at + // early diffusion steps where the signal is stronger. Not meant to be used together. - // Sample diffusion time double t = randFlat_.fire(); + if (tFocusFraction > 0.0) { + // Focus mode: exactly tFocusFraction of samples land in [tFocusLow, tFocusHigh); + // the remainder are drawn uniformly from the COMPLEMENT [0,tFocusLow) U [tFocusHigh,1) + // so the realized in-window fraction equals tFocusFraction independent of window width. + // (tLowBound / biasLowSigma are the non-focus sampling modes and do not apply here.) + if (randFlat_.fire() < tFocusFraction) { + t = tFocusLow + (tFocusHigh - tFocusLow) * t; // uniform inside the window + } else { + double comp = 1.0 - (tFocusHigh - tFocusLow); // length of the complement + double u = t * comp; // uniform in [0, comp) + t = (u < tFocusLow) ? u : u + (tFocusHigh - tFocusLow); // skip over the window + } + } else { + if (tLowBound > 0.0) { + t = tLowBound + (1.0 - tLowBound) * t; // scale t to [tLowBound, 1] + } + if (biasLowSigma) { + t = t * t; // focus on smaller t value (smaller sigma) + } + } + + // Select the training example for this draw and compute its importance weight. + // Default: the next entry from the per-epoch shuffled full-data index list (wIS=1). + // Peak sampling: per-window draw probability g_eff_k(sigma(t)); one cumulative draw + // selects a window pool or the out-of-window pool, each via a without-replacement + // cursor. wIS reweights the loss so each window is unbiased at alpha_k=1 and a + // deliberate up-weight for alpha_k<1, with the continuum (Q) kept unbiased. + double wIS = 1.0; + size_t chosenIdx; + bool inPeak = false; // true if this draw landed in a peak window (peak-window loss metric) + if (peakSampling) { + double s_t = sigma(t); + double sumGeff = 0.0; + for (size_t k = 0; k < pk.size(); ++k) { + geff[k] = std::max(f[k], pk[k].gMax * std::exp(-(s_t * s_t) / (2.0 * pk[k].sigma0 * pk[k].sigma0))); + sumGeff += geff[k]; + } + // sum gMax < 1 is validated up front and the exp factor is in (0,1], so this can + // only trip when some window has f[k] > gMax[k] (warned at setup). If it does, the + // out-of-window pool Q becomes unreachable and the importance weights stop being a + // valid estimator, so fail loudly rather than train on a silently biased sample. + if (sumGeff >= 1.0) + throw cet::exception("ScoreBasedDiffusionModel::train") + << "Peak sampling sum g_eff (" << sumGeff << ") >= 1 at sigma=" << s_t + << "; a window's empirical fraction exceeds its gMax. Raise gMax above f " + << "or narrow the window(s)."; + double u = randFlat_.fire(); + int sel = -1; // -1 = out-of-window pool Q + double cum = 0.0; + for (size_t k = 0; k < pk.size(); ++k) { + cum += geff[k]; + if (u < cum) { sel = static_cast(k); break; } + } + if (sel >= 0) { + inPeak = true; + chosenIdx = P[sel][pP[sel]++]; + if (pP[sel] >= P[sel].size()) { shuffleVec(P[sel]); pP[sel] = 0; } + wIS = std::pow(f[sel] / geff[sel], pk[sel].alpha); + } else { + // Out-of-window pool (reachable because sum gMax_k < 1 keeps 1 - sumGeff > 0). + chosenIdx = Q[pQ++]; + if (pQ >= Q.size()) { shuffleVec(Q); pQ = 0; } + wIS = (1.0 - sumF) / (1.0 - sumGeff); // alpha = 1: unbiased continuum + } + } else { + // Draw the next entry from the per-epoch shuffled full-data list via a + // wrapping cursor. When activeN <= N this is a plain subset (distinct + // samples). When activeN > N the cursor wraps and the list is reshuffled, + // so the requested count is drawn while reusing samples (fresh noise per + // draw); pIdx tracks the position independently of the epoch draw counter. + if (pIdx >= N) { shuffleVec(indices); pIdx = 0; } + chosenIdx = indices[pIdx++]; + } + const auto& sample = data[chosenIdx]; + // container for noise vector eps std::vector eps; auto xt = addNoise(sample.x,t,eps); - double s = sigma(t); + double s = std::max(sigma(t), eps_safe); // add eps_safe to prevent division by zero in case of very small sigma + double weight = std::pow(s, lossWeightPower_); // sigma(t)^power as the weight. Quadratic power good for convergence but missing details - // The target score is the negative of the noise scaled by the noise standard deviation, i.e., -eps/sigma(t). + // Target depends on predictionTarget_ (SCORE/EPS/V); see trainingTargetComponent. + // Epsilon/v prediction prevents 1/sigma blow-up at small t. + double a = std::sqrt(std::max(0.0, alphabar(t))); // signal coeff for v-prediction std::vector 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 + target[i] = trainingTargetComponent(eps[i], sample.x[i], s, a); } - // Prepare input vector for the network by concatenating the noisy sample xt with the diffusion time t. - std::vector input = xt; - input.insert(input.end(), sample.cond.begin(), sample.cond.end()); - input.push_back(t); + // Prepare input vector for the network from the noisy sample xt, condition, and time. + auto input = buildNetworkInput(xt, sample.cond, t); - // Check that input dimension matches expected dimension (dim_ + conditionDim_ + 1) + // Check that input dimension matches the network's expected input dimension if (input.size() != network_[0].W[0].size()) { throw cet::exception("ScoreBasedDiffusionModel::train") << "Training input dimension mismatch: got " << input.size() @@ -584,13 +1361,28 @@ namespace mu2e { auto score = forward(input); // Compute the loss (Mean Squared Error) per dimension between the predicted score and the target score. - double loss = computeLoss(score, target); - - // Compute the gradient of the loss w.r.t. the predicted score - // 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). + double loss = computeLoss(score, target, weight); + + // Compute the gradient of the loss w.r.t. the predicted score. + // dimWeights_[i] rescales each dimension's gradient; normalized to mean=1 so overall + // gradient scale is preserved. Accumulate raw squared residuals for the controller EMA. + // wIS (peak importance weight, 1.0 when peak sampling is off) reweights this + // sample's gradient. The per-dim controller accumulates the RAW squared residual + // (unweighted) so it still balances dimensions by their intrinsic difficulty. std::vector grad(dim_); + double sampleSqResid = 0.0; // sum_i residual_i^2 for this sample (peak-window metric) for (int i = 0; i < dim_; ++i) { - grad[i] = 2.0 * (score[i] - target[i]) / dim_; + double residual = score[i] - target[i]; + if (useDimWeightController_) + epochDimLoss[i] += residual * residual; + sampleSqResid += residual * residual; + grad[i] = 2.0 * weight * wIS * dimWeights_[i] * residual / dim_; + } + // Peak-window loss: accumulate the UNWEIGHTED mean squared residual (ignore wIS) + // for in-window draws, so the planner can track raw fit quality in the feature region. + if (inPeak) { + epochPeakLoss += sampleSqResid / dim_; + ++epochPeakCount; } // 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. @@ -618,10 +1410,11 @@ namespace mu2e { } else { updateWeights(learningRate_); } + if (useEMANetwork_) updateEMANetwork(); batchCounter = 0; } - epochLoss += loss; // Accumulate loss for monitoring. + epochLoss += loss * wIS; // Accumulate importance-weighted loss for monitoring (unbiased population loss at alpha=1). n++; } @@ -644,22 +1437,441 @@ namespace mu2e { } else { updateWeights(learningRate_); } + if (useEMANetwork_) updateEMANetwork(); + } + + // Update per-dimension EMA loss and recompute gradient weights + if (useDimWeightController_ && n > 0) { + double emaSum = 0.0; + for (int i = 0; i < dim_; ++i) { + dimLossEMA_[i] = dimWeightEMADecay_ * dimLossEMA_[i] + + (1.0 - dimWeightEMADecay_) * (epochDimLoss[i] / n); + emaSum += dimLossEMA_[i]; + } + double emaMean = emaSum / dim_; + if (emaMean > 0.0) { + // The raw ratio is unbounded and self-reinforcing, and a dimension + // driven far below 1 stops receiving useful gradient entirely, so the + // result is bounded (see clampDimWeights / the header note). No warning + // here: the controller hitting its bounds during a phase transient is + // expected, and warning every epoch would be noise. + for (int i = 0; i < dim_; ++i) + dimWeights_[i] = std::clamp(dimLossEMA_[i] / emaMean, + kDimWeightMin_, kDimWeightMax_); + } + std::ostringstream woss; + woss << "Epoch " << e << " dimWeights: ["; + for (int i = 0; i < dim_; ++i) { + woss << dimWeights_[i]; + if (i < dim_ - 1) woss << ", "; + } + woss << "]"; + mf::LogInfo("ScoreBasedDiffusionModel::train") << woss.str(); } epochLoss /= n; - mf::LogInfo("ScoreBasedDiffusionModel::train") << "Epoch " << e << " Loss=" << epochLoss; + double clipRatio = (totalClipChecks_ > 0) ? double(clipCount_) / totalClipChecks_ : 0.0; + double avgClipScale = (totalClipChecks_ > 0) ? clipScaleAccum_ / totalClipChecks_ : 1.0; + + // Finalize the peak-window loss for this epoch (NaN if no in-window draws occurred, + // e.g. peak sampling disabled). Transient per-run state read by the curriculum planner. + lastEpochPeakCount_ = epochPeakCount; + lastEpochPeakLoss_ = (epochPeakCount > 0) + ? epochPeakLoss / static_cast(epochPeakCount) + : std::numeric_limits::quiet_NaN(); + + mf::LogInfo("ScoreBasedDiffusionModel::train") << "Epoch " << e << " Loss=" << epochLoss + << " ClipRatio=" << clipRatio << " AvgClipScale=" << avgClipScale + << " PeakLoss=" << lastEpochPeakLoss_ << " PeakCount=" << lastEpochPeakCount_; epochLosses_.push_back(epochLoss); } } + double ScoreBasedDiffusionModel::evaluateAverageLoss( + const std::vector& data, + int subsetSize, + bool biasLowSigma, + double tLowBound, + double tFocusLow, + double tFocusHigh, + double tFocusFraction) + { + const size_t N = data.size(); + if (N == 0) return std::numeric_limits::quiet_NaN(); + + // Evaluate over a subset (first `subsetSize` after a shuffle) or the full set. + const size_t activeN = (subsetSize > 0 && static_cast(subsetSize) < N) + ? static_cast(subsetSize) : N; + std::vector indices(N); + for (size_t i = 0; i < N; ++i) indices[i] = i; + for (size_t i = N - 1; i > 0; --i) { + size_t j = static_cast(randFlat_.fire() * (i + 1)); + j = std::min(j, i); + std::swap(indices[i], indices[j]); + } + + const double eps_safe = 1e-12; + double sumLoss = 0.0; + int n = 0; + for (size_t idx = 0; idx < activeN; ++idx) { + const auto& sample = data[indices[idx]]; + // Same t sampling as train() (focus window / complement, or tLowBound / biasLowSigma). + double t = randFlat_.fire(); + if (tFocusFraction > 0.0) { + if (randFlat_.fire() < tFocusFraction) { + t = tFocusLow + (tFocusHigh - tFocusLow) * t; + } else { + double comp = 1.0 - (tFocusHigh - tFocusLow); + double u = t * comp; + t = (u < tFocusLow) ? u : u + (tFocusHigh - tFocusLow); + } + } else { + if (tLowBound > 0.0) t = tLowBound + (1.0 - tLowBound) * t; + if (biasLowSigma) t = t * t; + } + + std::vector eps; + auto xt = addNoise(sample.x, t, eps); + double s = std::max(sigma(t), eps_safe); + double weight = std::pow(s, lossWeightPower_); + + double a = std::sqrt(std::max(0.0, alphabar(t))); + std::vector target(dim_); + for (int i = 0; i < dim_; ++i) + target[i] = trainingTargetComponent(eps[i], sample.x[i], s, a); + + auto input = buildNetworkInput(xt, sample.cond, t); + auto score = forward(input); // forward only — no backward / optimizer / EMA / dimLossEMA + sumLoss += computeLoss(score, target, weight); + ++n; + } + return (n > 0) ? sumLoss / n : std::numeric_limits::quiet_NaN(); + } + + SBDMEpsLossSample ScoreBasedDiffusionModel::evalEpsLossSample( + const std::vector& xNorm, + const std::vector& condition, + bool useEMANetworkIfAvailable) + { + if (xNorm.size() != static_cast(dim_)) + throw cet::exception("ScoreBasedDiffusionModel::evalEpsLossSample") + << "State dimension mismatch: got " << xNorm.size() << ", expected " << dim_; + if (condition.size() != static_cast(conditionDim_)) + throw cet::exception("ScoreBasedDiffusionModel::evalEpsLossSample") + << "Conditioning dimension mismatch: got " << condition.size() << ", expected " << conditionDim_; + + const double eps_safe = 1e-12; + // Draw t uniformly, clamped away from the endpoints (sigma is ill-conditioned at t->0,1). + double t = std::min(1.0 - 1e-3, std::max(1e-3, randFlat_.fire())); + double s = std::max(sigma(t), eps_safe); + + std::vector eps; + auto xt = addNoise(xNorm, t, eps); + auto input = buildNetworkInput(xt, condition, t); + const auto& net = (useEMANetworkIfAvailable && useEMANetwork_) ? emaNetwork_ : network_; + auto out = forwardInference(input, net); + + SBDMEpsLossSample res; + res.t = t; + res.sigma = s; + res.perDimLoss.resize(dim_); + res.perDimNativeLoss.resize(dim_); + double a = std::sqrt(std::max(0.0, alphabar(t))); + for (int i = 0; i < dim_; ++i) { + // Eps-style loss: recover eps_hat in all modes so the lens is common across targets. + double epsHat = epsHatFromOutput(out[i], xt[i], s, a); + double dEps = epsHat - eps[i]; + res.perDimLoss[i] = dEps * dEps; + // Native-target loss: prediction vs target on the model's own training target. For SCORE + // the prediction is scoreFromOutput (matching how the sampler consumes the output); for + // EPS/V the prediction is the raw output. The target reuses trainingTargetComponent, so + // EPS's native loss equals the eps-style loss above. + double nativePred = (predictionTarget_ == PredictionTarget::SCORE) + ? scoreFromOutput(out[i], xt[i], s, a) : out[i]; + double nativeTgt = trainingTargetComponent(eps[i], xNorm[i], s, a); + double dNat = nativePred - nativeTgt; + res.perDimNativeLoss[i] = dNat * dNat; + } + return res; + } + + std::vector ScoreBasedDiffusionModel::firstLayerBlockMagnitudes( + const std::vector& data, + int nSamples, + std::vector& perDimLossOut) + { + // Build the input-feature-block layout, matching buildNetworkInput's column order: + // [raw x (dim_ cols, 1 per coord)] [Fourier(x): per coord j, inputEmbeddingDims_[j] cols] + // [raw cond (conditionDim_ cols)] [Fourier(cond): per coord j, conditionEmbeddingDims_[j] cols] + // [time: 1 + timeEmbeddingDim_ cols] + std::vector blocks; + std::vector blockStart; // input column where each block begins (parallel to blocks) + int col = 0; + auto pushBlock = [&](const std::string& name, int kind, int coord, int nCols) { + if (nCols <= 0) return; // skip dims with no Fourier columns / empty blocks + SBDMFeatureBlockMagnitude b; + b.name = name; b.kind = kind; b.coord = coord; b.nCols = nCols; + blocks.push_back(b); + blockStart.push_back(col); + col += nCols; + }; + for (int j = 0; j < dim_; ++j) pushBlock("raw_state[" + std::to_string(j) + "]", 0, j, 1); + for (int j = 0; j < dim_; ++j) pushBlock("fourier_state[" + std::to_string(j) + "]", 1, j, inputEmbeddingDims_[j]); + for (int j = 0; j < conditionDim_; ++j) pushBlock("raw_cond[" + std::to_string(j) + "]", 2, j, 1); + for (int j = 0; j < conditionDim_; ++j) pushBlock("fourier_cond[" + std::to_string(j) + "]", 3, j, conditionEmbeddingDims_[j]); + // timeEmbed() emits raw t (1 col) followed by timeEmbeddingDim_ sin/cos features; split them + // into separate blocks so the time Fourier embedding magnitude is visible on its own. + pushBlock("raw_time", 4, -1, 1); + pushBlock("fourier_time", 5, -1, timeEmbeddingDim_); + + const int inputSize = static_cast(network_[0].W[0].size()); + if (col != inputSize) + throw cet::exception("ScoreBasedDiffusionModel::firstLayerBlockMagnitudes") + << "Feature-block layout (" << col << " cols) does not match first-layer input size (" + << inputSize << ")"; + + // Weight L2 per block (over all output rows of the first layer). + const auto& W0 = network_[0].W; + for (size_t b = 0; b < blocks.size(); ++b) { + double sw = 0.0; + for (const auto& row : W0) + for (int c = blockStart[b]; c < blockStart[b] + blocks[b].nCols; ++c) + sw += row[c] * row[c]; + blocks[b].weightL2 = std::sqrt(sw); + } + + // Accumulate the training gradient over nSamples draws with NO optimizer step. + for (auto& layer : network_) { + for (auto& row : layer.gradW) std::fill(row.begin(), row.end(), 0.0); + std::fill(layer.gradb.begin(), layer.gradb.end(), 0.0); + } + perDimLossOut.assign(dim_, 0.0); + + const size_t N = data.size(); + const double eps_safe = 1e-12; + const size_t nUse = (nSamples > 0 && static_cast(nSamples) < N) + ? static_cast(nSamples) : N; + const size_t stride = std::max(1, N / std::max(1, nUse)); + size_t used = 0; + for (size_t k = 0; k < N && used < nUse; k += stride, ++used) { + const auto& sample = data[k]; + double t = std::min(1.0 - 1e-3, std::max(1e-3, randFlat_.fire())); + double s = std::max(sigma(t), eps_safe); + double weight = std::pow(s, lossWeightPower_); + + std::vector eps; + auto xt = addNoise(sample.x, t, eps); + double a = std::sqrt(std::max(0.0, alphabar(t))); + std::vector target(dim_); + for (int i = 0; i < dim_; ++i) + target[i] = trainingTargetComponent(eps[i], sample.x[i], s, a); + + auto input = buildNetworkInput(xt, sample.cond, t); + auto score = forward(input); + std::vector grad(dim_); + for (int i = 0; i < dim_; ++i) { + double residual = score[i] - target[i]; + perDimLossOut[i] += residual * residual; // fresh, unweighted per-output-dim loss + grad[i] = 2.0 * weight * dimWeights_[i] * residual / dim_; // same gradient train() would form + } + backward(grad); + } + const double invUsed = (used > 0) ? 1.0 / static_cast(used) : 0.0; + for (int i = 0; i < dim_; ++i) perDimLossOut[i] *= invUsed; + + // Gradient L2 per block (mean per sample), then leave the gradient buffers zeroed. + const auto& G0 = network_[0].gradW; + for (size_t b = 0; b < blocks.size(); ++b) { + double sg = 0.0; + for (const auto& row : G0) + for (int c = blockStart[b]; c < blockStart[b] + blocks[b].nCols; ++c) + sg += row[c] * row[c]; + blocks[b].gradL2 = std::sqrt(sg) * invUsed; + } + for (auto& layer : network_) { + for (auto& row : layer.gradW) std::fill(row.begin(), row.end(), 0.0); + std::fill(layer.gradb.begin(), layer.gradb.end(), 0.0); + } + return blocks; + } + void ScoreBasedDiffusionModel::saveModel( + const std::string& filename, + int basisTag + ) + { + // Write binary + std::ofstream out(filename, std::ios::binary); + if (!out) { + throw cet::exception("ScoreBasedDiffusionModel::saveModel") + << "Cannot open file " << filename; + } + + auto wI32 = [&](int32_t v){ out.write(reinterpret_cast(&v), 4); }; + auto wU32 = [&](uint32_t v){ out.write(reinterpret_cast(&v), 4); }; + auto wU64 = [&](uint64_t v){ out.write(reinterpret_cast(&v), 8); }; + auto wF64 = [&](double v){ out.write(reinterpret_cast(&v), 8); }; + auto wVec = [&](const std::vector& v){ + wU64(static_cast(v.size())); + if (!v.empty()) out.write(reinterpret_cast(v.data()), + static_cast(v.size() * sizeof(double))); + }; + auto wMat = [&](const std::vector>& m){ + // Writes a matrix as outSize rows, each row preceded by its column count. + // Row dimensions may differ (though in practice they don't). + wU64(static_cast(m.size())); + for (const auto& row : m) wVec(row); + }; + + // Magic + version + // Version history: 1 = original format; 2 = adds inputEmbeddingDim and + // conditionEmbeddingDim after timeEmbeddingDim; 3 = appends a 4-byte "ENDM" + // end-of-stream sentinel that the loader verifies to detect truncation; + // 4 = stores the batch-size-independent EMA decay *base* (emaNetworkDecayBase_) + // in the emaNetworkDecay field instead of the already-rescaled per-step + // effective value. Versions 1-3 stored the effective value, which the loader + // re-rescaled on construction, compounding the batch-size exponent on every + // round-trip and driving the decay toward zero; + // 5 = replaces the two scalar embedding fields with count-prefixed per-dimension + // vectors (inputEmbeddingDims_ of length dim_, conditionEmbeddingDims_ of length + // conditionDim_). Versions 2-4 stored a single scalar broadcast to all dims. + // 6 = stores the prediction target as an int32 enum (0=SCORE,1=EPS,2=V) in the + // field that versions 1-5 used for the epsPrediction bool (0/1). The enum values + // 0=SCORE,1=EPS were chosen to coincide with the old bool false/true, so a v<=6 + // loader maps the legacy bool directly with no behavior change. + // 7 = appends an opaque int32 basisTag (after diffusionSteps_, before the network + // weights). This class never interprets it; it is an application-level marker + // (see saveModel/basisTag()). v<=6 files have no such field and load it as 0. + // 8 = appends an int32 categoricalConditionDim (immediately after basisTag): the + // index of the one condition coordinate holding a class label, or -1 for none. + // v<=7 files predate it and load as -1, i.e. every condition dim z-scored, which + // is what those models were trained with. + // 9 = appends an int32 pdgId and a count-prefixed list of (int32 id, double value) + // build constants (both immediately after categoricalConditionDim). Like basisTag + // these are opaque here; the caller assigns the ids and decides what a disagreement + // means. v<=8 files load pdgId as 0 ("not recorded") and the list as empty. + out.write("SBDM", 4); + wU32(9u); + + // Model architecture & hyper-parameters + wI32(static_cast(dim_)); + wI32(static_cast(conditionDim_)); + wI32(static_cast(timeEmbeddingDim_)); + // Per-dimension Fourier embedding depths, each count-prefixed (version >= 5). + wI32(static_cast(inputEmbeddingDims_.size())); + for (int e : inputEmbeddingDims_) wI32(static_cast(e)); + wI32(static_cast(conditionEmbeddingDims_.size())); + for (int e : conditionEmbeddingDims_) wI32(static_cast(e)); + wI32(static_cast(hidden_)); + wI32(static_cast(layers_)); + wI32(optimizerType_ == OptimizerType::ADAM ? 0 : 1); + wF64(adamBeta1_); wF64(adamBeta2_); wF64(adamEps_); + int32_t schedIdx = (noiseScheduleType_ == NoiseScheduleType::LINEAR) ? 0 + : (noiseScheduleType_ == NoiseScheduleType::COSINE) ? 1 : 2; + wI32(schedIdx); + wF64(betaMin_); wF64(betaMax_); wF64(cosineOffset_); + wF64(logSigMin_); wF64(logSigMax_); + wI32(static_cast(predictionTarget_)); // 0=SCORE,1=EPS,2=V (was epsPrediction bool) + wF64(lossWeightPower_); + wI32(static_cast(batchSize_)); + wF64(gradientClipThreshold_); + wF64(learningRate_); + wI32(useDimWeightController_ ? 1 : 0); + wF64(dimWeightEMADecay_); + wI32(useEMANetwork_ ? 1 : 0); + // Store the batch-size-independent base (see version-4 note above), NOT the + // rescaled per-step effective decay. loadModel re-applies the batch-size + // rescaling exactly once when it reconstructs the model. + wF64(emaNetworkDecayBase_); + wI32(static_cast(diffusionSteps_)); + + // Opaque application-level basis tag (format v7+). Persisted verbatim; this + // class assigns it no meaning. Also remember it on the in-memory object so a + // subsequent basisTag() query is consistent with what was just written. + wI32(static_cast(basisTag)); + basisTag_ = basisTag; + + // Index of the class-label condition dim, or -1 for none (format v8+). Unlike + // basisTag this is not a caller argument: it is model state that normalizeData() + // has already acted on, so it must be persisted for the generate side to feed the + // label at the same scale the network was trained on. + wI32(static_cast(categoricalConditionDim_)); + + // Opaque caller-owned identity and build constants (format v9+). Set through + // setPdgId/setBuildConstants before saving; 0 and an empty list mean "not recorded". + // Count-prefixed like the embedding-dim vectors above. + wI32(static_cast(pdgId_)); + wU32(static_cast(buildConstants_.size())); + for (const auto& kv : buildConstants_) { + wI32(static_cast(kv.first)); + wF64(kv.second); + } + + // Network weights + wU32(static_cast(network_.size())); + for (const auto& layer : network_) { + wU32(static_cast(layer.W.size())); + wU32(static_cast(layer.W[0].size())); + for (const auto& row : layer.W) + out.write(reinterpret_cast(row.data()), + static_cast(row.size() * sizeof(double))); + out.write(reinterpret_cast(layer.b.data()), + static_cast(layer.b.size() * sizeof(double))); + } + + // Data normalisation + wVec(dataMean_); wVec(dataStdev_); wVec(normMin_); wVec(normMax_); + + // Training history + wU64(static_cast(epochLosses_.size())); + wU64(static_cast(trainingSampleSize_)); + for (double v : epochLosses_) wF64(v); + + // Optimizer state + wI32(static_cast(adamStep_)); + for (const auto& layer : network_) { + wMat(layer.mW); wMat(layer.vW); + wVec(layer.mb); wVec(layer.vb); + } + wVec(dimLossEMA_); wVec(dimWeights_); + + // EMA network (optional section flagged by a boolean) + wI32(useEMANetwork_ ? 1 : 0); + if (useEMANetwork_) { + wU32(static_cast(emaNetwork_.size())); + for (const auto& layer : emaNetwork_) { + wU32(static_cast(layer.W.size())); + wU32(static_cast(layer.W[0].size())); + for (const auto& row : layer.W) + out.write(reinterpret_cast(row.data()), + static_cast(row.size() * sizeof(double))); + out.write(reinterpret_cast(layer.b.data()), + static_cast(layer.b.size() * sizeof(double))); + } + } + + // End-of-stream sentinel (format version >= 3). A complete file ends in "ENDM"; + // the loader reads and verifies it, so a file truncated after otherwise-plausible + // contents is rejected instead of silently loading partial data. + out.write("ENDM", 4); + + // Flush and check: a write that failed part-way (a full disk or an exhausted grid + // quota) otherwise leaves a truncated checkpoint behind and the job exits 0. + out.flush(); + if (!out.good()) + throw cet::exception("ScoreBasedDiffusionModel::saveModel") + << "Failed while writing " << filename + << "; the file is incomplete. Check available space and quota."; + } + + void ScoreBasedDiffusionModel::saveModelCsv( const std::string& filename ) { std::ofstream out(filename); if (!out) { - throw cet::exception("ScoreBasedDiffusionModel::saveModel") + throw cet::exception("ScoreBasedDiffusionModel::saveModelCsv") << "Cannot open file " << filename; } @@ -669,6 +1881,11 @@ namespace mu2e { // Model architecture out << "dim," << dim_ << "\n"; out << "conditionDim," << conditionDim_ << "\n"; + out << "timeEmbeddingDim," << timeEmbeddingDim_ << "\n"; + // Per-dimension Fourier depths as bare comma-separated values (no brackets) so the CSV + // splitter parses them; an empty vector writes just the key with no values. + out << "inputEmbeddingDims"; for (int e : inputEmbeddingDims_) out << "," << e; out << "\n"; + out << "conditionEmbeddingDims"; for (int e : conditionEmbeddingDims_) out << "," << e; out << "\n"; out << "hidden," << hidden_ << "\n"; out << "layers," << layers_ << "\n"; // Optimizer configuration @@ -677,16 +1894,41 @@ namespace mu2e { out << "adamBeta2," << adamBeta2_ << "\n"; out << "adamEps," << adamEps_ << "\n"; // Noise schedule configuration - out << "noiseScheduleType," << (noiseScheduleType_ == NoiseScheduleType::COSINE ? "COSINE" : "LINEAR") << "\n"; + std::string schedName; + if (noiseScheduleType_ == NoiseScheduleType::COSINE) schedName = "COSINE"; + else if (noiseScheduleType_ == NoiseScheduleType::LOGSIG) schedName = "LOGSIG"; + else schedName = "LINEAR"; + out << "noiseScheduleType," << schedName << "\n"; out << "betaMin," << betaMin_ << "\n"; out << "betaMax," << betaMax_ << "\n"; out << "cosineOffset," << cosineOffset_ << "\n"; + out << "logSigMin," << logSigMin_ << "\n"; + out << "logSigMax," << logSigMax_ << "\n"; // Training configuration + out << "predictionTarget," << predictionTargetName(predictionTarget_) << "\n"; + out << "lossWeightPower," << lossWeightPower_ << "\n"; out << "batchSize," << batchSize_ << "\n"; out << "gradientClipThreshold," << gradientClipThreshold_ << "\n"; out << "learningRate," << learningRate_ << "\n"; + // Dimensional weight controller + out << "useDimWeightController," << (useDimWeightController_ ? "1" : "0") << "\n"; + out << "dimWeightEMADecay," << dimWeightEMADecay_ << "\n"; + // EMA network + out << "useEMANetwork," << (useEMANetwork_ ? "1" : "0") << "\n"; + out << "emaNetworkDecayBase," << emaNetworkDecayBase_ << "\n"; // batch-size-independent base + out << "emaNetworkDecay," << emaNetworkDecay_ << "\n"; // per-step effective (base^(batchSize/ref)) // Diffusion process configuration out << "diffusionSteps," << diffusionSteps_ << "\n"; + out << "basisTag," << basisTag_ << "\n"; // opaque app-level tag; see saveModel/basisTag() + // Opaque caller-owned markers; see pdgId()/buildConstants(). A CSV predating these + // keys reads back as 0 / empty, matching the binary v<=8 default. The constants are + // one "buildConstant,," row each so the flat key-value shape is kept. + out << "pdgId," << pdgId_ << "\n"; + for (const auto& kv : buildConstants_) + out << "buildConstant," << kv.first << "," << kv.second << "\n"; + // Index of the class-label condition dim, or -1 for none. A CSV predating this key + // reads back as -1, matching the binary v<=7 default. + out << "categoricalConditionDim," << categoricalConditionDim_ << "\n"; // Write network architecture header out << "\n[NETWORK_PARAMETERS]\n"; @@ -720,6 +1962,34 @@ namespace mu2e { out << "\n"; } + // Write data normalization parameters + out << "\n[DATA_NORMALIZATION]\n"; + out << "numDimensions," << dataMean_.size() << "\n"; + out << "dataMean\n"; + for (size_t i = 0; i < dataMean_.size(); ++i) { + out << dataMean_[i]; + if (i < dataMean_.size() - 1) out << ","; + } + out << "\n"; + out << "dataStdev\n"; + for (size_t i = 0; i < dataStdev_.size(); ++i) { + out << dataStdev_[i]; + if (i < dataStdev_.size() - 1) out << ","; + } + out << "\n"; + out << "normMin\n"; + for (size_t i = 0; i < normMin_.size(); ++i) { + out << normMin_[i]; + if (i < normMin_.size() - 1) out << ","; + } + out << "\n"; + out << "normMax\n"; + for (size_t i = 0; i < normMax_.size(); ++i) { + out << normMax_[i]; + if (i < normMax_.size() - 1) out << ","; + } + out << "\n"; + // Write training history out << "\n[TRAINING_HISTORY]\n"; out << "numEpochs," << epochLosses_.size() << "\n"; @@ -728,290 +1998,1157 @@ namespace mu2e { for (size_t i = 0; i < epochLosses_.size(); ++i) { out << i << "," << epochLosses_[i] << "\n"; } - } - - ScoreBasedDiffusionModel ScoreBasedDiffusionModel::loadModel( - CLHEP::RandFlat& randFlat, - CLHEP::RandGaussQ& randGaussQ, - 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 loadedNetwork; - std::vector epochLosses; - - // Helper lambda to split CSV - auto split = [](const std::string& s) { - std::vector 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)); - }; + // Write Adam optimizer state so training can be resumed from this checkpoint + out << "\n[OPTIMIZER_STATE]\n"; + out << "adamStep," << adamStep_ << "\n"; + for (size_t layerIdx = 0; layerIdx < network_.size(); ++layerIdx) { + auto& layer = network_[layerIdx]; - // Parse file - std::string section; + out << "\nLayer" << layerIdx << "_mW\n"; + for (const auto& row : layer.mW) { + for (size_t j = 0; j < row.size(); ++j) { + out << row[j]; + if (j < row.size() - 1) out << ","; + } + out << "\n"; + } + + out << "Layer" << layerIdx << "_vW\n"; + for (const auto& row : layer.vW) { + for (size_t j = 0; j < row.size(); ++j) { + out << row[j]; + if (j < row.size() - 1) out << ","; + } + out << "\n"; + } - while (std::getline(in, line)) { - if (line.empty()) continue; + out << "Layer" << layerIdx << "_mb\n"; + for (size_t j = 0; j < layer.mb.size(); ++j) { + out << layer.mb[j]; + if (j < layer.mb.size() - 1) out << ","; + } + out << "\n"; - // Detect section headers - if (line[0] == '[') { - section = line; - continue; + out << "Layer" << layerIdx << "_vb\n"; + for (size_t j = 0; j < layer.vb.size(); ++j) { + out << layer.vb[j]; + if (j < layer.vb.size() - 1) out << ","; + } + out << "\n"; + } + + // Dim weight controller state + out << "dimLossEMA\n"; + for (int i = 0; i < dim_; ++i) { + out << dimLossEMA_[i]; + if (i < dim_ - 1) out << ","; + } + out << "\n"; + out << "dimWeights\n"; + for (int i = 0; i < dim_; ++i) { + out << dimWeights_[i]; + if (i < dim_ - 1) out << ","; + } + out << "\n"; + + // EMA network weights (used for inference) + if (useEMANetwork_) { + out << "\n[EMA_NETWORK]\n"; + out << "numLayers," << emaNetwork_.size() << "\n"; + for (size_t layerIdx = 0; layerIdx < emaNetwork_.size(); ++layerIdx) { + auto& layer = emaNetwork_[layerIdx]; + out << "\nLayer" << layerIdx << "_OutSize," << layer.W.size() << "\n"; + out << "Layer" << layerIdx << "_InSize," << layer.W[0].size() << "\n"; + + 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"; + } + + 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"; + } + } + + // See saveModel: an unchecked write leaves a truncated file and a zero exit code. + out.flush(); + if (!out.good()) + throw cet::exception("ScoreBasedDiffusionModel::saveModelCsv") + << "Failed while writing " << filename + << "; the file is incomplete. Check available space and quota."; + } + + ScoreBasedDiffusionModel ScoreBasedDiffusionModel::loadModel( + CLHEP::RandFlat& randFlat, + CLHEP::RandGaussQ& randGaussQ, + const std::string& filename + ) + { + // Dispatch based on file extension. ".dat" is the current binary extension; ".bin" is its + // legacy spelling and stays accepted so existing checkpoints keep loading. The human- + // readable dump remains ".csv". + auto hasExt = [&filename](const char* ext) { + return filename.size() >= 4 && filename.compare(filename.size() - 4, 4, ext) == 0; + }; + bool isBin = hasExt(".dat") || hasExt(".bin"); + bool isCsv = hasExt(".csv"); + if (isBin) { + std::ifstream bin(filename, std::ios::binary); + if (!bin) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Cannot open binary file " << filename; } - // 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"; + // Fail with an attributable message on a short read, instead of letting + // an uninitialized/garbage value propagate into a vector allocation (which + // surfaces as an opaque std::length_error / std::bad_alloc). + auto fail = [&](const char* ctx) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Truncated or corrupt binary checkpoint " << filename + << ": unexpected EOF or read error while reading " << ctx << "."; + }; + auto rI32 = [&]() -> int32_t { int32_t v = 0; bin.read(reinterpret_cast(&v), 4); if (!bin) fail("int32"); return v; }; + auto rU32 = [&]() -> uint32_t { uint32_t v = 0; bin.read(reinterpret_cast(&v), 4); if (!bin) fail("uint32"); return v; }; + auto rU64 = [&]() -> uint64_t { uint64_t v = 0; bin.read(reinterpret_cast(&v), 8); if (!bin) fail("uint64"); return v; }; + auto rF64 = [&]() -> double { double v = 0; bin.read(reinterpret_cast(&v), 8); if (!bin) fail("double"); return v; }; + // Sanity-bound a size/count field read from the file BEFORE it is used to size + // a container. A misaligned-but-not-yet-EOF stream yields a self-consistent + // header followed by garbage counts; bounding them here converts that into a + // clear error rather than an allocation that throws length_error. + constexpr uint64_t kMaxBinElems = 100000000ull; // 1e8 doubles (~800 MB) per dimension + auto checkCount = [&](uint64_t v, const char* ctx) -> uint64_t { + if (v > kMaxBinElems) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Corrupt binary checkpoint " << filename << ": implausible " << ctx + << " count " << v << " (max " << kMaxBinElems + << "). File is likely truncated or damaged."; + return v; + }; + auto rVec = [&](size_t n) -> std::vector { + uint64_t stored = rU64(); + if (stored != n) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Binary vector size mismatch: expected " << n << ", got " << stored; + std::vector v(n); + if (n) { + bin.read(reinterpret_cast(v.data()), static_cast(n * sizeof(double))); + if (!bin) fail("vector payload"); + } + return v; + }; + auto rMat = [&](size_t rows, size_t cols) -> std::vector> { + uint64_t storedRows = rU64(); + if (storedRows != rows) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Binary matrix row count mismatch: expected " << rows << ", got " << storedRows; + std::vector> m(rows); + for (auto& row : m) { + uint64_t storedCols = rU64(); + if (storedCols != cols) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Binary matrix col count mismatch: expected " << cols << ", got " << storedCols; + row.resize(cols); + bin.read(reinterpret_cast(row.data()), static_cast(cols * sizeof(double))); + if (!bin) fail("matrix payload"); + } + return m; + }; + + // Magic + version + char magic[4]; bin.read(magic, 4); + if (std::string(magic, 4) != "SBDM") + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Invalid magic bytes in binary file " << filename; + uint32_t version = rU32(); + if (version < 1 || version > 9) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Unsupported binary file version " << version; + + // Model parameters + int dim = rI32(); + int conditionDim = rI32(); + int timeEmbeddingDim = rI32(); + // Fourier embedding depths. Version 1 predates them (none). Versions 2-4 store a + // single scalar each, broadcast to all dims. Version 5+ stores count-prefixed + // per-dimension vectors. resolveEmbeddingDims (in the constructor) maps {} / {k} / + // full-length the same way, so we hand it the raw form read here. + std::vector inputEmbeddingDims, conditionEmbeddingDims; + if (version >= 5) { + int nIn = rI32(); + for (int i = 0; i < nIn; ++i) inputEmbeddingDims.push_back(rI32()); + int nCond = rI32(); + for (int i = 0; i < nCond; ++i) conditionEmbeddingDims.push_back(rI32()); + } else if (version >= 2) { + inputEmbeddingDims = { rI32() }; // single scalar -> broadcast + conditionEmbeddingDims = { rI32() }; + } // version 1: leave empty -> no embedding + int hidden = rI32(); + int layers = rI32(); + int optType = rI32(); + OptimizerType optimizerType = (optType == 0) ? OptimizerType::ADAM : OptimizerType::SGD; + double adamBeta1 = rF64(), adamBeta2 = rF64(), adamEps = rF64(); + int schedIdx = rI32(); + NoiseScheduleType scheduleType = (schedIdx == 0) ? NoiseScheduleType::LINEAR + : (schedIdx == 1) ? NoiseScheduleType::COSINE + : NoiseScheduleType::LOGSIG; + double betaMin = rF64(), betaMax = rF64(), cosineOffset = rF64(); + double logSigMin = rF64(), logSigMax = rF64(); + // Versions 1-5 stored an epsPrediction bool (0/1); version 6+ stores the + // PredictionTarget enum (0=SCORE,1=EPS,2=V). The enum reuses 0/1 for SCORE/EPS, + // so the legacy bool maps directly with no behavior change. + int32_t ptRaw = rI32(); + PredictionTarget predictionTarget; + if (version >= 6) { + if (ptRaw < 0 || ptRaw > 2) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Invalid predictionTarget enum value " << ptRaw << " in " << filename; + predictionTarget = static_cast(ptRaw); + } else { + predictionTarget = (ptRaw != 0) ? PredictionTarget::EPS : PredictionTarget::SCORE; + } + double lossWeightPower = rF64(); + int batchSize = rI32(); + double gradientClipThreshold = rF64(); + double learningRate = rF64(); + bool useDimWeightController = (rI32() != 0); + double dimWeightEMADecay = rF64(); + bool useEMANetwork = (rI32() != 0); + // Version 4 stores the batch-size-independent EMA decay base; versions 1-3 + // stored the already-rescaled per-step effective value. The constructor + // below treats this as the base and rescales it by batchSize/kEMABatchSizeRef_, + // which is correct for version >= 4 but a spurious second rescaling for + // version <= 3 (corrected after construction). + double emaNetworkDecayStored = rF64(); + int diffusionSteps = rI32(); + + // Opaque application-level basis tag (format v7+). v<=6 files predate it, + // so default to 0. Set on the constructed model below. + int basisTag = (version >= 7) ? rI32() : 0; + + // Index of the class-label condition dim (format v8+), or -1 for none. v<=7 + // files predate it and load as -1: every condition dim z-scored, which is how + // those models were trained. Applied to the constructed model below. + int categoricalConditionDim = (version >= 8) ? rI32() : -1; + + // Opaque caller-owned identity and build constants (format v9+). v<=8 files + // predate both: pdgId loads as 0 and the list as empty, which the caller reads + // as "not recorded" and handles with its own fallback. Applied below. + int pdgId = 0; + std::vector> buildConstants; + if (version >= 9) { + pdgId = rI32(); + const uint32_t nConst = + static_cast(checkCount(rU32(), "build constant")); + buildConstants.reserve(nConst); + for (uint32_t k = 0; k < nConst; ++k) { + const int id = rI32(); + buildConstants.emplace_back(id, rF64()); + } + } + + // Network weights. The stored count must match the header's `layers`: the copy + // loop below iterates over the constructed model's network_ (sized from layers) + // while indexing loadedNetwork, so a smaller stored count reads out of bounds. + // The CSV path performs the same check. + uint32_t numLayers = static_cast(checkCount(rU32(), "network layer")); + if (numLayers != static_cast(layers)) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Binary checkpoint " << filename << " declares " << layers + << " layers in its header but stores " << numLayers << " network layer(s)."; + std::vector loadedNetwork(numLayers); + for (auto& layer : loadedNetwork) { + uint32_t outSize = static_cast(checkCount(rU32(), "layer outSize")); + uint32_t inSize = static_cast(checkCount(rU32(), "layer inSize")); + layer.W.resize(outSize, std::vector(inSize)); + for (auto& row : layer.W) { + bin.read(reinterpret_cast(row.data()), + static_cast(inSize * sizeof(double))); + if (!bin) fail("network weight row"); + } + layer.b.resize(outSize); + bin.read(reinterpret_cast(layer.b.data()), + static_cast(outSize * sizeof(double))); + if (!bin) fail("network biases"); + } + + // Data normalisation. saveModel writes each of the four vectors with its own + // u64 size prefix (wVec), so they must be read back symmetrically — one size + + // payload per vector. (An earlier reader read a single size then four raw + // payloads, swallowing the intervening size prefixes and desynchronising the + // stream for every field that followed.) + auto readNorm = [&](const char* ctx) -> std::vector { + uint64_t n = checkCount(rU64(), ctx); + std::vector v(n); + if (n) { + bin.read(reinterpret_cast(v.data()), static_cast(n * sizeof(double))); + if (!bin) fail(ctx); + } + return v; + }; + std::vector dataMean = readNorm("dataMean"); + std::vector dataStdev = readNorm("dataStdev"); + std::vector normMin = readNorm("normMin"); + std::vector normMax = readNorm("normMax"); + uint64_t normDim = dataMean.size(); + + // Training history + uint64_t numEpochs = checkCount(rU64(), "epoch-loss"); + uint64_t trainingSampleSz = rU64(); + std::vector epochLosses(numEpochs); + for (auto& v : epochLosses) v = rF64(); + + // Optimizer state + int loadedAdamStep = rI32(); + std::vector>> loadedMW(numLayers), loadedVW(numLayers); + std::vector> loadedMb(numLayers), loadedVb(numLayers); + for (size_t l = 0; l < numLayers; ++l) { + size_t outSize = loadedNetwork[l].W.size(); + size_t inSize = loadedNetwork[l].W[0].size(); + loadedMW[l] = rMat(outSize, inSize); + loadedVW[l] = rMat(outSize, inSize); + loadedMb[l] = rVec(outSize); + loadedVb[l] = rVec(outSize); + } + std::vector loadedDimLossEMA = rVec(static_cast(dim)); + std::vector loadedDimWeights = rVec(static_cast(dim)); + + // EMA network + bool hasEMA = (rI32() != 0); + std::vector loadedEmaNetwork; + if (hasEMA) { + uint32_t emaLayers = static_cast(checkCount(rU32(), "EMA layer")); + // Same reasoning as the main network above: the EMA copy loop is bounded by + // emaNetwork_.size(), not by what the file stored. + if (emaLayers != static_cast(layers)) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Binary checkpoint " << filename << " declares " << layers + << " layers in its header but stores " << emaLayers << " EMA layer(s)."; + loadedEmaNetwork.resize(emaLayers); + for (auto& layer : loadedEmaNetwork) { + uint32_t outSize = static_cast(checkCount(rU32(), "EMA layer outSize")); + uint32_t inSize = static_cast(checkCount(rU32(), "EMA layer inSize")); + layer.W.resize(outSize, std::vector(inSize)); + for (auto& row : layer.W) { + bin.read(reinterpret_cast(row.data()), + static_cast(inSize * sizeof(double))); + if (!bin) fail("EMA weight row"); } + layer.b.resize(outSize); + bin.read(reinterpret_cast(layer.b.data()), + static_cast(outSize * sizeof(double))); + if (!bin) fail("EMA biases"); + } + } + + if (!bin) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Read error or unexpected EOF in binary file " << filename; + } + + // End-of-stream sentinel (version >= 3 only; versions 1-2 predate it, so they + // are accepted without it for back-compatibility). A complete version-3 file + // ends in "ENDM"; its absence means the file was truncated after otherwise- + // plausible contents. + if (version >= 3) { + char endTag[4] = {0}; + bin.read(endTag, 4); + if (!bin || std::string(endTag, 4) != "ENDM") + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Missing or invalid end-of-stream sentinel in binary file " << filename + << "; the file is truncated or corrupt."; + } + + // Validate + if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Invalid model parameters in binary file"; + if ((int)normDim != dim + conditionDim) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Normalisation dimension mismatch in binary file"; + + // Reconstruct model + // Argument order follows the temporary constructor signature; see its declaration. + ScoreBasedDiffusionModel model( + randFlat, randGaussQ, dim, conditionDim, hidden, layers, + optimizerType, adamBeta1, adamBeta2, adamEps, scheduleType, + betaMin, betaMax, cosineOffset, + batchSize, gradientClipThreshold, learningRate, + diffusionSteps, false, + logSigMin, logSigMax, predictionTarget, lossWeightPower, + useDimWeightController, dimWeightEMADecay, useEMANetwork, emaNetworkDecayStored, + timeEmbeddingDim, inputEmbeddingDims, conditionEmbeddingDims + ); + model.setBasisTag(basisTag); // opaque tag (0 for v<=6); see saveModel/basisTag() + // Opaque caller-owned markers (0 / empty for v<=8); see pdgId()/buildConstants(). + model.setPdgId(pdgId); + model.setBuildConstants(buildConstants); + // Class-label condition dim (-1 for v<=7). Set BEFORE the normalization arrays + // are restored below so normalizeCondition/dimStats agree with them immediately; + // the setter also re-checks the dim against this file's own conditionDim and + // condition embedding depths, so a hand-edited or mismatched checkpoint is + // rejected here rather than silently mis-scaling the label at generation time. + model.setCategoricalConditionDim(categoricalConditionDim); + + // EMA decay semantics fix-up. The constructor rescaled emaNetworkDecayStored + // as if it were the batch-size-independent base. That is correct for version-4 + // checkpoints. Versions 1-3, however, stored the already-rescaled per-step + // effective decay, so the constructor applied the batch-size exponent a second + // time. Undo that here: treat the stored value as the effective per-step decay + // and back out a consistent base for subsequent saves. + if (useEMANetwork && version <= 3) { + model.emaNetworkDecay_ = emaNetworkDecayStored; + model.emaNetworkDecayBase_ = + (batchSize > 0) + ? std::pow(emaNetworkDecayStored, + (double)kEMABatchSizeRef_ / batchSize) + : emaNetworkDecayStored; + } + + 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.dataMean_ = dataMean; + model.dataStdev_ = dataStdev; + model.normMin_ = normMin; + model.normMax_ = normMax; + model.epochLosses_ = epochLosses; + model.trainingSampleSize_ = static_cast(trainingSampleSz); + + // Restore optimizer state + model.adamStep_ = loadedAdamStep; + for (size_t l = 0; l < model.network_.size(); ++l) { + model.network_[l].mW = loadedMW[l]; + model.network_[l].vW = loadedVW[l]; + model.network_[l].mb = loadedMb[l]; + model.network_[l].vb = loadedVb[l]; + } + model.dimLossEMA_ = loadedDimLossEMA; + model.dimWeights_ = loadedDimWeights; + // Checkpoints written before the dim-weight bounds existed can carry + // arbitrarily skewed weights; enforce the invariant on restore. + model.clampDimWeights("Binary checkpoint restore"); + + // Restore EMA network + if (hasEMA && useEMANetwork) { + for (size_t l = 0; l < model.emaNetwork_.size(); ++l) { + model.emaNetwork_[l].W = loadedEmaNetwork[l].W; + model.emaNetwork_[l].b = loadedEmaNetwork[l].b; + } + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") + << "EMA network weights restored from binary checkpoint."; + } else if (useEMANetwork) { + for (size_t l = 0; l < model.emaNetwork_.size(); ++l) { + model.emaNetwork_[l].W = model.network_[l].W; + model.emaNetwork_[l].b = model.network_[l].b; + } + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") + << "No EMA network in binary checkpoint — initialized from network weights."; + } + + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") + << "Binary model loaded from " << filename + << " (format version " << version << ", basisTag " << basisTag + << ", categoricalConditionDim=" << categoricalConditionDim + << ", adamStep=" << loadedAdamStep << ", epochs=" << numEpochs << ")"; + return model; + } else if (isCsv) { + + 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, timeEmbeddingDim = 0, hidden = 0, layers = 0; + // Raw Fourier embedding specs as read from the file: {} (legacy files predating the + // feature, or an empty list), {k} (legacy singular key, broadcast), or a full per-dim + // list (plural key). resolveEmbeddingDims() normalizes all three at construction. + std::vector inputEmbeddingDims, conditionEmbeddingDims; + 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; + double logSigMin = 1e-5, logSigMax = 1.0; + PredictionTarget predictionTarget = PredictionTarget::SCORE; + double lossWeightPower = 2.0; + int batchSize = 1, diffusionSteps = 1; + int basisTag = 0; // opaque app-level tag; absent in CSVs predating the feature + // Class-label condition dim; absent in CSVs predating it, hence -1 = none. + int categoricalConditionDim = -1; + // Opaque caller-owned markers; absent in CSVs predating them, hence 0 / empty. + int pdgId = 0; + std::vector> buildConstants; + double gradientClipThreshold = 0.0, learningRate = 0.0; + + std::vector dataMean, dataStdev, normMin, normMax; + + std::vector loadedNetwork; + std::vector epochLosses; + + // Optimizer state (optional — absent in files saved before this feature was added) + int loadedAdamStep = 0; + bool optimizerStateLoaded = false; + std::vector>> loadedMW, loadedVW; + std::vector> loadedMb, loadedVb; + std::vector loadedDimLossEMA, loadedDimWeights; + + // New controller / EMA network parameters (default to constructor defaults if absent) + bool useDimWeightController = false; + double dimWeightEMADecay = 0.99; + bool useEMANetwork = true; + double emaNetworkDecay = 0.9999; // legacy field: per-step effective decay + double emaNetworkDecayBase = 0.9999; // batch-size-independent base + bool emaNetworkDecayBasePresent = false; + + // EMA network weights (optional) + std::vector loadedEmaNetwork; + bool emaNetworkLoaded = false; + + // Helper lambda to split CSV + auto split = [](const std::string& s) { + std::vector 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; } - // Layer sizes - if (line.find("_OutSize") != std::string::npos) { + // MODEL PARAMETERS + if (section == "[MODEL_PARAMETERS]") { + if (line == "Parameter,Value") continue; + 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); + if (tokens.empty()) continue; + + // Fourier embedding depths are variable-length, so handle them before the + // strict 2-token check. Accept the plural keys (per-dim comma list, possibly + // empty) and the legacy singular keys (single scalar, broadcast at construction). + if (tokens[0] == "inputEmbeddingDims" || tokens[0] == "inputEmbeddingDim") { + inputEmbeddingDims.clear(); + for (size_t i = 1; i < tokens.size(); ++i) inputEmbeddingDims.push_back(std::stoi(tokens[i])); + continue; + } + if (tokens[0] == "conditionEmbeddingDims" || tokens[0] == "conditionEmbeddingDim") { + conditionEmbeddingDims.clear(); + for (size_t i = 1; i < tokens.size(); ++i) conditionEmbeddingDims.push_back(std::stoi(tokens[i])); + continue; + } + + // Opaque build constants: one "buildConstant,," row each, so + // handled here rather than by the two-token key/value path below. + if (tokens[0] == "buildConstant") { + if (tokens.size() != 3) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "invalid buildConstant line in [MODEL_PARAMETERS]: " << line; + buildConstants.emplace_back(std::stoi(tokens[1]), std::stod(tokens[2])); + continue; + } + + 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 == "timeEmbeddingDim") timeEmbeddingDim = 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") { + if (val == "COSINE") scheduleType = NoiseScheduleType::COSINE; + else if (val == "LOGSIG") scheduleType = NoiseScheduleType::LOGSIG; + else scheduleType = 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 == "logSigMin") logSigMin = std::stod(val); + else if (key == "logSigMax") logSigMax = std::stod(val); + // Legacy CSV key (versions before the enum): map the old bool. + else if (key == "epsPrediction") + predictionTarget = (val == "1") ? PredictionTarget::EPS : PredictionTarget::SCORE; + // New CSV key: SCORE / EPS / V (string, matching the writer). + else if (key == "predictionTarget") + predictionTarget = predictionTargetFromName(val); + else if (key == "lossWeightPower") lossWeightPower = 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 == "useDimWeightController") useDimWeightController = (val == "1"); + else if (key == "dimWeightEMADecay") dimWeightEMADecay = std::stod(val); + else if (key == "useEMANetwork") useEMANetwork = (val == "1"); + else if (key == "emaNetworkDecayBase") { emaNetworkDecayBase = std::stod(val); emaNetworkDecayBasePresent = true; } + else if (key == "emaNetworkDecay") emaNetworkDecay = std::stod(val); + else if (key == "diffusionSteps") diffusionSteps = std::stoi(val); + else if (key == "basisTag") basisTag = std::stoi(val); + else if (key == "categoricalConditionDim") categoricalConditionDim = std::stoi(val); + else if (key == "pdgId") pdgId = std::stoi(val); } - 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 << ")"; + + // 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; } - if (loadedNetwork[layerIdx].W.empty()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "InSize encountered before OutSize for layer " << layerIdx; + + // 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); } - 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; + else if (line.find("_InSize") != std::string::npos) { + auto tokens = split(line); + int inSize = std::stoi(tokens[1]); + int layerIdx = getLayerIdx(tokens[0]); + // Resolve the raw embedding specs to full per-dim vectors (same rules as + // the constructor) so the expected input width uses their summed depths. + auto resIn = resolveEmbeddingDims(inputEmbeddingDims, dim, "inputEmbeddingDims"); + auto resCond = resolveEmbeddingDims(conditionEmbeddingDims, conditionDim, "conditionEmbeddingDims"); + int sumIn = 0; for (int e : resIn) sumIn += e; + int sumCo = 0; for (int e : resCond) sumCo += e; + int expectedInSize = dim + sumIn + conditionDim + sumCo + 1 + timeEmbeddingDim; + if (layerIdx == 0 && inSize != expectedInSize) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Input layer input size mismatch (expected " + << expectedInSize << ", 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); } - if (loadedNetwork[layerIdx].W.empty() || loadedNetwork[layerIdx].W[0].empty()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Weights encountered before layer size definition"; + // 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]); + } } - for (auto& row : loadedNetwork[layerIdx].W) { + // 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 weights"; + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading biases"; } auto vals = split(line); - if (vals.size() != row.size()) { + if (vals.size() != loadedNetwork[layerIdx].b.size()) { throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Weight row size mismatch for layer " << layerIdx; + << "Bias size mismatch for layer " << layerIdx; } for (size_t j = 0; j < vals.size(); ++j) - row[j] = std::stod(vals[j]); + loadedNetwork[layerIdx].b[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; + + // DATA NORMALIZATION + else if (section == "[DATA_NORMALIZATION]") { + if (line.rfind("numDimensions", 0) == 0) { + int normalizationDim = std::stoi(split(line)[1]); + dataMean.resize(normalizationDim); + dataStdev.resize(normalizationDim); + normMin.resize(normalizationDim); + normMax.resize(normalizationDim); + continue; } - if (!std::getline(in, line)) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading biases"; + if (line == "dataMean") { + if (!std::getline(in, line)) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading dataMean"; + } + auto vals = split(line); + if (vals.size() != dataMean.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Data mean size mismatch"; + } + for (size_t j = 0; j < vals.size(); ++j) { + dataMean[j] = std::stod(vals[j]); + } + continue; } - auto vals = split(line); - if (vals.size() != loadedNetwork[layerIdx].b.size()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Bias size mismatch for layer " << layerIdx; + if (line == "dataStdev") { + if (!std::getline(in, line)) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading dataStdev"; + } + auto vals = split(line); + if (vals.size() != dataStdev.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Data stdev size mismatch"; + } + for (size_t j = 0; j < vals.size(); ++j) { + dataStdev[j] = std::stod(vals[j]); + } + continue; + } + if (line == "normMin") { + if (!std::getline(in, line)) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading normMin"; + } + auto vals = split(line); + if (vals.size() != normMin.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Norm min size mismatch"; + } + for (size_t j = 0; j < vals.size(); ++j) { + normMin[j] = std::stod(vals[j]); + } + continue; } - 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. - // If needed later, parse as size_t to avoid truncation, e.g. size_t trainingSampleSize = std::stoull(tokens[1]); - // and apply checked conversion only when assigning to narrower types. - }else { - epochLosses.push_back(std::stod(tokens[1])); + if (line == "normMax") { + if (!std::getline(in, line)) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF while reading normMax"; + } + auto vals = split(line); + if (vals.size() != normMax.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Norm max size mismatch"; + } + for (size_t j = 0; j < vals.size(); ++j) { + normMax[j] = std::stod(vals[j]); + } + continue; + } + } + + // 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. + // If needed later, parse as size_t to avoid truncation, e.g. size_t trainingSampleSize = std::stoull(tokens[1]); + // and apply checked conversion only when assigning to narrower types. + }else { + epochLosses.push_back(std::stod(tokens[1])); + } + } + } + + // OPTIMIZER STATE + else if (section == "[OPTIMIZER_STATE]") { + if (line.rfind("adamStep,", 0) == 0) { + loadedAdamStep = std::stoi(split(line)[1]); + optimizerStateLoaded = true; + continue; + } + if (line.find("_mW") != std::string::npos && line.rfind("Layer", 0) == 0) { + int layerIdx = getLayerIdx(line); + if (layerIdx >= (int)loadedNetwork.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "mW layer index out of range: " << layerIdx; + } + if (layerIdx >= (int)loadedMW.size()) loadedMW.resize(layerIdx + 1); + auto& outRows = loadedMW[layerIdx]; + outRows.resize(loadedNetwork[layerIdx].W.size()); + for (auto& row : outRows) { + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading mW"; + auto vals = split(line); + row.resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) row[j] = std::stod(vals[j]); + } + } + else if (line.find("_vW") != std::string::npos && line.rfind("Layer", 0) == 0) { + int layerIdx = getLayerIdx(line); + if (layerIdx >= (int)loadedNetwork.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "vW layer index out of range: " << layerIdx; + } + if (layerIdx >= (int)loadedVW.size()) loadedVW.resize(layerIdx + 1); + auto& outRows = loadedVW[layerIdx]; + outRows.resize(loadedNetwork[layerIdx].W.size()); + for (auto& row : outRows) { + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading vW"; + auto vals = split(line); + row.resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) row[j] = std::stod(vals[j]); + } + } + else if (line.find("_mb") != std::string::npos && line.rfind("Layer", 0) == 0) { + int layerIdx = getLayerIdx(line); + if (layerIdx >= (int)loadedNetwork.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "mb layer index out of range: " << layerIdx; + } + if (layerIdx >= (int)loadedMb.size()) loadedMb.resize(layerIdx + 1); + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading mb"; + auto vals = split(line); + loadedMb[layerIdx].resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) loadedMb[layerIdx][j] = std::stod(vals[j]); + } + else if (line.find("_vb") != std::string::npos && line.rfind("Layer", 0) == 0) { + int layerIdx = getLayerIdx(line); + if (layerIdx >= (int)loadedNetwork.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "vb layer index out of range: " << layerIdx; + } + if (layerIdx >= (int)loadedVb.size()) loadedVb.resize(layerIdx + 1); + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading vb"; + auto vals = split(line); + loadedVb[layerIdx].resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) loadedVb[layerIdx][j] = std::stod(vals[j]); + } + else if (line == "dimLossEMA") { + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading dimLossEMA"; + auto vals = split(line); + loadedDimLossEMA.resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) loadedDimLossEMA[j] = std::stod(vals[j]); + } + else if (line == "dimWeights") { + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading dimWeights"; + auto vals = split(line); + loadedDimWeights.resize(vals.size()); + for (size_t j = 0; j < vals.size(); ++j) loadedDimWeights[j] = std::stod(vals[j]); + } + } + + // EMA NETWORK + else if (section == "[EMA_NETWORK]") { + if (line.rfind("numLayers", 0) == 0) { + int numLayers = std::stoi(split(line)[1]); + loadedEmaNetwork.resize(numLayers); + emaNetworkLoaded = true; + continue; + } + if (line.find("_OutSize") != std::string::npos) { + auto tokens = split(line); + int outSize = std::stoi(tokens[1]); + int layerIdx = getLayerIdx(tokens[0]); + loadedEmaNetwork[layerIdx].W.resize(outSize); + loadedEmaNetwork[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 (loadedEmaNetwork[layerIdx].W.empty()) + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "EMA InSize before OutSize for layer " << layerIdx; + for (auto& row : loadedEmaNetwork[layerIdx].W) row.resize(inSize); + } + else if (line.find("_Weights") != std::string::npos) { + int layerIdx = getLayerIdx(line); + for (auto& row : loadedEmaNetwork[layerIdx].W) { + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading EMA weights"; + auto vals = split(line); + if (vals.size() != row.size()) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "EMA weight row size mismatch layer " << layerIdx; + for (size_t j = 0; j < vals.size(); ++j) row[j] = std::stod(vals[j]); + } + } + else if (line.find("_Biases") != std::string::npos) { + int layerIdx = getLayerIdx(line); + if (!std::getline(in, line)) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Unexpected EOF reading EMA biases"; + auto vals = split(line); + if (vals.size() != loadedEmaNetwork[layerIdx].b.size()) throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "EMA bias size mismatch layer " << layerIdx; + for (size_t j = 0; j < vals.size(); ++j) loadedEmaNetwork[layerIdx].b[j] = std::stod(vals[j]); } } } - } - 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 without random weight initialization - ScoreBasedDiffusionModel model( - randFlat, - randGaussQ, - dim, - conditionDim, - hidden, - layers, - optimizerType, - adamBeta1, - adamBeta2, - adamEps, - scheduleType, - betaMin, - betaMax, - cosineOffset, - batchSize, - gradientClipThreshold, - learningRate, - diffusionSteps, - false - ); - - // Validate loaded network shape and dimensions match the initialized model before overwriting. - if (loadedNetwork.size() != model.network_.size()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Layer count mismatch: loaded " << loadedNetwork.size() - << " layers, expected " << model.network_.size() - << ". The model file may be missing or malformed in [NETWORK_PARAMETERS] (e.g. numLayers)."; - } - for (size_t l = 0; l < loadedNetwork.size(); ++l) { - if (loadedNetwork[l].W.empty() || loadedNetwork[l].W[0].empty() || loadedNetwork[l].b.empty()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Loaded layer " << l - << " is not fully sized before overwrite (empty weights/biases)."; + if (dim <= 0 || conditionDim < 0 || hidden <= 0 || layers <= 0) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Invalid model parameters in file"; } - // Check that loaded layer dimensions match the model's initialized dimensions - if (loadedNetwork[l].W.size() != model.network_[l].W.size() || - loadedNetwork[l].W[0].size() != model.network_[l].W[0].size()) { - throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Layer " << l << " weight dimension mismatch: loaded W[" - << loadedNetwork[l].W.size() << "][" - << loadedNetwork[l].W[0].size() << "], expected W[" - << model.network_[l].W.size() << "][" - << model.network_[l].W[0].size() << "]."; + 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"; + } } - if (loadedNetwork[l].b.size() != model.network_[l].b.size()) { + + std::ostringstream logMsg; + logMsg << "Model loaded successfully from " << filename + << " (format CSV (unversioned), basisTag " << basisTag + << ", categoricalConditionDim " << categoricalConditionDim << ")\n"; + if (optimizerStateLoaded) { + logMsg << " Optimizer state (Adam moments, step=" << loadedAdamStep << ") restored — training can be resumed.\n"; + } else { + logMsg << " No optimizer state found — fresh optimizer state will be used.\n"; + } + logMsg << " DimWeightController: " << (useDimWeightController ? "enabled" : "disabled"); + if (useDimWeightController) logMsg << " (EMADecay=" << dimWeightEMADecay << ")"; + logMsg << "\n"; + logMsg << " EMANetwork: " << (useEMANetwork ? "enabled" : "disabled"); + if (useEMANetwork) logMsg << " (decay=" << emaNetworkDecay << ", " << (emaNetworkLoaded ? "weights from checkpoint" : "initialized from network") << ")"; + logMsg << "\n"; + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") << logMsg.str(); + + // Reconstruct model without random weight initialization + // Argument order follows the temporary constructor signature; see its declaration. + ScoreBasedDiffusionModel model( + randFlat, + randGaussQ, + dim, + conditionDim, + hidden, + layers, + optimizerType, + adamBeta1, + adamBeta2, + adamEps, + scheduleType, + betaMin, + betaMax, + cosineOffset, + batchSize, + gradientClipThreshold, + learningRate, + diffusionSteps, + false, // initializeRandomWeights + logSigMin, + logSigMax, + predictionTarget, + lossWeightPower, + useDimWeightController, + dimWeightEMADecay, + useEMANetwork, + // Newer CSVs carry the batch-size-independent base, which the constructor + // rescales correctly. Older CSVs only have the effective decay; pass it as + // a placeholder base and fix it up below. + emaNetworkDecayBasePresent ? emaNetworkDecayBase : emaNetworkDecay, + timeEmbeddingDim, + inputEmbeddingDims, + conditionEmbeddingDims + ); + model.setBasisTag(basisTag); // opaque tag (0 if absent); see saveModel/basisTag() + // Opaque caller-owned markers (0 / empty if absent); see pdgId()/buildConstants(). + model.setPdgId(pdgId); + model.setBuildConstants(buildConstants); + // Class-label condition dim (-1 if absent); see the binary path for why this is + // applied before the normalization arrays are restored. + model.setCategoricalConditionDim(categoricalConditionDim); + + // Legacy-CSV EMA decay fix-up (mirrors the binary version <= 3 path). When the + // file predates emaNetworkDecayBase, the stored emaNetworkDecay is the per-step + // effective value; adopt it directly and back out a consistent base so future + // saves and any later batch-size change behave correctly. + if (useEMANetwork && !emaNetworkDecayBasePresent) { + model.emaNetworkDecay_ = emaNetworkDecay; + model.emaNetworkDecayBase_ = + (batchSize > 0) + ? std::pow(emaNetworkDecay, (double)kEMABatchSizeRef_ / batchSize) + : emaNetworkDecay; + } + + // Validate loaded network shape and dimensions match the initialized model before overwriting. + if (loadedNetwork.size() != model.network_.size()) { throw cet::exception("ScoreBasedDiffusionModel::loadModel") - << "Layer " << l << " bias dimension mismatch: loaded b[" - << loadedNetwork[l].b.size() << "], expected b[" - << model.network_[l].b.size() << "]."; + << "Layer count mismatch: loaded " << loadedNetwork.size() + << " layers, expected " << model.network_.size() + << ". The model file may be missing or malformed in [NETWORK_PARAMETERS] (e.g. numLayers)."; + } + for (size_t l = 0; l < loadedNetwork.size(); ++l) { + if (loadedNetwork[l].W.empty() || loadedNetwork[l].W[0].empty() || loadedNetwork[l].b.empty()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Loaded layer " << l + << " is not fully sized before overwrite (empty weights/biases)."; + } + // Check that loaded layer dimensions match the model's initialized dimensions + if (loadedNetwork[l].W.size() != model.network_[l].W.size() || + loadedNetwork[l].W[0].size() != model.network_[l].W[0].size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Layer " << l << " weight dimension mismatch: loaded W[" + << loadedNetwork[l].W.size() << "][" + << loadedNetwork[l].W[0].size() << "], expected W[" + << model.network_[l].W.size() << "][" + << model.network_[l].W[0].size() << "]."; + } + if (loadedNetwork[l].b.size() != model.network_[l].b.size()) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Layer " << l << " bias dimension mismatch: loaded b[" + << loadedNetwork[l].b.size() << "], expected b[" + << model.network_[l].b.size() << "]."; + } } - } - // Allocate 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; + if ((int)dataMean.size() != (dim + conditionDim)) { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") << "Normalization parameter size mismatch"; + } // dataStdev, normMin, normMax have same dimensions + + // Restore network weights + 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.dataMean_ = dataMean; + model.dataStdev_ = dataStdev; + model.normMin_ = normMin; + model.normMax_ = normMax; + model.epochLosses_ = epochLosses; + + // Restore Adam optimizer state if present, enabling seamless training resumption + if (optimizerStateLoaded) { + bool momentSizeOk = (loadedMW.size() == model.network_.size() && + loadedVW.size() == model.network_.size() && + loadedMb.size() == model.network_.size() && + loadedVb.size() == model.network_.size()); + if (momentSizeOk) { + for (size_t l = 0; l < model.network_.size() && momentSizeOk; ++l) { + if (loadedMW[l].size() != model.network_[l].mW.size() || + loadedVW[l].size() != model.network_[l].vW.size() || + loadedMb[l].size() != model.network_[l].mb.size() || + loadedVb[l].size() != model.network_[l].vb.size()) { + momentSizeOk = false; + break; + } + for (size_t i = 0; i < loadedMW[l].size() && momentSizeOk; ++i) { + if (loadedMW[l][i].size() != model.network_[l].mW[i].size() || + loadedVW[l][i].size() != model.network_[l].vW[i].size()) { + momentSizeOk = false; + } + } + } + } + if (momentSizeOk) { + model.adamStep_ = loadedAdamStep; + for (size_t l = 0; l < model.network_.size(); ++l) { + model.network_[l].mW = loadedMW[l]; + model.network_[l].vW = loadedVW[l]; + model.network_[l].mb = loadedMb[l]; + model.network_[l].vb = loadedVb[l]; + } + } else { + mf::LogWarning("ScoreBasedDiffusionModel::loadModel") + << "Optimizer state dimensions do not match the reconstructed model — " + << "falling back to fresh optimizer state. Training can continue but will not resume from the saved checkpoint."; + } + + // Restore dim weight controller state + if (loadedDimLossEMA.size() == static_cast(dim) && + loadedDimWeights.size() == static_cast(dim)) { + model.dimLossEMA_ = loadedDimLossEMA; + model.dimWeights_ = loadedDimWeights; + // Checkpoints written before the dim-weight bounds existed can carry + // arbitrarily skewed weights; enforce the invariant on restore. + model.clampDimWeights("CSV checkpoint restore"); + } else if (!loadedDimLossEMA.empty()) { + mf::LogWarning("ScoreBasedDiffusionModel::loadModel") + << "Dim weight controller state size mismatch — resetting to defaults (all-zeros EMA, all-ones weights)."; + } + } - return model; + // Restore EMA network weights + if (emaNetworkLoaded && useEMANetwork) { + bool emaOk = (loadedEmaNetwork.size() == model.emaNetwork_.size()); + for (size_t l = 0; l < loadedEmaNetwork.size() && emaOk; ++l) { + if (loadedEmaNetwork[l].W.size() != model.emaNetwork_[l].W.size() || + (!loadedEmaNetwork[l].W.empty() && loadedEmaNetwork[l].W[0].size() != model.emaNetwork_[l].W[0].size()) || + loadedEmaNetwork[l].b.size() != model.emaNetwork_[l].b.size()) + emaOk = false; + } + if (emaOk) { + for (size_t l = 0; l < model.emaNetwork_.size(); ++l) { + model.emaNetwork_[l].W = loadedEmaNetwork[l].W; + model.emaNetwork_[l].b = loadedEmaNetwork[l].b; + } + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") << "EMA network weights restored from checkpoint."; + } else { + mf::LogWarning("ScoreBasedDiffusionModel::loadModel") + << "EMA network dimensions mismatch — re-initializing EMA from loaded network weights."; + for (size_t l = 0; l < model.emaNetwork_.size(); ++l) { + model.emaNetwork_[l].W = model.network_[l].W; + model.emaNetwork_[l].b = model.network_[l].b; + } + } + } else if (useEMANetwork) { + // Old checkpoint without [EMA_NETWORK] — seed EMA from the loaded network weights + for (size_t l = 0; l < model.emaNetwork_.size(); ++l) { + model.emaNetwork_[l].W = model.network_[l].W; + model.emaNetwork_[l].b = model.network_[l].b; + } + mf::LogInfo("ScoreBasedDiffusionModel::loadModel") + << "No [EMA_NETWORK] section found — EMA network initialized from loaded network weights."; + } + + return model; + } else { + throw cet::exception("ScoreBasedDiffusionModel::loadModel") + << "Unrecognized file extension in \"" << filename + << "\"; expected \".dat\" (or legacy \".bin\") for binary, or \".csv\" for text."; + } } - std::vector ScoreBasedDiffusionModel::generateSample( + SBDMGeneratedSample ScoreBasedDiffusionModel::generateSample( const std::vector& condition, + bool useEMANetworkIfAvailable, bool useHeun, - int diffusionSteps + bool useSDE, + int diffusionSteps, + double sdeToOdeSigmaThreshold ) { if (condition.size() != static_cast(conditionDim_)) { @@ -1030,66 +3167,245 @@ namespace mu2e { x[i] = randGaussQ_.fire(); } + return reverseDiffuseFrom(std::move(x), condition, steps, + useEMANetworkIfAvailable, useHeun, useSDE, + steps, sdeToOdeSigmaThreshold); + } + + SBDMGeneratedSample ScoreBasedDiffusionModel::reverseDiffuseFrom( + std::vector x, + const std::vector& condition, + int stepStart, + bool useEMANetworkIfAvailable, + bool useHeun, + bool useSDE, + int steps, + double sdeToOdeSigmaThreshold + ) + { + double sigma_safe = 1e-12; // small constant to prevent division by zero in case of very small sigma + // Reverse diffusion process - for (int step = steps - 1; step >= 0; --step) { + for (int step = stepStart - 1; step >= 0; --step) { - double t = (double)step/steps; + double t = ((double)step + 1.0)/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. + double beta_val = beta(t); + bool effectiveSDE = useSDE && + (sdeToOdeSigmaThreshold < 0.0 || sigma(t) >= sdeToOdeSigmaThreshold); + const auto& inferNet = (useEMANetworkIfAvailable && useEMANetwork_) ? emaNetwork_ : network_; if (!useHeun) { // Euler method (1st order) - std::vector input = x; - input.insert(input.end(), condition.begin(), condition.end()); - input.push_back(t); + auto input = buildNetworkInput(x, condition, t); - auto score = forward(input); + auto score = forwardInference(input, inferNet); + { + // Convert network output to the score for any prediction target. For SCORE + // this is the identity (epsHat=-out*s -> -epsHat/s = out). + // s = sigma(t); as long as diffusionSteps_ is not too large the smallest t on + // the grid keeps s above sigma_safe, so the 1/s in scoreFromOutput stays well-behaved. + double s = std::max(sigma(t), sigma_safe); + double a = std::sqrt(std::max(0.0, alphabar(t))); + for (int i = 0; i < dim_; ++i) + score[i] = scoreFromOutput(score[i], x[i], s, a); + } for (int i = 0; i < dim_; ++i) { - x[i] += -s * s * score[i] * dt; + if (effectiveSDE) { + // SDE solver: + double drift = 0.5 * beta_val * x[i] + beta_val * score[i]; + double noise = std::sqrt(beta_val * dt) * randGaussQ_.fire(); + x[i] += drift * dt + noise; + } else { + // ODE solver: + double drift = 0.5 * beta_val * x[i] + 0.5 * beta_val * score[i]; + x[i] += drift * dt; + } } } else { - // Heun's method (2nd order) + // Heun's method (2nd order) Only ODE solver, no noise added + + // sahred noise vector + std::vector dw(dim_); + double noiseScale = std::sqrt(beta_val * dt); + for (int i = 0; i < dim_; ++i) { + dw[i] = noiseScale * randGaussQ_.fire(); + } + // k1 = f(x,t) - std::vector input = x; - input.insert(input.end(), condition.begin(), condition.end()); - input.push_back(t); - auto score_k1 = forward(input); + auto input = buildNetworkInput(x, condition, t); + auto score_k1 = forwardInference(input, inferNet); + { + double s = std::max(sigma(t), sigma_safe); + double a = std::sqrt(std::max(0.0, alphabar(t))); + for (int i = 0; i < dim_; ++i) + score_k1[i] = scoreFromOutput(score_k1[i], x[i], s, a); + } std::vector k1(dim_); for (int i = 0; i < dim_; ++i) { - k1[i] = -s * s * score_k1[i]; + if (effectiveSDE) { + // SDE solver: + k1[i] = 0.5 * beta_val * x[i] + beta_val * score_k1[i]; + } else { + // ODE solver: + k1[i] = 0.5 * beta_val * x[i] + 0.5 * beta_val * score_k1[i]; + } } // predictor std::vector x_pred(dim_); for (int i = 0; i < dim_; ++i) { x_pred[i] = x[i] + k1[i] * dt; + if (effectiveSDE) { + x_pred[i] += dw[i]; // add noise + } } // next time double t_next = std::max(0.0, t - dt); - double s_next = sigma(t_next); + double b_next = beta(t_next); // k2 = f(x_pred,t_next) - std::vector 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); + auto input_next = buildNetworkInput(x_pred, condition, t_next); + auto score_k2 = forwardInference(input_next, inferNet); + { + double s_next = std::max(sigma(t_next), sigma_safe); + double a_next = std::sqrt(std::max(0.0, alphabar(t_next))); + for (int i = 0; i < dim_; ++i) + score_k2[i] = scoreFromOutput(score_k2[i], x_pred[i], s_next, a_next); + } std::vector k2(dim_); for (int i = 0; i < dim_; ++i) { - k2[i] = -s_next * s_next * score_k2[i]; + if (effectiveSDE) { + // SDE solver: + k2[i] = 0.5 * b_next * x_pred[i] + b_next * score_k2[i]; + } else { + // ODE solver: + k2[i] = 0.5 * b_next * x_pred[i] + 0.5 * b_next * score_k2[i]; + } } // trapezoidal update for (int i = 0; i < dim_; ++i) { x[i] += 0.5 * (k1[i] + k2[i]) * dt; + if (effectiveSDE) { + x[i] += dw[i]; // add noise + } } } } - return x; + SBDMGeneratedSample generatedSample; + generatedSample.zscore = x; + generatedSample.value.resize(dim_); + for (int i = 0; i < dim_; ++i) { + generatedSample.value[i] = x[i]*dataStdev_[i] + dataMean_[i]; + // note the order in the stdev and mean is always x then cond + } + return generatedSample; + } + + // Partial-reverse diagnostic: noise a normalized data sample at t0 (snapped to the + // sampler's time grid so the first network evaluation matches the noising time), then + // integrate the same reverse process generateSample() uses from t0 down to 0. + SBDMGeneratedSample ScoreBasedDiffusionModel::partialReverseSample( + const std::vector& xNorm, + const std::vector& condition, + double t0, + bool useEMANetworkIfAvailable, + bool useHeun, + bool useSDE, + int diffusionSteps, + double sdeToOdeSigmaThreshold, + std::vector* noisedZscoreOut + ) + { + if (xNorm.size() != static_cast(dim_)) { + throw cet::exception("ScoreBasedDiffusionModel::partialReverseSample") + << "State dimension mismatch: got " << xNorm.size() << ", expected " << dim_; + } + if (condition.size() != static_cast(conditionDim_)) { + throw cet::exception("ScoreBasedDiffusionModel::partialReverseSample") + << "Conditioning dimension mismatch: got " << condition.size() + << ", expected " << conditionDim_; + } + if (t0 <= 0.0 || t0 > 1.0) { + throw cet::exception("ScoreBasedDiffusionModel::partialReverseSample") + << "Invalid diffusion time t0=" << t0 << ": must be in (0,1]"; + } + + // Use provided diffusionSteps if positive, otherwise use the model's configured value + int steps = (diffusionSteps > 0) ? diffusionSteps : diffusionSteps_; + + // Snap t0 to the sampler's grid t = stepStart/steps; the reverse loop's first + // evaluation is at exactly that time, so the noised state sits on the grid. + int stepStart = std::min(steps, std::max(1, static_cast(std::lround(t0 * steps)))); + double tGrid = static_cast(stepStart) / steps; + + std::vector eps; + auto x = addNoise(xNorm, tGrid, eps); + if (noisedZscoreOut) *noisedZscoreOut = x; // the noised starting state (z-score), before reverse + + return reverseDiffuseFrom(std::move(x), condition, stepStart, + useEMANetworkIfAvailable, useHeun, useSDE, + steps, sdeToOdeSigmaThreshold); + } + + // One-step denoising diagnostic: noise a normalized sample at fixed t, run one forward + // pass, and reconstruct x0_hat via Tweedie's formula for the VP forward process + // x_t = sqrt(alphabar(t)) * x0 + sigma(t) * eps => x0_hat = (x_t - sigma * eps_hat) / sqrt(alphabar) + SBDMGeneratedSample ScoreBasedDiffusionModel::denoiseOneStep( + const std::vector& xNorm, + const std::vector& condition, + double t, + bool useEMANetworkIfAvailable, + std::vector* noisedZscoreOut + ) + { + if (xNorm.size() != static_cast(dim_)) { + throw cet::exception("ScoreBasedDiffusionModel::denoiseOneStep") + << "State dimension mismatch: got " << xNorm.size() << ", expected " << dim_; + } + if (condition.size() != static_cast(conditionDim_)) { + throw cet::exception("ScoreBasedDiffusionModel::denoiseOneStep") + << "Conditioning dimension mismatch: got " << condition.size() + << ", expected " << conditionDim_; + } + if (t <= 0.0 || t >= 1.0) { + throw cet::exception("ScoreBasedDiffusionModel::denoiseOneStep") + << "Invalid diffusion time t=" << t << ": must be in (0,1)"; + } + + const double eps_safe = 1e-12; + + std::vector eps; + auto xt = addNoise(xNorm, t, eps); + if (noisedZscoreOut) *noisedZscoreOut = xt; // the noised input (z-score) the network sees + double s = std::max(sigma(t), eps_safe); + double ab = std::sqrt(std::max(alphabar(t), eps_safe)); + + auto input = buildNetworkInput(xt, condition, t); + const auto& inferNet = (useEMANetworkIfAvailable && useEMANetwork_) ? emaNetwork_ : network_; + auto out = forwardInference(input, inferNet); + + // Recover eps_hat from the network output for any prediction target (SCORE/EPS/V), + // then x0_hat = (xt - sigma*eps_hat)/alpha. + std::vector x0hat(dim_); + for (int i = 0; i < dim_; ++i) { + double epsHat = epsHatFromOutput(out[i], xt[i], s, ab); + x0hat[i] = (xt[i] - s * epsHat) / ab; + } + + SBDMGeneratedSample result; + result.zscore = x0hat; + result.value.resize(dim_); + for (int i = 0; i < dim_; ++i) { + result.value[i] = x0hat[i] * dataStdev_[i] + dataMean_[i]; + } + return result; } } // namespace mu2e From 2aafe5967ab2ddf8e6ed3645b6d01f2d2a86c70b Mon Sep 17 00:00:00 2001 From: YongyiBWu <47263079+YongyiBWu@users.noreply.github.com> Date: Sat, 5 Sep 2026 06:40:58 -0500 Subject: [PATCH 2/2] VDResampler: coordinate transforms, naming convention, pTotal resampler --- STMMC/inc/VDResamplerNameHelper.hh | 194 ++++++ STMMC/inc/VDResamplerPtotResampler.hh | 297 ++++++++ STMMC/inc/VDResamplerTransforms.hh | 962 +++++++++++++++++++++++--- 3 files changed, 1351 insertions(+), 102 deletions(-) create mode 100644 STMMC/inc/VDResamplerNameHelper.hh create mode 100644 STMMC/inc/VDResamplerPtotResampler.hh diff --git a/STMMC/inc/VDResamplerNameHelper.hh b/STMMC/inc/VDResamplerNameHelper.hh new file mode 100644 index 0000000000..0aaf4fd04d --- /dev/null +++ b/STMMC/inc/VDResamplerNameHelper.hh @@ -0,0 +1,194 @@ +#pragma once + +// File-name convention for the VD resampler pipeline (single source of truth). +// +// Names follow the Mu2e dataset convention .mu2e...., +// where the sequencer is _: a 6-digit run number and the 8-digit enumeration of the +// data source. Every artifact embeds versionTag, run number and dataSourceTag so several campaigns +// can coexist in one directory AND the generator can recover all three from a summary file name: +// summary : etc.mu2e.STMVDResamplerConfigure_VD_hitSummary.._.txt +// model : nts.mu2e.STMVDResamplerModel_VD_pdg_.._.dat +// (role: stage1/stage2/allAtOnce) +// ROOT : nts.mu2e.STMVDResamplerConfigure_VD_hitDump.._.root +// (recommended TFileService name; module-external) +// is 'm' for negative pdgIds; is the VD id; is dataSourceIndex() of the +// dataSourceTag. The summary keeps its CSV *content* (comma-separated) under the .txt extension; +// the model keeps its binary content under .dat. +// +// versionTag MUST NOT contain '.' — the parser splits the name on dots, so a dotted version tag +// would be ambiguous. dataSourceTag MUST be one of the enumerated sources (see kDataSourceNames); +// an unrecognized tag is a hard error at encode time rather than a name the generator cannot parse. +// +// The producer (VDResamplerConfigure) and the consumer (VDResamplerGenerateMix) BOTH call these +// helpers, so a change here can never desynchronize the two sides. VDResamplerGenerateFromModel +// takes explicit model-file paths, so it does not use these builders. +// +// Header-only, dependency-free (just the standard library, plus cetlib_except for the encode-time +// errors). Yongyi Wu, Jul. 2026 + +#include "cetlib_except/exception.h" + +#include +#include +#include + +namespace mu2e { +namespace VDResampler { + +// Replace every character that is not [A-Za-z0-9_] with '_', so a tag is safe in a file name. +inline std::string sanitizeTag(const std::string& tag) { + std::string s = tag; + for (char& c : s) { + const bool ok = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_'; + if (!ok) c = '_'; + } + return s; +} + +// 'm' for negative pdgIds (bare number otherwise), matching the model-file convention. +inline std::string pdgFileToken(int pdgId) { + return (pdgId < 0) ? "m" + std::to_string(-pdgId) : std::to_string(pdgId); +} + +// The enumerated data sources, in the order that defines their numeric encoding. The INDEX of a +// tag in this list is what lands in the file name, so entries may be APPENDED but never reordered +// or removed — doing so would silently re-point every existing name at a different source. +inline const std::vector& dataSourceNames() { + static const std::vector kNames = { "EleBeam", "MuBeam", "TargetStops1809", "Neutrals" }; + return kNames; +} + +// Index of dataSourceTag in the enumeration, or -1 if it is not an enumerated source. +inline int dataSourceIndex(const std::string& dataSourceTag) { + const std::vector& names = dataSourceNames(); + for (std::size_t i = 0; i < names.size(); ++i) + if (names[i] == dataSourceTag) return static_cast(i); + return -1; +} + +// Inverse of dataSourceIndex; empty string if the index is out of range. +inline std::string dataSourceFromIndex(int index) { + const std::vector& names = dataSourceNames(); + if (index < 0 || static_cast(index) >= names.size()) return std::string(); + return names[index]; +} + +// Zero-pad a non-negative integer to `width` digits. Values that already exceed the width are +// emitted in full rather than truncated — a wrong-but-complete number beats a silently mangled one. +inline std::string zeroPad(int value, std::size_t width) { + std::string s = std::to_string(value); + return (s.size() >= width) ? s : std::string(width - s.size(), '0') + s; +} + +// The _ sequencer field shared by every name. Throws if dataSourceTag is not one of the +// enumerated sources, or if runNumber is negative. +inline std::string sequencerField(int runNumber, const std::string& dataSourceTag) { + if (runNumber < 0) + throw cet::exception("VDResamplerNameHelper") + << "runNumber must be non-negative (got " << runNumber << ")."; + const int srcIndex = dataSourceIndex(dataSourceTag); + if (srcIndex < 0) { + cet::exception e("VDResamplerNameHelper"); + e << "dataSourceTag \"" << dataSourceTag << "\" is not an enumerated data source. Expected one of:"; + for (const std::string& n : dataSourceNames()) e << " " << n; + e << ". Add it to dataSourceNames() (APPEND only) if it is a new source."; + throw e; + } + return zeroPad(runNumber, 6) + "_" + zeroPad(srcIndex, 8); +} + +// versionTag is embedded between two dots, so a dotted tag would break parseSummaryFileName. +inline std::string checkedVersionTag(const std::string& versionTag) { + if (versionTag.find('.') != std::string::npos) + throw cet::exception("VDResamplerNameHelper") + << "versionTag \"" << versionTag << "\" must not contain '.'."; + return sanitizeTag(versionTag); +} + +inline std::string summaryFileName(const std::string& versionTag, int virtualDetectorID, + const std::string& dataSourceTag, int runNumber) { + return "etc.mu2e.STMVDResamplerConfigure_VD" + std::to_string(virtualDetectorID) + "_hitSummary." + + checkedVersionTag(versionTag) + "." + sequencerField(runNumber, dataSourceTag) + ".txt"; +} + +// role is "stage1" | "stage2" | "allAtOnce". Extension is .dat (binary, full double precision). +inline std::string modelFileName(const std::string& role, const std::string& versionTag, + int virtualDetectorID, const std::string& dataSourceTag, + int pdgId, int runNumber) { + return "nts.mu2e.STMVDResamplerModel_VD" + std::to_string(virtualDetectorID) + + "_pdg" + pdgFileToken(pdgId) + "_" + role + "." + + checkedVersionTag(versionTag) + "." + sequencerField(runNumber, dataSourceTag) + ".dat"; +} + +inline std::string rootDumpFileName(const std::string& versionTag, int virtualDetectorID, + const std::string& dataSourceTag, int runNumber) { + return "nts.mu2e.STMVDResamplerConfigure_VD" + std::to_string(virtualDetectorID) + "_hitDump." + + checkedVersionTag(versionTag) + "." + sequencerField(runNumber, dataSourceTag) + ".root"; +} + +// Recover (versionTag, virtualDetectorID, dataSourceTag, runNumber) from a summary file name of the +// form ".../etc.mu2e.STMVDResamplerConfigure_VD_hitSummary.._.txt". Strips any +// directory prefix and returns false if the pattern does not match. +// +// The name is split on '.' into exactly five fields: tier, "mu2e", description, version, sequencer +// (the ".txt" extension is stripped first). That is unambiguous because versionTag is dot-free (see +// checkedVersionTag) — but the description field may itself contain underscores, so within it we +// anchor on the fixed "STMVDResamplerConfigure_VD" prefix and "_hitSummary" suffix. The source is +// recovered by decoding the 8-digit index back through the enumeration, so an index written by a +// newer build with extra sources will fail here rather than resolve to the wrong source. +inline bool parseSummaryFileName(const std::string& path, std::string& versionTag, + int& virtualDetectorID, std::string& dataSourceTag, + int& runNumber) { + const std::size_t slash = path.find_last_of("/\\"); + std::string base = (slash == std::string::npos) ? path : path.substr(slash + 1); + + const std::string ext = ".txt"; + if (base.size() <= ext.size()) return false; + if (base.compare(base.size() - ext.size(), ext.size(), ext) != 0) return false; + base = base.substr(0, base.size() - ext.size()); + + // Split the remainder on '.' — expect exactly {"etc", "mu2e", , , }. + std::vector fields; + std::size_t start = 0; + while (true) { + const std::size_t dot = base.find('.', start); + if (dot == std::string::npos) { fields.push_back(base.substr(start)); break; } + fields.push_back(base.substr(start, dot - start)); + start = dot + 1; + } + if (fields.size() != 5) return false; + if (fields[0] != "etc" || fields[1] != "mu2e") return false; + + // = "STMVDResamplerConfigure_VD_hitSummary" + const std::string pre = "STMVDResamplerConfigure_VD"; + const std::string suf = "_hitSummary"; + const std::string& desc = fields[2]; + if (desc.size() <= pre.size() + suf.size()) return false; + if (desc.compare(0, pre.size(), pre) != 0) return false; + if (desc.compare(desc.size() - suf.size(), suf.size(), suf) != 0) return false; + const std::string idStr = desc.substr(pre.size(), desc.size() - pre.size() - suf.size()); + if (idStr.empty() || idStr.find_first_not_of("0123456789") != std::string::npos) return false; + + const std::string& ver = fields[3]; + if (ver.empty()) return false; + + // = "_" + const std::string& seq = fields[4]; + const std::size_t us = seq.find('_'); + if (us == std::string::npos) return false; + const std::string runStr = seq.substr(0, us); + const std::string srcStr = seq.substr(us + 1); + if (runStr.empty() || runStr.find_first_not_of("0123456789") != std::string::npos) return false; + if (srcStr.empty() || srcStr.find_first_not_of("0123456789") != std::string::npos) return false; + const std::string src = dataSourceFromIndex(std::stoi(srcStr)); + if (src.empty()) return false; + + versionTag = ver; + virtualDetectorID = std::stoi(idStr); + dataSourceTag = src; + runNumber = std::stoi(runStr); + return true; +} + +} // namespace VDResampler +} // namespace mu2e diff --git a/STMMC/inc/VDResamplerPtotResampler.hh b/STMMC/inc/VDResamplerPtotResampler.hh new file mode 100644 index 0000000000..792513ad86 --- /dev/null +++ b/STMMC/inc/VDResamplerPtotResampler.hh @@ -0,0 +1,297 @@ +#pragma once + +// Non-diffusion 1-D resampler for the total momentum |p| (pTotal, MeV/c), used as +// the V2 two-stage "stage 1". The empirical distribution IS the source data, so this +// needs no training and no saved model: at generation time the source ROOT file is +// read once, physical pTotal is extracted per accepted hit, and samples are drawn +// from the empirical distribution. Operating on PHYSICAL pTotal (not the transformed +// log(pTotal/p0)) is where the fine energy structure is defined and where sub-keV +// resolution is meaningful; the caller applies log/z-score afterward only to build +// the stage-2 conditioning input, while the RAW drawn pTotal is carried directly to +// the V2 inversion (no log/exp round-trip). +// +// Header-only; depends on ROOT (TFile/TTree, TSpline3) and CLHEP randoms. +// Yongyi Wu, Jun. 2026 + +#include +#include +#include +#include +#include +#include + +#include "CLHEP/Random/RandFlat.h" +#include "CLHEP/Random/RandGaussQ.h" +#include "cetlib_except/exception.h" +#include "messagefacility/MessageLogger/MessageLogger.h" + +#include "TFile.h" +#include "TTree.h" +#include "TSpline.h" + +namespace mu2e { +namespace VDResampler { + +// --------------------------------------------------------------------------- +// Stage-1 generation method for pTotal. +// DIFFUSION : NOT handled here — a trained 1-D SBDM is used instead. +// INVERSE_CDF : sorted order-statistics inverse-CDF (no binning; sub-keV). Default. +// SPLINE_CDF : monotone cubic (TSpline3) through the empirical CDF, inverted. +// KDE : Gaussian-kernel smoothing (blurs sharp lines; not recommended +// for sub-keV structure, provided as an option). +// --------------------------------------------------------------------------- +enum class Stage1Method { DIFFUSION = 0, INVERSE_CDF = 1, SPLINE_CDF = 2, KDE = 3 }; + +inline Stage1Method parseStage1Method(const std::string& m, const std::string& moduleName) { + if (m == "DIFFUSION") return Stage1Method::DIFFUSION; + if (m == "INVERSE_CDF") return Stage1Method::INVERSE_CDF; + if (m == "SPLINE_CDF") return Stage1Method::SPLINE_CDF; + if (m == "KDE") return Stage1Method::KDE; + throw cet::exception(moduleName) + << "Unrecognized SBDMstage1Method value \"" << m + << "\" (expected DIFFUSION / INVERSE_CDF / SPLINE_CDF / KDE)."; +} + +// --------------------------------------------------------------------------- +// setBranchAddressChecked — SetBranchAddress that throws instead of returning a code. +// SetBranchAddress returns kMissingBranch and binds NOTHING when the branch is absent, +// so GetEntry then leaves the local untouched and the caller silently reads whatever +// was there. Naming the missing branch turns a wrong-tree mistake into a setup error. +// --------------------------------------------------------------------------- +template +inline void setBranchAddressChecked(TTree* ttree, const char* branch, T* target, + const std::string& moduleName) { + if (ttree->SetBranchAddress(branch, target) < 0) + throw cet::exception(moduleName) + << "TTree '" << ttree->GetName() << "' has no branch '" << branch + << "'. Check that the input file is the expected dump."; +} + +// --------------------------------------------------------------------------- +// forEachAcceptedHitRoot — open `file`/`tree` and invoke cb(x,y,z,t,px,py,pz) for +// every hit passing the standard VD-resampler selection (matching VD id, pdgId +// (0 = any), pz>0). Mirrors the read loop in VDResamplerTrainFromRoot so the +// resampler sees the SAME accepted-hit set the trainer did. Branch schema: +// double time,x,y,z,px,py,pz; int pdgId; ULong64_t virtualdetectorId. +// --------------------------------------------------------------------------- +inline void forEachAcceptedHitRoot( + const std::string& file, const std::string& tree, + unsigned long virtualDetectorID, int pdgID, const std::string& moduleName, + const std::function& cb) +{ + // TFile::Open is used to handle xroot:// paths. + auto fin = std::unique_ptr{TFile::Open(file.c_str(), "READ")}; + if (!fin || fin->IsZombie()) + throw cet::exception(moduleName) << "Cannot open ROOT file: " << file; + TTree* ttree = dynamic_cast(fin->Get(tree.c_str())); + if (!ttree) + throw cet::exception(moduleName) << "Cannot find TTree: " << tree; + + double time = 0., x = 0., y = 0., z = 0., px = 0., py = 0., pz = 0.; + int stepPdgId = 0; + ULong64_t vdId = 0; + setBranchAddressChecked(ttree, "time", &time, moduleName); + setBranchAddressChecked(ttree, "x", &x, moduleName); + setBranchAddressChecked(ttree, "y", &y, moduleName); + setBranchAddressChecked(ttree, "z", &z, moduleName); + setBranchAddressChecked(ttree, "px", &px, moduleName); + setBranchAddressChecked(ttree, "py", &py, moduleName); + setBranchAddressChecked(ttree, "pz", &pz, moduleName); + setBranchAddressChecked(ttree, "pdgId", &stepPdgId, moduleName); + setBranchAddressChecked(ttree, "virtualdetectorId", &vdId, moduleName); + + for (Long64_t i = 0; i < ttree->GetEntries(); ++i) { + ttree->GetEntry(i); + if (vdId != virtualDetectorID || (stepPdgId != pdgID && pdgID != 0) || pz <= 0) + continue; + cb(x, y, z, time, px, py, pz); + } + fin->Close(); +} + +// --------------------------------------------------------------------------- +// PtotResampler — holds the sorted physical pTotal samples and draws from them. +// Built in memory each generation job (no serialization). For SPLINE_CDF/KDE the +// same sorted array is kept and the spline / kernel sum is evaluated on demand. +// --------------------------------------------------------------------------- +class PtotResampler { +public: + // Below this, the empirical distribution is sparsely sampled — a generic + // small-dataset warning (statistics of any method will be noisy), independent of + // the basis/transform resolution. + static constexpr std::size_t kSmallSampleWarn = 1000; + + // Bound on redraws when SPLINE_CDF / KDE produce a non-positive pTotal (physical + // magnitude must be > 0). Exceeding this points at a malformed spline / oversized + // KDE bandwidth rather than an unlucky draw, so we throw instead of looping forever. + static constexpr int kMaxNonPositiveRedraws = 100; + + // Populate from a source ROOT file using the standard selection. After this the + // sorted physical pTotal vector is ready; method-specific structures are built lazily. + void buildFromRoot(const std::string& file, const std::string& tree, + unsigned long virtualDetectorID, int pdgID, + Stage1Method method, const std::string& moduleName) + { + method_ = method; + ptot_.clear(); + forEachAcceptedHitRoot(file, tree, virtualDetectorID, pdgID, moduleName, + [&](double, double, double, double, double px, double py, double pz) { + ptot_.push_back(std::sqrt(px * px + py * py + pz * pz)); + }); + if (ptot_.empty()) + throw cet::exception(moduleName) + << "PtotResampler: no accepted hits in " << file << ":" << tree + << " (VDid=" << virtualDetectorID << ", pdgId=" << pdgID << ", pz>0)."; + std::sort(ptot_.begin(), ptot_.end()); + + // Report the accumulated sample size and, when small, an estimate of the + // achievable resolution. Median adjacent-sample spacing is a rough proxy for + // the inverse-CDF's local resolution (a draw interpolates between neighbors). + const std::size_t N = ptot_.size(); + mf::LogInfo(moduleName) + << "PtotResampler: accumulated " << N << " physical pTotal sample(s) from " + << file << ":" << tree << " (range [" << ptot_.front() << ", " << ptot_.back() + << "] MeV/c)."; + if (N < kSmallSampleWarn) { + mf::LogWarning(moduleName) + << "PtotResampler: only " << N << " samples (< " << kSmallSampleWarn + << ") — a small source dataset; resampled statistics will be noisy. " + << "Estimated median adjacent-sample spacing ~ " << medianSpacing() + << " MeV/c (roughly the finest structure INVERSE_CDF can resolve). " + << "Consider a larger source file."; + } + + if (method_ == Stage1Method::SPLINE_CDF) buildSpline(moduleName); + // KDE bandwidth: Silverman's rule on the sorted data (used only for KDE draws). + if (method_ == Stage1Method::KDE) bandwidth_ = silvermanBandwidth(); + } + + std::size_t size() const { return ptot_.size(); } + + // Draw one physical pTotal (MeV/c) via the configured method. + double draw(CLHEP::RandFlat& rf, CLHEP::RandGaussQ& rg) const { + switch (method_) { + case Stage1Method::SPLINE_CDF: return drawSpline(rf); + case Stage1Method::KDE: return drawKde(rf, rg); + case Stage1Method::INVERSE_CDF: return drawInverseCdf(rf); + case Stage1Method::DIFFUSION: + default: + throw cet::exception("PtotResampler") + << "draw() called with DIFFUSION method (use the stage-1 diffusion model instead)."; + } + } + +private: + Stage1Method method_ = Stage1Method::INVERSE_CDF; + std::vector ptot_; // sorted physical pTotal (MeV/c) + std::unique_ptr cdfSpline_; // SPLINE_CDF: pTotal as a function of CDF value + double bandwidth_ = 0.0; // KDE bandwidth + + // Quantile (pTotal value) at CDF level c in [0,1] by continuous-rank linear + // interpolation between order statistics. With rank q = c*(N-1), c=0 -> ptot_[0] + // and c=1 -> ptot_[N-1]; intermediate c interpolates the two bracketing samples. + // No grid quantization, so local resolution = inter-sample spacing. + double quantile(double c) const { + const std::size_t N = ptot_.size(); + if (N == 1) return ptot_[0]; + const double q = c * static_cast(N - 1); + const std::size_t i = static_cast(std::floor(q)); + if (i + 1 >= N) return ptot_.back(); // only reached at c==1 exactly; no pileup for c<1 + const double frac = q - static_cast(i); + return ptot_[i] + frac * (ptot_[i + 1] - ptot_[i]); + } + + // INVERSE_CDF: draw u~U[0,1) (RandFlat::fire is half-open, so c<1 and the c==1 + // branch in quantile() never fires -> no artificial pileup at the max). + double drawInverseCdf(CLHEP::RandFlat& rf) const { + return quantile(rf.fire()); + } + + // Knot count for the CDF spline. The spline is only a smoothing convenience (the + // sub-keV-faithful path is INVERSE_CDF, which has NO binning); knots merely need to + // be dense enough that the spline tracks the empirical CDF's shape without + // overfitting per-sample noise. Heuristic: ~1 knot per kSamplesPerKnot samples + // (so each spline segment averages over that many points), floored for a smooth + // curve on small datasets and capped to bound the spline size. + int splineKnotCount() const { + constexpr std::size_t kSamplesPerKnot = 100; // ~averaging window per segment + constexpr int kMinKnots = 50; // floor: smoothness on small N + constexpr int kMaxKnots = 100000; // cap: bound TSpline3 size + const std::size_t N = ptot_.size(); + const int byDensity = static_cast(N / kSamplesPerKnot); + int n = std::max(kMinKnots, std::min(kMaxKnots, byDensity)); + return std::min(n, static_cast(N)); // never more knots than samples + } + + // SPLINE_CDF: invert a monotone-cubic fit of pTotal vs the empirical CDF level. + // Knots at evenly spaced CDF levels; each value is the corresponding quantile. + void buildSpline(const std::string& moduleName) { + const int nKnots = splineKnotCount(); + mf::LogInfo(moduleName) + << "PtotResampler: SPLINE_CDF using " << nKnots << " knots over " + << ptot_.size() << " samples."; + std::vector cdf(nKnots), val(nKnots); + for (int k = 0; k < nKnots; ++k) { + const double c = (nKnots == 1) ? 0.5 : static_cast(k) / (nKnots - 1); + cdf[k] = c; + val[k] = quantile(c); + } + cdfSpline_ = std::make_unique("ptotCdfSpline", cdf.data(), val.data(), nKnots); + } + + // Cubic-spline inversion can extrapolate below 0 near the CDF ends; pTotal is a + // physical magnitude, so reject a non-positive draw and try again (bounded). + double drawSpline(CLHEP::RandFlat& rf) const { + for (int attempt = 0; attempt < kMaxNonPositiveRedraws; ++attempt) { + const double p = cdfSpline_->Eval(rf.fire()); + if (p > 0.0) return p; + mf::LogWarning("PtotResampler") + << "SPLINE_CDF drew non-positive pTotal " << p << " MeV/c (spline extrapolation); redrawing."; + } + throw cet::exception("PtotResampler") + << "SPLINE_CDF failed to draw a positive pTotal after " << kMaxNonPositiveRedraws + << " attempts; the spline fit is likely producing negative values."; + } + + // KDE: pick a sample at random and jitter by a Gaussian of width bandwidth_. The + // Gaussian jitter can push a small-pTotal sample below 0; reject and redraw (bounded). + double drawKde(CLHEP::RandFlat& rf, CLHEP::RandGaussQ& rg) const { + for (int attempt = 0; attempt < kMaxNonPositiveRedraws; ++attempt) { + const std::size_t i = static_cast(rf.fireInt(static_cast(ptot_.size()))); + const double p = ptot_[i] + bandwidth_ * rg.fire(); + if (p > 0.0) return p; + mf::LogWarning("PtotResampler") + << "KDE drew non-positive pTotal " << p << " MeV/c (Gaussian jitter); redrawing."; + } + throw cet::exception("PtotResampler") + << "KDE failed to draw a positive pTotal after " << kMaxNonPositiveRedraws + << " attempts; the bandwidth is likely too large for the smallest pTotal samples."; + } + + // Median spacing between adjacent sorted samples — a rough resolution proxy. + double medianSpacing() const { + const std::size_t N = ptot_.size(); + if (N < 2) return 0.0; + std::vector gaps; + gaps.reserve(N - 1); + for (std::size_t i = 1; i < N; ++i) gaps.push_back(ptot_[i] - ptot_[i - 1]); + std::nth_element(gaps.begin(), gaps.begin() + gaps.size() / 2, gaps.end()); + return gaps[gaps.size() / 2]; + } + + double silvermanBandwidth() const { + const std::size_t N = ptot_.size(); + if (N < 2) return 0.0; + double mean = 0.0; + for (double v : ptot_) mean += v; + mean /= static_cast(N); + double var = 0.0; + for (double v : ptot_) var += (v - mean) * (v - mean); + var /= static_cast(N - 1); + const double sigma = std::sqrt(var); + return 1.06 * sigma * std::pow(static_cast(N), -0.2); + } +}; + +} // namespace VDResampler +} // namespace mu2e diff --git a/STMMC/inc/VDResamplerTransforms.hh b/STMMC/inc/VDResamplerTransforms.hh index b6e4178b07..25cc44453c 100644 --- a/STMMC/inc/VDResamplerTransforms.hh +++ b/STMMC/inc/VDResamplerTransforms.hh @@ -2,6 +2,11 @@ #include #include +#include +#include +#include + +#include "cetlib_except/exception.h" namespace mu2e { namespace VDResampler { @@ -15,127 +20,597 @@ namespace mu2e { // tunable momentum scale constexpr double kP0 = 1.0; // MeV/c - // safety constants for numerical stability in the forward and inverse transforms + // safety constants for numerical stability in the forward and inverse transforms. + // NOTE kRadiusSafetyEpsilon has several unrelated uses (an r>eps test in mm before + // dividing by r to build the local polar frame, and a floor on pTot in MeV before + // log), so it is deliberately NOT reused as the radial clamp -- see + // kRhoClampEpsilon below, which is dimensionless and guards u = r/VDr. constexpr double kRadiusSafetyEpsilon = 1e-6; constexpr double kMinSafeTime = 0.1; - // Transform detector-space quantities to model-space variables used for training. - inline void forwardTransformSample( - const double x, - const double y, - const double z, - const double t, - const double px, - const double py, - const double pz, - const double x0, - const double y0, - const double t0, - const double tScale, - const double p0, - const double VDr, - const double VDz0, - double& xTrans, - double& yTrans, - double& tTrans, - double& prTrans, - double& pphiTrans, - double& pzTrans - ) { - // 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. + // Radial clamp: rho = r/VDr is capped at 1-kRhoClampEpsilon before the radial map, + // so an event exactly on (or numerically past) the rim cannot produce u = inf. This + // is a pure numerical guard. It is NOT a fix for the radial mis-modelling described + // under PositionBasis: no generated event ever reached the clamp (0/100000 + // measured), so that was not rim saturation, and tightening or loosening eps does + // not address it. Raising it only truncates real area near the rim (the truncated + // fraction is ~2 eps for a uniform disc), so it is kept at the historical value. + // + // Both PositionBasis maps are atanh-based and so diverge at rho = 1; the clamp is + // what keeps u finite there. + constexpr double kRhoClampEpsilon = 1e-6; + // Independent floor for pz in divisions (slope basis). Distinct from + // kRadiusSafetyEpsilon because it guards a different quantity (longitudinal + // momentum, not transverse radius). If it is ever used, the input pz was + // pathologically small or non-positive (pz>0 is expected from the selection). + constexpr double kPzSafetyEpsilon = 1e-9; + // asinh slope scales for the V2 asinh variant (slopes are near 0 for a forward + // beam, so asinh ~ identity there; this only compresses large-|u| tails). + // when additional kSlopeScales < 1 are used, the regions near u=0 are stretched + // and the tails are more compressed. + // Separate scales for the radial (ur=pr/pz) and azimuthal (uphi=pphi/pz) slopes + // so each can be tuned to its own spread. + constexpr double kUrSlopeScale = 0.05; + constexpr double kUphiSlopeScale = 0.05; + + // Tail-taming of the log-time coordinate for the V3 basis. After the shared + // forwardTime (ln(t/t0)/tScale), V3 applies asinh((base - kTBulkCenter)/kTTailScale) + // to compress a heavy but physical right tail (~1e-6 of events reach z-score ~100 + // after the log alone). kTBulkCenter is the empirical bulk center of ln(t/t0)/tScale + // (measured ~-2.25 across particle species) so the bulk lands at asinh(0)=0 where the + // map is locally linear/symmetric and only the far tail is compressed. asinh is smooth, + // monotone, and exactly invertible with no hard cutoff (unlike asinh(pz/p0), which was + // rejected for its wall at 0). kTTailScale sets the asinh width; smaller => tails + // compressed harder, bulk near center stretched more. + constexpr double kTBulkCenter = -2.25; + constexpr double kTTailScale = 1.0; + + // ------------------------------------------------------------------------ + // PzFallbackStats — accumulates pz-fallback occurrences across a transform + // loop so the CALLER can emit a SINGLE summary warning at the end (instead + // of one line per hit). This struct performs NO I/O so the header stays + // free of messagefacility/iostream: it only records. The caller (a module + // with mf::LogWarning available) inspects count/firstValues after its loop. + // ------------------------------------------------------------------------ + struct PzFallbackStats { + static constexpr std::size_t kMaxSamples = 50; + std::size_t count = 0; // total fallbacks + std::vector firstValues; // first kMaxSamples offending pz values + + void record(double pz) { + ++count; + if (firstValues.size() < kMaxSamples) firstValues.push_back(pz); + } + + // Hits whose extrapolated radius reached the rim, i.e. rho >= 1 before the clamp in + // forwardPosition. Distinct from the pz fallback above: a near-zero pz inflates the + // extrapolation and lands here too, but so does any hit that simply extrapolates + // outside VDr. Recorded separately so the caller can say which happened. + std::size_t clampCount = 0; + std::vector firstRhos; // first kMaxSamples offending rho values + + void recordRhoClamp(double rho) { + ++clampCount; + if (firstRhos.size() < kMaxSamples) firstRhos.push_back(rho); + } + }; + + // ------------------------------------------------------------------------ + // MomentumBasis — selects the momentum transform math used by the ALL-AT-ONCE + // 6-vector (t,x,y, m0,m1,m2). It defines what the three momentum slots mean + // and therefore which forward/inverse math applies. (For two-stage models the + // per-stage layout is given by ModelLayout below; the V2 stages reuse the V2 + // slope/pTotal math.) + // V1_CylindricalTransformed : original basis. A per-component transform of + // the LOCAL CYLINDRICAL momentum (pr, pphi, pz): + // (asinh(pr/p0), asinh(pphi/p0), log(pz/p0)). NOTE this is NOT an exact + // direction/magnitude factorization — it transforms the cylindrical + // components independently. Default, for backward compatibility. + // Momentum slot order: m0=asinh(pr/p0), m1=asinh(pphi/p0), m2=log(pz/p0). + // V2_PtotSlopes : pTotal carries the energy structure; ur,uphi are transverse + // slopes (theta=0 regular, unbounded). Inversion is exact and enforces + // |p|=pTotal, pz>0. Momentum slot order: m0=log(pTotal/p0), m1=ur=pr/pz, + // m2=uphi=pphi/pz (pTotal FIRST). + // V2_PtotSlopesAsinh : as V2_PtotSlopes but slopes (m1,m2) wrapped in + // asinh(ur/kUrSlopeScale), asinh(uphi/kUphiSlopeScale) to tame heavy wide-angle tails. + // V3_PtotSlopesAsinhTimeAsinh : identical momentum treatment to V2_PtotSlopesAsinh, + // plus asinh tail-taming of the log-time coordinate (see kTBulkCenter/kTTailScale). + // Use V2_PtotSlopesAsinh when the extra time transform is NOT wanted. + // In the all-at-once vector t,x,y occupy slots 0,1,2 and momentum occupies 3,4,5. + // ------------------------------------------------------------------------ + enum class MomentumBasis { + V1_CylindricalTransformed = 0, + V2_PtotSlopes = 1, + V2_PtotSlopesAsinh = 2, + V3_PtotSlopesAsinhTimeAsinh = 3 + }; + + // ------------------------------------------------------------------------ + // PositionBasis - the radial boundary-removing map u(rho), rho = r/VDr in [0,1), + // selected PER SPECIES and orthogonal to MomentumBasis. Every variant keeps the + // SAME polar encoding (xTrans,yTrans) = u*(cos theta, sin theta) with + // theta = atan2(dy,dx), so the angular treatment is untouched; only u(rho) + // differs. All are monotone, send rho -> [0,inf), and invert in closed form. + // + // Three effects decide whether a map suits a species. They are easy to conflate, so + // state them explicitly: + // + // EFFECT 1 - RESOLUTION, drho/du. How far a fixed error in the network's output u + // displaces the physical radius. Where drho/du is large, a given u-error costs + // more in rho. This is about precision, and it is the one that matters when a + // species' structure is concentrated somewhere specific. + // + // EFFECT 2 - DENSITY, p(u) = p(rho) drho/du. What the training distribution + // actually looks like in the coordinate the network fits. A density that is flat + // and compactly supported is easy; one with a thin tail is not, because the tail + // is sparsely sampled and a score model generically over-populates it. This is + // about how many samples land where. + // + // For a uniform disc p(rho) ~ rho, so BOTH maps give p(u) -> 0 at the origin: + // that is inherited from the geometry, not a property of either map, and cannot + // be fixed by choosing between them. The maps differ in their TAILS. + // + // EFFECT 3 - CORE WIDTH AFTER NORMALIZATION. The model z-scores each coordinate, + // so what it fits is u divided by the stdev of u over the whole training set. A + // map whose u grows steeply toward the rim inflates that stdev, and dividing by + // it squeezes the populated core into a narrow band -- which shows up directly + // as a generated peak that is too narrow. This is what ruled out the + // rho/(1-rho) style maps, whose u reaches 99 at rho=0.99. It is invisible in u + // itself and only appears after normalization, so it is easily missed when maps + // are compared on their drho/du alone. + // + // A caution from an earlier investigation. A radial mis-modelling seen under + // V1_Atanh (generated/mother ratio ~0.7 at u~0, peaking ~1.38 at u~1.7, back to + // ~0.7 by u~2.8 -- a deficit at BOTH ends with an excess in the middle) looked + // like a coordinate problem but was not: that shape is a distribution contracted + // toward its own mode, the signature of a weak score field. Those runs were later + // found to have stopped with a training loss still above 0.99, i.e. undertrained. + // So before reaching for this enum to explain a radial discrepancy, check the + // training loss and the planner's stop reason first. + // + // V1_Atanh : u = atanh(rho); rho = tanh(u); drho/du = 1-rho^2. + // Effect 1: drho/du = 1 at rho=0, 0.75 at rho=0.5, 0.19 at rho=0.9. Coarsest + // at the centre, finest at the rim. Since u ~ rho + O(rho^3) near the + // origin, it gives no extra ABSOLUTE resolution where the events actually + // are for a centrally-peaked species. + // Effect 2: for a uniform disc p(u) = tanh(u) sech^2(u), peaking near u~0.66 + // with an e^{-2u} tail -- thin, and hard to terminate correctly. + // V2_AtanhSqrt : u = atanh(sqrt(rho)); rho = tanh^2(u); + // drho/du = 2 sqrt(rho) (1-rho). + // Effect 1: the sqrt sends drho/du -> 0 as rho -> 0, so du/drho DIVERGES at + // the centre. Where atanh spends a fixed amount of u range per unit rho near + // the origin, this spends an unbounded amount: the inner region is STRETCHED + // rather than compressed. Away from the centre it converges back toward + // atanh -- 0.71 at rho=0.5, 0.19 at rho=0.9 -- so the rim keeps the same + // tanh saturation and the same class of tail. + // Effect 3: this is the reason to prefer it. Output range stays modest + // (u = 1.47 at rho=0.9, 2.65 at rho=0.99), so the stdev is not inflated by a + // steep rim and the core survives the z-score at a usable width. + // V3_AtanhSq : u = atanh(rho^2); rho = sqrt(tanh(u)); + // drho/du = (1-rho^4)/(2 rho). + // The mirror image of V2: the square sends du/drho -> 0 at the centre, so the + // inner region is COMPRESSED and the u range is spent further out. Fraction of + // the u range (below the rim) landing in each rho band: + // rho atanh(rho) atanh(sqrt) atanh(rho^2) + // 0.0-0.2 0.203 0.481 0.040 + // 0.2-0.4 0.221 0.264 0.121 + // 0.4-0.6 0.270 0.286 0.216 + // 0.6-0.8 0.406 0.412 0.381 + // Use it for a species whose events sit at LARGE rho, where V2 would spend most + // of its resolution on a sparsely populated centre. Its cost is the mirror of + // V2's benefit: the z-scored core comes out narrower (measured coreFrac 0.76 + // against V2's 0.68 on a rim-weighted species), so it trades Effect 3 for + // Effect 1 where the mass is. + // + // All three saturate identically at the rim -- the atanh sets that, not the inner + // power (u-range on rho 0.9-1 is 5.78 for every one of them), so none of them buys + // resolution against the boundary itself. + // ------------------------------------------------------------------------ + enum class PositionBasis { + V1_Atanh = 0, + V2_AtanhSqrt = 1, + V3_AtanhSq = 2 + }; + + // ------------------------------------------------------------------------ + // PeakTag — one monoenergetic line, named by its position on the RAW physical pTotal + // axis. A hit is "in" the line when |pTot - center| <= halfWidth, both in MeV/c. + // + // WHY A DISCRETE TAG. The stage-2 condition is m0 = log(pTotal/p0), z-scored across the + // whole spectrum. For a line at 1.809 MeV a +/-1 keV window is Delta log p ~ 5.5e-4, + // which after the z-score is well under 1e-3 sigma. Making p(t | pTot) change character + // across a gap that narrow would demand an effective slope of order 1e3 in the condition; + // a Fourier condition embedding could only reach that with a base frequency ~1e3, which + // would alias everywhere else in the spectrum. The conditional distribution really does + // differ — a prompt line photon has a different arrival-time shape from the degraded + // continuum at a neighbouring energy — so the clean fix is to hand the network the class + // membership directly instead of asking it to infer it from a near-discontinuity. + // + // The tag is a DETERMINISTIC FUNCTION of pTotal, which is what makes it safe: pTotal is + // already drawn (by stage 1 or the resampler) before stage 2 runs, so the tag is available + // at generation time with no extra sampling, adds no new marginal to model, and cannot + // desynchronize from the condition it was derived from. + // + // CAVEAT worth knowing when reading generated output. The tag only fires correctly if the + // pTotal feeding it reproduces the line. A DIFFUSION stage-1 over log(pTotal/p0) smears a + // 1 keV line, so the tagged fraction will not match the truth. The empirical pTotal + // resampler (INVERSE_CDF / SPLINE_CDF) preserves the line exactly and is the configuration + // this feature is meant for. + // ------------------------------------------------------------------------ + struct PeakTag { + double center = 0.0; // MeV/c, RAW pTotal (not transformed) + double halfWidth = 0.0; // MeV/c, inclusive half-window + }; + + // Index of the class label within the V2 stage-2 condition vector, which is + // {log(pTotal/p0)} or {log(pTotal/p0), label}. The label goes LAST so the existing + // coordinate keeps index 0 and every normalizeCondition(0, ...) call means what it always + // meant. Defined here so the train and generate sides index the same slot by name rather + // than by a repeated literal. + constexpr int kPeakTagCondIdx = 1; + + // Class label for one hit: 0 = in no configured line (continuum), k+1 = inside tags[k]. + // Zero-based tag index k maps to label k+1 so that 0 always means "none", which keeps the + // label meaningful when no tags are configured at all (every hit is 0). + // + // Windows are expected to be disjoint; if they do overlap the FIRST match wins, so the + // label stays single-valued and the ordering of the configured list decides. Callers that + // care validate disjointness up front (see assemblePeakTags). + inline int peakTagFor(double pTot, const std::vector& tags) { + for (std::size_t k = 0; k < tags.size(); ++k) + if (std::abs(pTot - tags[k].center) <= tags[k].halfWidth) + return static_cast(k) + 1; + return 0; + } + + // ------------------------------------------------------------------------ + // ModelLayout — the role/vector layout of a single saved model, so the + // generate side knows how to interpret and invert it from the file alone. + // AllAtOnce6D : one model over (t,x,y, m0,m1,m2) per MomentumBasis. + // TwoStageStage1Ptot1D : V2 two-stage stage 1, a 1-D model over (log(pTotal/p0)). + // TwoStageStage2_5D : V2 two-stage stage 2, a 5-D model over + // (t,x,y,ur,uphi) CONDITIONED on log(pTotal/p0). + // ------------------------------------------------------------------------ + enum class ModelLayout { + AllAtOnce6D = 0, + TwoStageStage1Ptot1D = 1, + TwoStageStage2_5D = 2 + }; + + // ------------------------------------------------------------------------ + // Opaque basis-tag packing. The SBDM stores a single int32 it never interprets; + // VDResampler packs (ModelLayout, PositionBasis, MomentumBasis) into it and + // unpacks on load. Encoding: peakTagged*1000 + layout*100 + position*10 + momentum, + // one decade per field, which caps ModelLayout, PositionBasis and MomentumBasis at + // 10 entries each. Widen the strides if that is ever reached. + // + // This stays backward compatible with the pre-PositionBasis encoding + // (layout*100 + momentum): an old tag has momentum < 10 in the units place and + // nothing in the tens, so it decodes to the same layout and momentum with + // position = 0 = V1_Atanh, which IS the map those models were trained with. A + // pre-tag (v<=6) model still loads as tag 0 = (AllAtOnce6D, V1_Atanh, + // V1_CylindricalTransformed), preserving the original behaviour. + // + // The thousands digit is the peak-tag flag: 1 when the model carries the extra + // categorical condition dim described by PeakTag, 0 otherwise. Every tag written + // before that feature has nothing in the thousands place and so decodes to 0 = not + // tagged, again matching how those models were trained. The flag is what lets the + // generate side decide, FROM THE CHECKPOINT ALONE, whether stage 2 wants a one- or + // two-element condition vector, so a tagged model and an untagged one cannot be fed + // the wrong shape. + // ------------------------------------------------------------------------ + constexpr int kBasisTagPeakStride = 1000; + constexpr int kBasisTagLayoutStride = 100; + constexpr int kBasisTagPositionStride = 10; + + inline int packBasisTag(ModelLayout layout, PositionBasis position, MomentumBasis basis, + bool peakTagged = false) { + return (peakTagged ? kBasisTagPeakStride : 0) + + static_cast(layout) * kBasisTagLayoutStride + + static_cast(position) * kBasisTagPositionStride + + static_cast(basis); + } + // Highest defined enumerator of each field, so unpacking a tag written by a NEWER build + // is rejected instead of static_cast'ing to a value no switch handles. Bump these when + // an enumerator is added. + constexpr int kMaxModelLayout = static_cast(ModelLayout::TwoStageStage2_5D); + constexpr int kMaxPositionBasis = static_cast(PositionBasis::V3_AtanhSq); + constexpr int kMaxMomentumBasis = static_cast(MomentumBasis::V3_PtotSlopesAsinhTimeAsinh); + + inline ModelLayout unpackModelLayout(int tag) { + const int v = (tag % kBasisTagPeakStride) / kBasisTagLayoutStride; + if (v < 0 || v > kMaxModelLayout) + throw cet::exception("VDResamplerTransforms") + << "basisTag " << tag << " decodes to ModelLayout " << v + << ", which this build does not define (max " << kMaxModelLayout + << "). The checkpoint was written by a newer build."; + return static_cast(v); + } + inline PositionBasis unpackPositionBasis(int tag) { + const int v = (tag % kBasisTagLayoutStride) / kBasisTagPositionStride; + if (v < 0 || v > kMaxPositionBasis) + throw cet::exception("VDResamplerTransforms") + << "basisTag " << tag << " decodes to PositionBasis " << v + << ", which this build does not define (max " << kMaxPositionBasis + << "). The checkpoint was written by a newer build."; + return static_cast(v); + } + inline MomentumBasis unpackMomentumBasis(int tag) { + const int v = tag % kBasisTagPositionStride; + if (v < 0 || v > kMaxMomentumBasis) + throw cet::exception("VDResamplerTransforms") + << "basisTag " << tag << " decodes to MomentumBasis " << v + << ", which this build does not define (max " << kMaxMomentumBasis + << "). The checkpoint was written by a newer build."; + return static_cast(v); + } + // True when the model carries the peak-tag condition dim (see PeakTag). + inline bool unpackPeakTagged(int tag) { + return (tag / kBasisTagPeakStride) != 0; + } + + // ------------------------------------------------------------------------ + // Build constants — the map-defining numeric parameters, recorded in every + // checkpoint so a retune is detected on load rather than silently changing + // what a stored model means. The basis enums say which map; these say what + // it was tuned to. + // + // The id is the contract, not the position: ids are explicit, never reused, + // and new ones are appended. Pure numerical guards (kRadiusSafetyEpsilon, + // kMinSafeTime, kRhoClampEpsilon, kPzSafetyEpsilon) are excluded — they do + // not define the map. + // ------------------------------------------------------------------------ + enum class BuildConstantId : int { + kVDr = 0, // position, every PositionBasis + kVDz0 = 1, // position, every PositionBasis + kP0 = 2, // momentum, every MomentumBasis + kT0 = 3, // time, every MomentumBasis + kTScale = 4, // time, every MomentumBasis + kUrSlopeScale = 5, // momentum, MomentumBasis V2_PtotSlopesAsinh + V3 + kUphiSlopeScale = 6, // momentum, MomentumBasis V2_PtotSlopesAsinh + V3 + kTBulkCenter = 7, // time, MomentumBasis V3 only + kTTailScale = 8 // time, MomentumBasis V3 only + }; + + inline const char* buildConstantName(BuildConstantId id) { + switch (id) { + case BuildConstantId::kVDr: return "VDr"; + case BuildConstantId::kVDz0: return "VDz0"; + case BuildConstantId::kP0: return "kP0"; + case BuildConstantId::kT0: return "kT0"; + case BuildConstantId::kTScale: return "kTScale"; + case BuildConstantId::kUrSlopeScale: return "kUrSlopeScale"; + case BuildConstantId::kUphiSlopeScale: return "kUphiSlopeScale"; + case BuildConstantId::kTBulkCenter: return "kTBulkCenter"; + case BuildConstantId::kTTailScale: return "kTTailScale"; + } + return "unknown"; + } + + // This job's values. VDr/VDz0 are per-job (from the training plan) so they are + // passed in; the rest are this build's compile-time constants. All are written + // regardless of basis — an unused one costs nothing and keeps the list uniform. + inline std::vector> currentBuildConstants(double VDr, double VDz0) { + return { + {static_cast(BuildConstantId::kVDr), VDr}, + {static_cast(BuildConstantId::kVDz0), VDz0}, + {static_cast(BuildConstantId::kP0), kP0}, + {static_cast(BuildConstantId::kT0), kT0}, + {static_cast(BuildConstantId::kTScale), kTScale}, + {static_cast(BuildConstantId::kUrSlopeScale), kUrSlopeScale}, + {static_cast(BuildConstantId::kUphiSlopeScale), kUphiSlopeScale}, + {static_cast(BuildConstantId::kTBulkCenter), kTBulkCenter}, + {static_cast(BuildConstantId::kTTailScale), kTTailScale} + }; + } + + // Compare a checkpoint's stored constants against this job's, throwing on the first + // disagreement. An empty list means a pre-v9 checkpoint: current values are assumed + // and the check is skipped. A non-empty list must be complete. + inline void checkBuildConstants(const std::vector>& stored, + double VDr, double VDz0, + const std::string& what, const std::string& moduleName) { + if (stored.empty()) return; + + for (const auto& expected : currentBuildConstants(VDr, VDz0)) { + const auto it = std::find_if(stored.begin(), stored.end(), + [&](const std::pair& kv) { + return kv.first == expected.first; + }); + const auto id = static_cast(expected.first); + if (it == stored.end()) + throw cet::exception(moduleName) + << what << " records build constants but is missing '" << buildConstantName(id) + << "' (id " << expected.first << ")."; + if (it->second != expected.second) + throw cet::exception(moduleName) + << what << " was built with " << buildConstantName(id) << " = " << it->second + << ", but this job uses " << expected.second << ". Re-train, or restore the value."; + } + } + + // Cross-check a checkpoint's recorded pdgId against the one inferred from the model + // file name. A stored 0 means a pre-v9 checkpoint, so the inferred value stands. + inline void checkPdgId(int stored, int inferred, + const std::string& what, const std::string& moduleName) { + if (stored == 0) return; + if (stored != inferred) + throw cet::exception(moduleName) + << what << " was trained for pdgId " << stored << ", but this job inferred " + << inferred << " from the model file name."; + } + + // Human-readable enum names and a full basisTag decode, shared by the train and + // generate modules so an opaque tag is never printed raw. Pure string helpers — + // they add no I/O dependency, keeping this header messagefacility-free. + inline const char* modelLayoutName(ModelLayout l) { + switch (l) { + case ModelLayout::AllAtOnce6D: return "AllAtOnce6D"; + case ModelLayout::TwoStageStage1Ptot1D: return "TwoStageStage1Ptot1D"; + case ModelLayout::TwoStageStage2_5D: return "TwoStageStage2_5D"; + } + return "unknown"; + } + inline const char* positionBasisName(PositionBasis p) { + switch (p) { + case PositionBasis::V1_Atanh: return "V1_Atanh"; + case PositionBasis::V2_AtanhSqrt: return "V2_AtanhSqrt"; + case PositionBasis::V3_AtanhSq: return "V3_AtanhSq"; + } + return "unknown"; + } + inline const char* momentumBasisName(MomentumBasis b) { + switch (b) { + case MomentumBasis::V1_CylindricalTransformed: return "V1_CylindricalTransformed"; + case MomentumBasis::V2_PtotSlopes: return "V2_PtotSlopes"; + case MomentumBasis::V2_PtotSlopesAsinh: return "V2_PtotSlopesAsinh"; + case MomentumBasis::V3_PtotSlopesAsinhTimeAsinh: return "V3_PtotSlopesAsinhTimeAsinh"; + } + return "unknown"; + } + // Decode an opaque basisTag (peakTagged*1000 + layout*100 + position*10 + momentum) + // into "layout=<...>, basis=<...> (tag )" for logs and error messages. + // Decodes the raw fields rather than calling the unpack* helpers: this is used INSIDE the + // messages those helpers throw, so it must render an out-of-range tag rather than throw + // again. The name helpers return "unknown" for a value they do not cover. + inline std::string basisTagToString(int tag) { + std::ostringstream os; + os << "layout=" << modelLayoutName(static_cast( + (tag % kBasisTagPeakStride) / kBasisTagLayoutStride)) + << ", position=" << positionBasisName(static_cast( + (tag % kBasisTagLayoutStride) / kBasisTagPositionStride)) + << ", basis=" << momentumBasisName(static_cast( + tag % kBasisTagPositionStride)) + << ", peakTagged=" << (unpackPeakTagged(tag) ? "yes" : "no") + << " (tag " << tag << ")"; + return os.str(); + } + + // ------------------------------------------------------------------------ + // Shared geometry helpers (basis-independent) + // ------------------------------------------------------------------------ + + // Extrapolate (x,y) along the momentum to the nominal VDz0 plane, then shift to + // detector-centered coordinates. Outputs dx, dy and the in-plane radius r. + inline void extrapolateAndCenter( + double x, double y, double z, double px, double py, double pz, + double x0, double y0, double VDz0, + double& dx, double& dy, double& r) + { const double extrapolationFactor = (VDz0 - z) / pz; const double xExtrapolated = x + extrapolationFactor * px; const double yExtrapolated = y + extrapolationFactor * py; + dx = xExtrapolated - x0; + dy = yExtrapolated - y0; + r = std::sqrt(dx * dx + dy * dy); + } - // 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') + // The radial map u(rho) and its inverse, selected by PositionBasis. Split out from + // forwardPosition/invertPosition so the two directions cannot drift apart: they are + // exact inverses for every enum value, which invertPosition relies on. Also used by + // the validation plots to build the transformed radial coordinate. + inline double radialForward(double rho, PositionBasis basis) { + switch (basis) { + case PositionBasis::V2_AtanhSqrt: { + const double s = std::sqrt(rho); + return 0.5 * std::log((1.0 + s) / (1.0 - s)); // atanh(sqrt(rho)) + } + case PositionBasis::V3_AtanhSq: { + const double q = rho * rho; + return 0.5 * std::log((1.0 + q) / (1.0 - q)); // atanh(rho^2) + } + case PositionBasis::V1_Atanh: + return 0.5 * std::log((1.0 + rho) / (1.0 - rho)); // atanh(rho) + } + // No default arm, so -Wswitch flags a new enumerator at compile time instead of + // letting it fall through to the V1 map. Unreachable for a valid enum value. + throw cet::exception("VDResamplerTransforms") + << "radialForward: unhandled PositionBasis " << static_cast(basis); + } + inline double radialInverse(double u, PositionBasis basis) { + switch (basis) { + case PositionBasis::V2_AtanhSqrt: { + const double th = std::tanh(u); + return th * th; // tanh^2(u) + } + case PositionBasis::V3_AtanhSq: + // tanh(u) >= 0 for u >= 0, which radialForward always produces. + return std::sqrt(std::tanh(u)); // sqrt(tanh(u)) + case PositionBasis::V1_Atanh: + return std::tanh(u); + } + throw cet::exception("VDResamplerTransforms") + << "radialInverse: unhandled PositionBasis " << static_cast(basis); + } - // Shift to detector-centered coordinates and map to transformed position variables. - const double dx = xExtrapolated - x0; - const double dy = yExtrapolated - y0; - // polar coordinates - const double r = std::sqrt(dx * dx + dy * dy); + // Forward position transform: (dx,dy,r) -> (xTrans,yTrans) = u(rho)*(cos,sin theta) + // with rho = r/VDr. The angular part is identical for every PositionBasis; only the + // radial map differs (see PositionBasis for which to use and why). + inline void forwardPosition(double dx, double dy, double r, double VDr, + PositionBasis basis, + double& xTrans, double& yTrans, + PzFallbackStats* stats = nullptr) + { double rho = r / VDr; - // numerical safety (avoid rho >= 1) - rho = std::min(rho, 1.0 - kRadiusSafetyEpsilon); - // boundary-removing transform - // u = atanh(r/R) - const double u = 0.5 * std::log((1.0 + rho) / (1.0 - rho)); - // angle + // Two distinct cases share this clamp: the documented rho=1 guard (so u cannot be + // inf; see kRhoClampEpsilon) and a hit that extrapolated outside VDr entirely, which + // is silently relocated to the rim. Only the latter is worth reporting, so record + // rho >= 1 rather than every application of the clamp. + if (rho >= 1.0 && stats) stats->recordRhoClamp(rho); + rho = std::min(rho, 1.0 - kRhoClampEpsilon); + const double u = radialForward(rho, basis); const double theta = std::atan2(dy, dx); - // map back to Cartesian-like coordinates xTrans = u * std::cos(theta); yTrans = 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 > kRadiusSafetyEpsilon) { // avoid division by zero, if r is very small, we can approximate pr ~ px and pphi ~ py - const double rx = dx / r; // unit vector in the radial direction - const double ry = dy / r; - const double phix = -ry; // unit vector in the angular direction - const double phiy = rx; - pr = px * rx + py * ry; // radial momentum component - pphi = px * phix + py * phiy; // angular momentum component - } else { - pr = px; - pphi = py; - } - // momentum scaling - prTrans = std::asinh(pr / p0); // tunable scale where I want best resolution - pphiTrans = std::asinh(pphi / p0); - pzTrans = std::asinh(pz / p0); - - // time transform - const double tSafe = (t > kMinSafeTime) ? t : kMinSafeTime; // avoid log(0) - tTrans = std::log(tSafe / t0) / tScale; } - // Invert transformed sample coordinates/momenta back to detector-space quantities. - inline void invertGeneratedSample( - const double xTrans, - const double yTrans, - const double tTrans, - const double prTrans, - const double pphiTrans, - const double pzTrans, - const double x0, - const double y0, - const double t0, - const double tScale, - const double p0, - const double VDr, - const double VDz0, - double& x, - double& y, - double& z, - double& t, - double& px, - double& py, - double& pz - ) { + // Inverse position transform: (xTrans,yTrans) -> detector-space (x,y) and the local + // frame (dx,dy,r) needed to rotate momentum back. MUST be given the same + // PositionBasis the forward transform used; radialInverse is its exact inverse. + // + // Deliberately NOT clamped to the forward map's kRhoClampEpsilon range: every + // radialInverse already returns rho<1 by construction, and a generated u above the + // forward u_max is better left to land smoothly within the last fraction of a mm + // than snapped onto a hard edge, which would build a rim spike. The generated-vs- + // source radial ratio near rho=1 is the diagnostic for whether that tail matters. + inline void invertPosition(double xTrans, double yTrans, double x0, double y0, + double VDr, PositionBasis basis, + double& x, double& y, double& dx, double& dy, double& r) + { 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 = dx + x0; - y = dy + y0; - z = VDz0; - - t = t0 * std::exp(tTrans * tScale); + const double rho = radialInverse(u, basis); + r = rho * VDr; + dx = r * std::cos(theta); + dy = r * std::sin(theta); + x = dx + x0; + y = dy + y0; + } - const double pr = p0 * std::sinh(prTrans); - const double pphi = p0 * std::sinh(pphiTrans); - pz = p0 * std::sinh(pzTrans); + // Project Cartesian (px,py) onto the LOCAL cylindrical frame defined by (dx,dy). + // Returns radial (pr) and azimuthal (pphi) momentum. Universal across bases. + inline void cartesianToLocalPolar(double px, double py, double dx, double dy, double r, + double& pr, double& pphi) + { + if (r > kRadiusSafetyEpsilon) { // else r~0: pr~px, pphi~py + const double rx = dx / r, ry = dy / r; // radial unit vector + const double phix = -ry, phiy = rx; // azimuthal unit vector + pr = px * rx + py * ry; + pphi = px * phix + py * phiy; + } else { + pr = px; + pphi = py; + } + } + // Rotate local cylindrical (pr,pphi) back to Cartesian (px,py) using (dx,dy). + // Inverse of cartesianToLocalPolar. Universal across bases. + inline void localPolarToCartesian(double pr, double pphi, double dx, double dy, double r, + double& px, double& py) + { if (r > kRadiusSafetyEpsilon) { - const double rx = dx / r; - const double ry = dy / r; - const double phix = -ry; - const double phiy = rx; + const double rx = dx / r, ry = dy / r; + const double phix = -ry, phiy = rx; px = pr * rx + pphi * phix; py = pr * ry + pphi * phiy; } else { @@ -144,5 +619,288 @@ namespace mu2e { } } + // Forward time transform (shared). + inline double forwardTime(double t, double t0, double tScale) { + const double tSafe = (t > kMinSafeTime) ? t : kMinSafeTime; // avoid log(0) + return std::log(tSafe / t0) / tScale; + } + + // Inverse time transform (shared). + inline double invertTime(double tTrans, double t0, double tScale) { + return t0 * std::exp(tTrans * tScale); + } + + // V3 time transform: asinh tail-taming layered on top of the shared log-time. + // Centered at kTBulkCenter so the bulk maps near asinh(0)=0. Exactly invertible. + inline double forwardTimeAsinh(double t, double t0, double tScale) { + const double base = forwardTime(t, t0, tScale); // ln(tSafe/t0)/tScale + return std::asinh((base - kTBulkCenter) / kTTailScale); + } + inline double invertTimeAsinh(double tTrans, double t0, double tScale) { + const double base = kTBulkCenter + kTTailScale * std::sinh(tTrans); + return t0 * std::exp(base * tScale); + } + + // Basis feature predicates — single source of truth for "which knobs does this + // basis turn on", so callers never enumerate basis values themselves. A basis is + // added here once and every dispatch site follows. + inline bool basisUsesAsinhSlopes(MomentumBasis b) { + return b == MomentumBasis::V2_PtotSlopesAsinh + || b == MomentumBasis::V3_PtotSlopesAsinhTimeAsinh; + } + inline bool basisUsesAsinhTime(MomentumBasis b) { + return b == MomentumBasis::V3_PtotSlopesAsinhTimeAsinh; + } + + // Basis-aware time transforms: dispatch to the asinh variant for bases that + // request it, else the plain log-time. Use these wherever a bare forwardTime/ + // invertTime would otherwise be called against a known basis. + inline double forwardTimeForBasis(double t, double t0, double tScale, MomentumBasis b) { + return basisUsesAsinhTime(b) ? forwardTimeAsinh(t, t0, tScale) + : forwardTime(t, t0, tScale); + } + inline double invertTimeForBasis(double tTrans, double t0, double tScale, MomentumBasis b) { + return basisUsesAsinhTime(b) ? invertTimeAsinh(tTrans, t0, tScale) + : invertTime(tTrans, t0, tScale); + } + + // Shared momentum reconstruction for V2 given a RAW physical pTotal (MeV/c) and + // the (already de-asinh'd) slopes ur,uphi. Lets the resampler path pass its raw + // drawn pTotal directly (no log/exp round-trip); also used by invertGeneratedSampleV2. + inline void invertMomentumV2FromPtot( + double pTot, double ur, double uphi, double dx, double dy, double r, + double& px, double& py, double& pz) + { + pz = pTot / std::sqrt(1.0 + ur * ur + uphi * uphi); // >0; denom >= 1, no guard needed + const double pr = ur * pz; + const double pphi = uphi * pz; + localPolarToCartesian(pr, pphi, dx, dy, r, px, py); + } + + // ======================================================================== + // V1_CylindricalTransformed (original basis; default) + // momentum slots: m0=asinh(pr/p0), m1=asinh(pphi/p0), m2=log(pz/p0) + // ======================================================================== + inline void forwardTransformSampleV1( + const double x, const double y, const double z, const double t, + const double px, const double py, const double pz, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& xTrans, double& yTrans, double& tTrans, + double& prTrans, double& pphiTrans, double& pzTrans, + const PositionBasis posBasis = PositionBasis::V1_Atanh, + PzFallbackStats* pzStats = nullptr) + { + // Floored for the extrapolation's divide only. Unlike V2 the momentum slot here is + // log(pz/p0), which stays finite for a tiny pz, so pz itself is left untouched below. + double pzSafe = pz; + if (std::abs(pz) < kPzSafetyEpsilon) { + if (pzStats) pzStats->record(pz); + pzSafe = kPzSafetyEpsilon; + } + + double dx, dy, r; + extrapolateAndCenter(x, y, z, px, py, pzSafe, x0, y0, VDz0, dx, dy, r); + forwardPosition(dx, dy, r, VDr, posBasis, xTrans, yTrans, pzStats); + + double pr, pphi; + cartesianToLocalPolar(px, py, dx, dy, r, pr, pphi); + prTrans = std::asinh(pr / p0); + pphiTrans = std::asinh(pphi / p0); + pzTrans = std::log(pz / p0); // tried asinh(pz/p0) but hard cutoff at 0 was unfriendly for DM + + tTrans = forwardTime(t, t0, tScale); + } + + inline void invertGeneratedSampleV1( + const double xTrans, const double yTrans, const double tTrans, + const double prTrans, const double pphiTrans, const double pzTrans, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& x, double& y, double& z, double& t, + double& px, double& py, double& pz, + const PositionBasis posBasis = PositionBasis::V1_Atanh) + { + double dx, dy, r; + invertPosition(xTrans, yTrans, x0, y0, VDr, posBasis, x, y, dx, dy, r); + z = VDz0; + t = invertTime(tTrans, t0, tScale); + + const double pr = p0 * std::sinh(prTrans); + const double pphi = p0 * std::sinh(pphiTrans); + pz = p0 * std::exp(pzTrans); + localPolarToCartesian(pr, pphi, dx, dy, r, px, py); + } + + // ======================================================================== + // V2_PtotSlopes / V2_PtotSlopesAsinh (all-at-once 6-vector) + // momentum slots: m0=log(pTotal/p0), m1=ur=pr/pz, m2=uphi=pphi/pz + // asinh variant wraps the slopes: m1=asinh(ur/kUrSlopeScale), m2=asinh(uphi/kUphiSlopeScale) + // Inversion: pz = pTotal/sqrt(1+ur^2+uphi^2) (>0), pr=ur*pz, pphi=uphi*pz. + // pzStats (optional): records pz-fallback occurrences for a single summary + // warning by the caller (see PzFallbackStats). The inverse needs no such + // guard — its denominator sqrt(1+ur^2+uphi^2) >= 1. + // ======================================================================== + inline void forwardTransformSampleV2( + const double x, const double y, const double z, const double t, + const double px, const double py, const double pz, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& xTrans, double& yTrans, double& tTrans, + double& pTotTrans, double& urTrans, double& uphiTrans, + const bool asinhSlopes, const bool asinhTime = false, + PzFallbackStats* pzStats = nullptr, + const PositionBasis posBasis = PositionBasis::V1_Atanh) + { + // Floored BEFORE the extrapolation, which also divides by pz: an unfloored pz~1e-30 + // sends the extrapolation factor (and the radius with it) to ~1e30, and the hit is + // then silently relocated to the rim by the clamp in forwardPosition. + double pzSafe = pz; + if (std::abs(pz) < kPzSafetyEpsilon) { + if (pzStats) pzStats->record(pz); + pzSafe = kPzSafetyEpsilon; + } + + double dx, dy, r; + extrapolateAndCenter(x, y, z, px, py, pzSafe, x0, y0, VDz0, dx, dy, r); + forwardPosition(dx, dy, r, VDr, posBasis, xTrans, yTrans, pzStats); + + double pr, pphi; + cartesianToLocalPolar(px, py, dx, dy, r, pr, pphi); + + double ur = pr / pzSafe; + double uphi = pphi / pzSafe; + if (asinhSlopes) { + ur = std::asinh(ur / kUrSlopeScale); + uphi = std::asinh(uphi / kUphiSlopeScale); + } + + const double pTot = std::sqrt(px * px + py * py + pz * pz); + pTotTrans = std::log(std::max(pTot, kRadiusSafetyEpsilon) / p0); + urTrans = ur; + uphiTrans = uphi; + + tTrans = asinhTime ? forwardTimeAsinh(t, t0, tScale) : forwardTime(t, t0, tScale); + } + + inline void invertGeneratedSampleV2( + const double xTrans, const double yTrans, const double tTrans, + const double pTotTrans, const double urTrans, const double uphiTrans, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& x, double& y, double& z, double& t, + double& px, double& py, double& pz, + const bool asinhSlopes, const bool asinhTime = false, + const PositionBasis posBasis = PositionBasis::V1_Atanh) + { + double dx, dy, r; + invertPosition(xTrans, yTrans, x0, y0, VDr, posBasis, x, y, dx, dy, r); + z = VDz0; + t = asinhTime ? invertTimeAsinh(tTrans, t0, tScale) : invertTime(tTrans, t0, tScale); + + double ur = urTrans, uphi = uphiTrans; + if (asinhSlopes) { + ur = kUrSlopeScale * std::sinh(ur); + uphi = kUphiSlopeScale * std::sinh(uphi); + } + const double pTot = p0 * std::exp(pTotTrans); + invertMomentumV2FromPtot(pTot, ur, uphi, dx, dy, r, px, py, pz); + } + + // ------------------------------------------------------------------------ + // V2 slope helpers (de-asinh / asinh of the slope pair) — used by the two-stage + // generate path which assembles (pTotal, ur, uphi) from stage1 + stage2 and the + // raw resampled pTotal, then calls invertMomentumV2FromPtot directly. + // ------------------------------------------------------------------------ + inline void v2DecodeSlopes(double m1, double m2, bool asinhSlopes, + double& ur, double& uphi) + { + if (asinhSlopes) { ur = kUrSlopeScale * std::sinh(m1); uphi = kUphiSlopeScale * std::sinh(m2); } + else { ur = m1; uphi = m2; } + } + + // ======================================================================== + // Dispatching wrappers for the ALL-AT-ONCE 6-vector. Default basis = + // V1_CylindricalTransformed so existing callers (which omit the basis argument) + // are unchanged / backward compatible. Momentum out-params are basis-neutral + // here (m0,m1,m2). (Two-stage models assemble the V2 vector explicitly and use + // invertMomentumV2FromPtot / v2DecodeSlopes rather than these wrappers.) + // ======================================================================== + inline void forwardTransformSample( + const double x, const double y, const double z, const double t, + const double px, const double py, const double pz, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& xTrans, double& yTrans, double& tTrans, + double& m0, double& m1, double& m2, + const MomentumBasis basis = MomentumBasis::V1_CylindricalTransformed, + PzFallbackStats* pzStats = nullptr, + const PositionBasis posBasis = PositionBasis::V1_Atanh) + { + switch (basis) { + case MomentumBasis::V2_PtotSlopes: + forwardTransformSampleV2(x, y, z, t, px, py, pz, x0, y0, t0, tScale, p0, VDr, VDz0, + xTrans, yTrans, tTrans, m0, m1, m2, + /*asinhSlopes=*/false, /*asinhTime=*/false, pzStats, posBasis); + return; + case MomentumBasis::V2_PtotSlopesAsinh: + forwardTransformSampleV2(x, y, z, t, px, py, pz, x0, y0, t0, tScale, p0, VDr, VDz0, + xTrans, yTrans, tTrans, m0, m1, m2, + /*asinhSlopes=*/true, /*asinhTime=*/false, pzStats, posBasis); + return; + case MomentumBasis::V3_PtotSlopesAsinhTimeAsinh: + forwardTransformSampleV2(x, y, z, t, px, py, pz, x0, y0, t0, tScale, p0, VDr, VDz0, + xTrans, yTrans, tTrans, m0, m1, m2, + /*asinhSlopes=*/true, /*asinhTime=*/true, pzStats, posBasis); + return; + case MomentumBasis::V1_CylindricalTransformed: + forwardTransformSampleV1(x, y, z, t, px, py, pz, x0, y0, t0, tScale, p0, VDr, VDz0, + xTrans, yTrans, tTrans, m0, m1, m2, posBasis, pzStats); + return; + } + // No default arm, so -Wswitch flags a new enumerator at compile time. Each case + // returns and this throws, which also tells the compiler every path either assigns + // the outputs or leaves — without it callers get -Wmaybe-uninitialized. + throw cet::exception("VDResamplerTransforms") + << "forwardTransformSample: unhandled MomentumBasis " << static_cast(basis); + } + + inline void invertGeneratedSample( + const double xTrans, const double yTrans, const double tTrans, + const double m0, const double m1, const double m2, + const double x0, const double y0, const double t0, + const double tScale, const double p0, const double VDr, const double VDz0, + double& x, double& y, double& z, double& t, + double& px, double& py, double& pz, + const MomentumBasis basis = MomentumBasis::V1_CylindricalTransformed, + const PositionBasis posBasis = PositionBasis::V1_Atanh) + { + switch (basis) { + case MomentumBasis::V2_PtotSlopes: + invertGeneratedSampleV2(xTrans, yTrans, tTrans, m0, m1, m2, x0, y0, t0, tScale, p0, VDr, VDz0, + x, y, z, t, px, py, pz, /*asinhSlopes=*/false, /*asinhTime=*/false, + posBasis); + return; + case MomentumBasis::V2_PtotSlopesAsinh: + invertGeneratedSampleV2(xTrans, yTrans, tTrans, m0, m1, m2, x0, y0, t0, tScale, p0, VDr, VDz0, + x, y, z, t, px, py, pz, /*asinhSlopes=*/true, /*asinhTime=*/false, + posBasis); + return; + case MomentumBasis::V3_PtotSlopesAsinhTimeAsinh: + invertGeneratedSampleV2(xTrans, yTrans, tTrans, m0, m1, m2, x0, y0, t0, tScale, p0, VDr, VDz0, + x, y, z, t, px, py, pz, /*asinhSlopes=*/true, /*asinhTime=*/true, + posBasis); + return; + case MomentumBasis::V1_CylindricalTransformed: + invertGeneratedSampleV1(xTrans, yTrans, tTrans, m0, m1, m2, x0, y0, t0, tScale, p0, VDr, VDz0, + x, y, z, t, px, py, pz, posBasis); + return; + } + // See forwardTransformSample: each case returns and this throws, so the compiler can + // see every path assigns the outputs or leaves. + throw cet::exception("VDResamplerTransforms") + << "invertGeneratedSample: unhandled MomentumBasis " << static_cast(basis); + } + } // namespace VDResampler } // namespace mu2e