def tgv_update_w(w, p, q, tau_p):
"""Update w
Args:
w (gpuarray): w.
p (gpuarray): p.
q (gpuarray): q.
tau_p (float): tau_p.
"""
tgv_update_w_func(w,
p,
q,
np.float32(tau_p),
np.int32(w.shape[0]),
np.int32(w.shape[1]),
block=block,
grid=get_grid(w))
def tgv_update_u(u, p, f, tau_p):
"""Update u
Args:
u (gpuarray): u.
p (gpuarray): p.
f (gpuarray): f.
tau_p (float): tau_p.
"""
tgv_update_u_func(u,
p,
f,
np.float32(tau_p),
np.int32(u.shape[0]),
np.int32(u.shape[1]),
block=block,
grid=get_grid(u))
def tgv_update_q(w, q, tau_d, alpha):
"""Update q
Args:
w (gpuarray): w.
q (gpuarray): q.
tau_d (float): tau_d.
alpha (float): alpha.
"""
tgv_update_q_func(w,
q,
np.float32(1.0 / alpha),
np.float32(tau_d),
np.int32(w.shape[0]),
np.int32(w.shape[1]),
block=block,
grid=get_grid(w))
def tgv_update_p(u, w, p, tau_d, alpha):
"""Update p
Args:
u (gpuarray): u.
w (gpuarray): w.
p (gpuarray): p.
tau_d (float): tau_d.
alpha (float): alpha.
"""
tgv_update_p_func(u,
w,
p,
np.float32(1.0 / alpha),
np.float32(tau_d),
np.int32(u.shape[0]),
np.int32(u.shape[1]),
block=block,
grid=get_grid(u))
def inner_cg_rhs(rhs, u, v, EHs, tau):
"""Compute right hand side for inner CG method
``rhs = u^n + tau * (div_h v^{n+1} + EHs)``
Args:
rhs (gpuarray): Right hand side.
u (gpuarray): u.
v (gpuarray): v.
EHs (gpuarray): EHs.
tau (float): tau.
"""
inner_cg_rhs_func(rhs,
u,
v,
EHs,
np.float32(tau),
np.int32(u.shape[0]),
np.int32(u.shape[1]),
np.int32(1),
block=block,
grid=get_grid(u))