[pypy-commit] pypy default: (jterrace) merge numpy-identity branch, adding numpy.identity function

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r50484:3b81c36d2bd1
Date: 2011-12-14 01:43 +0200
http://bitbucket.org/pypy/pypy/changeset/3b81c36d2bd1/

Log:	(jterrace) merge numpy-identity branch, adding numpy.identity
	function

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -91,6 +91,7 @@
         'mean': 'app_numpy.mean',
         'sum': 'app_numpy.sum',
         'min': 'app_numpy.min',
+        'identity': 'app_numpy.identity',
         'max': 'app_numpy.max',
         'inf': 'app_numpy.inf',
         'e': 'app_numpy.e',
diff --git a/pypy/module/micronumpy/app_numpy.py b/pypy/module/micronumpy/app_numpy.py
--- a/pypy/module/micronumpy/app_numpy.py
+++ b/pypy/module/micronumpy/app_numpy.py
@@ -13,6 +13,11 @@
     # weighting, just the average part!
     return mean(a)
 
+def identity(n, dtype=None):
+    a = numpypy.zeros((n,n), dtype=dtype)
+    for i in range(n):
+        a[i][i] = 1
+    return a
 
 def mean(a):
     if not hasattr(a, "mean"):
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -727,6 +727,26 @@
         a = array([True] * 5, bool)
         assert a.sum() == 5
 
+    def test_identity(self):
+        from numpypy import identity, array
+        from numpypy import int32, float64, dtype
+        a = identity(0)
+        assert len(a) == 0
+        assert a.dtype == dtype('float64')
+        assert a.shape == (0,0)
+        b = identity(1, dtype=int32)
+        assert len(b) == 1
+        assert b[0][0] == 1
+        assert b.shape == (1,1)
+        assert b.dtype == dtype('int32')
+        c = identity(2)
+        assert c.shape == (2,2)
+        assert (c == [[1,0],[0,1]]).all()
+        d = identity(3, dtype='int32')
+        assert d.shape == (3,3)
+        assert d.dtype == dtype('int32')
+        assert (d == [[1,0,0],[0,1,0],[0,0,1]]).all()
+
     def test_prod(self):
         from numpypy import array
         a = array(range(1, 6))