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