Make comm first argument in localarray load_* functions.

Author: Kurt Smith

distarray/dist/context.py

@@ -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'):
         """
@@ -476,7 +475,7 @@ def load_hdf5(self, filename, distribution, key='buffer'):
         def _local_load_hdf5(filename, ddpr, comm, key):
             from distarray.local import load_hdf5
             dim_data = ddpr[comm.Get_rank()]
-            return proxyize(load_hdf5(filename, dim_data, comm, key))
+            return proxyize(load_hdf5(comm, filename, dim_data, key))
 
         ddpr = distribution.get_dim_data_per_rank()
 

distarray/dist/distarray.py

@@ -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/localarray.py

@@ -198,7 +198,7 @@ def compatibility_hash(self):
     #-------------------------------------------------------------------------
 
     @classmethod
-    def from_distarray(cls, obj, comm):
+    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
@@ -625,7 +625,7 @@ def save_dnpy(file, arr):
             fid.close()
 
 
-def load_dnpy(file, comm):
+def load_dnpy(comm, file):
     """
     Load a LocalArray from a ``.dnpy`` file.
 
@@ -649,7 +649,7 @@ def load_dnpy(file, comm):
 
     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:
@@ -741,7 +741,7 @@ def unstructured_index(dd):
     return tuple(index)
 
 
-def load_hdf5(filename, dim_data, comm, key='buffer'):
+def load_hdf5(comm, filename, dim_data, key='buffer'):
     """
     Load a LocalArray from an ``.hdf5`` file.
 
@@ -789,7 +789,7 @@ def load_hdf5(filename, dim_data, comm, key='buffer'):
     return LocalArray(distribution=distribution, dtype=dtype, buf=buf)
 
 
-def load_npy(filename, dim_data, comm):
+def load_npy(comm, filename, dim_data):
     """
     Load a LocalArray from a ``.npy`` file.
 
@@ -891,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

distarray/local/tests/paralleltest_distributed_array_protocol.py

@@ -43,12 +43,11 @@ def assert_round_trip_equality(self, larr):
         assert_array_equal(larr.ndarray, self.larr.ndarray)
 
     def test_round_trip_equality_from_object(self):
-        larr = LocalArray.from_distarray(self.larr, comm=self.comm)
+        larr = LocalArray.from_distarray(comm=self.comm, obj=self.larr)
         self.assert_round_trip_equality(larr)
 
     def test_round_trip_equality_from_dict(self):
-        larr = LocalArray.from_distarray(self.larr.__distarray__(),
-                                         comm=self.comm)
+        larr = LocalArray.from_distarray(comm=self.comm, obj=self.larr.__distarray__())
         self.assert_round_trip_equality(larr)
 
 
@@ -65,7 +64,7 @@ def test_with_validator(self):
         validate_distbuffer(self.larr.__distarray__())
 
     def test_round_trip_elements(self):
-        larr = LocalArray.from_distarray(self.larr, comm=self.comm)
+        larr = LocalArray.from_distarray(comm=self.comm, obj=self.larr)
         if self.comm.Get_rank() == 0:
             idx = (0,) * larr.ndarray.ndim
             larr.ndarray[idx] = 99
@@ -183,7 +182,7 @@ def test_values(self):
         elif self.comm.Get_rank() == 1:
             assert_array_equal(np.arange(30), self.larr.ndarray)
 
-        larr = LocalArray.from_distarray(self.larr, comm=self.comm)
+        larr = LocalArray.from_distarray(comm=self.comm, obj=self.larr)
         if self.comm.Get_rank() == 0:
             assert_array_equal(np.arange(20), larr.ndarray)
         elif self.comm.Get_rank() == 1:

distarray/local/tests/paralleltest_io.py

@@ -47,14 +47,14 @@ def test_flat_file_save_with_file_object(self):
 
     def test_flat_file_save_load_with_filename(self):
         save_dnpy(self.output_path, self.larr0)
-        larr1 = load_dnpy(self.output_path, comm=self.comm)
+        larr1 = load_dnpy(comm=self.comm, file=self.output_path)
         self.assertTrue(isinstance(larr1, LocalArray))
         assert_allclose(self.larr0, larr1)
 
     def test_flat_file_save_load_with_file_object(self):
         save_dnpy(self.output_path, self.larr0)
         with open(self.output_path, 'rb') as fp:
-            larr1 = load_dnpy(fp, comm=self.comm)
+            larr1 = load_dnpy(comm=self.comm, file=fp)
         self.assertTrue(isinstance(larr1, LocalArray))
         assert_allclose(self.larr0, larr1)
 
@@ -162,25 +162,25 @@ def tearDown(self):
 
     def test_load_bn(self):
         dim_data_per_rank = bn_test_data
-        la = load_npy(self.output_path, dim_data_per_rank[self.rank],
-                      comm=self.comm)
+        la = load_npy(comm=self.comm, filename=self.output_path,
+                      dim_data=dim_data_per_rank[self.rank])
         assert_equal(la, self.expected[numpy.newaxis, self.rank])
 
     def test_load_nc(self):
         dim_data_per_rank = nc_test_data
         expected_slices = [(slice(None), slice(0, None, 2)),
                            (slice(None), slice(1, None, 2))]
 
-        la = load_npy(self.output_path, dim_data_per_rank[self.rank],
-                      comm=self.comm)
+        la = load_npy(comm=self.comm, filename=self.output_path,
+                      dim_data=dim_data_per_rank[self.rank])
         assert_equal(la, self.expected[expected_slices[self.rank]])
 
     def test_load_nu(self):
         dim_data_per_rank = nu_test_data
         expected_indices = [dd[1]['indices'] for dd in dim_data_per_rank]
 
-        la = load_npy(self.output_path, dim_data_per_rank[self.rank],
-                      comm=self.comm)
+        la = load_npy(comm=self.comm, filename=self.output_path,
+                      dim_data=dim_data_per_rank[self.rank])
         assert_equal(la, self.expected[:, expected_indices[self.rank]])
 
 
@@ -257,8 +257,9 @@ def tearDown(self):
 
     def test_load_bn(self):
         dim_data_per_rank = bn_test_data
-        la = load_hdf5(self.output_path, dim_data_per_rank[self.rank],
-                       key=self.key, comm=self.comm)
+        la = load_hdf5(comm=self.comm, filename=self.output_path,
+                       dim_data=dim_data_per_rank[self.rank],
+                       key=self.key)
         with self.h5py.File(self.output_path, 'r', driver='mpio',
                             comm=self.comm) as fp:
             assert_equal(la, self.expected[numpy.newaxis, self.rank])
@@ -267,8 +268,9 @@ def test_load_nc(self):
         dim_data_per_rank = nc_test_data
         expected_slices = [(slice(None), slice(0, None, 2)),
                            (slice(None), slice(1, None, 2))]
-        la = load_hdf5(self.output_path, dim_data_per_rank[self.rank],
-                       key=self.key, comm=self.comm)
+        la = load_hdf5(comm=self.comm, filename=self.output_path,
+                       dim_data=dim_data_per_rank[self.rank],
+                       key=self.key)
         with self.h5py.File(self.output_path, 'r', driver='mpio',
                             comm=self.comm) as fp:
             expected_slice = expected_slices[self.rank]
@@ -277,8 +279,9 @@ def test_load_nc(self):
     def test_load_nu(self):
         dim_data_per_rank = nu_test_data
         expected_indices = [dd[1]['indices'] for dd in dim_data_per_rank]
-        la = load_hdf5(self.output_path, dim_data_per_rank[self.rank],
-                       key=self.key, comm=self.comm)
+        la = load_hdf5(comm=self.comm, filename=self.output_path,
+                       dim_data=dim_data_per_rank[self.rank],
+                       key=self.key)
         with self.h5py.File(self.output_path, 'r', driver='mpio',
                             comm=self.comm) as fp:
             assert_equal(la, self.expected[:, expected_indices[self.rank]])