def inv(dst, src):
# verbosity
verbose = g.default.is_verbose("split")
if len(src) % nparallel != 0:
raise Exception(
f"Cannot divide {len(src)} global vectors into {nparallel} groups"
)
t0 = g.time()
src_split = g.split(src, matrix_split.grid[1], cache)
dst_split = g.split(dst, matrix_split.grid[0], cache)
t1 = g.time()
operation_split(dst_split, src_split)
t2 = g.time()
g.unsplit(dst, dst_split, cache)
t3 = g.time()
if verbose:
g.message(
f"Split {len(src)} global vectors to {len(src_split)} local vectors\n"
+ f"Timing: {t1-t0} s (split), {t2-t1} s (operation), {t3-t2} s (unsplit)"
)
################################################################################
# Test split grid
################################################################################
for src in [l, l_rb]:
split_grid = src[0].grid.split(
g.default.get_ivec("--mpi_split", None, 4), src[0].grid.fdimensions
)
# perform this test without cache and with cache (to fill and then to test)
cache = {}
for it in range(3):
for group_policy in [
g.split_group_policy.separate,
g.split_group_policy.together,
]:
g.message(f"Iteration {it}, group_policy = {group_policy.__name__}")
src_unsplit = [g.lattice(x) for x in src]
t0 = g.time()
src_split = g.split(
src, split_grid, None if it == 0 else cache, group_policy
)
t1 = g.time()
g.unsplit(src_unsplit, src_split, None if it == 0 else cache, group_policy)
t2 = g.time()
g.message(f"Timing: {t1-t0}s for split and {t2-t1}s for unsplit")
for i in range(len(src)):
eps2 = g.norm2(src_unsplit[i] - src[i]) / g.norm2(src[i])
g.message(f"Split test {i} / {len(src)}: {eps2}")
assert eps2 == 0.0
# create rng if needed
rng = None if p_rng_seed is None else g.random(p_rng_seed)
# load source
U = g.load(p_source)
# split in time
Nt = U[0].grid.gdimensions[3]
g.message(f"Separate {Nt} time slices")
Usep = [g.separate(u, 3) for u in U[0:3]]
Vt = [g.mcolor(Usep[0][0].grid) for t in range(Nt)]
cache = {}
split_grid = Usep[0][0].grid.split(p_mpi_split, Usep[0][0].grid.fdimensions)
g.message("Split grid")
Usep_split = [g.split(Usep[mu], split_grid, cache) for mu in range(3)]
Vt_split = g.split(Vt, split_grid, cache)
# optimizer
opt = g.algorithms.optimize
cg = opt.non_linear_cg(
maxiter=p_maxiter_cg,
eps=p_eps,
step=p_step,
line_search=opt.line_search_quadratic,
beta=opt.polak_ribiere,
max_abs_step=p_max_abs_step,
)
gd = opt.gradient_descent(maxiter=p_maxiter_gd, eps=p_eps, step=p_gd_step)
# Coulomb functional on each time-slice
k = i * 4 + j
assert g.norm2(l[k][1, 2, 0, 0] - m[i][1, j, 2, 0, 0]) == 0.0
test = g.separate(m, 1)
assert len(test) == Nlat
for i in range(len(l)):
assert g.norm2(l[i] - test[i]) == 0.0
# default arguments should be compatible
test = g.separate(g.merge(l))
for i in range(len(l)):
assert g.norm2(l[i] - test[i]) == 0.0
################################################################################
# Test split grid
################################################################################
for src in [l, l_rb]:
split_grid = src[0].grid.split(g.default.get_ivec("--mpi_split", None, 4),
src[0].grid.fdimensions)
src_unsplit = [g.lattice(x) for x in src]
# perform this test twice, second time cached plan will be used
for it in range(2):
src_split = g.split(src, split_grid=split_grid)
g.unsplit(src_unsplit, src_split)
for i in range(len(src)):
eps2 = g.norm2(src_unsplit[i] - src[i]) / g.norm2(src[i])
g.message(f"Split test {i} / {len(src)}: {eps2}")
assert eps2 == 0.0
