Add convex analysis ufuncs to cupyx.scipy.special

cupyx/scipy/special/__init__.py

@@ -19,3 +19,9 @@
 from cupyx.scipy.special.erf import erfcx  # NOQA
 from cupyx.scipy.special.erf import erfinv  # NOQA
 from cupyx.scipy.special.erf import erfcinv  # NOQA
+
+from cupyx.scipy.special.convex_analysis import entr  # NOQA
+from cupyx.scipy.special.convex_analysis import huber  # NOQA
+from cupyx.scipy.special.convex_analysis import kl_div  # NOQA
+from cupyx.scipy.special.convex_analysis import pseudo_huber  # NOQA
+from cupyx.scipy.special.convex_analysis import rel_entr  # NOQA

cupyx/scipy/special/convex_analysis.py

@@ -0,0 +1,129 @@
+from cupy import core
+
+
+_float_preamble = '''
+#ifndef NAN
+#define NAN __int_as_float(0x7fffffff)
+#endif
+
+#ifndef INF
+#define INF __int_as_float(0x7f800000)
+#endif
+
+
+double __device__ entr(double x) {
+    if(isnan(x)) {
+        return NAN;
+    }
+    else if(x > 0){
+        return -x * log(x);
+    } else if(x == 0){
+        return 0;
+    } else {
+        return -INF;
+    }
+}
+
+double __device__ kl_div(double x, double y) {
+    if (isnan(x) | isnan(y)) {
+        return NAN;
+    } else if (x > 0 & y > 0) {
+        return x * log(x / y) - x + y;
+    } else if (x == 0 & y >= 0) {
+        return y;
+    } else {
+        return INF;
+    }
+}
+
+double __device__ rel_entr(double x, double y) {
+    if (isnan(x) | isnan(y)) {
+        return out0_type(NAN);
+    } else if (x > 0 & y > 0) {
+        return x * log(x / y);
+    } else if (x == 0 & y >= 0) {
+        return 0;
+    } else {
+        return INF;
+    }
+}
+
+double __device__ huber(double delta, double r) {
+    if (delta < 0) {
+        return INF;
+    } else if (abs(r) <= delta) {
+        return 0.5 * r * r;
+    } else {
+        return delta * (abs(r) - 0.5 * delta);
+    }
+}
+
+double __device__ pseudo_huber(double delta, double r) {
+    double u, v;
+    if (delta < 0) {
+        return INF;
+    } else if (delta == 0 | r == 0) {
+        return 0;
+    } else {
+        u = delta;
+        v = r / delta;
+        return u * u * (sqrt(1 + v * v) - 1);
+    }
+}
+
+'''
+
+
+entr = core.create_ufunc(
+    'cupyx_scipy_entr', ('f->f', 'd->d'),
+    'out0 = out0_type(entr(in0));',
+    preamble=_float_preamble,
+    doc='''Elementwise function for computing entropy.
+
+    .. seealso:: :meth:`scipy.special.entr`
+
+    ''')
+
+
+kl_div = core.create_ufunc(
+    'cupyx_scipy_kl_div', ('ff->f', 'dd->d'),
+    'out0 = out0_type(kl_div(in0, in1));',
+    preamble=_float_preamble,
+    doc='''Elementwise function for computing Kullback-Leibler divergence.
+
+    .. seealso:: :meth:`scipy.special.kl_div`
+
+    ''')
+
+
+rel_entr = core.create_ufunc(
+    'cupyx_scipy_rel_entr', ('ff->f', 'dd->d'),
+    'out0 = out0_type(rel_entr(in0, in1));',
+    preamble=_float_preamble,
+    doc='''Elementwise function for computing relative entropy.
+
+    .. seealso:: :meth:`scipy.special.rel_entr`
+
+    ''')
+
+
+huber = core.create_ufunc(
+    'cupyx_scipy_huber', ('ff->f', 'dd->d'),
+    'out0 = out0_type(huber(in0, in1));',
+    preamble=_float_preamble,
+    doc='''Elementwise function for computing the Huber loss.
+
+    .. seealso:: :meth:`scipy.special.huber`
+
+    ''')
+
+
+pseudo_huber = core.create_ufunc(
+    'cupyx_scipy_pseudo_huber', ('ff->f', 'dd->d'),
+    'out0 = out0_type(pseudo_huber(in0, in1));',
+    preamble=_float_preamble,
+    doc='''Elementwise function for computing the Pseudo-Huber loss.
+
+    .. seealso:: :meth:`scipy.special.pseudo_huber`
+
+    ''')

tests/cupyx_tests/scipy_tests/special_tests/test_convex_analysis.py

@@ -0,0 +1,58 @@
+import itertools
+import unittest
+
+import cupy
+import numpy
+from cupy import testing
+import scipy.special  # NOQA
+
+import cupyx.scipy.special
+
+
+@testing.gpu
+@testing.with_requires('scipy')
+class TestSpecialConvex(unittest.TestCase):
+
+    def test_huber_basic(self):
+        assert cupyx.scipy.special.huber(-1, 1.5) == cupy.inf
+        testing.assert_allclose(cupyx.scipy.special.huber(2, 1.5),
+                                0.5 * 1.5**2)
+        testing.assert_allclose(cupyx.scipy.special.huber(2, 2.5),
+                                2 * (2.5 - 0.5 * 2))
+
+    @testing.for_dtypes(['e', 'f', 'd'])
+    @testing.numpy_cupy_allclose(scipy_name="scp")
+    def test_huber(self, xp, scp, dtype):
+        z = testing.shaped_random((10, 2), xp=xp, dtype=dtype)
+        return scp.special.huber(z[:, 0], z[:, 1])
+
+    @testing.for_dtypes(['e', 'f', 'd'])
+    @testing.numpy_cupy_allclose(scipy_name="scp")
+    def test_entr(self, xp, scp, dtype):
+        values = (0, 0.5, 1.0, cupy.inf)
+        signs = [-1, 1]
+        arr = []
+        for sgn, v in itertools.product(signs, values):
+            arr.append(sgn * v)
+        z = xp.asarray(arr, dtype=dtype)
+        return scp.special.entr(z)
+
+    @testing.for_dtypes(['e', 'f', 'd'])
+    @testing.numpy_cupy_allclose(scipy_name="scp")
+    def test_rel_entr(self, xp, scp, dtype):
+        values = (0, 0.5, 1.0)
+        signs = [-1, 1]
+        arr = []
+        arr = []
+        for sgna, va, sgnb, vb in itertools.product(signs, values, signs,
+                                                    values):
+            arr.append((sgna*va, sgnb*vb))
+        z = xp.asarray(numpy.array(arr, dtype=dtype))
+        return scp.special.kl_div(z[:, 0], z[:, 1])
+
+    @testing.for_dtypes(['e', 'f', 'd'])
+    @testing.numpy_cupy_allclose(scipy_name="scp")
+    def test_pseudo_huber(self, xp, scp, dtype):
+        z = testing.shaped_random((10, 2), xp=numpy, dtype=dtype).tolist()
+        z = xp.asarray(z + [[0, 0.5], [0.5, 0]], dtype=dtype)
+        return scp.special.pseudo_huber(z[:, 0], z[:, 1])