test_complex.py

import math
import itertools

import numpy as np

from numba.cuda.testing import unittest, CUDATestCase
from numba.core import types
from numba import cuda
from numba.tests.complex_usecases import (real_usecase, imag_usecase,
                                          conjugate_usecase, phase_usecase,
                                          polar_as_complex_usecase,
                                          rect_usecase, isnan_usecase,
                                          isinf_usecase, isfinite_usecase,
                                          exp_usecase, log_usecase,
                                          log_base_usecase, log10_usecase,
                                          sqrt_usecase, asin_usecase,
                                          acos_usecase, atan_usecase,
                                          cos_usecase, sin_usecase,
                                          tan_usecase, acosh_usecase,
                                          asinh_usecase, atanh_usecase,
                                          cosh_usecase, sinh_usecase,
                                          tanh_usecase)
from numba.np import numpy_support


def compile_scalar_func(pyfunc, argtypes, restype):
    # First compile a scalar device function
    assert not any(isinstance(tp, types.Array) for tp in argtypes)
    assert not isinstance(restype, types.Array)
    device_func = cuda.jit(restype(*argtypes), device=True)(pyfunc)

    kernel_types = [types.Array(tp, 1, "C")
                    for tp in [restype] + list(argtypes)]

    if len(argtypes) == 1:
        def kernel_func(out, a):
            i = cuda.grid(1)
            if i < out.shape[0]:
                out[i] = device_func(a[i])
    elif len(argtypes) == 2:
        def kernel_func(out, a, b):
            i = cuda.grid(1)
            if i < out.shape[0]:
                out[i] = device_func(a[i], b[i])
    else:
        assert 0

    kernel = cuda.jit(tuple(kernel_types))(kernel_func)

    def kernel_wrapper(values):
        n = len(values)
        inputs = [np.empty(n, dtype=numpy_support.as_dtype(tp))
                  for tp in argtypes]
        output = np.empty(n, dtype=numpy_support.as_dtype(restype))
        for i, vs in enumerate(values):
            for v, inp in zip(vs, inputs):
                inp[i] = v
        args = [output] + inputs
        kernel[int(math.ceil(n / 256)), 256](*args)
        return list(output)
    return kernel_wrapper


class BaseComplexTest(CUDATestCase):

    def basic_values(self):
        reals = [-0.0, +0.0, 1, -1, +1.5, -3.5,
                 float('-inf'), float('+inf'), float('nan')]
        return [complex(x, y) for x, y in itertools.product(reals, reals)]

    def more_values(self):
        reals = [0.0, +0.0, 1, -1, -math.pi, +math.pi,
                 float('-inf'), float('+inf'), float('nan')]
        return [complex(x, y) for x, y in itertools.product(reals, reals)]

    def non_nan_values(self):
        reals = [-0.0, +0.0, 1, -1, -math.pi, +math.pi,
                 float('inf'), float('-inf')]
        return [complex(x, y) for x, y in itertools.product(reals, reals)]

    def run_func(self, pyfunc, sigs, values, ulps=1, ignore_sign_on_zero=False):
        for sig in sigs:
            if isinstance(sig, types.Type):
                sig = sig,
            if isinstance(sig, tuple):
                # Assume return type is the type of first argument
                sig = sig[0](*sig)
            prec = ('single'
                    if sig.args[0] in (types.float32, types.complex64)
                    else 'double')
            cudafunc = compile_scalar_func(pyfunc, sig.args, sig.return_type)
            ok_values = []
            expected_list = []
            for args in values:
                if not isinstance(args, (list, tuple)):
                    args = args,
                try:
                    expected_list.append(pyfunc(*args))
                    ok_values.append(args)
                except ValueError as e:
                    self.assertIn("math domain error", str(e))
                    continue
            got_list = cudafunc(ok_values)
            for got, expected, args in zip(got_list, expected_list, ok_values):
                msg = 'for input %r with prec %r' % (args, prec)
                self.assertPreciseEqual(got, expected, prec=prec,
                                        ulps=ulps,
                                        ignore_sign_on_zero=ignore_sign_on_zero,
                                        msg=msg)

    run_unary = run_func
    run_binary = run_func


class TestComplex(BaseComplexTest):

    def check_real_image(self, pyfunc):
        values = self.basic_values()
        self.run_unary(pyfunc,
                       [tp.underlying_float(tp)
                        for tp in (types.complex64, types.complex128)],
                       values)

    def test_real(self):
        self.check_real_image(real_usecase)

    def test_imag(self):
        self.check_real_image(imag_usecase)

    def test_conjugate(self):
        pyfunc = conjugate_usecase
        values = self.basic_values()
        self.run_unary(pyfunc,
                       [types.complex64, types.complex128],
                       values)


class TestCMath(BaseComplexTest):
    """
    Tests for cmath module support.
    """

    def check_predicate_func(self, pyfunc):
        self.run_unary(pyfunc,
                       [types.boolean(tp)
                        for tp in (types.complex128, types.complex64)],
                       self.basic_values())

    def check_unary_func(self, pyfunc, ulps=1, values=None,
                         returns_float=False, ignore_sign_on_zero=False):
        if returns_float:
            def sig(tp):
                return tp.underlying_float(tp)
        else:
            def sig(tp):
                return tp(tp)
        self.run_unary(pyfunc, [sig(types.complex128)],
                       values or self.more_values(), ulps=ulps,
                       ignore_sign_on_zero=ignore_sign_on_zero)
        # Avoid discontinuities around pi when in single precision.
        self.run_unary(pyfunc, [sig(types.complex64)],
                       values or self.basic_values(), ulps=ulps,
                       ignore_sign_on_zero=ignore_sign_on_zero)

    # Conversions

    def test_phase(self):
        self.check_unary_func(phase_usecase, returns_float=True)

    def test_polar(self):
        self.check_unary_func(polar_as_complex_usecase)

    def test_rect(self):
        def do_test(tp, seed_values):
            values = [(z.real, z.imag) for z in seed_values
                      if not math.isinf(z.imag) or z.real == 0]
            float_type = tp.underlying_float
            self.run_binary(rect_usecase, [tp(float_type, float_type)],
                            values)
        do_test(types.complex128, self.more_values())
        # Avoid discontinuities around pi when in single precision.
        do_test(types.complex64, self.basic_values())

    # Classification

    def test_isnan(self):
        self.check_predicate_func(isnan_usecase)

    def test_isinf(self):
        self.check_predicate_func(isinf_usecase)

    def test_isfinite(self):
        self.check_predicate_func(isfinite_usecase)

    # Power and logarithms

    def test_exp(self):
        self.check_unary_func(exp_usecase, ulps=2)

    def test_log(self):
        self.check_unary_func(log_usecase)

    def test_log_base(self):
        values = list(itertools.product(self.more_values(), self.more_values()))
        value_types = [(types.complex128, types.complex128),
                       (types.complex64, types.complex64)]
        self.run_binary(log_base_usecase, value_types, values,
                        ulps=3)

    def test_log10(self):
        self.check_unary_func(log10_usecase)

    def test_sqrt(self):
        self.check_unary_func(sqrt_usecase)

    # Trigonometric functions

    def test_acos(self):
        self.check_unary_func(acos_usecase, ulps=2)

    def test_asin(self):
        self.check_unary_func(asin_usecase, ulps=2)

    def test_atan(self):
        self.check_unary_func(atan_usecase, ulps=2,
                              values=self.non_nan_values())

    def test_cos(self):
        self.check_unary_func(cos_usecase, ulps=2)

    def test_sin(self):
        # See test_sinh.
        self.check_unary_func(sin_usecase, ulps=2)

    def test_tan(self):
        self.check_unary_func(tan_usecase, ulps=2,
                              ignore_sign_on_zero=True)

    # Hyperbolic functions

    def test_acosh(self):
        self.check_unary_func(acosh_usecase)

    def test_asinh(self):
        self.check_unary_func(asinh_usecase, ulps=2)

    def test_atanh(self):
        self.check_unary_func(atanh_usecase, ulps=2,
                              ignore_sign_on_zero=True)

    def test_cosh(self):
        self.check_unary_func(cosh_usecase, ulps=2)

    def test_sinh(self):
        self.check_unary_func(sinh_usecase, ulps=2)

    def test_tanh(self):
        self.check_unary_func(tanh_usecase, ulps=2,
                              ignore_sign_on_zero=True)


class TestAtomicOnComplexComponents(CUDATestCase):
    # Based on the reproducer from Issue #8309. array.real and array.imag could
    # not be used because they required returning an array from a generated
    # function, and even if this was permitted, they could not be resolved from
    # the atomic lowering when they were overloads.
    #
    # See https://github.com/numba/numba/issues/8309

    def test_atomic_on_real(self):
        @cuda.jit
        def atomic_add_one(values):
            i = cuda.grid(1)
            cuda.atomic.add(values.real, i, 1)

        N = 32
        arr1 = np.arange(N) + np.arange(N) * 1j
        arr2 = arr1.copy()
        atomic_add_one[1, N](arr2)
        np.testing.assert_equal(arr1 + 1, arr2)

    def test_atomic_on_imag(self):
        @cuda.jit
        def atomic_add_one_j(values):
            i = cuda.grid(1)
            cuda.atomic.add(values.imag, i, 1)

        N = 32
        arr1 = np.arange(N) + np.arange(N) * 1j
        arr2 = arr1.copy()
        atomic_add_one_j[1, N](arr2)
        np.testing.assert_equal(arr1 + 1j, arr2)


if __name__ == '__main__':
    unittest.main()