Fix and test 0d arrays

distarray/dist/context.py

@@ -216,8 +216,11 @@ def _create_local(self, local_call, distribution, dtype):
         """Creates LocalArrays with the method named in `local_call`."""
         def create_local(local_call, ddpr, dtype, comm):
             from distarray.local.maps import Distribution
+            if len(ddpr) == 0:
+                dim_data = ()
+            else:
+                dim_data = ddpr[comm.Get_rank()]
             local_call = eval(local_call)
-            dim_data = ddpr[comm.Get_rank()]
             distribution = Distribution(comm=comm, dim_data=dim_data)
             rval = local_call(distribution=distribution, dtype=dtype)
             return proxyize(rval)
@@ -441,7 +444,7 @@ def load_npy(self, filename, distribution):
         result : DistArray
             A DistArray encapsulating the file loaded.
         """
-        
+
         def _local_load_npy(filename, ddpr, comm):
             from distarray.local import load_npy
             dim_data = ddpr[comm.Get_rank()]

distarray/dist/distarray.py

@@ -21,7 +21,7 @@
 
 import numpy as np
 
-import distarray
+import distarray.local
 from distarray.metadata_utils import sanitize_indices
 from distarray.dist.maps import Distribution
 from distarray.utils import _raise_nie
@@ -215,7 +215,10 @@ def _set_view(self, index, value):
         def set_view(arr, index, value, ddpr, comm):
             from distarray.local.localarray import LocalArray
             from distarray.local.maps import Distribution
-            dim_data = ddpr[comm.Get_rank()]
+            if len(ddpr) == 0:
+                dim_data = ()
+            else:
+                dim_data = ddpr[comm.Get_rank()]
             dist = Distribution(comm=comm, dim_data=dim_data)
             if isinstance(value, LocalArray):
                 arr.global_index[index] = value.ndarray

distarray/dist/tests/test_distarray.py

@@ -279,21 +279,25 @@ def test_all_ellipsis(self):
         assert_array_equal(arr[..., ..., ..., ...].toarray(),
                            expected[..., ..., ..., ...])
 
-    @unittest.skip("Waiting on 0d-array support.")
-    def test_0d_ellipsis(self):
-        shape = ()
-        expected = numpy.random.randint(10, size=shape)
-        arr = self.context.fromarray(expected)
-        assert_array_equal(arr[...].toarray(),
-                           expected[...])
-
     def test_resulting_slice(self):
         dist = Distribution.from_shape(self.context, (10, 20))
         da = self.context.ones(dist)
         db = da[:5, :10]
         dc = db * 2
         assert_array_equal(dc.toarray(), numpy.ones(dc.shape) * 2)
 
+    def test_0d_empty_tuple(self):
+        shape = ()
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[()], expected[()])
+
+    def test_0d_ellipsis(self):
+        shape = ()
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[...].toarray(), expected[...])
+
 
 class TestSetItemSlicing(ContextTestCase):
 
@@ -421,6 +425,18 @@ def test_set_DistArray_slice(self):
         db = self.context.zeros(dist)
         da[...] = db
 
+    def test_set_0d_slice(self):
+        expected = numpy.array(33)
+        arr = self.context.fromarray(expected)
+
+        val0 = 55
+        arr[...] = val0
+        assert_array_equal(arr[...].toarray(), numpy.array(val0))
+
+        val1 = 99
+        arr[()] = val1
+        assert_array_equal(arr[...].toarray(), numpy.array(val1))
+
 
 class TestDistArrayCreationFromGlobalDimData(ContextTestCase):
 
@@ -634,15 +650,36 @@ def test_zeros(self):
         zero_ndarray = numpy.zeros(shape)
         assert_array_equal(zero_distarray.tondarray(), zero_ndarray)
 
+    def test_zeros_0d(self):
+        shape = ()
+        distribution = Distribution.from_shape(self.context, shape)
+        zero_distarray = self.context.zeros(distribution)
+        zero_ndarray = numpy.zeros(shape)
+        assert_array_equal(zero_distarray.tondarray(), zero_ndarray)
+
     def test_ones(self):
         shape = (16, 16)
         distribution = Distribution.from_shape(self.context, shape)
-        one_distarray = self.context.ones(distribution)
-        one_ndarray = numpy.ones(shape)
-        assert_array_equal(one_distarray.tondarray(), one_ndarray)
+        ones_distarray = self.context.ones(distribution)
+        ones_ndarray = numpy.ones(shape)
+        assert_array_equal(ones_distarray.tondarray(), ones_ndarray)
+
+    def test_ones_0d(self):
+        shape = ()
+        distribution = Distribution.from_shape(self.context, shape)
+        ones_distarray = self.context.ones(distribution)
+        ones_ndarray = numpy.ones(shape)
+        assert_array_equal(ones_distarray.tondarray(), ones_ndarray)
 
     def test_empty(self):
-        distribution = Distribution.from_shape(self.context, (16, 16))
+        shape = (16, 16)
+        distribution = Distribution.from_shape(self.context, shape)
+        empty_distarray = self.context.empty(distribution)
+        self.assertEqual(empty_distarray.shape, distribution.shape)
+
+    def test_empty_0d(self):
+        shape = ()
+        distribution = Distribution.from_shape(self.context, shape)
         empty_distarray = self.context.empty(distribution)
         self.assertEqual(empty_distarray.shape, distribution.shape)
 
@@ -652,6 +689,11 @@ def test_fromndarray(self):
         for (i, j), val in numpy.ndenumerate(ndarr):
             self.assertEqual(distarr[i, j], ndarr[i, j])
 
+    def test_fromndarray_0d(self):
+        ndarr = numpy.array(42)
+        distarr = self.context.fromarray(ndarr)
+        assert_array_equal(ndarr, distarr.toarray())
+
     def test_grid_rank(self):
         # regression test for issue #235
         d = Distribution.from_shape(self.context, (4, 4, 4),

distarray/local/localarray.py

@@ -111,8 +111,10 @@ def __init__(self, distribution, dtype=None, buf=None):
         if buf is None:
             self.ndarray = np.empty(self.local_shape, dtype=dtype)
         else:
-            mv = memoryview(buf)
-            self.ndarray = np.asarray(mv, dtype=dtype)
+            self.ndarray = np.asarray(buf, dtype=dtype)
+            if distribution.local_shape != self.ndarray.shape:
+                msg = "distribution shape must equal buf shape."
+                raise RuntimeError(msg)
 
         # We pass a view of self.ndarray because we want the
         # GlobalIndex object to be able to change the LocalArray

distarray/local/tests/paralleltest_distributed_array_protocol.py

@@ -166,15 +166,32 @@ class TestDapLopsided(DapValidatorMixin, MpiTestCase):
     comm_size = 2
 
     def setUp(self):
+        global_size = 50
         if self.comm.Get_rank() == 0:
-            arr = np.arange(20)
+            local_size = 20
+            arr = np.arange(local_size)
+            dim_data = ({
+                'dist_type': 'b',
+                'size': global_size,
+                'proc_grid_size': 2,
+                'proc_grid_rank': 0,
+                'start': 0,
+                'stop': local_size,
+            },)
         elif self.comm.Get_rank() == 1:
-            arr = np.arange(30)
-
-        d = Distribution.from_shape(comm=self.comm, shape=(50,),
-                                    dist={0: 'b', 1: 'n'}, grid_shape=(2,))
-
-        self.larr = LocalArray(d, dtype='float64', buf=arr)
+            local_size = 30
+            arr = np.arange(local_size)
+            dim_data = ({
+                'dist_type': 'b',
+                'size': global_size,
+                'proc_grid_size': 2,
+                'proc_grid_rank': 1,
+                'start': 20,
+                'stop': global_size,
+            },)
+
+        d = Distribution(comm=self.comm, dim_data=dim_data)
+        self.larr = LocalArray(d, buf=arr)
 
     def test_values(self):
         if self.comm.Get_rank() == 0: