Allow and test 1's in grid_shape

distarray/local/localarray.py

@@ -222,8 +222,8 @@ def _init(self, dim_data, grid_shape, dtype=None, buf=None, comm=None):
         # object's data.
         self.global_index = GlobalIndex(self.maps, self.local_array.view())
 
-        self.base = None
-        self.ctypes = None
+        self.base = None  # mimic numpy.ndarray.base
+        self.ctypes = None  # mimic numpy.ndarray.ctypes
 
     def _init_grid_shape(self, grid_shape):
 

distarray/local/tests/paralleltest_localarray.py

@@ -230,6 +230,8 @@ def test_grid_shape(self):
         self.larr = LocalArray((100, 10, 300), dist=('b', 'n', 'c'), comm=self.comm)
         self.assertEqual(self.larr.grid_shape, (2, 1, 6))
         self.larr = LocalArray((100, 50, 300), dist='b', comm=self.comm)
+        self.assertEqual(self.larr.grid_shape, (2, 1, 6))
+        self.larr = LocalArray((100, 100, 150), dist='b', comm=self.comm)
         self.assertEqual(self.larr.grid_shape, (2, 2, 3))
 
     def test_ones_in_grid_shape(self):

distarray/metadata_utils.py

@@ -24,16 +24,6 @@ class GridShapeError(Exception):
     pass
 
 
-def make_grid_shape(global_shape, dist, comm_size):
-    """ Generate a `grid_shape` from `global_shape` tuple and `dist` tuple.
-
-    Does not assume that `dim_data` has `proc_grid_size` set for each
-    dimension.
-    """
-    distdims = tuple(i for (i, v) in enumerate(dist) if v != 'n')
-    return optimize_grid_shape(global_shape, distdims, comm_size)
-
-
 def normalize_grid_shape(grid_shape, ndims):
     """ Adds 1's to grid_shape so it has `ndims` dimensions.
     """
@@ -58,8 +48,13 @@ def validate_grid_shape(grid_shape, dist, comm_size):
     return grid_shape
 
 
-def optimize_grid_shape(shape, distdims, comm_size):
-    ''' Attempts to allocate processes optimally for distributed dimensions.
+def make_grid_shape(shape, dist, comm_size):
+    """ Generate a `grid_shape` from `shape` tuple and `dist` tuple.
+
+    Does not assume that `dim_data` has `proc_grid_size` set for each
+    dimension.
+
+    Attempts to allocate processes optimally for distributed dimensions.
 
     Parameters
     ----------
@@ -78,7 +73,8 @@ def optimize_grid_shape(shape, distdims, comm_size):
     ------
         GridShapeError if not possible to distribute `comm_size` processes over
         number of dimensions.
-    '''
+    """
+    distdims = tuple(i for (i, v) in enumerate(dist) if v != 'n')
     ndistdim = len(distdims)
 
     if ndistdim == 1:

distarray/tests/test_context.py

@@ -105,5 +105,30 @@ def test_purge_and_dump_keys(self):
         self.assertEqual(num_keys2, num_keys0)
 
 
+class TestPrimeCluster(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.context = Context(targets=range(3))
+
+    def test_1D(self):
+        a = self.context.empty((3,))
+        self.assertEqual(a.grid_shape, (3,))
+
+    def test_2D(self):
+        a = self.context.empty((3, 3))
+        b = self.context.empty((3, 3), dist=('n', 'b'))
+        self.assertEqual(a.grid_shape, (3, 1))
+        self.assertEqual(b.grid_shape, (1, 3))
+
+    def test_3D(self):
+        a = self.context.empty((3, 3, 3))
+        b = self.context.empty((3, 3, 3), dist=('n', 'b', 'n'))
+        c = self.context.empty((3, 3, 3), dist=('n', 'n', 'b'))
+        self.assertEqual(a.grid_shape, (3, 1, 1))
+        self.assertEqual(b.grid_shape, (1, 3, 1))
+        self.assertEqual(c.grid_shape, (1, 1, 3))
+
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)

distarray/tests/test_utils.py

@@ -16,14 +16,18 @@ class TestMultPartitions(unittest.TestCase):
     Test the multiplicative parition code.
     """
 
-    def test_both_methods(self):
-        """
-        Do the two methods of computing the multiplicative partitions agree?
-        """
-        for s in [2, 3]:
-            for n in range(2, 512):
-                self.assertEqual(utils.mult_partitions(n, s),
-                                 utils.create_factors(n, s))
+    def test_mult_partitions(self):
+        self.assertEqual(utils.mult_partitions(1, 2), [(1, 1)])
+        self.assertEqual(utils.mult_partitions(2, 2), [(1, 2)])
+        self.assertEqual(utils.mult_partitions(1, 7), [(1,) * 7])
+        self.assertEqual(utils.mult_partitions(7, 2), [(1, 7)])
+        self.assertEqual(utils.mult_partitions(7, 3), [(1, 1, 7)])
+        self.assertEqual(utils.mult_partitions(16, 4), [(1, 1, 1, 16),
+                                                        (1, 1, 2, 8),
+                                                        (1, 1, 4, 4),
+                                                        (1, 2, 2, 4),
+                                                        (2, 2, 2, 2)])
+        self.assertEqual(utils.mult_partitions(6, 3), [(1, 1, 6), (1, 2, 3)])
 
 
 class TestSanitizeIndices(unittest.TestCase):

distarray/utils.py

@@ -11,13 +11,6 @@
 from distarray.externals.six import next
 
 
-def divisors(n):
-    i = 2
-    while i<n:
-        if n % i == 0:
-            yield i
-        i += 1
-
 
 def multi_for(iterables):
     if not iterables:
@@ -28,22 +21,6 @@ def multi_for(iterables):
                 yield (item,) + rest_tuple
 
 
-def create_factors(n, size=2):
-    divs = list(divisors(n))
-    factors = []
-    for indices in multi_for( [range(p) for p in size*[len(divs)]] ):
-        total = 1
-        for i in indices:
-            total = total*divs[i]
-        if n == total:
-            factor = [divs[i] for i in indices]
-            factor.sort()
-            factor = tuple(factor)
-            if factor not in factors:
-                factors.append(factor)
-    return factors
-
-
 def divisors_minmax(n, dmin, dmax):
     """Find the divisors of n in the interval (dmin,dmax]."""
     i = dmin+1
@@ -75,10 +52,10 @@ def mult_partitions(n, s):
     >>> mult_partitions(52,2)
     [(2, 26), (4, 13)]
     """
-    return [tuple(flatten(p)) for p in mult_partitions_recurs(n,s)]
+    return [tuple(flatten(p)) for p in mult_partitions_recurs(n, s)]
 
 
-def mult_partitions_recurs(n, s, pd=1):
+def mult_partitions_recurs(n, s, pd=0):
     if s == 1:
         return [n]
     divs = divisors_minmax(n, pd, int(sqrt(n)))
@@ -107,18 +84,6 @@ def mirror_sort(seq, ref_seq):
     return newseq
 
 
-def outer_zip(seqa, seqb):
-    """An outer product, but using zip rather than multiplication.
-
-    >>> a = 2*range(4)
-    >>> b = range(8)
-    >>> outer_zip(a,b)
-    [[(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7)], [(1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (1, 7)], [(2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2, 7)], [(3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (3, 7)], [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7)], [(1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (1, 7)], [(2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2, 7)], [(3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (3, 7)]]
-
-    """
-    return [[(i,j) for j in seqb] for i in seqa]
-
-
 def _raise_nie():
     msg = "This has not yet been implemented for distributed arrays"
     raise NotImplementedError(msg)