Commits

Taavi Burns  committed 93fbec9

Adding logaddexp and logaddexp2 functions to numpypy

  • Participants
  • Parent commits 5b3e480

Comments (0)

Files changed (4)

File pypy/module/micronumpy/__init__.py

         ('log1p', 'log1p'),
         ('power', 'power'),
         ('floor_divide', 'floor_divide'),
+        ('logaddexp', 'logaddexp'),
+        ('logaddexp2', 'logaddexp2'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 

File pypy/module/micronumpy/interp_ufuncs.py

             ("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)
 

File pypy/module/micronumpy/test/test_ufuncs.py

         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 prob12 == -113.87649168120691
+        assert math.exp(prob12) == 3.5000000000000057e-50
+
+        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 prob12 == -164.28904982231052
+
+        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')

File pypy/module/micronumpy/types.py

         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 = 2**v1
+        except OverflowError:
+            v1e = rfloat.INFINITY
+        try:
+            v2e = 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