Commit b218f4cb authored by Dennis Noll's avatar Dennis Noll
Browse files

[processor] util: completes tf driven neutrino reco

parent 4ecd37f5
......@@ -335,70 +335,111 @@ 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
+ tf.sqrt(wmass ** 2 / 2 * (lepx ** 2 + lepy ** 2) * (2 * lepx * vx + wmass ** 2 / 2))
+ (wmass ** 2 / 2 * (lepx ** 2 + lepy ** 2) * (2 * lepx * vx + wmass ** 2 / 2)) ** 0.5
) / lepx ** 2
# neutrino reconstruction
@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
wmass = W_MASS.nominal
k = wmass ** 2 / 2 + lepx * vx + lepy * vy
leppt = np.sqrt(lepx ** 2 + lepy ** 2)
h = leppt * np.sqrt(vx ** 2 + vy ** 2) / k
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
vz1 = k / leppt ** 2 * (lepz + lepe * np.sqrt(1 - h ** 2))
vzp = k / leppt ** 2 * (lepz + lepe * (1 - h ** 2) ** 0.5)
# negative solution
vz2 = k / leppt ** 2 * (lepz - lepe * np.sqrt(1 - h ** 2))
vzn = k / leppt ** 2 * (lepz - lepe * (1 - h ** 2) ** 0.5)
# smaller momentum is correct in 70% of cases
vz = vz1 * (abs(vz1) <= abs(vz2)) + vz2 * (abs(vz2) < abs(vz1))
vz = tf.where(tf.abs(vzp) <= tf.abs(vzn), vzp, vzn)
return vx, vy, vz
def neutrino_reco_complex(lepe, lepx, lepy, lepz, metx, mety):
@tf.function
def vx_constrain(vx, lepx):
inf = tf.constant(999999.0) # TODO: 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
wmass = W_MASS.nominal
cond = lepx >= 0
constr_val = np.divide(-(wmass ** 2), 4 * lepx, out=np.zeros_like(lepx), where=lepx != 0)
constr = lambda x: cond * tf.clip_by_value(x, constr_val, np.infty) + (
~cond
) * tf.clip_by_value(x, -np.infty, constr_val)
vx = tf.Variable(metx, constraint=constr)
opt = tf.keras.optimizers.Adam()
for i in range(100):
opt.minimize(lambda: d2(vx, metx, mety, lepx, lepy), var_list=[vx])
vx = vx.numpy()
vy = calc_vy(vx, metx, mety, lepx, lepy).numpy()
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
def reconstruct_neutrino(lepe, lepx, lepy, lepz, metx, mety):
@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_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
wmass = W_MASS.nominal
k = wmass ** 2 / 2 + lepx * metx + lepy * mety
leppt = np.sqrt(lepx ** 2 + lepy ** 2)
h = leppt * np.sqrt(metx ** 2 + mety ** 2) / 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))
vx = np.where(h <= 1, vx_real, vx_complex)
vy = np.where(h <= 1, vy_real, vy_complex)
vz = np.where(h <= 1, vz_real, vz_complex)
neutrino = TLorentzVectorArray.from_xyzm(vx, vy, vz, 0 * vx)
@tf.function
def call(self, inputs, training=False):
self.reset_optimizer()
lepe, lepx, lepy, lepz, metx, mety = (inputs[:, i] for i in range(6))
return neutrino
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):
......
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