def from_shape(cls, context, shape, dist=None, grid_shape=None):
self = cls.__new__(cls)
self.context = context
self.shape = shape
self.ndim = len(shape)
if dist is None:
dist = {0: 'b'}
self.dist = normalize_dist(dist, self.ndim)
if grid_shape is None: # Make a new grid_shape if not provided.
self.grid_shape = make_grid_shape(self.shape, self.dist,
len(context.targets))
else: # Otherwise normalize the one passed in.
self.grid_shape = normalize_grid_shape(grid_shape, self.ndim)
# In either case, validate.
validate_grid_shape(self.grid_shape, self.dist, len(context.targets))
# TODO: FIXME: assert that self.rank_from_coords is valid and conforms
# to how MPI does it.
nelts = reduce(operator.mul, self.grid_shape)
self.rank_from_coords = np.arange(nelts).reshape(*self.grid_shape)
# List of `ClientMap` objects, one per dimension.
self.maps = [map_from_sizes(*args)
for args in zip(self.shape, self.dist, self.grid_shape)]
return self
def from_shape(cls, comm, shape, dist=None, grid_shape=None):
"""Create a Distribution from a `shape` and optional arguments."""
dist = {0: 'b'} if dist is None else dist
ndim = len(shape)
dist_tuple = normalize_dist(dist, ndim)
base_comm = construct.init_base_comm(comm)
comm_size = base_comm.Get_size()
if grid_shape is None: # Make a new grid_shape if not provided.
grid_shape = make_grid_shape(shape, dist_tuple, comm_size)
grid_shape = normalize_grid_shape(grid_shape, ndim,
dist_tuple, comm_size)
comm = construct.init_comm(base_comm, grid_shape)
grid_coords = comm.Get_coords(comm.Get_rank())
dim_data = []
for dist, size, grid_rank, grid_size in zip(dist_tuple, shape,
grid_coords, grid_shape):
dim_dict = dict(dist_type=dist,
size=size,
proc_grid_rank=grid_rank,
proc_grid_size=grid_size)
distribute_indices(dim_dict)
dim_data.append(dim_dict)
return cls(comm=base_comm, dim_data=dim_data)
def from_shape(cls, comm, shape, dist=None, grid_shape=None):
"""Create a Distribution from a `shape` and optional arguments."""
dist = {0: 'b'} if dist is None else dist
ndim = len(shape)
dist_tuple = normalize_dist(dist, ndim)
base_comm = construct.init_base_comm(comm)
comm_size = base_comm.Get_size()
if grid_shape is None: # Make a new grid_shape if not provided.
grid_shape = make_grid_shape(shape, dist_tuple, comm_size)
grid_shape = normalize_grid_shape(grid_shape, shape, dist_tuple,
comm_size)
comm = construct.init_comm(base_comm, grid_shape)
grid_coords = comm.Get_coords(comm.Get_rank())
dim_data = []
for dist, size, grid_rank, grid_size in zip(dist_tuple, shape,
grid_coords, grid_shape):
dim_dict = dict(dist_type=dist,
size=size,
proc_grid_rank=grid_rank,
proc_grid_size=grid_size)
distribute_indices(dim_dict)
dim_data.append(dim_dict)
return cls(comm=base_comm, dim_data=dim_data)
def __new__(cls, context, shape, dist=None, grid_shape=None, targets=None):
"""Create a Distribution from a `shape` and other optional args.
Parameters
----------
context : Context object
shape : tuple of int
Shape of the resulting Distribution, one integer per dimension.
dist : str, list, tuple, or dict, optional
Shorthand data structure representing the distribution type for
every dimension. Default: {0: 'b'}, with all other dimensions 'n'.
grid_shape : tuple of int
targets : Sequence of int, optional
Sequence of engine target numbers. Default: all available
Returns
-------
Distribution
"""
# special case when dist is all 'n's.
if (dist is not None) and all(d == 'n' for d in dist):
if (targets is not None) and (len(targets) != 1):
raise ValueError('target dist conflict')
elif targets is None:
targets = [context.targets[0]]
else:
# then targets is set correctly
pass
ndim = len(shape)
dist = dist or {0: 'b'}
dist = normalize_dist(dist, ndim)
targets = sorted(targets or context.targets)
grid_shape = grid_shape or make_grid_shape(shape, dist, len(targets))
grid_shape = normalize_grid_shape(grid_shape, shape, dist,
len(targets))
# choose targets from grid_shape
ntargets = reduce(operator.mul, grid_shape, 1)
targets = targets[:ntargets]
# list of `ClientMap` objects, one per dimension.
maps = [map_from_sizes(*args) for args in zip(shape, dist, grid_shape)]
self = cls.from_maps(context=context, maps=maps, targets=targets)
# TODO: FIXME: this is a workaround. The reason we slice here is to
# return a distribution with no empty local shapes. The `from_maps()`
# classmethod should be fixed to ensure no empty local arrays are
# created in the first place. That will remove the need to slice the
# distribution to remove empty localshapes.
if all(d in ('n', 'b') for d in self.dist):
self = self.slice((slice(None), ) * self.ndim)
return self
def __new__(cls, context, shape, dist=None, grid_shape=None, targets=None):
"""Create a Distribution from a `shape` and other optional args.
Parameters
----------
context : Context object
shape : tuple of int
Shape of the resulting Distribution, one integer per dimension.
dist : str, list, tuple, or dict, optional
Shorthand data structure representing the distribution type for
every dimension. Default: {0: 'b'}, with all other dimensions 'n'.
grid_shape : tuple of int
targets : Sequence of int, optional
Sequence of engine target numbers. Default: all available
Returns
-------
Distribution
"""
# special case when dist is all 'n's.
if (dist is not None) and all(d == 'n' for d in dist):
if (targets is not None) and (len(targets) != 1):
raise ValueError('target dist conflict')
elif targets is None:
targets = [context.targets[0]]
else:
# then targets is set correctly
pass
ndim = len(shape)
dist = dist or {0: 'b'}
dist = normalize_dist(dist, ndim)
targets = sorted(targets or context.targets)
grid_shape = grid_shape or make_grid_shape(shape, dist, len(targets))
grid_shape = normalize_grid_shape(grid_shape, shape, dist, len(targets))
# choose targets from grid_shape
ntargets = reduce(operator.mul, grid_shape, 1)
targets = targets[:ntargets]
# list of `ClientMap` objects, one per dimension.
maps = [map_from_sizes(*args) for args in zip(shape, dist, grid_shape)]
self = cls.from_maps(context=context, maps=maps, targets=targets)
# TODO: FIXME: this is a workaround. The reason we slice here is to
# return a distribution with no empty local shapes. The `from_maps()`
# classmethod should be fixed to ensure no empty local arrays are
# created in the first place. That will remove the need to slice the
# distribution to remove empty localshapes.
if all(d in ('n', 'b') for d in self.dist):
self = self.slice((slice(None),)*self.ndim)
return self
def test_make_grid_shape(self):
grid_shape = metadata_utils.make_grid_shape((20, 20), ('b', 'b'), 12)
self.assertEqual(grid_shape, (3, 4))
