def cg_loop(x, dx, dy, residual, iterations):
dx_dy = math.sum(dx * dy, axis=non_batch_dims, keepdims=True)
step_size = math.divide_no_nan(math.sum(dx * residual, axis=non_batch_dims, keepdims=True), dx_dy)
x += step_size * dx
residual -= step_size * dy
dx = residual - math.divide_no_nan(math.sum(residual * dy, axis=non_batch_dims, keepdims=True) * dx, dx_dy)
dy = function(dx)
return [x, dx, dy, residual, iterations + 1]
def fourier_poisson(tensor, times=1):
""" Inverse operation to `fourier_laplace`. """
frequencies = math.fft(math.to_complex(tensor))
k = fftfreq(math.staticshape(tensor)[1:-1], mode='square')
fft_laplace = -(2 * np.pi)**2 * k
fft_laplace[(0, ) * math.ndims(k)] = np.inf
return math.cast(
math.real(
math.ifft(math.divide_no_nan(frequencies, fft_laplace**times))),
math.dtype(tensor))
def loop_body(pressure, momentum, A_times_momentum, residual, loop_index):
"""
iteratively solve for:
x : pressure
momentum : momentum
laplace_momentum : A_times_momentum
residual : residual
"""
tmp = math.sum(momentum * A_times_momentum,
axis=non_batch_dims,
keepdims=True) # t = sum(mAm)
a = math.divide_no_nan(
math.sum(momentum * residual, axis=non_batch_dims, keepdims=True),
tmp) # a = sum(mr)/sum(mAm)
pressure += a * momentum # p += am
residual -= a * A_times_momentum # r -= aAm
momentum = residual - math.divide_no_nan(
math.sum(residual * A_times_momentum,
axis=non_batch_dims,
keepdims=True) * momentum,
tmp) # m = r-sum(rAm)*m/t = r-sum(rAm)*m/sum(mAm)
A_times_momentum = apply_A(momentum) # Am = A*m
return [pressure, momentum, A_times_momentum, residual, loop_index + 1]