Merge pull request #412 from enthought/bugfix/144

Fixes #144.

Author: bgrant

distarray/dist/context.py

@@ -218,7 +218,7 @@ def _create_local(self, local_call, distribution, dtype):
         ddpr = distribution.get_dim_data_per_rank()
         ddpr_name, dtype_name =  self._key_and_push(ddpr, dtype)
         cmd = ('{da_key} = {local_call}(distarray.local.maps.Distribution('
-               '{ddpr_name}[{comm_name}.Get_rank()], comm={comm_name}), '
+               'comm={comm_name}, dim_data={ddpr_name}[{comm_name}.Get_rank()]), '
                'dtype={dtype_name})')
         self._execute(cmd.format(**locals()), targets=distribution.targets)
         return DistArray.from_localarrays(da_key, distribution=distribution,
@@ -364,20 +364,18 @@ def load_dnpy(self, name):
         da_key = self._generate_key()
 
         if isinstance(name, six.string_types):
-            subs = (da_key,) + self._key_and_push(name) + (self.comm,
-                    self.comm)
+            subs = (da_key,) + (self.comm,) + self._key_and_push(name) + (self.comm,)
             self._execute(
-                '%s = distarray.local.load_dnpy(%s + "_" + str(%s.Get_rank()) + ".dnpy", %s)' % subs,
+                '%s = distarray.local.load_dnpy(%s, %s + "_" + str(%s.Get_rank()) + ".dnpy")' % subs,
                 targets=self.targets
             )
         elif isinstance(name, collections.Sequence):
             if len(name) != len(self.targets):
                 errmsg = "`name` must be the same length as `self.targets`."
                 raise TypeError(errmsg)
-            subs = (da_key,) + self._key_and_push(name) + (self.comm,
-                    self.comm)
+            subs = (da_key,) + (self.comm,) + self._key_and_push(name) + (self.comm,)
             self._execute(
-                '%s = distarray.local.load_dnpy(%s[%s.Get_rank()], %s)' % subs,
+                '%s = distarray.local.load_dnpy(%s, %s[%s.Get_rank()])' % subs,
                 targets=self.targets
             )
         else:
@@ -438,16 +436,17 @@ def load_npy(self, filename, distribution):
         result : DistArray
             A DistArray encapsulating the file loaded.
         """
-        da_key = self._generate_key()
+        
+        def _local_load_npy(filename, ddpr, comm):
+            from distarray.local import load_npy
+            dim_data = ddpr[comm.Get_rank()]
+            return proxyize(load_npy(comm, filename, dim_data))
+
         ddpr = distribution.get_dim_data_per_rank()
-        subs = ((da_key,) + self._key_and_push(filename, ddpr) +
-                (distribution.comm,) + (distribution.comm,))
 
-        self._execute(
-            '%s = distarray.local.load_npy(%s, %s[%s.Get_rank()], %s)' % subs,
-            targets=distribution.targets
-        )
-        return DistArray.from_localarrays(da_key, distribution=distribution)
+        da_key = self.apply(_local_load_npy, (filename, ddpr, distribution.comm),
+                            targets=distribution.targets)
+        return DistArray.from_localarrays(da_key[0], distribution=distribution)
 
     def load_hdf5(self, filename, distribution, key='buffer'):
         """
@@ -473,16 +472,17 @@ def load_hdf5(self, filename, distribution, key='buffer'):
             errmsg = "An MPI-enabled h5py must be available to use load_hdf5."
             raise ImportError(errmsg)
 
-        da_key = self._generate_key()
+        def _local_load_hdf5(filename, ddpr, comm, key):
+            from distarray.local import load_hdf5
+            dim_data = ddpr[comm.Get_rank()]
+            return proxyize(load_hdf5(comm, filename, dim_data, key))
+
         ddpr = distribution.get_dim_data_per_rank()
-        subs = ((da_key,) + self._key_and_push(filename, ddpr) +
-                (distribution.comm,) + self._key_and_push(key) + (distribution.comm,))
 
-        self._execute(
-            '%s = distarray.local.load_hdf5(%s, %s[%s.Get_rank()], %s, %s)' % subs,
-            targets=distribution.targets
-        )
-        return DistArray.from_localarrays(da_key, distribution=distribution)
+        da_key = self.apply(_local_load_hdf5, (filename, ddpr, distribution.comm, key),
+                   targets=distribution.targets)
+
+        return DistArray.from_localarrays(da_key[0], distribution=distribution)
 
     def fromndarray(self, arr, distribution=None):
         """Create a DistArray from an ndarray.
@@ -530,7 +530,7 @@ def fromfunction(self, function, shape, **kwargs):
         comm_name = distribution.comm
         cmd = ('{da_name} = distarray.local.fromfunction({function_name}, '
                'distarray.local.maps.Distribution('
-               '{ddpr_name}[{comm_name}.Get_rank()], comm={comm_name}),'
+               'comm={comm_name}, dim_data={ddpr_name}[{comm_name}.Get_rank()]),'
                '**{kwargs_name})')
         self._execute(cmd.format(**locals()), targets=distribution.targets)
         return DistArray.from_localarrays(da_name, distribution=distribution)

distarray/dist/distarray.py

@@ -49,8 +49,8 @@ def __init__(self, distribution, dtype=float):
         ddpr_name, dtype_name = ctx._key_and_push(ddpr, dtype)
         cmd = ('{da_key} = distarray.local.empty('
                'distarray.local.maps.Distribution('
-               '{ddpr_name}[{comm_name}.Get_rank()], '
-               '{comm_name}), {dtype_name})')
+               'comm={comm_name}, dim_data={ddpr_name}[{comm_name}.Get_rank()]), '
+               '{dtype_name})')
         ctx._execute(cmd.format(**locals()), targets=distribution.targets)
         self.distribution = distribution
         self.key = da_key
@@ -286,7 +286,7 @@ def _reduce(self, local_reduce_name, axes=None, dtype=None, out=None):
         def _local_reduce(local_name, larr, out_comm, ddpr, dtype, axes):
             import distarray.local.localarray as la
             local_reducer = getattr(la, local_name)
-            res = proxyize(la.local_reduction(local_reducer, out_comm, larr,  # noqa
+            res = proxyize(la.local_reduction(out_comm, local_reducer, larr,  # noqa
                                               ddpr, dtype, axes))
             return res
 

distarray/local/construct.py

@@ -8,7 +8,6 @@
 
 from distarray.local.mpiutils import MPI
 from distarray.local.error import NullCommError, InvalidBaseCommError
-from distarray.local import mpiutils
 
 
 # ---------------------------------------------------------------------------
@@ -27,8 +26,6 @@ def init_base_comm(comm):
     """Sanitize an MPI.comm instance or create one."""
     if comm == MPI.COMM_NULL:
         raise NullCommError("Cannot create a LocalArray with COMM_NULL")
-    elif comm is None:
-        return mpiutils.COMM_PRIVATE
     elif isinstance(comm, MPI.Comm):
         return comm
     else:

distarray/local/localarray.py

@@ -10,7 +10,6 @@
 # ---------------------------------------------------------------------------
 # Imports
 # ---------------------------------------------------------------------------
-import math
 from collections import Mapping
 from numbers import Integral
 
@@ -19,7 +18,6 @@
 from distarray.externals import six
 from distarray.externals.six.moves import zip
 
-import distarray.local
 from distarray.local.mpiutils import MPI
 from distarray.utils import _raise_nie
 from distarray.local import format, maps
@@ -200,7 +198,7 @@ def compatibility_hash(self):
     #-------------------------------------------------------------------------
 
     @classmethod
-    def from_distarray(cls, obj, comm=None):
+    def from_distarray(cls, comm, obj):
         """Make a LocalArray from Distributed Array Protocol data structure.
 
         An object that supports the Distributed Array Protocol will have
@@ -228,7 +226,7 @@ def from_distarray(cls, obj, comm=None):
         buf = np.asarray(distbuffer['buffer'])
         dim_data = distbuffer['dim_data']
 
-        distribution = maps.Distribution(dim_data=dim_data, comm=comm)
+        distribution = maps.Distribution(comm=comm, dim_data=dim_data)
         return cls(distribution=distribution, buf=buf)
 
     def __distarray__(self):
@@ -627,7 +625,7 @@ def save_dnpy(file, arr):
             fid.close()
 
 
-def load_dnpy(file, comm=None):
+def load_dnpy(comm, file):
     """
     Load a LocalArray from a ``.dnpy`` file.
 
@@ -651,7 +649,7 @@ def load_dnpy(file, comm=None):
 
     try:
         distbuffer = format.read_localarray(fid)
-        return LocalArray.from_distarray(distbuffer, comm=comm)
+        return LocalArray.from_distarray(comm=comm, obj=distbuffer)
 
     finally:
         if own_fid:
@@ -743,7 +741,7 @@ def unstructured_index(dd):
     return tuple(index)
 
 
-def load_hdf5(filename, dim_data, key='buffer', comm=None):
+def load_hdf5(comm, filename, dim_data, key='buffer'):
     """
     Load a LocalArray from an ``.hdf5`` file.
 
@@ -756,10 +754,10 @@ def load_hdf5(filename, dim_data, key='buffer', comm=None):
         https://github.com/enthought/distributed-array-protocol, describing
         which portions of the HDF5 file to load into this LocalArray, and with
         what metadata.
+    comm : MPI comm object
     key : str, optional
         The identifier for the group to load the LocalArray from (the default
         is 'buffer').
-    comm : MPI comm object, optional
 
     Returns
     -------
@@ -787,11 +785,11 @@ def load_hdf5(filename, dim_data, key='buffer', comm=None):
         buf = dset[index]
         dtype = dset.dtype
 
-    distribution = maps.Distribution(dim_data=dim_data, comm=comm)
+    distribution = maps.Distribution(comm=comm, dim_data=dim_data)
     return LocalArray(distribution=distribution, dtype=dtype, buf=buf)
 
 
-def load_npy(filename, dim_data, comm=None):
+def load_npy(comm, filename, dim_data):
     """
     Load a LocalArray from a ``.npy`` file.
 
@@ -804,7 +802,7 @@ def load_npy(filename, dim_data, comm=None):
         https://github.com/enthought/distributed-array-protocol, describing
         which portions of the HDF5 file to load into this LocalArray, and with
         what metadata.
-    comm : MPI comm object, optional
+    comm : MPI comm object
 
     Returns
     -------
@@ -823,7 +821,7 @@ def load_npy(filename, dim_data, comm=None):
     # http://stackoverflow.com/questions/6397495/unmap-of-numpy-memmap
 
     #data._mmap.close()
-    distribution = maps.Distribution(dim_data=dim_data, comm=comm)
+    distribution = maps.Distribution(comm=comm, dim_data=dim_data)
     return LocalArray(distribution=distribution, dtype=data.dtype, buf=buf)
 
 
@@ -893,7 +891,7 @@ def get_printoptions():
 # Reduction functions
 # ---------------------------------------------------------------------------
 
-def local_reduction(reducer, out_comm, larr, ddpr, dtype, axes):
+def local_reduction(out_comm, reducer, larr, ddpr, dtype, axes):
     """ Entry point for reductions on local arrays.
 
     Parameters
@@ -923,8 +921,8 @@ def local_reduction(reducer, out_comm, larr, ddpr, dtype, axes):
         out = None
     else:
         dim_data = ddpr[out_comm.Get_rank()] if ddpr else ()
-        dist = distarray.local.maps.Distribution(dim_data, out_comm)
-        out = distarray.local.empty(dist, dtype)
+        dist = maps.Distribution(comm=out_comm, dim_data=dim_data)
+        out = empty(dist, dtype)
 
     remaining_dims = [False] * larr.ndim
     for axis in axes:

distarray/local/maps.py

@@ -39,14 +39,14 @@ class Distribution(object):
     Manages one or more one-dimensional map classes.
     """
 
-    def __init__(self, dim_data, comm=None):
+    def __init__(self, comm, dim_data):
         """Create a Distribution from a `dim_data` structure."""
         self._maps = tuple(map_from_dim_dict(dim_dict) for dim_dict in dim_data)
         self.base_comm = construct.init_base_comm(comm)
         self.comm = construct.init_comm(self.base_comm, self.grid_shape)
 
     @classmethod
-    def from_shape(cls, shape, dist=None, grid_shape=None, comm=None):
+    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)
@@ -72,7 +72,7 @@ def from_shape(cls, shape, dist=None, grid_shape=None, comm=None):
             distribute_indices(dim_dict)
             dim_data.append(dim_dict)
 
-        return cls(dim_data, comm=base_comm)
+        return cls(comm=base_comm, dim_data=dim_data)
 
     def __getitem__(self, idx):
         return self._maps[idx]