[pypy-commit] pypy numpy-identity: Added numpy.identity function

[jterrace]

Author: Jeff Terrace <jterrace at gmail.com>
Branch: numpy-identity
Changeset: r50482:2ff54fec76e7
Date: 2011-12-13 17:15 -0500
http://bitbucket.org/pypy/pypy/changeset/2ff54fec76e7/

Log:	Added 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
@@ -81,6 +81,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
@@ -722,6 +722,37 @@
         a = array([True] * 5, bool)
         assert a.sum() == 5
 
+    def test_identity(self):
+        from numpypy import identity, array
+        from numpypy import int32, float64
+        a = identity(0)
+        assert len(a) == 0
+        assert repr(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 repr(b.dtype) == "dtype('int32')"
+        c = identity(2)
+        assert c.shape == (2,2)
+        assert c[0][0] == 1.0
+        assert c[0][1] == 0.0
+        assert c[1][0] == 0.0
+        assert c[1][1] == 1.0
+        d = identity(3, dtype='int32')
+        assert d.shape == (3,3)
+        assert repr(d.dtype) == "dtype('int32')"
+        assert d[0][0] == 1.0
+        assert d[0][1] == 0.0
+        assert d[0][2] == 0.0
+        assert d[1][0] == 0.0
+        assert d[1][1] == 1.0
+        assert d[1][2] == 0.0
+        assert d[2][0] == 0.0
+        assert d[2][1] == 0.0
+        assert d[2][2] == 1.0
+
     def test_prod(self):
         from numpypy import array
         a = array(range(1, 6))