def forward_thomas(inf, diag, sup, rhs, out):
with gt.region(iteration=gt.FORWARD, k_interval=(0, 1)):
sup = sup / diag
rhs = rhs / diag
with gt.region(iteration=gt.FORWARD, k_interval=(1, None)):
sup = sup / (diag - sup[0, 0, -1] * inf)
rhs = (rhs - inf * rhs[0, 0, -1]) / (diag - sup[0, 0, -1] * inf)
def set_stage_id(out):
with gt.region(iteration=gt.PARALLEL, k_interval=(0, 1)):
out = 1
with gt.region(iteration=gt.PARALLEL, k_interval=(1, -1)):
out = 2
with gt.region(iteration=gt.PARALLEL, k_interval=(-1, None)):
out = 3
def u_backward_function(utens_stage, u_pos, dtr_stage, ccol, dcol, data_col):
with gt.region(iteration=gt.BACKWARD, k_interval=(0, -1)):
datacol = dcol[0, 0, 0] - ccol[0, 0, 0] * data_col[0, 0, 1]
data_col = datacol
utens_stage = dtr_stage[0, 0, 0] * (datacol - u_pos[0, 0, 0])
with gt.region(iteration=gt.BACKWARD, k_interval=(-1, None)):
datacol = dcol[0, 0, 0]
data_col = datacol
utens_stage = dtr_stage[0, 0, 0] * (datacol - u_pos[0, 0, 0])
def u_forward_function(utens_stage, wcon, u_stage, u_pos, utens, dtr_stage,
acol, bcol, ccol, dcol):
with gt.region(iteration=gt.FORWARD, k_interval=(0, 1)):
gcv = 0.25 * (wcon[1, 0, 1] + wcon[0, 0, 1])
cs = gcv * BET_M
ccol = gcv * BET_P
bcol = dtr_stage[0, 0, 0] - ccol[0, 0, 0]
# update the d column
correction_term = -cs * (u_stage[0, 0, 1] - u_stage[0, 0, 0])
dcol = (dtr_stage[0, 0, 0] * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / bcol[0, 0, 0]
ccol = ccol[0, 0, 0] * divided
dcol = dcol[0, 0, 0] * divided
with gt.region(iteration=gt.FORWARD, k_interval=(1, -1)):
gav = -0.25 * (wcon[1, 0, 0] + wcon[0, 0, 0])
gcv = 0.25 * (wcon[1, 0, 1] + wcon[0, 0, 1])
as_ = gav * BET_M
cs = gcv * BET_M
acol = gav * BET_P
ccol = gcv * BET_P
bcol = dtr_stage[0, 0, 0] - acol[0, 0, 0] - ccol[0, 0, 0]
# update the d column
correction_term = -as_ * (u_stage[0, 0, -1] - u_stage[
0, 0, 0]) - cs * (u_stage[0, 0, 1] - u_stage[0, 0, 0])
dcol = (dtr_stage[0, 0, 0] * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / (bcol[0, 0, 0] - ccol[0, 0, -1] * acol[0, 0, 0])
ccol = ccol[0, 0, 0] * divided
dcol = (dcol[0, 0, 0] - (dcol[0, 0, -1]) * acol[0, 0, 0]) * divided
with gt.region(iteration=gt.FORWARD, k_interval=(-1, None)):
gav = -0.25 * (wcon[1, 0, 0] + wcon[0, 0, 0])
as_ = gav * BET_M
acol = gav * BET_P
bcol = dtr_stage[0, 0, 0] - acol[0, 0, 0]
# update the d column
correction_term = -as_ * (u_stage[0, 0, -1] - u_stage[0, 0, 0])
dcol = (dtr_stage[0, 0, 0] * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / (bcol[0, 0, 0] - ccol[0, 0, -1] * acol[0, 0, 0])
dcol = (dcol[0, 0, 0] - (dcol[0, 0, -1]) * acol[0, 0, 0]) * divided
def definition_func(inf, diag, sup, rhs, out):
with gt.region(iteration=gt.FORWARD, k_interval=(0, 1)):
sup = sup / diag
rhs = rhs / diag
with gt.region(iteration=gt.FORWARD, k_interval=(1, None)):
sup = sup / (diag - sup[0, 0, -1] * inf)
rhs = (rhs - inf * rhs[0, 0, -1]) / (diag - sup[0, 0, -1] * inf)
with gt.region(iteration=gt.BACKWARD, k_interval=(0, -1)):
out = rhs - sup * out[0, 0, 1]
with gt.region(iteration=gt.BACKWARD, k_interval=(-1, None)):
out = rhs
def definition_func(utens_stage, u_stage, wcon, u_pos, utens, *,
dtr_stage):
with gt.region(iteration=gt.FORWARD, k_interval=(0, 1)):
gcv = 0.25 * (wcon[1, 0, 1] + wcon[0, 0, 1])
cs = gcv * BET_M
ccol = gcv * BET_P
bcol = dtr_stage - ccol[0, 0, 0]
# update the d column
correction_term = -cs * (u_stage[0, 0, 1] - u_stage[0, 0, 0])
dcol = (dtr_stage * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / bcol[0, 0, 0]
ccol = ccol[0, 0, 0] * divided
dcol = dcol[0, 0, 0] * divided
with gt.region(iteration=gt.FORWARD, k_interval=(1, -1)):
gav = -0.25 * (wcon[1, 0, 0] + wcon[0, 0, 0])
gcv = 0.25 * (wcon[1, 0, 1] + wcon[0, 0, 1])
as_ = gav * BET_M
cs = gcv * BET_M
acol = gav * BET_P
ccol = gcv * BET_P
bcol = dtr_stage - acol[0, 0, 0] - ccol[0, 0, 0]
# update the d column
correction_term = -as_ * (u_stage[0, 0, -1] - u_stage[
0, 0, 0]) - cs * (u_stage[0, 0, 1] - u_stage[0, 0, 0])
dcol = (dtr_stage * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / (bcol[0, 0, 0] - ccol[0, 0, -1] * acol[0, 0, 0])
ccol = ccol[0, 0, 0] * divided
dcol = (dcol[0, 0, 0] - (dcol[0, 0, -1]) * acol[0, 0, 0]) * divided
with gt.region(iteration=gt.FORWARD, k_interval=(-1, None)):
gav = -0.25 * (wcon[1, 0, 0] + wcon[0, 0, 0])
as_ = gav * BET_M
acol = gav * BET_P
bcol = dtr_stage - acol[0, 0, 0]
# update the d column
correction_term = -as_ * (u_stage[0, 0, -1] - u_stage[0, 0, 0])
dcol = (dtr_stage * u_pos[0, 0, 0] + utens[0, 0, 0] +
utens_stage[0, 0, 0] + correction_term)
# Thomas forward
divided = 1.0 / (bcol[0, 0, 0] - ccol[0, 0, -1] * acol[0, 0, 0])
dcol = (dcol[0, 0, 0] - (dcol[0, 0, -1]) * acol[0, 0, 0]) * divided
# with gt.region(iteration=gt.BACKWARD, k_interval=(0, -1)):
# datacol = dcol[0, 0, 0] - ccol[0, 0, 0] * data_col[0, 0, 1]
# data_col = datacol
# utens_stage = dtr_stage * (datacol) - u_pos[0, 0, 0]
with gt.region(iteration=gt.BACKWARD, k_interval=(-1, None)):
datacol = dcol[0, 0, 0]
data_col = datacol
utens_stage = dtr_stage * (datacol - u_pos[0, 0, 0])
with gt.region(iteration=gt.BACKWARD, k_interval=(0, -1)):
datacol = dcol[0, 0, 0] - ccol[0, 0, 0] * data_col[0, 0, 1]
data_col = datacol
utens_stage = dtr_stage * (datacol - u_pos[0, 0, 0])
def backward_thomas(inf, diag, sup, rhs, out):
with gt.region(iteration=gt.BACKWARD, k_interval=(0, -1)):
out = rhs - sup * out[0, 0, 1]
with gt.region(iteration=gt.BACKWARD, k_interval=(-1, None)):
out = rhs
