[pypy-commit] pypy numpy-ufuncs2: Adding logaddexp and logaddexp2

[taavi_burns]

Author: Taavi Burns <taavi.burns at gmail.com>
Branch: numpy-ufuncs2
Changeset: r53438:36652d68f0ce
Date: 2012-03-13 01:19 -0700
http://bitbucket.org/pypy/pypy/changeset/36652d68f0ce/

Log:	Adding logaddexp and logaddexp2

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
@@ -116,6 +116,8 @@
         ('log1p', 'log1p'),
         ('power', 'power'),
         ('floor_divide', 'floor_divide'),
+        ('logaddexp', 'logaddexp'),
+        ('logaddexp2', 'logaddexp2'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -450,6 +450,8 @@
             ("log2", "log2", 1, {"promote_to_float": True}),
             ("log10", "log10", 1, {"promote_to_float": True}),
             ("log1p", "log1p", 1, {"promote_to_float": True}),
+            ("logaddexp", "logaddexp", 2, {"promote_to_float": True}),
+            ("logaddexp2", "logaddexp2", 2, {"promote_to_float": True}),
         ]:
             self.add_ufunc(space, *ufunc_def)
 
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -564,3 +564,50 @@
         b = floor_divide(a, 2.5)
         for i in range(len(a)):
             assert b[i] == a[i] // 2.5
+
+    def test_logaddexp(self):
+        import math
+        from _numpypy import logaddexp
+
+        # From the numpy documentation
+        prob1 = math.log(1e-50)
+        prob2 = math.log(2.5e-50)
+        prob12 = logaddexp(prob1, prob2)
+        assert math.fabs(-113.87649168120691 - prob12) < 0.000000000001
+
+        assert logaddexp(0, 0) == math.log(2)
+        assert logaddexp(float('-inf'), 0) == 0
+        assert logaddexp(12345678, 12345678) == float('inf')
+
+        assert math.isnan(logaddexp(float('nan'), 1))
+        assert math.isnan(logaddexp(1, float('nan')))
+        assert math.isnan(logaddexp(float('nan'), float('inf')))
+        assert math.isnan(logaddexp(float('inf'), float('nan')))
+        assert logaddexp(float('-inf'), float('-inf')) == float('-inf')
+        assert logaddexp(float('-inf'), float('inf')) == float('inf')
+        assert logaddexp(float('inf'), float('-inf')) == float('inf')
+        assert logaddexp(float('inf'), float('inf')) == float('inf')
+
+    def test_logaddexp2(self):
+        import math
+        from _numpypy import logaddexp2
+        log2 = math.log(2)
+
+        # From the numpy documentation
+        prob1 = math.log(1e-50) / log2
+        prob2 = math.log(2.5e-50) / log2
+        prob12 = logaddexp2(prob1, prob2)
+        assert math.fabs(-164.28904982231052 - prob12) < 0.000000000001
+
+        assert logaddexp2(0, 0) == 1
+        assert logaddexp2(float('-inf'), 0) == 0
+        assert logaddexp2(12345678, 12345678) == float('inf')
+
+        assert math.isnan(logaddexp2(float('nan'), 1))
+        assert math.isnan(logaddexp2(1, float('nan')))
+        assert math.isnan(logaddexp2(float('nan'), float('inf')))
+        assert math.isnan(logaddexp2(float('inf'), float('nan')))
+        assert logaddexp2(float('-inf'), float('-inf')) == float('-inf')
+        assert logaddexp2(float('-inf'), float('inf')) == float('inf')
+        assert logaddexp2(float('inf'), float('-inf')) == float('inf')
+        assert logaddexp2(float('inf'), float('inf')) == float('inf')
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -601,6 +601,48 @@
         except ValueError:
             return rfloat.NAN
 
+    @simple_binary_op
+    def logaddexp(self, v1, v2):
+        try:
+            v1e = math.exp(v1)
+        except OverflowError:
+            v1e = rfloat.INFINITY
+        try:
+            v2e = math.exp(v2)
+        except OverflowError:
+            v2e = rfloat.INFINITY
+
+        v12e = v1e + v2e
+        try:
+            return math.log(v12e)
+        except ValueError:
+            if v12e == 0.0:
+                # CPython raises ValueError here, so we have to check
+                # the value to find the correct numpy return value
+                return -rfloat.INFINITY
+            return rfloat.NAN
+
+    @simple_binary_op
+    def logaddexp2(self, v1, v2):
+        try:
+            v1e = math.pow(2, v1)
+        except OverflowError:
+            v1e = rfloat.INFINITY
+        try:
+            v2e = math.pow(2, v2)
+        except OverflowError:
+            v2e = rfloat.INFINITY
+
+        v12e = v1e + v2e
+        try:
+            return math.log(v12e) / log2
+        except ValueError:
+            if v12e == 0.0:
+                # CPython raises ValueError here, so we have to check
+                # the value to find the correct numpy return value
+                return -rfloat.INFINITY
+            return rfloat.NAN
+
 
 class Float32(BaseType, Float):
     T = rffi.FLOAT