def get_ranks(arr):
"""
Given a distarray arr, return a distarray with the same shape, but
with the elements equal to the rank of the process the element is
on.
"""
from distarray.localapi import LocalArray
out = LocalArray(distribution=arr.distribution, dtype=int)
out.fill(arr.comm_rank)
return out
def test_save_2d(self):
d = Distribution.from_shape(comm=self.comm, shape=(11, 15))
la = LocalArray(d)
np_arr = numpy.random.random(la.local_shape)
la.ndarray = np_arr
save_hdf5(self.output_path, la, key=self.key, mode='w')
with self.h5py.File(self.output_path, 'r', driver='mpio',
comm=self.comm) as fp:
for i, v in ndenumerate(la):
self.assertEqual(v, fp[self.key][i])
def test_save_2d(self):
d = Distribution.from_shape(comm=self.comm, shape=(11, 15))
la = LocalArray(d)
np_arr = numpy.random.random(la.local_shape)
la.ndarray = np_arr
save_hdf5(self.output_path, la, key=self.key, mode='w')
with self.h5py.File(self.output_path,
'r',
driver='mpio',
comm=self.comm) as fp:
for i, v in ndenumerate(la):
self.assertEqual(v, fp[self.key][i])
def local_filter_avg3(la):
''' Filter a local array via 3-point average over z axis. '''
def filter_avg3(a):
''' Filter a local array via 3-point average over z axis.
A 2-point average is used at the ends. '''
from numpy import empty_like
b = empty_like(a)
b[:, :, 0] = (a[:, :, 0] + a[:, :, 1]) / 2.0
b[:, :, 1:-1] = (a[:, :, :-2] + a[:, :, 1:-1] + a[:, :, 2:]) / 3.0
b[:, :, -1] = (a[:, :, -2] + a[:, :, -1]) / 2.0
return b
from distarray.localapi import LocalArray
a = la.ndarray
b = filter_avg3(a)
res = LocalArray(la.distribution, buf=b)
rtn = proxyize(res)
return rtn
def local_filter_max3(la):
''' Filter a local array via 3-point average over z axis. '''
def filter_max3(a):
''' Filter a numpy array via a 3-element window maximum. '''
from numpy import empty_like
b = empty_like(a)
shape = a.shape
for k in xrange(shape[2]):
k0 = max(k - 1, 0)
k1 = min(k + 1, shape[2] - 1)
b[:, :, k] = a[:, :, k0:k1 + 1].max(axis=2)
return b
from distarray.localapi import LocalArray
a = la.ndarray
b = filter_max3(a)
res = LocalArray(la.distribution, buf=b)
rtn = proxyize(res)
return rtn
def local_julia_calc(la, c, z_max, n_max, kernel):
"""Calculate the number of iterations for the point to escape.
Parameters
----------
la : LocalArray
LocalArray of complex values whose iterations we will count.
c : complex
Complex number to add at each iteration.
z_max : float
Magnitude of complex value that we assume goes to infinity.
n_max : int
Maximum number of iterations.
kernel : function
Kernel to use for computation of the Julia set. Options are 'fancy',
'numpy', or 'cython'.
"""
from distarray.localapi import LocalArray
counts = kernel(la, c, z_max, n_max)
res = LocalArray(la.distribution, buf=counts)
return proxyize(res) # noqa
def setUp(self):
d = Distribution.from_shape(comm=self.comm, shape=(7, ))
self.larr0 = LocalArray(d)
# a different file on every engine
self.output_path = temp_filepath(extension='.dnpy')
def local_process(frames, output, params, kernel):
from distarray.localapi import LocalArray
recon = kernel(frames, output, params)
res = LocalArray(output.distribution, buf=recon)
return proxyize(res) # noqa