Modify Random.randn to use new from_localarrays

enthought · May 1, 2014 · 7272d21 · 7272d21
1 parent e9b3432
commit 7272d21
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/distarray/random.py b/distarray/random.py
@@ -241,14 +241,17 @@ def randn(self, size=None, dist=None, grid_shape=None):
         if dist is None:
             dist = {0: 'b'}
         da_key = self.context._generate_key()
-        size_key, dist_key, grid_shape_key = \
-            self.context._key_and_push(size, dist, grid_shape)
-        comm_key = self.context._comm_key
+
+        distribution = Distribution.from_shape(context=self.context,
+                                               shape=size,
+                                               dist=dist,
+                                               grid_shape=grid_shape)
+        ddpr = distribution.get_dim_data_per_rank()
+        ddpr_name = self.context._key_and_push(ddpr)[0]
+        comm_name = self.context._comm_key
         self.context._execute(
             '{da_key} = distarray.local.random.randn('
-            'distribution=distarray.local.maps.Distribution.from_shape('
-            'shape={size_key}, dist={dist_key},'
-            'grid_shape={grid_shape_key}, comm={comm_key}'
-            '))'.format(**locals())
-        )
-        return DistArray.from_localarrays(da_key, self.context)
+            'distribution=distarray.local.maps.Distribution('
+            'dim_data={ddpr_name}[{comm_name}.Get_rank()], '
+            'comm={comm_name}))'.format(**locals()))
+        return DistArray.from_localarrays(da_key, distribution=distribution)