Negative indexing in distarrays.

distarray/cleanup.py

@@ -8,30 +8,13 @@
 
 from distarray.ipython_utils import IPythonClient
 
+def cleanup(view, module_name, context_name):
+    """ Delete Context object with the given name from the given module"""
+    def _cleanup(module_name, context_name):
+        ns = __import__(module_name)
+        delattr(ns, context_name)
 
-def engine_cleanup(module_name, prefix):
-    """ Remove variables with ``prefix`` prefix from the namespace of the
-    module with ``module_name``. 
-    """
-    mod = __import__(module_name)
-    ns = mod.__dict__
-    keys = tuple(ns.keys())
-    deleted = 0
-    for k in keys:
-        if k.startswith(prefix):
-            del ns[k]
-            deleted += 1
-    count = 0
-    for k in ns:
-        if k.startswith(prefix):
-            count += 1
-    return (deleted, count)
-
-
-def cleanup(view, module_name, prefix):
-    """ Delete keys with prefix from client's engines. """
-    remaining = view.apply_async(engine_cleanup, module_name, prefix).get_dict()
-    return remaining
+    view.apply_async(_cleanup, module_name, context_name)
 
 
 def cleanup_all(module_name, prefix):

distarray/client_map.py

@@ -33,6 +33,7 @@
 from distarray.metadata_utils import (normalize_dist,
                                       normalize_grid_shape,
                                       make_grid_shape,
+                                      positivify,
                                       validate_grid_shape,
                                       _start_stop_block)
 
@@ -469,6 +470,7 @@ def owning_ranks(self, idxs):
         If the `idxs` tuple is out of bounds, raises `IndexError`.
 
         """
+        idxs = map(positivify, idxs, self.shape) # positivify and check
         dim_coord_hits = [m.owners(idx) for (m, idx) in zip(self.maps, idxs)]
         all_coords = product(*dim_coord_hits)
         ranks = [self.rank_from_coords[c] for c in all_coords]

distarray/context.py

@@ -82,8 +82,10 @@ def _setup_context_key(self):
         Create a dict on the engines which will hold everything from
         this context.
         """
-        context_key = self.uid()
-        cmd = '%s = {}' % (context_key)
+        context_key = DISTARRAY_BASE_NAME + self.uid()
+        cmd = ("import types, sys;"
+               "%s = types.ModuleType('%s');")
+        cmd %= (context_key, context_key)
         self._execute(cmd, targets=range(len(self.view)))
         return context_key
 
@@ -151,7 +153,7 @@ def uid():
 
     def _generate_key(self):
         """ Generate a unique key name for this context. """
-        key = "%s['%s']" % (self.context_key, self.uid())
+        key = "%s.%s" % (self.context_key, 'key_' + self.uid())
         return key
 
     def _key_and_push(self, *values):
@@ -166,7 +168,7 @@ def delete_key(self, key):
 
     def cleanup(self):
         """ Delete keys that this context created from all the engines. """
-        cleanup.cleanup(view=self.view, module_name='__main__', prefix=self._key_prefix())
+        cleanup.cleanup(view=self.view, module_name='__main__', context_name=self.context_key)
 
     def close(self):
         self.cleanup()

distarray/local/maps.py

@@ -29,7 +29,7 @@
 from distarray.local import construct
 from distarray.metadata_utils import (validate_grid_shape, make_grid_shape,
                                       normalize_grid_shape, normalize_dist,
-                                      distribute_indices)
+                                      distribute_indices, positivify)
 
 
 class Distribution(object):
@@ -136,6 +136,7 @@ def rank_from_coords(self, coords):
 
     def local_from_global(self, *global_ind):
         """ Given `global_ind` indices, translate into local indices."""
+        global_ind = tuple(map(positivify, global_ind, self.global_shape))
         return tuple(self._maps[dim].local_from_global(global_ind[dim])
                      for dim in range(self.ndim))
 

distarray/metadata_utils.py

@@ -192,3 +192,12 @@ def distribute_indices(dd):
     except KeyError:
         msg = "dist_type %r not supported."
         raise TypeError(msg % dist_type)
+
+
+def positivify(index, size):
+    if 0 <= index < size:
+        return index
+    elif -size <= index < 0:
+        return size + index
+    else:
+        raise IndexError("Index %s out of bounds" % index)

distarray/tests/test_client.py

@@ -41,6 +41,10 @@ def test_set_and_getitem_block_dist(self):
         for val in range(size):
             self.assertEqual(dap[val], val)
 
+        for i in range(1, size + 1):
+            dap[-i] = i
+            self.assertEqual(dap[-i], i)
+
     def test_set_and_getitem_nd_block_dist(self):
         size = 5
         dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'})
@@ -49,31 +53,39 @@ def test_set_and_getitem_nd_block_dist(self):
             for col in range(size):
                 val = size*row + col
                 dap[row, col] = val
-
-        for row in range(size):
-            for col in range(size):
-                val = size*row + col
                 self.assertEqual(dap[row, col], val)
 
+        for row in range(1 ,size + 1):
+            for col in range(1, size + 1):
+                dap[-row, -col] = row + col
+                self.assertEqual(dap[-row, -col], row + col)
+
+
     def test_set_and_getitem_cyclic_dist(self):
         size = 10
         dap = self.dac.empty((size,), dist={0: 'c'})
 
         for val in range(size):
             dap[val] = val
-
-        for val in range(size):
             self.assertEqual(dap[val], val)
 
+        for i in range(1, size + 1):
+            dap[-i] = i
+            self.assertEqual(dap[-i], i)
+
     def test_get_index_error(self):
         dap = self.dac.empty((10,), dist={0: 'c'})
         with self.assertRaises(IndexError):
             dap[11]
+        with self.assertRaises(IndexError):
+            dap[-11]
 
     def test_set_index_error(self):
         dap = self.dac.empty((10,), dist={0: 'c'})
         with self.assertRaises(IndexError):
             dap[11] = 55
+        with self.assertRaises(IndexError):
+            dap[-11] = 55
 
     def test_iteration(self):
         size = 10