def from_maps(cls, context, maps, targets=None):
"""Create a Distribution from a sequence of `Map`\s.
Parameters
----------
context : Context object
maps : Sequence of Map objects
targets : Sequence of int, optional
Sequence of engine target numbers. Default: all available
Returns
-------
Distribution
"""
# This constructor is called by all the others
self = super(Distribution, cls).__new__(cls)
self.context = context
self.targets = sorted(targets or context.targets)
self._comm = None
self.maps = maps
self.shape = tuple(m.size for m in self.maps)
self.ndim = len(self.maps)
self.dist = tuple(m.dist for m in self.maps)
self.grid_shape = tuple(m.grid_size for m in self.maps)
self.grid_shape = normalize_grid_shape(self.grid_shape, self.shape,
self.dist, len(self.targets))
nelts = reduce(operator.mul, self.grid_shape, 1)
self.rank_from_coords = np.arange(nelts).reshape(self.grid_shape)
return self
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_maps(cls, context, maps, targets=None):
"""Create a Distribution from a sequence of `Map`\s.
Parameters
----------
context : Context object
maps : Sequence of Map objects
targets : Sequence of int, optional
Sequence of engine target numbers. Default: all available
Returns
-------
Distribution
"""
# This constructor is called by all the others
self = super(Distribution, cls).__new__(cls)
self.context = context
self.targets = sorted(targets or context.targets)
self._comm = None
self.maps = maps
self.shape = tuple(m.size for m in self.maps)
self.ndim = len(self.maps)
self.dist = tuple(m.dist for m in self.maps)
self.grid_shape = tuple(m.grid_size for m in self.maps)
self.grid_shape = normalize_grid_shape(self.grid_shape, self.shape,
self.dist, len(self.targets))
nelts = reduce(operator.mul, self.grid_shape, 1)
self.rank_from_coords = np.arange(nelts).reshape(self.grid_shape)
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 __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 comm_union(self, *dists):
"""
Make a communicator that includes the union of all targets in `dists`.
Parameters
----------
dists: sequence of distribution objects.
Returns
-------
tuple
First element is encompassing communicator proxy; second is a
sequence of all targets in `dists`.
"""
dist_targets = [d.targets for d in dists]
all_targets = sorted(reduce(set.union, dist_targets, set(self.targets)))
return self.context.make_subcomm(all_targets), all_targets
def comm_union(self, *dists):
"""
Make a communicator that includes the union of all targets in `dists`.
Parameters
----------
dists: sequence of distribution objects.
Returns
-------
tuple
First element is encompassing communicator proxy; second is a
sequence of all targets in `dists`.
"""
dist_targets = [d.targets for d in dists]
all_targets = sorted(reduce(set.union, dist_targets,
set(self.targets)))
return self.context.make_subcomm(all_targets), all_targets
def test_ones_in_grid_shape(self):
"""Test not-distributed dimensions in grid_shape."""
dist = ('n', 'b', 'n', 'c', 'n')
glb_shape = (2, 6, 2, 8, 2)
grid_shape = (1, 3, 1, 4, 1)
dist_5d = Distribution.from_shape(comm=self.comm, shape=glb_shape,
grid_shape=grid_shape, dist=dist)
self.assertEqual(dist_5d.global_shape, glb_shape)
self.assertEqual(dist_5d.grid_shape, grid_shape)
self.assertEqual(dist_5d.base_comm, self.comm)
self.assertEqual(dist_5d.comm_size, 12)
self.assertTrue(dist_5d.comm_rank in range(12))
self.assertEqual(dist_5d.comm.Get_topo(),
(list(dist_5d.grid_shape),
[0]*5, list(dist_5d.cart_coords)))
self.assertEqual(len(dist_5d), 5)
self.assertEqual(dist_5d.global_shape, glb_shape)
self.assertEqual(dist_5d.local_shape, (2, 2, 2, 2, 2))
self.assertEqual(dist_5d.local_size,
reduce(int.__mul__, glb_shape) // self.comm_size)
def test_ones_in_grid_shape(self):
"""Test not-distributed dimensions in grid_shape."""
dist = ('n', 'b', 'n', 'c', 'n')
glb_shape = (2, 6, 2, 8, 2)
grid_shape = (1, 3, 1, 4, 1)
dist_5d = Distribution.from_shape(comm=self.comm,
shape=glb_shape,
grid_shape=grid_shape,
dist=dist)
self.assertEqual(dist_5d.global_shape, glb_shape)
self.assertEqual(dist_5d.grid_shape, grid_shape)
self.assertEqual(dist_5d.base_comm, self.comm)
self.assertEqual(dist_5d.comm_size, 12)
self.assertTrue(dist_5d.comm_rank in range(12))
self.assertEqual(
dist_5d.comm.Get_topo(),
(list(dist_5d.grid_shape), [0] * 5, list(dist_5d.cart_coords)))
self.assertEqual(len(dist_5d), 5)
self.assertEqual(dist_5d.global_shape, glb_shape)
self.assertEqual(dist_5d.local_shape, (2, 2, 2, 2, 2))
self.assertEqual(dist_5d.local_size,
reduce(int.__mul__, glb_shape) // self.comm_size)
def __init__(self, context, global_dim_data):
"""Make a Distribution from a global_dim_data structure.
Parameters
----------
global_dim_data : tuple of dict
A global dimension dictionary per dimension. See following `Note`
section.
Returns
-------
result : Distribution
An empty DistArray of the specified size, dimensionality, and
distribution.
Note
----
The `global_dim_data` tuple is a simple, straightforward data structure
that allows full control over all aspects of a DistArray's distribution
information. It does not contain any of the array's *data*, only the
*metadata* needed to specify how the array is to be distributed. Each
dimension of the array is represented by corresponding dictionary in
the tuple, one per dimension. All dictionaries have a `dist_type` key
that specifies whether the array is block, cyclic, or unstructured.
The other keys in the dictionary are dependent on the `dist_type` key.
**Block**
* ``dist_type`` is ``'b'``.
* ``bounds`` is a sequence of integers, at least two elements.
The ``bounds`` sequence always starts with 0 and ends with the global
``size`` of the array. The other elements indicate the local array
global index boundaries, such that successive pairs of elements from
``bounds`` indicates the ``start`` and ``stop`` indices of the
corresponding local array.
* ``comm_padding`` integer, greater than or equal to zero.
* ``boundary_padding`` integer, greater than or equal to zero.
These integer values indicate the communication or boundary padding,
respectively, for the local arrays. Currently only a single value for
both ``boundary_padding`` and ``comm_padding`` is allowed for the
entire dimension.
**Cyclic**
* ``dist_type`` is ``'c'``
* ``proc_grid_size`` integer, greater than or equal to one.
The size of the process grid in this dimension. Equivalent to the
number of local arrays in this dimension and determines the number of
array sections.
* ``size`` integer, greater than or equal to zero.
The global size of the array in this dimension.
* ``block_size`` integer, optional. Greater than or equal to one.
If not present, equivalent to being present with value of one.
**Unstructured**
* ``dist_type`` is ``'u'``
* ``indices`` sequence of one-dimensional numpy integer arrays or
buffers.
The ``len(indices)`` is the number of local unstructured arrays in
this dimension.
To compute the global size of the array in this dimension, compute
``sum(len(ii) for ii in indices)``.
**Not-distributed**
The ``'n'`` distribution type is a convenience to specify that an array
is not distributed along this dimension.
* ``dist_type`` is ``'n'``
* ``size`` integer, greater than or equal to zero.
The global size of the array in this dimension.
"""
self.context = context
self.maps = [map_from_global_dim_dict(gdd) for gdd in global_dim_data]
self.shape = tuple(m.size for m in self.maps)
self.ndim = len(self.maps)
self.dist = tuple(m.dist for m in self.maps)
self.grid_shape = tuple(m.grid_size for m in self.maps)
validate_grid_shape(self.grid_shape, self.dist, len(context.targets))
nelts = reduce(operator.mul, self.grid_shape)
self.rank_from_coords = np.arange(nelts).reshape(*self.grid_shape)