File indexing completed on 2026-09-01 09:34:06
0001 """Shared helpers used by the standalone energy-response reproduction."""
0002
0003 import os
0004 import sys
0005 from dataclasses import dataclass
0006 from typing import Callable, Optional
0007
0008 import matplotlib.pyplot as plt
0009 import numpy as np
0010 import pandas as pd
0011 from scipy.optimize import curve_fit
0012
0013 LOAD_CHUNK_SIZE = 1_000_000
0014 nETruthBins = 30
0015 nERecoBins = 150
0016
0017
0018 def applyPlotStyle():
0019 """Apply 2× scaling to matplotlib font and line sizes for readability."""
0020 plt.rcParams.update({
0021 'axes.titlesize': 24, 'axes.labelsize': 22,
0022 'xtick.labelsize': 20, 'ytick.labelsize': 20,
0023 'legend.fontsize': 16, 'figure.titlesize': 18,
0024 'lines.linewidth': 1.3, 'lines.markersize': 7.8,
0025 })
0026
0027
0028 def crystal_ball(x, A, mu, sigma, alpha, n):
0029 """Evaluate the un-normalized Crystal Ball model."""
0030 x = np.array(x)
0031 t = (x - mu) / sigma
0032 abs_alpha = np.abs(alpha)
0033 a = (n / abs_alpha) ** n * np.exp(-0.5 * abs_alpha ** 2)
0034 b = n / abs_alpha - abs_alpha
0035 result = np.zeros_like(t)
0036 gaussian = t < abs_alpha
0037 result[gaussian] = np.exp(-0.5 * t[gaussian] ** 2)
0038 result[~gaussian] = a * (b + t[~gaussian]) ** (-n)
0039 return A * result
0040
0041
0042 def loadData(path, pdg, columns, chunksize=LOAD_CHUNK_SIZE):
0043 """Yield filtered DataFrame chunks from a space-separated data file.
0044
0045 Each chunk is filtered to status >= 0 and the requested PDG value; empty
0046 chunks are skipped so downstream histogram accumulation can stream safely.
0047 """
0048 reader = pd.read_csv(path, sep=' ', names=columns, chunksize=chunksize)
0049 for raw_chunk in reader:
0050 chunk = raw_chunk[(raw_chunk['status'] >= 0) & (raw_chunk['PDG'] == pdg)]
0051 if len(chunk):
0052 yield chunk
0053
0054
0055 def buildSliceHistogramFromChunks(chunkIter, xcol, ycol, xbins, ybins,
0056 xTransform=None, yTransform=None):
0057 """Stream-accumulate a 2D histogram from DataFrame chunks.
0058
0059 Optional transforms replace the named columns and are useful for derived
0060 quantities. Returns centers, histogram, and Poisson-style errors.
0061 """
0062 hist = np.zeros((len(xbins) - 1, len(ybins) - 1))
0063 for chunk in chunkIter:
0064 x = xTransform(chunk) if xTransform is not None else chunk[xcol].to_numpy()
0065 y = yTransform(chunk) if yTransform is not None else chunk[ycol].to_numpy()
0066 h, _, _ = np.histogram2d(x, y, bins=(xbins, ybins))
0067 hist += h
0068 err = hist / np.sqrt(hist)
0069 xcenters = 0.5 * (xbins[1:] + xbins[:-1])
0070 ycenters = 0.5 * (ybins[1:] + ybins[:-1])
0071 return xcenters, ycenters, hist, err
0072
0073
0074 def plotThetaDistribution(ax, thetaCounts, bins):
0075 """Draw the theta-distribution side panel from accumulated bin counts."""
0076 ax.stairs(thetaCounts, bins)
0077 ax.set_xlabel(r'$\theta$ (deg)')
0078
0079
0080 def p0CrystalBall(ycenters, histSlice):
0081 """Build robust initial Crystal Ball parameters for one histogram slice."""
0082 cdf = np.cumsum(histSlice)
0083 cdf = cdf / cdf[-1]
0084 mu0 = float(np.interp(0.5, cdf, ycenters))
0085 y_low = float(np.interp(0.1587, cdf, ycenters))
0086 y_high = float(np.interp(0.8413, cdf, ycenters))
0087 sigma0 = 0.5 * (y_high - y_low)
0088 if sigma0 <= 0 or np.isnan(sigma0):
0089 yrange = ycenters[-1] - ycenters[0]
0090 sigma0 = yrange / 6.0 if yrange > 0 else 1.0
0091 return [np.max(histSlice), mu0, sigma0, 1.5, 2.0]
0092
0093
0094 def boundsCrystalBallPositive(ycenters, histSlice, p0):
0095 """Return positive-peak Crystal Ball parameter bounds."""
0096 return (
0097 [0.1 * p0[0], 0.5 * p0[1], 0.1 * p0[2], 0.1, 1.01],
0098 [5 * p0[0], max(p0[1], ycenters[-1]), 5 * p0[2], 10, 100],
0099 )
0100
0101
0102 def maskPositive(xcenter, ycenters, histSlice):
0103 """Keep positive-count bins."""
0104 return histSlice > 0
0105
0106
0107 def maskPositiveAboveThreshold(thresholdFrac, lowXCutoff):
0108 """Make the historical positive-response slice mask."""
0109 def _mask(xcenter, ycenters, histSlice):
0110 if xcenter < lowXCutoff:
0111 return histSlice > 0
0112 threshold = thresholdFrac * max(ycenters)
0113 return (histSlice > 0) & (ycenters > threshold)
0114 return _mask
0115
0116
0117 @dataclass
0118 class SliceResult:
0119 """Per-slice center and width estimates returned by a reducer."""
0120 center: float
0121 centerErr: float
0122 width: float
0123 widthErr: float
0124 render: Optional[Callable] = None
0125 method: str = 'fit'
0126
0127
0128 def reduceSlices(xcenters, ycenters, hist, err, reducer, sliceMask=None, outDir=None):
0129 """Reduce every x-slice and return centers, errors, widths, and errors.
0130
0131 Slices with too few usable bins receive NaNs. If `outDir` is supplied, a
0132 reducer's optional render callback writes one diagnostic image per slice.
0133 """
0134 if sliceMask is None:
0135 sliceMask = maskPositive
0136 centers, center_errors, widths, width_errors = [], [], [], []
0137 for i, (hist_slice, err_slice) in enumerate(zip(hist, err)):
0138 selected = sliceMask(xcenters[i], ycenters, hist_slice)
0139 hm = hist_slice[selected]
0140 em = err_slice[selected]
0141 ym = ycenters[selected]
0142 result = reducer(ym, hm, em, ycenters) if len(hm) > 3 and np.sum(hm) > 0 else None
0143 if result is None:
0144 centers.append(np.nan); center_errors.append(np.nan)
0145 widths.append(np.nan); width_errors.append(np.nan)
0146 continue
0147 centers.append(result.center); center_errors.append(result.centerErr)
0148 widths.append(result.width); width_errors.append(result.widthErr)
0149 if outDir is not None and result.render is not None:
0150 fig, ax = plt.subplots(figsize=(7, 5))
0151 result.render(ax)
0152 fig.tight_layout()
0153 fig.savefig(os.path.join(outDir, 'img_%d.png' % i))
0154 plt.close(fig)
0155 return (np.array(centers), np.array(center_errors),
0156 np.array(widths), np.array(width_errors))
0157
0158
0159 def _mode_and_crossing(ycenters, histSlice, heightFrac):
0160 """Estimate histogram mode and half-width at a chosen height fraction."""
0161 ycenters, histSlice = np.array(ycenters), np.array(histSlice)
0162 peak_index = np.argmax(histSlice)
0163 mode = float(ycenters[peak_index])
0164 height = heightFrac * float(histSlice[peak_index])
0165 left = np.nan
0166 for i in range(peak_index - 1, -1, -1):
0167 if histSlice[i] <= height:
0168 denom = histSlice[i + 1] - histSlice[i]
0169 frac = (height - histSlice[i]) / denom if denom else 0.5
0170 left = ycenters[i] + np.clip(frac, 0, 1) * (ycenters[i + 1] - ycenters[i])
0171 break
0172 right = np.nan
0173 for i in range(peak_index + 1, len(histSlice)):
0174 if histSlice[i] <= height:
0175 denom = histSlice[i] - histSlice[i - 1]
0176 frac = (height - histSlice[i - 1]) / denom if denom else 0.5
0177 right = ycenters[i - 1] + np.clip(frac, 0, 1) * (ycenters[i] - ycenters[i - 1])
0178 break
0179 width = (right - left) / 2 if np.isfinite(left) and np.isfinite(right) else np.nan
0180 return mode, width, left, right, height
0181
0182
0183 def _bootstrap_mode_crossing(ycenters, histSlice, n=200):
0184 """Estimate mode/HWHM uncertainties with deterministic multinomial bootstrap."""
0185 total = int(round(np.sum(histSlice)))
0186 if total < 5:
0187 return np.nan, np.nan
0188 rng = np.random.default_rng(12345)
0189 probabilities = histSlice / histSlice.sum()
0190 centers, widths = [], []
0191 for _ in range(n):
0192 result = _mode_and_crossing(ycenters, rng.multinomial(total, probabilities), 0.5)
0193 centers.append(result[0])
0194 if np.isfinite(result[1]):
0195 widths.append(result[1])
0196 center_error = np.std(centers, ddof=1) if len(centers) >= 2 else np.nan
0197 width_error = np.std(widths, ddof=1) if len(widths) >= 10 else np.nan
0198 return center_error, width_error
0199
0200
0201 def curveFitReducer(fitFunc, p0Func, boundsFunc=None):
0202 """Create a per-slice Crystal Ball fit reducer with historical fallback.
0203
0204 Fits use scipy's `curve_fit`; low-statistics, failed, or physically
0205 implausible fits fall back to a mode plus HWHM estimate with bootstrap
0206 uncertainties, matching the parent pipeline's behavior.
0207 """
0208 def reduce_one(ycMasked, histMasked, errMasked, ycFull):
0209 total = np.sum(histMasked)
0210 if total < 20:
0211 return fallback(ycMasked, histMasked, errMasked, ycFull, 'low_stats')
0212 p0 = p0Func(ycMasked, histMasked)
0213 kwargs = dict(sigma=errMasked, absolute_sigma=True, p0=p0, maxfev=5000)
0214 if boundsFunc is not None:
0215 kwargs['bounds'] = boundsFunc(ycFull, histMasked, p0)
0216 try:
0217 popt, pcov = curve_fit(fitFunc, ycMasked, histMasked, **kwargs)
0218 except Exception as exc:
0219 return fallback(ycMasked, histMasked, errMasked, ycFull, f'fit exception {type(exc).__name__}')
0220 yrange = ycFull[-1] - ycFull[0]
0221 diag = np.diag(pcov)
0222 reason = None
0223 if popt[2] <= 0 or popt[2] > yrange:
0224 reason = 'invalid sigma'
0225 elif popt[1] < ycFull[0] or popt[1] > ycFull[-1]:
0226 reason = 'invalid mu'
0227 elif np.any(~np.isfinite(diag)) or np.any(diag < 0):
0228 reason = 'invalid covariance'
0229 elif pcov[1, 1] > yrange ** 2 or pcov[2, 2] > yrange ** 2:
0230 reason = 'uncertain fit'
0231 if reason:
0232 return fallback(ycMasked, histMasked, errMasked, ycFull, reason)
0233 return SliceResult(float(popt[1]), float(np.sqrt(pcov[1, 1])),
0234 float(popt[2]), float(np.sqrt(pcov[2, 2])), method=fitFunc.__name__)
0235
0236 def fallback(yc, histSlice, errSlice, ycFull, reason):
0237 """Return the historical mode/HWHM fallback result."""
0238 print(f'[curveFitReducer] slice fallback -> mode+HWHM: {reason}', file=sys.stderr)
0239 mode, width, left, right, height = _mode_and_crossing(yc, histSlice, 0.5)
0240 center_error, width_error = _bootstrap_mode_crossing(yc, histSlice)
0241 return SliceResult(float(mode), center_error, float(width) if np.isfinite(width) else np.nan,
0242 width_error, method='mode_hwhm')
0243
0244 return reduce_one
