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
