def coordinates(o, order="lexicographic"): if type(o) == gpt.grid and o.cb.n == 1: return coordinates((o, gpt.none), order=order) elif type(o) == tuple and type(o[0]) == gpt.grid and len(o) == 2: dim = len(o[0].ldimensions) cb = o[1].tag checker_dim_mask = o[0].cb.cb_mask cbf = [o[0].fdimensions[i] // o[0].gdimensions[i] for i in range(dim)] top = [ o[0].processor_coor[i] * o[0].ldimensions[i] * cbf[i] for i in range(dim) ] bottom = [top[i] + o[0].ldimensions[i] * cbf[i] for i in range(dim)] # cache tag = f"{top}-{bottom}-{checker_dim_mask}-{cb}-{order}" if tag in _coordinates_cache: return _coordinates_cache[tag] val = cgpt.coordinates_from_cartesian_view(top, bottom, checker_dim_mask, cb, order) _coordinates_cache[tag] = val return val elif type(o) == gpt.lattice: return coordinates((o.grid, o.checkerboard()), order=order) elif type(o) == gpt.cartesian_view: return cgpt.coordinates_from_cartesian_view(o.top, o.bottom, o.checker_dim_mask, o.cb, order) else: assert 0
def coor(self, grid, tag=None): # and pos -> core/block/ coor = numpy.zeros((self.grid.gsites, self.grid.nd), dtype=numpy.int32) n = 0 # global offset of the local lattice ofs = [ grid.processor_coor[mu] * grid.ldimensions[mu] for mu in range(grid.nd) ] for ib in range(self.nb): bc = index2coor(ib, self.bdl) _eo = int(numpy.sum(bc) % 2) if _eo == self.eo: sl = slice(n * self.bv, (n + 1) * self.bv) n += 1 top = [ ofs[mu] + bc[mu] * self.bs[mu] for mu in range(len(ofs)) ] bottom = [top[mu] + self.bs[mu] for mu in range(len(ofs))] pos = cgpt.coordinates_from_cartesian_view( top, bottom, grid.cb.cb_mask, tag, "lexicographic") coor[sl, :] = pos assert n * 2 == self.nb return coor
def map_tidx_and_shape(l, key): # create shape of combined lattices shapes = [x.otype.shape for x in l] assert all([shapes[0][1:] == s[1:] for s in shapes[1:]]) shape = (sum([s[0] for s in shapes]), ) + shapes[0][1:] nd = len(shape) # if key is None, numpy array of all indices if key is None: tidx = cgpt.coordinates_from_cartesian_view([0] * nd, list(shape), [0] * nd, gpt.none.tag, "reverse_lexicographic") return tidx, shape # if key is a list, convert to numpy array if type(key) == list: key = numpy.array(key, dtype=numpy.int32) # if key is numpy array, no further processing needed if isinstance(key, numpy.ndarray): # Need to decide how to index tensor indices. With lexicographical return key, (len(key), ) # if not, we expect a tuple of either coordinates or slices assert type(key) == tuple # slices key = tuple([k if type(k) == slice else slice(k, k + 1) for k in key]) assert all([k.step is None for k in key]) top = [0 if k.start is None else k.start for i, k in enumerate(key)] bottom = [ shape[i] if k.stop is None else k.stop for i, k in enumerate(key) ] assert all([ 0 <= top[i] and top[i] <= bottom[i] and bottom[i] <= shape[i] for i in range(nd) ]) tidx = cgpt.coordinates_from_cartesian_view(top, bottom, [0] * nd, gpt.none.tag, "reverse_lexicographic") shape = tuple([bottom[i] - top[i] for i in range(nd)]) return tidx, shape
def coordinates(o, order="lexicographic", margin=None): if type(o) == gpt.grid and o.cb.n == 1: return coordinates((o, gpt.none), order=order, margin=margin) elif type(o) == tuple and type(o[0]) == gpt.grid and len(o) == 2: dim = len(o[0].ldimensions) cb = o[1].tag checker_dim_mask = o[0].cb.cb_mask cbf = [o[0].fdimensions[i] // o[0].gdimensions[i] for i in range(dim)] top = [ o[0].processor_coor[i] * o[0].ldimensions[i] * cbf[i] for i in range(dim) ] bottom = [top[i] + o[0].ldimensions[i] * cbf[i] for i in range(dim)] if margin is not None: top = [t - m for t, m in zip(top, margin)] bottom = [b + m for b, m in zip(bottom, margin)] x = cgpt.coordinates_from_cartesian_view(top, bottom, checker_dim_mask, cb, order) if margin is None: x = x.view(local_coordinates) else: L = numpy.array(o[0].gdimensions, dtype=numpy.int32) x = numpy.mod(x, L) return x elif type(o) == gpt.lattice: return coordinates((o.grid, o.checkerboard()), order=order, margin=margin) elif type(o) == gpt.cartesian_view: assert margin is None return cgpt.coordinates_from_cartesian_view(o.top, o.bottom, o.checker_dim_mask, o.cb, order) else: assert 0
def pos(self, tag=None): self.pos = numpy.zeros((self.grid.gsites, self.grid.nd), dtype=numpy.int32) Nd = len(self.bs) ofs = [0] * Nd for mu in range(Nd): if self.ebs[mu] > 1: ofs[mu] = 1 for n in range(self.nb // 2): sl = slice(n * self.bv, (n + 1) * self.bv) top = [ofs[mu] * n * self.bs[mu] for mu in range(Nd)] bottom = [top[mu] + self.bs[mu] for mu in range(Nd)] _pos = cgpt.coordinates_from_cartesian_view( top, bottom, self.grid.cb.cb_mask, tag, "lexicographic") self.pos[sl, :] = _pos
def map_pos(grid, cb, key): # if list, convert to numpy array if type(key) == list: key = numpy.array(key, dtype=numpy.int32) # if key is numpy array, no further processing needed if isinstance(key, numpy.ndarray): return key # if not, we expect a tuple of slices assert type(key) == tuple # slices without specified start/stop corresponds to memory view limitation for this rank if all([k == slice(None, None, None) for k in key]): # go through gpt.coordinates to use its caching feature return gpt.coordinates((grid, cb), order="lexicographic") nd = grid.nd key = tuple([k if type(k) == slice else slice(k, k + 1) for k in key]) assert all([k.step is None for k in key]) top = [ grid.fdimensions[i] // grid.mpi[i] * grid.processor_coor[i] if k.start is None else k.start for i, k in enumerate(key) ] bottom = [ grid.fdimensions[i] // grid.mpi[i] * (1 + grid.processor_coor[i]) if k.stop is None else k.stop for i, k in enumerate(key) ] assert all([ 0 <= top[i] and top[i] <= bottom[i] and bottom[i] <= grid.fdimensions[i] for i in range(nd) ]) return cgpt.coordinates_from_cartesian_view(top, bottom, grid.cb.cb_mask, cb.tag, "lexicographic")
def __init__(self, grid, block_size, parity): self.parity = parity self.block_size = block_size self.block_volume = int(numpy.prod(block_size)) self.grid = grid assert grid.cb.n == 1 # for now only full grids are supported self.promote_plan = {} self.project_plan = {} # general blocking info nd = len(block_size) assert nd == grid.nd assert numpy.all(numpy.mod(grid.ldimensions, block_size) == 0) self.local_blocks_per_dimension = [ int(d // b) for d, b in zip(grid.ldimensions, block_size) ] self.blocks_per_dimension = [ int(d // b) for d, b in zip(grid.gdimensions, block_size) ] self.number_of_local_blocks = int( numpy.prod(self.local_blocks_per_dimension)) assert self.number_of_local_blocks % 2 == 0 # find checkerboarding dimension self.extended_local_blocks_per_dimension = [1] * nd for mu in reversed(range(nd)): if self.local_blocks_per_dimension[mu] > 1: self.extended_local_blocks_per_dimension[mu] = ( self.number_of_local_blocks // 2) break extended_block_size = [ block_size[mu] * self.extended_local_blocks_per_dimension[mu] for mu in range(nd) ] # block grid local to the node self.local_grid = grid.split([1] * nd, extended_block_size) # map local positions in one-dimensional stack of blocks self.lcoor = numpy.zeros((self.local_grid.gsites, nd), dtype=numpy.int32) offset = [0] * nd for mu in range(nd): if self.extended_local_blocks_per_dimension[mu] > 1: offset[mu] = 1 for n in range(self.number_of_local_blocks // 2): sl = slice(n * self.block_volume, (n + 1) * self.block_volume) top = [offset[mu] * n * self.block_size[mu] for mu in range(nd)] bottom = [top[mu] + self.block_size[mu] for mu in range(nd)] _pos = cgpt.coordinates_from_cartesian_view( top, bottom, self.local_grid.cb.cb_mask, None, "lexicographic") self.lcoor[sl, :] = _pos # map corresponding global positions self.gcoor = numpy.zeros((self.local_grid.gsites, nd), dtype=numpy.int32) n = 0 offset = [ grid.processor_coor[mu] * grid.ldimensions[mu] for mu in range(nd) ] for ib in range(self.number_of_local_blocks): block_coordinate = index_to_coordinate( ib, self.local_blocks_per_dimension) _eo = int(numpy.sum(block_coordinate) % 2) if _eo == parity.tag: sl = slice(n * self.block_volume, (n + 1) * self.block_volume) n += 1 top = [ offset[mu] + block_coordinate[mu] * block_size[mu] for mu in range(len(offset)) ] bottom = [ top[mu] + block_size[mu] for mu in range(len(offset)) ] pos = cgpt.coordinates_from_cartesian_view( top, bottom, grid.cb.cb_mask, None, "lexicographic") self.gcoor[sl, :] = pos assert n * 2 == self.number_of_local_blocks