def cal_monentum(m, ma, mb):
mabp = ma + mb
mabm = ma - mb
s = m * m
p2 = (s - mabp * mabp) * (s - mabm * mabm) / 4 / s
zeros = tf.zeros_like(s)
p_p = tf.complex(tf.sqrt(tf.abs(p2)), zeros)
p_m = tf.complex(zeros, tf.sqrt(tf.abs(p2)))
return tf.where(p2 > 0, p_p, p_m)
def poly_i(i):
tmp = zeros
for j in range(i - 1, i + 3):
if j < 0 or j > N - 1:
continue
r = ones
for k in range(j - 1, j + 3):
if k == i or k < 0 or k > N:
continue
r = r * (x - xi[k]) / (xi[i] - xi[k])
r = tf.where((x >= xi[j]) & (x < xi[j + 1]), r, zeros)
tmp = tmp + r
return tmp
def poly_i(i, xi):
tmp = zeros
for j in range(i - 1, i + 1):
if j < 0 or j > self.interp_N - 1:
continue
r = ones
for k in range(j, j + 2):
if k == i:
continue
r = r * (m - xi[k]) / (xi[i] - xi[k])
r = tf.where((m >= xi[j]) & (m < xi[j + 1]), r, zeros)
tmp = tmp + r
return tmp
def poly_i(i):
tmp = zeros
x_i = (xi[i] + xi[i - 1]) / 2
for j in range(i - 1, i + 3):
if j < 0 or j > N - 1:
continue
r = ones
for k in range(j - 1, j + 3):
if k == i or k < 1 or k > N:
continue
x_k = (xi[k] + xi[k - 1]) / 2
r = r * (x - x_k) / (x_i - x_k)
r = tf.where(
(x >= (xi[j] + xi[j - 1]) / 2) & (x < (xi[j] + xi[j + 1]) / 2),
r,
zeros,
)
tmp = tmp + r
return tmp
def add_f(x, bl, br, pl, pr):
return tf.where(
(x > bl) & (x <= br),
(x - bl) / (br - bl) * (pr - pl) + pl,
zeros,
)
def add_f(x, bl, br):
return tf.where((x > bl) & (x <= br), ones, zeros)
def poly_i(i):
cut = (x >= xi[i]) & (x < xi[i + 1])
return tf.where(cut, x_p, zeros)