def spline_matrix(x, xi, yi, bc_type="not-a-knot"):
"""calculate spline interpolation"""
xi_m = spline_xi_matrix(xi) # (N_range, 4, N_yi)
x_m = spline_x_matrix(x, xi) # (..., N_range, 4)
x_m = tf.expand_dims(x_m, axis=-1)
m = tf.reduce_sum(xi_m * x_m, axis=[-3, -2])
return tf.reduce_sum(tf.cast(m, yi.dtype) * yi, axis=-1)
def interp(self, m):
zeros = tf.zeros_like(m)
p = self.point_value()
p_r = tf.math.real(p)
p_i = tf.math.imag(p)
xi_m = self.h_matrix
x_m = spline_x_matrix(m, self.points)
x_m = tf.expand_dims(x_m, axis=-1)
m_xi = tf.reduce_sum(xi_m * x_m, axis=[-3, -2])
m_xi = tf.stop_gradient(m_xi)
ret_r = tf.reduce_sum(tf.cast(m_xi, p_r.dtype) * p_r, axis=-1)
ret_i = tf.reduce_sum(tf.cast(m_xi, p_i.dtype) * p_i, axis=-1)
return tf.complex(ret_r, ret_i)
def spline_x_matrix(x, xi):
"""build matrix of x for spline interpolation"""
ones = tf.ones_like(x)
x2 = x * x
x3 = x2 * x
x_p = tf.stack([ones, x, x2, x3], axis=-1)
x = tf.expand_dims(x, axis=-1)
zeros = tf.zeros_like(x)
def poly_i(i):
cut = (x >= xi[i]) & (x < xi[i + 1])
return tf.where(cut, x_p, zeros)
xs = [poly_i(i) for i in range(len(xi) - 1)]
return tf.stack(xs, axis=-2)