From 731dabdc297ce91ca3fb904ee01e06ecf87fde65 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 14:28:22 -0500 Subject: [PATCH 1/8] Add failing test for fromfunction. --- distarray/tests/test_client.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/distarray/tests/test_client.py b/distarray/tests/test_client.py index 042d7326..f2cc0308 100644 --- a/distarray/tests/test_client.py +++ b/distarray/tests/test_client.py @@ -342,6 +342,13 @@ def test_grid_rank(self): grid_shape=(1, 1, 4)) self.assertEqual(a.grid_shape, (1, 1, 4)) + def test_fromfunction(self): + fn = lambda i, j: i + j + shape = (3, 3) + expected = numpy.fromfunction(fn, shape, dtype=int) + result = self.context.fromfunction(fn, shape, dtype=int) + assert_array_equal(expected, result.tondarray()) + class TestReduceMethods(unittest.TestCase): """Test reduction methods""" From fdf97fd6aedb0cf789c7383675545b3fd5476bdd Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 14:34:01 -0500 Subject: [PATCH 2/8] Fix fromfunction. --- distarray/context.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/distarray/context.py b/distarray/context.py index 2f3c27e8..f1c85048 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -475,11 +475,10 @@ def fromndarray(self, arr, dist=None, grid_shape=None): fromarray = fromndarray def fromfunction(self, function, shape, **kwargs): - func_key = self._generate_key() - self.view.push_function({func_key: function}, targets=self.targets, - block=True) - keys = self._key_and_push(shape, kwargs) + keys = self._key_and_push(function, shape, kwargs) new_key = self._generate_key() - subs = (new_key, func_key) + keys - self._execute('%s = distarray.local.fromfunction(%s,%s,**%s)' % subs) + subs = (new_key,) + keys + cmd = ('%s = distarray.local.fromfunction(%s,' + 'distarray.local.maps.Distribution.from_shape(%s),**%s)') + self._execute(cmd % subs) return DistArray.from_localarrays(new_key, context=self) From 550bb4b0058a9989c645d4ce3445c835027ccaba Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 14:41:07 -0500 Subject: [PATCH 3/8] Pass dtype to from_localarrays if available. --- distarray/context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/distarray/context.py b/distarray/context.py index f1c85048..8436e49c 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -475,10 +475,11 @@ def fromndarray(self, arr, dist=None, grid_shape=None): fromarray = fromndarray def fromfunction(self, function, shape, **kwargs): + dtype = kwargs.get('dtype', None) keys = self._key_and_push(function, shape, kwargs) new_key = self._generate_key() subs = (new_key,) + keys cmd = ('%s = distarray.local.fromfunction(%s,' 'distarray.local.maps.Distribution.from_shape(%s),**%s)') self._execute(cmd % subs) - return DistArray.from_localarrays(new_key, context=self) + return DistArray.from_localarrays(new_key, context=self, dtype=dtype) From 85d9b8e6c069b0d8ec2fa711d456fd30a084c3a4 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 15:28:38 -0500 Subject: [PATCH 4/8] Use a weirder test shape. --- distarray/tests/test_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distarray/tests/test_client.py b/distarray/tests/test_client.py index f2cc0308..d73c46ca 100644 --- a/distarray/tests/test_client.py +++ b/distarray/tests/test_client.py @@ -344,7 +344,7 @@ def test_grid_rank(self): def test_fromfunction(self): fn = lambda i, j: i + j - shape = (3, 3) + shape = (7, 9) expected = numpy.fromfunction(fn, shape, dtype=int) result = self.context.fromfunction(fn, shape, dtype=int) assert_array_equal(expected, result.tondarray()) From 25c2e62a765fa9d42ca7a93c3dd4567fb4b63dd7 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 15:28:59 -0500 Subject: [PATCH 5/8] Modify fromfunction to use the newer from_localarrays --- distarray/context.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/distarray/context.py b/distarray/context.py index 8436e49c..a52aecb4 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -476,10 +476,20 @@ def fromndarray(self, arr, dist=None, grid_shape=None): def fromfunction(self, function, shape, **kwargs): dtype = kwargs.get('dtype', None) - keys = self._key_and_push(function, shape, kwargs) - new_key = self._generate_key() - subs = (new_key,) + keys - cmd = ('%s = distarray.local.fromfunction(%s,' - 'distarray.local.maps.Distribution.from_shape(%s),**%s)') - self._execute(cmd % subs) - return DistArray.from_localarrays(new_key, context=self, dtype=dtype) + dist = kwargs.get('dist', None) + grid_shape = kwargs.get('grid_shape', None) + distribution = Distribution.from_shape(context=self, + shape=shape, dist=dist, + grid_shape=grid_shape) + ddpr = distribution.get_dim_data_per_rank() + function_name, ddpr_name, kwargs_name = \ + self._key_and_push(function, ddpr, kwargs) + da_name = self._generate_key() + comm_name = self._comm_key + cmd = ('{da_name} = distarray.local.fromfunction({function_name}, ' + 'distarray.local.maps.Distribution(' + '{ddpr_name}[{comm_name}.Get_rank()], comm={comm_name}),' + '**{kwargs_name})') + self._execute(cmd.format(**locals())) + return DistArray.from_localarrays(da_name, distribution=distribution, + dtype=dtype) \ No newline at end of file From 2d2ba771be5104214ebc5b83790c20fa72730e41 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 15:30:22 -0500 Subject: [PATCH 6/8] Add a simple docstring. --- distarray/context.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/distarray/context.py b/distarray/context.py index a52aecb4..4a52a35a 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -475,6 +475,10 @@ def fromndarray(self, arr, dist=None, grid_shape=None): fromarray = fromndarray def fromfunction(self, function, shape, **kwargs): + """Create a DistArray from a function over global indices. + + See np.fromfunction for details. + """ dtype = kwargs.get('dtype', None) dist = kwargs.get('dist', None) grid_shape = kwargs.get('grid_shape', None) From 04544d5aab4b7b80029be602a32118d248479a76 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 15:47:42 -0500 Subject: [PATCH 7/8] Don't restrict output dtype to input dtype. --- distarray/context.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/distarray/context.py b/distarray/context.py index 4a52a35a..a2dd0a12 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -495,5 +495,4 @@ def fromfunction(self, function, shape, **kwargs): '{ddpr_name}[{comm_name}.Get_rank()], comm={comm_name}),' '**{kwargs_name})') self._execute(cmd.format(**locals())) - return DistArray.from_localarrays(da_name, distribution=distribution, - dtype=dtype) \ No newline at end of file + return DistArray.from_localarrays(da_name, distribution=distribution) From 575ba3a70fe597d7f6198bb5d24f924923bf42a7 Mon Sep 17 00:00:00 2001 From: Robert David Grant Date: Thu, 1 May 2014 15:52:55 -0500 Subject: [PATCH 8/8] Write a better docstring. --- distarray/context.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/distarray/context.py b/distarray/context.py index a2dd0a12..de039f67 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -477,7 +477,11 @@ def fromndarray(self, arr, dist=None, grid_shape=None): def fromfunction(self, function, shape, **kwargs): """Create a DistArray from a function over global indices. - See np.fromfunction for details. + Unlike numpy's `fromfunction`, the result of distarray's + `fromfunction` is restricted to the same Distribution as the + index array generated from `shape`. + + See numpy.fromfunction for more details. """ dtype = kwargs.get('dtype', None) dist = kwargs.get('dist', None)