Skip to content

Commit

Permalink
Merge pull request #342 from bendavid/calibrationvars
Browse files Browse the repository at this point in the history 
Update muon response maps and scale/resolution weights
  • Loading branch information
@bendavid
bendavid authored Nov 18, 2023
2 parents 78ffbdc + 00dc72e commit 270c3c9
Show file tree
Hide file tree
Showing 4 changed files with 173 additions and 88 deletions.
169 changes: 90 additions & 79 deletions scripts/corrections/make_muon_response_maps.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,11 @@
mpl.rcParams['figure.dpi'] = 300


infile = "w_z_muonresponse_scetlib_dyturboCorr_NonClosureCorl.hdf5"
infile = "w_z_muonresponse_scetlib_dyturboCorr_maxFiles_m1.hdf5"

hist_response = None
hist_response_scaled = None
hist_response_smeared = None

procs = []
procs.append("ZmumuPostVFP")
Expand All @@ -30,16 +32,31 @@
results = narf.ioutils.pickle_load_h5py(f["results"])
for proc in procs:
hist_response_proc = results[proc]["output"]["hist_qopr"].get()
hist_response_scaled_proc = results[proc]["output"]["hist_qopr_shifted"].get()
hist_response_smeared_proc = results[proc]["output"]["hist_qopr_smearedmulti"].get()
if hist_response is None:
hist_response = hist_response_proc
hist_response_scaled = hist_response_scaled_proc
hist_response_smeared = hist_response_smeared_proc
else:
hist_response += hist_response_proc
hist_response_scaled += hist_response_scaled_proc
hist_response_smeared += hist_response_smeared_proc


print(hist_response)

dscale = hist_response_scaled.metadata["scalerel"]
dsigma = hist_response_smeared.metadata["sigmarel"]
dsigmasq = dsigma**2

dscale = tf.constant(dscale, tf.float64)
dsigmasq = tf.constant(dsigmasq, tf.float64)


hist_response = hist_response.project("genCharge", "qopr", "genPt", "genEta")
hist_response_scaled = hist_response_scaled.project("genCharge", "qopr", "genPt", "genEta")
hist_response_smeared = hist_response_smeared.project("genCharge", "qopr", "genPt", "genEta")

print(hist_response)

Expand All @@ -57,14 +74,22 @@
print("quant_cdfvals", quant_cdfvals)


quants, errs = narf.fitutils.hist_to_quantiles(hist_response, quant_cdfvals, axis = 1)
quants, _ = narf.fitutils.hist_to_quantiles(hist_response, quant_cdfvals, axis = 1)
quants_scaled, _ = narf.fitutils.hist_to_quantiles(hist_response_scaled, quant_cdfvals, axis = 1)
quants_smeared, _ = narf.fitutils.hist_to_quantiles(hist_response_smeared, quant_cdfvals, axis = 1)


# print("quants", quants, quants_scaled, quants_smeared)
dquants = np.sum((quants_scaled-quants)**2)
print("dquants", dquants)

print("non-finite quants:", np.count_nonzero(np.invert(np.isfinite(quants))))



quants = tf.constant(quants, tf.float64)
quants_scaled = tf.constant(quants_scaled, tf.float64)
quants_smeared = tf.constant(quants_smeared, tf.float64)

grid_points = [tf.constant(axis.centers) for axis in hist_response.axes]
grid_points = grid_points[2:]
Expand Down Expand Up @@ -92,74 +117,54 @@
print("pt bounds:", pt_low, pt_high)
print("eta bounds:", eta_low, eta_high)

# first use a cubic spline to estimate the gradients at the knot points
knot_slopes = np.zeros_like(quants)
for i0 in range(knot_slopes.shape[0]):
for i2 in range(knot_slopes.shape[2]):
for i3 in range(knot_slopes.shape[3]):
spline = scipy.interpolate.make_interp_spline(x = quants[i0,:,i2,i3], y = quant_cdfvals_interp, k=3, bc_type="natural")
grad = spline(quants[i0,:,i2,i3], nu=1)
knot_slopes[i0, :, i2, i3] = grad

knot_slopes = tf.constant(knot_slopes, tf.float64)

# now the final interpolation in tensorflow, where the C2 continuity of the cubic spline guarantees continuous second derivatives needed for the resolution variation weights

# n.b. in an ideal world we would be able to do this directly in one step with a C3 continuous monotonic spline
# but an implementation of this doesn't currently exist in tensorflow

def interp_pdf(genPt, genEta, genCharge, qopr):
def interp_cdf(quants_sel, genPt, genEta, genCharge, qopr):
chargeIdx = tf.where(genCharge > 0., 1, 0)
quants_charge = quants[chargeIdx]

qopr = tf.reshape(qopr, [-1])

quants_charge = quants_sel[chargeIdx]

x = tf.stack([genPt, genEta], axis=0)
x = x[None, :]
quants_interp = tfp.math.batch_interp_rectilinear_nd_grid(x, x_grid_points = grid_points, y_ref = quants_charge, axis = 1)
quants_interp = tf.reshape(quants_interp, [-1])

knot_slopes_charge = knot_slopes[chargeIdx]
knot_slopes_interp = tfp.math.batch_interp_rectilinear_nd_grid(x, x_grid_points = grid_points, y_ref = knot_slopes_charge, axis = 1)
knot_slopes_interp = tf.reshape(knot_slopes_interp, [-1])

pdf = narf.fitutils.cubic_spline_interpolate(xi = quants_interp[None, :], yi = knot_slopes_interp[None, :], x = qopr[None, :])
pdf = tf.reshape(pdf, [])
quants_interp = tf.reshape(quants_interp, [-1])
quant_cdfvals_interp = tf.reshape(quant_cdfvals, [-1])

qopr = tf.clip_by_value(qopr, qopr_low, qopr_high)

# neither monotonicity of the initial spline estimating the gradients, nor positivity of the second spline is
# otherwise enforced so we need to explicitly clip the pdf to prevent negative values
pdf = tf.maximum (pdf, tf.constant(0., tf.float64))
qopr = tf.reshape(qopr, [-1])

return pdf
# cdf = narf.fitutils.pchip_interpolate(xi = quants_interp, yi = quant_cdfvals_interp, x = qopr)
cdf = narf.fitutils.cubic_spline_interpolate(xi = quants_interp[..., None], yi = quant_cdfvals_interp[..., None], x = qopr[..., None], axis=0)
cdf = cdf[..., 0]

return cdf

def interp_grads(genPt, genEta, genCharge, qopr):
def interp_dpdf(quants_sel, genPt, genEta, genCharge, qopr):
with tf.GradientTape() as t0:
t0.watch(qopr)
with tf.GradientTape() as t1:
t1.watch(qopr)
pdf = interp_pdf(genPt, genEta, genCharge, qopr)
dpdf = t1.gradient(pdf, qopr)
d2pdf = t0.gradient(dpdf, qopr)
cdf = interp_cdf(quants_sel, genPt, genEta, genCharge, qopr)
pdf = t1.gradient(cdf, qopr)
dpdf = t0.gradient(pdf, qopr)

return pdf, dpdf, d2pdf
return cdf, pdf, dpdf

def interp_dweight(genPt, genEta, genCharge, qopr):
def interp_pdf(quants_sel, genPt, genEta, genCharge, qopr):

with tf.GradientTape() as t0:
t0.watch(qopr)
with tf.GradientTape() as t1:
t1.watch(qopr)
pdf = interp_pdf(genPt, genEta, genCharge, qopr)
dpdf = t1.gradient(pdf, qopr)
d2pdf = t0.gradient(dpdf, qopr)
cdf = interp_cdf(quants_sel, genPt, genEta, genCharge, qopr)
pdf = t0.gradient(cdf, qopr)

pdf, dpdf, d2pdf = interp_grads(genPt, genEta, genCharge, qopr)
return cdf, pdf

def interp_dweight(genPt, genEta, genCharge, qopr):
cdf, pdf, dpdf = interp_dpdf(quants, genPt, genEta, genCharge, qopr)
cdf_smeared, pdf_smeared= interp_pdf(quants_smeared, genPt, genEta, genCharge, qopr)

dweightdscale = -dpdf/pdf
dweightdsigmasq = 0.5*d2pdf/pdf
dweightdsigmasq = (pdf_smeared - pdf)/pdf/dsigmasq

in_range = (qopr > qopr_low) & (qopr < qopr_high) & (genPt > pt_low) & (genPt < pt_high) & (genEta > eta_low) & (genEta < eta_high)

Expand All @@ -171,16 +176,18 @@ def interp_dweight(genPt, genEta, genCharge, qopr):

return dweightdscale, dweightdsigmasq


genPt_test = tf.constant(25., tf.float64)
genEta_test = tf.constant(0.1, tf.float64)
genCharge_test = tf.constant(1., tf.float64)
qopr_test = tf.constant(1.002, tf.float64)

res = interp_pdf(genPt_test, genEta_test, genCharge_test, qopr_test)
res2 = interp_dweight(genPt_test, genEta_test, genCharge_test, qopr_test)
res = interp_cdf(quants, genPt_test, genEta_test, genCharge_test, qopr_test)
res2a, res2b = interp_dweight(genPt_test, genEta_test, genCharge_test, qopr_test)

print("res", res)
print("res2", res2)
print("res2a", res2a)
print("res2b", res2b)

scalar_spec = tf.TensorSpec([], tf.float64)
input_signature = 4*[scalar_spec]
Expand All @@ -192,39 +199,23 @@ def interp_dweight(genPt, genEta, genCharge, qopr):
with open(output_filename, 'wb') as f:
f.write(tflite_model)

#
# qoprvals = np.linspace(0., 2., 1000)
qoprvals = np.linspace(0.9, 1.1, 1000)
pttf = tf.constant(40., tf.float64)
etatf = tf.constant(0., tf.float64)
# pttf = tf.constant(9.5, tf.float64)
# etatf = tf.constant(2.35, tf.float64)
chargetf = tf.constant(1., tf.float64)

pdfs = []
dpdfs = []
d2pdfs = []
#this is just for plotting
def func_pdf(h):
dtype = tf.float64
xvals = [tf.constant(center, dtype=dtype) for center in h.axes.centers]
xedges = [tf.constant(edge, dtype=dtype) for edge in h.axes.edges]
axis=1

for qoprval in qoprvals:
qopr = tf.constant(qoprval, tf.float64)
pdf, dpdf, d2pdf = interp_grads(pttf,etatf, chargetf, qopr)

pdfs.append(pdf)
dpdfs.append(dpdf)
d2pdfs.append(d2pdf)
# cdf = narf.fitutils.pchip_interpolate(xi = quants, yi = quant_cdfvals, x = xedges[axis], axis=axis)
cdf = narf.fitutils.cubic_spline_interpolate(xi = quants, yi = quant_cdfvals, x = xedges[axis], axis=axis)

plt.figure()
plt.plot(qoprvals, pdfs)
plt.savefig("pdf.png")
pdf = cdf[:,1:] - cdf[:,:-1]
# pdf = tf.maximum(pdf, tf.zeros_like(pdf))

plt.figure()
plt.plot(qoprvals, dpdfs)
plt.savefig("dpdf.png")

plt.figure()
plt.plot(qoprvals, d2pdfs)
plt.savefig("d2pdf.png")
return pdf

qoprvals = np.linspace(0., 2., 1000)
# qoprvals = np.linspace(0.9, 1.1, 1000)

# etaidx = 47
# ptidx = 20
Expand Down Expand Up @@ -252,13 +243,21 @@ def interp_dweight(genPt, genEta, genCharge, qopr):


pdfvals_sel = []
dpdfvals_sel = []
d2pdfvals_sel = []

for qoprval in qoprvals:
testqopr = tf.constant(qoprval, tf.float64)
pdf = interp_pdf(testpt, testeta, testcharge, testqopr)
cdf, pdf, dpdf = interp_dpdf(quants, testpt, testeta, testcharge, testqopr)
cdf_smeared, pdf_smeared = interp_pdf(quants_smeared, testpt, testeta, testcharge, testqopr)
d2pdf = (pdf_smeared - pdf)/dsigmasq
pdfvals_sel.append(pdf.numpy())
dpdfvals_sel.append(dpdf.numpy())
d2pdfvals_sel.append(d2pdf.numpy())

pdfvals_sel = np.array(pdfvals_sel)
dpdfvals_sel = np.array(dpdfvals_sel)
d2pdfvals_sel = np.array(d2pdfvals_sel)


print("pdfvals_sel", pdfvals_sel)
Expand All @@ -279,3 +278,15 @@ def interp_dweight(genPt, genEta, genCharge, qopr):
plt.yscale("log")
plot.savefig("test_log.png")


plot = plt.figure()
plt.plot(qoprvals, dpdfvals_sel)
plt.xlim([0.9, 1.1])
plot.savefig("testd.png")


plot = plt.figure()
plt.plot(qoprvals, d2pdfvals_sel)
plt.xlim([0.9, 1.1])
plot.savefig("testd2.png")

35 changes: 27 additions & 8 deletions scripts/histmakers/w_z_muonresponse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,15 @@
smearing_helper, smearing_uncertainty_helper = (None, None) if args.noSmearing else muon_calibration.make_muon_smearing_helpers()
bias_helper = muon_calibration.make_muon_bias_helpers(args) if args.biasCalibration else None

sigmarel = 5e-3
scalerel = 5e-4
nreps = 100

smearing_helper_simple_multi = ROOT.wrem.SmearingHelperSimpleMulti[nreps](sigmarel)

if args.testHelpers:
response_helper = ROOT.wrem.SplinesDifferentialWeightsHelper(f"{wremnants.data_dir}/calibration/muon_response.tflite")

sigmarel = 5e-3
scalerel = 5e-4

smearing_helper_simple = ROOT.wrem.SmearingHelperSimple(sigmarel, ROOT.ROOT.GetThreadPoolSize())
smearing_helper_simple_weights = ROOT.wrem.SmearingHelperSimpleWeight(sigmarel)
smearing_helper_simple_transform = ROOT.wrem.SmearingHelperSimpleTransform(sigmarel)
Expand Down Expand Up @@ -119,6 +122,27 @@ def build_graph(df, dataset):
hist_qopr = df.HistoBoost("hist_qopr", response_axes, [*response_cols, "nominal_weight"])
results.append(hist_qopr)

df = df.Define("selMuons_shiftedqopr", f"(1. + {scalerel})*selMuons_qopr")

response_cols_shifted = ["selMuons_genPt", "selMuons_genEta", "selMuons_genCharge", "selMuons_shiftedqopr"]
hist_qopr_shifted = df.HistoBoost("hist_qopr_shifted", response_axes, [*response_cols_shifted, "nominal_weight"])
hist_qopr_shifted._hist.metadata = { "scalerel" : scalerel }
results.append(hist_qopr_shifted)

df = df.Define("selMuonsMulti_smearedmqop", smearing_helper_simple_multi, ["run", "luminosityBlock", "event", "selMuons_correctedPt", "selMuons_correctedEta", "selMuons_correctedCharge"])

df = df.Define("selMuonsMulti_genPt", f"wrem::replicate_rvec(selMuons_genPt, {nreps})")
df = df.Define("selMuonsMulti_genEta", f"wrem::replicate_rvec(selMuons_genEta, {nreps})")
df = df.Define("selMuonsMulti_genCharge", f"wrem::replicate_rvec(selMuons_genCharge, {nreps})")
df = df.Define("selMuonsMulti_genQop", f"wrem::replicate_rvec(selMuons_genQop, {nreps})")
df = df.Define("selMuonsMulti_smearedqopr", "selMuonsMulti_smearedmqop/selMuonsMulti_genQop")

response_cols_smeared_multi = ["selMuonsMulti_genPt", "selMuonsMulti_genEta", "selMuonsMulti_genCharge", "selMuonsMulti_smearedqopr"]
hist_qopr_smearedmulti = df.HistoBoost("hist_qopr_smearedmulti", response_axes, [*response_cols_smeared_multi, "nominal_weight"])
hist_qopr_smearedmulti._hist.metadata = { "sigmarel" : sigmarel }
results.append(hist_qopr_smearedmulti)


if args.testHelpers:

df = df.Define("selMuons_response_weight", response_helper, ["selMuons_correctedPt", "selMuons_correctedEta", "selMuons_correctedCharge", "selMuons_genPt", "selMuons_genEta", "selMuons_genCharge"])
Expand All @@ -135,7 +159,6 @@ def build_graph(df, dataset):
df = df.Define("selMuons_transformedqop", "selMuons_correctedCharge*1.0/(selMuons_transformedPt*cosh(selMuons_correctedEta))")
df = df.Define("selMuons_transformedqopr", "selMuons_transformedqop/selMuons_genQop")

df = df.Define("selMuons_shiftedqopr", f"(1. + {scalerel})*selMuons_qopr")
df = df.Define("weight_scale", scale_helper_simple_weights, ["selMuons_correctedPt", "selMuons_correctedEta", "selMuons_correctedCharge", "selMuons_response_weight", "nominal_weight"])

response_cols_smeared = ["selMuons_genPt", "selMuons_genEta", "selMuons_genCharge", "selMuons_smearedqopr"]
Expand All @@ -149,10 +172,6 @@ def build_graph(df, dataset):
hist_qopr_smeared_weight = df.HistoBoost("hist_qopr_smeared_weight", response_axes, [*response_cols, "weight_smear"])
results.append(hist_qopr_smeared_weight)

response_cols_shifted = ["selMuons_genPt", "selMuons_genEta", "selMuons_genCharge", "selMuons_shiftedqopr"]
hist_qopr_shifted = df.HistoBoost("hist_qopr_shifted", response_axes, [*response_cols_shifted, "nominal_weight"])
results.append(hist_qopr_shifted)

hist_qopr_scaled_weight = df.HistoBoost("hist_qopr_scaled_weight", response_axes, [*response_cols, "weight_scale"])
results.append(hist_qopr_scaled_weight)

Expand Down
2 changes: 1 addition & 1 deletion wremnants-data
Submodule wremnants-data updated from 7ad111 to acb3d1
Loading

0 comments on commit 270c3c9

Please sign in to comment.