def randn(size=None, dist={0:'b'}, grid_shape=None, comm=None):
if size is None:
return np.random.randn()
else:
base_comm = init_base_comm(comm)
comm_size = base_comm.Get_size()
local_shape = find_local_shape(size, dist=dist, grid_shape=grid_shape, comm_size=comm_size)
local_result = np.random.randn(*local_shape)
return densedistarray.DistArray(size, local_result.dtype, dist, grid_shape, comm, buf=local_result)
def rand(size=None, dist={0: 'b'}, grid_shape=None, comm=None):
if size is None:
return np.random.rand()
else:
base_comm = init_base_comm(comm)
comm_size = base_comm.Get_size()
local_shape = find_local_shape(size,
dist=dist,
grid_shape=grid_shape,
comm_size=comm_size)
local_result = np.random.rand(*local_shape)
return densedistarray.DistArray(size,
local_result.dtype,
dist,
grid_shape,
comm,
buf=local_result)
def __init__(self, shape, dtype=float, dist={0:'b'} , grid_shape=None,
comm=None, buf=None, offset=0):
"""Create a distributed memory array on a set of processors.
"""
if comm==MPI.COMM_NULL:
raise NullCommError("cannot create a DistArray with COMM_NULL")
self.shape = shape
self.ndim = len(shape)
self.dtype = np.dtype(dtype)
self.size = reduce(lambda x,y: x*y, shape)
self.itemsize = self.dtype.itemsize
self.nbytes = self.size*self.itemsize
self.data = None
self.base = None
self.ctypes = None
# This order is extremely important and is shown by the arguments passed on to
# subsequent _init_* methods. It is critical that these _init_* methods are free
# of side effects and stateless. This means that they cannot set or get class or
# instance attributes
self.base_comm = init_base_comm(comm)
self.comm_size = self.base_comm.Get_size()
self.comm_rank = self.base_comm.Get_rank()
self.dist = init_dist(dist, self.ndim)
self.distdims = init_distdims(self.dist, self.ndim)
self.ndistdim = len(self.distdims)
self.map_classes = init_map_classes(self.dist)
self.grid_shape = init_grid_shape(self.shape, grid_shape,
self.distdims, self.comm_size)
self.comm = init_comm(self.base_comm, self.grid_shape, self.ndistdim)
self.cart_coords = self.comm.Get_coords(self.comm_rank)
self.local_shape, self.maps = init_local_shape_and_maps(self.shape,
self.grid_shape, self.distdims, self.map_classes)
self.local_size = reduce(lambda x,y: x*y, self.local_shape)
def ones(shape, dtype=float, dist={0:'b'}, grid_shape=None, comm=None):
base_comm = init_base_comm(comm)
local_shape = find_local_shape(shape, dist, grid_shape, base_comm.Get_size())
local_ones = np.ones(local_shape, dtype=dtype)
return DistArray(shape, dtype, dist, grid_shape, comm, buf=local_ones)