Commit 025ec412 authored by Dennis Noll's avatar Dennis Noll
Browse files

[tasks] neutrino: moved neutrino reco + to task (outputs model)

parent fe3b0992
......@@ -14,7 +14,6 @@ from coffea.nanoaod.nanoevents import NanoCollection
from utils.util import parametrized
from utils.order import NanoAOD_version
from config.constants import W_MASS
logger = logging.getLogger(__name__)
......@@ -335,113 +334,6 @@ def nthlargest(arr, n=1):
return max
# neutrino reconstruction
@tf.function
def calc_vy(vx, metx, mety, lepx, lepy):
wmass = W_MASS.nominal
return (
lepx * lepy * vx
+ lepy * wmass ** 2 / 2
+ (wmass ** 2 / 2 * (lepx ** 2 + lepy ** 2) * (2 * lepx * vx + wmass ** 2 / 2)) ** 0.5
) / lepx ** 2
@tf.function
def d2(vx, metx, mety, lepx, lepy):
return (metx - vx) ** 2 + (mety - calc_vy(vx, metx, mety, lepx, lepy)) ** 2
@tf.function
def neutrino_reco_real(lepe, lepx, lepy, lepz, vx, vy):
# standard reconstruction
k = W_MASS.nominal ** 2 / 2 + lepx * vx + lepy * vy
leppt = (lepx ** 2 + lepy ** 2) ** 0.5
h = leppt * (vx ** 2 + vy ** 2) ** 0.5 / k
# positive solution
vzp = k / leppt ** 2 * (lepz + lepe * (1 - h ** 2) ** 0.5)
# negative solution
vzn = k / leppt ** 2 * (lepz - lepe * (1 - h ** 2) ** 0.5)
# smaller momentum is correct in 70% of cases
vz = tf.where(tf.abs(vzp) <= tf.abs(vzn), vzp, vzn)
return vx, vy, vz
@tf.function
def vx_constrain(vx, lepx):
inf = tf.constant(np.inf)
wmass = tf.constant(W_MASS.nominal)
constr_val = -(wmass ** 2) / (4 * lepx)
return tf.where(
lepx >= 0,
tf.clip_by_value(vx, constr_val, inf),
tf.clip_by_value(vx, -inf, constr_val),
)
@tf.function
def optimize(optimizer, vx, metx, mety, lepx, lepy):
batch_size = tf.shape(metx)[0]
for _ in range(100):
with tf.control_dependencies(
[optimizer.minimize(lambda: d2(vx[:batch_size], metx, mety, lepx, lepy), var_list=[vx])]
):
vx[:batch_size].assign(vx_constrain(vx[:batch_size], lepx))
return vx[:batch_size]
@tf.function
def neutrino_reco_complex(lepe, lepx, lepy, lepz, metx, mety, vx, opt):
# complex solution, fit using tensorflow
batch_size = tf.shape(metx)[0]
vx = optimize(opt, vx, metx, mety, lepx, lepy)
vy = calc_vy(vx, metx, mety, lepx, lepy)
vz = neutrino_reco_real(lepe, lepx, lepy, lepz, vx, vy)[2]
return vx, vy, vz
@tf.function
def reconstruct_nu(lepe, lepx, lepy, lepz, metx, mety, vx, opt):
# interleave standard reconstruction (1) and complex solution (2)
vx_real, vy_real, vz_real = neutrino_reco_real(lepe, lepx, lepy, lepz, metx, mety)
vx_complex, vy_complex, vz_complex = neutrino_reco_complex(
lepe, lepx, lepy, lepz, metx, mety, vx, opt
)
k = W_MASS.nominal ** 2 / 2 + lepx * metx + lepy * mety
leppt = (lepx ** 2 + lepy ** 2) ** 0.5
h = leppt * (metx ** 2 + mety ** 2) ** 0.5 / k
vx = tf.where(h <= 1, vx_real, vx_complex)
vy = tf.where(h <= 1, vy_real, vy_complex)
vz = tf.where(h <= 1, vz_real, vz_complex)
return vx, vy, vz
class NeutrinoReco(tf.keras.layers.Layer):
def __init__(self, batch_size=1000):
super(NeutrinoReco, self).__init__(name="")
self.optimizer = tf.keras.optimizers.Adam()
self.batch_size = batch_size
def build(self, input_shape):
# self.vx = tf.Variable(np.zeros(self.batch_size, dtype=np.float32), shape=(self.batch_size,))
self.vx = self.add_weight(shape=(self.batch_size,))
def reset_optimizer(self, *args, **kwargs):
# reset optimizer to initial values (all zeros)
for var in self.optimizer.variables():
var.assign(tf.zeros_like(var))
@tf.function
def call(self, inputs, training=False):
self.reset_optimizer()
lepe, lepx, lepy, lepz, metx, mety = (inputs[:, i] for i in range(6))
with tf.control_dependencies([self.vx[: tf.shape(metx)[0]].assign(metx)]):
vx, vy, vz = reconstruct_nu(lepe, lepx, lepy, lepz, metx, mety, self.vx, self.optimizer)
return vx, vy, vz
def linear_fit(x, y, eigen_decomp=False):
coeff, cov = np.polyfit(x, y, 1, cov="unscaled")
return linear_func(coeff, cov, eigen_decomp)
......
# coding: utf-8
import luigi
import tensorflow as tf
import numpy as np
from tasks.base import BaseTask
from config.constants import W_MASS
# neutrino reconstruction
@tf.function
def calc_vy(vx, metx, mety, lepx, lepy):
wmass = W_MASS.nominal
return (
lepx * lepy * vx
+ lepy * wmass ** 2 / 2
+ (wmass ** 2 / 2 * (lepx ** 2 + lepy ** 2) * (2 * lepx * vx + wmass ** 2 / 2)) ** 0.5
) / lepx ** 2
@tf.function
def d2(vx, metx, mety, lepx, lepy):
return (metx - vx) ** 2 + (mety - calc_vy(vx, metx, mety, lepx, lepy)) ** 2
@tf.function
def neutrino_reco_real(lepe, lepx, lepy, lepz, vx, vy):
# standard reconstruction
k = W_MASS.nominal ** 2 / 2 + lepx * vx + lepy * vy
leppt = (lepx ** 2 + lepy ** 2) ** 0.5
h = leppt * (vx ** 2 + vy ** 2) ** 0.5 / k
# positive solution
vzp = k / leppt ** 2 * (lepz + lepe * (1 - h ** 2) ** 0.5)
# negative solution
vzn = k / leppt ** 2 * (lepz - lepe * (1 - h ** 2) ** 0.5)
# smaller momentum is correct in 70% of cases
vz = tf.where(tf.abs(vzp) <= tf.abs(vzn), vzp, vzn)
return vx, vy, vz
@tf.function
def vx_constrain(vx, lepx):
inf = tf.constant(np.inf)
wmass = tf.constant(W_MASS.nominal)
constr_val = -(wmass ** 2) / (4 * lepx)
return tf.where(
lepx >= 0,
tf.clip_by_value(vx, constr_val, inf),
tf.clip_by_value(vx, -inf, constr_val),
)
@tf.function
def optimize(optimizer, vx, metx, mety, lepx, lepy):
batch_size = tf.shape(metx)[0]
for _ in range(100):
with tf.control_dependencies(
[optimizer.minimize(lambda: d2(vx[:batch_size], metx, mety, lepx, lepy), var_list=[vx])]
):
vx[:batch_size].assign(vx_constrain(vx[:batch_size], lepx))
return vx[:batch_size]
@tf.function
def neutrino_reco_complex(lepe, lepx, lepy, lepz, metx, mety, vx, opt):
# complex solution, fit using tensorflow
batch_size = tf.shape(metx)[0]
vx = optimize(opt, vx, metx, mety, lepx, lepy)
vy = calc_vy(vx, metx, mety, lepx, lepy)
vz = neutrino_reco_real(lepe, lepx, lepy, lepz, vx, vy)[2]
return vx, vy, vz
@tf.function
def reconstruct_nu(lepe, lepx, lepy, lepz, metx, mety, vx, opt):
# interleave standard reconstruction (1) and complex solution (2)
vx_real, vy_real, vz_real = neutrino_reco_real(lepe, lepx, lepy, lepz, metx, mety)
vx_complex, vy_complex, vz_complex = neutrino_reco_complex(
lepe, lepx, lepy, lepz, metx, mety, vx, opt
)
k = W_MASS.nominal ** 2 / 2 + lepx * metx + lepy * mety
leppt = (lepx ** 2 + lepy ** 2) ** 0.5
h = leppt * (metx ** 2 + mety ** 2) ** 0.5 / k
vx = tf.where(h <= 1, vx_real, vx_complex)
vy = tf.where(h <= 1, vy_real, vy_complex)
vz = tf.where(h <= 1, vz_real, vz_complex)
return vx, vy, vz
class NeutrinoReco(tf.keras.layers.Layer):
def __init__(self, batch_size=1000):
super(NeutrinoReco, self).__init__(name="")
self.optimizer = tf.keras.optimizers.Adam()
self.batch_size = batch_size
def build(self, input_shape):
self.vx = self.add_weight(shape=(self.batch_size,), name="vx")
def reset_optimizer(self, *args, **kwargs):
# reset optimizer to initial values (all zeros)
for var in self.optimizer.variables():
var.assign(tf.zeros_like(var))
@tf.function
def call(self, inputs, training=False):
self.reset_optimizer()
lepe, lepx, lepy, lepz, metx, mety = (inputs[:, i] for i in range(6))
with tf.control_dependencies([self.vx[: tf.shape(metx)[0]].assign(metx)]):
vx, vy, vz = reconstruct_nu(lepe, lepx, lepy, lepz, metx, mety, self.vx, self.optimizer)
return vx, vy, vz
class NeutrinoReconstructionFit1(BaseTask):
batch_size = luigi.IntParameter(default=150_000)
def output(self):
return self.local_target("model/.done")
def run(self):
inp = tf.keras.layers.Input(shape=(6,))
reco = NeutrinoReco(batch_size=self.batch_size)
model = tf.keras.models.Model(inputs=inp, outputs=reco(inp))
model.save(self.output().parent.path)
self.output().touch()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment