Exemplo n.º 1
0
 def wrapper(func):
     return numba.generated_jit(func,
                                nopython=nopython,
                                cache=cache,
                                parallel=parallel,
                                fastmath=fastmath,
                                error_model=error_model)
Exemplo n.º 2
0
 def test_generated_dtype(self):
     f = generated_jit(nopython=True)(dtype_generated_usecase)
     a = np.ones((10,), dtype=np.float32)
     b = np.ones((10,), dtype=np.float64)
     self.assertEqual(f(a, b).dtype, np.float64)
     self.assertEqual(f(a, b, dtype=np.dtype('int32')).dtype, np.int32)
     self.assertEqual(f(a, b, dtype=np.int32).dtype, np.int32)
Exemplo n.º 3
0
 def test_generated(self):
     f = generated_jit(nopython=True)(generated_usecase)
     self.assertEqual(f(8), 8 - 5)
     self.assertEqual(f(x=8), 8 - 5)
     self.assertEqual(f(x=8, y=4), 8 - 4)
     self.assertEqual(f(1j), 5 + 1j)
     self.assertEqual(f(1j, 42), 42 + 1j)
     self.assertEqual(f(x=1j, y=7), 7 + 1j)
Exemplo n.º 4
0
 def test_generated(self):
     f = generated_jit(nopython=True)(generated_usecase)
     self.assertEqual(f(8), 8 - 5)
     self.assertEqual(f(x=8), 8 - 5)
     self.assertEqual(f(x=8, y=4), 8 - 4)
     self.assertEqual(f(1j), 5 + 1j)
     self.assertEqual(f(1j, 42), 42 + 1j)
     self.assertEqual(f(x=1j, y=7), 7 + 1j)
Exemplo n.º 5
0
 def test_signature_errors(self):
     """
     Check error reporting when implementation signature doesn't match
     generating function signature.
     """
     f = generated_jit(nopython=True)(bad_generated_usecase)
     # Mismatching # of arguments
     with self.assertRaises(TypeError) as raises:
         f(1j)
     self.assertIn("should be compatible with signature '(x, y=5)', but has signature '(x)'",
                   str(raises.exception))
     # Mismatching defaults
     with self.assertRaises(TypeError) as raises:
         f(1)
     self.assertIn("should be compatible with signature '(x, y=5)', but has signature '(x, y=6)'",
                   str(raises.exception))
Exemplo n.º 6
0
"""
   File Name:jit_mode
   Description :  numba jit mode 定义(常用模式)
   pytorch 中也自带 jit 编译
   Email : [email protected]
   Date:18-2-1
"""

from numba import float32, float64, generated_jit, int16, int32, int64, jit, vectorize, void

# jit 模式的定义
jit_cpu = jit(nopython=True, parallel=True)
jit_gpu = jit(nopython=True, parallel=True, target='cuda')

# generated_jit 模式的定义
gjit_cpu = generated_jit(nopython=True, parallel=True)
gjit_gpu = generated_jit(nopython=True, parallel=True, target='cuda')

# ufunc example
vec_cpu = vectorize([
    void(int64, int64),
    int16(int16, int16),
    int32(int32, int32),
    int64(int64, int64),
    float32(float32, float32),
    float64(float64, float64)
])

vec_gpu = vectorize([
    void(int64, int64),
    int16(int16, int16),
Exemplo n.º 7
0
        for s in range(radec.shape[0]):
            l, m = lm[s]
            n = np.sqrt(1.0 - l**2 - m**2)

            radec[s, 1] = np.arcsin(m*cos_d0 + n*sin_d0)
            radec[s, 0] = pc_ra + np.arctan(l / (n*cos_d0 - m*sin_d0))

        return radec

    return _lm_to_radec_impl


# inspect.getargspec doesn't work on a numba dispatcher object
# so rtd fails.
if not on_rtd():
    jitter = numba.generated_jit(nopython=True, nogil=True, cache=True)
    lmn_to_radec = jitter(lmn_to_radec)
    lm_to_radec = jitter(lm_to_radec)
    radec_to_lmn = jitter(radec_to_lmn)
    radec_to_lm = jitter(radec_to_lm)


RADEC_TO_LMN_DOCS = DocstringTemplate(r"""
Converts Right-Ascension/Declination coordinates in radians
to a Direction Cosine lm coordinates, relative to the Phase Centre.

.. math::
    :nowrap:

    \begin{eqnarray}
        & l =& \, \cos \, \delta  \sin \, \Delta \alpha  \\
Exemplo n.º 8
0
    def execute_user_code(self, user_args, user_kwargs, tracing):
        # type: (list, dict, bool) -> (object, COMPSsException)
        """ Executes the user code.

        Disables the tracing hook if tracing is enabled. Restores it
        at the end of the user code execution.

        :param user_args: Function args.
        :param user_kwargs: Function kwargs.
        :param tracing: If tracing enabled.
        :return: The user function returns and the compss exception (if any).
        """
        # Tracing hook is disabled by default during the user code of the task.
        # The user can enable it with tracing_hook=True in @task decorator for
        # specific tasks or globally with the COMPSS_TRACING_HOOK=true
        # environment variable.
        restore_hook = False
        pro_f = None
        if tracing:
            global_tracing_hook = False
            if TRACING_HOOK_ENV_VAR in os.environ:
                hook_enabled = os.environ[TRACING_HOOK_ENV_VAR] == "true"
                global_tracing_hook = hook_enabled
            if self.decorator_arguments['tracing_hook'] or global_tracing_hook:
                # The user wants to keep the tracing hook
                pass
            else:
                # When Extrae library implements the function to disable,
                # use it, as:
                #     import pyextrae
                #     pro_f = pyextrae.shutdown()
                # Since it is not available yet, we manage the tracing hook
                # by ourselves
                pro_f = sys.getprofile()
                sys.setprofile(None)
                restore_hook = True

        user_returns = None
        compss_exception = None
        if self.decorator_arguments['numba']:
            # Import all supported functionalities
            from numba import jit
            from numba import njit
            from numba import generated_jit
            from numba import vectorize
            from numba import guvectorize
            from numba import stencil
            from numba import cfunc
            numba_mode = self.decorator_arguments['numba']
            numba_flags = self.decorator_arguments['numba_flags']
            if type(numba_mode) is dict:
                # Use the flags defined by the user
                numba_flags['cache'] = True  # Always force cache
                user_returns = jit(self.user_function,
                                   **numba_flags)(*user_args, **user_kwargs)
            elif numba_mode is True or numba_mode == 'jit':
                numba_flags['cache'] = True  # Always force cache
                user_returns = jit(self.user_function,
                                   **numba_flags)(*user_args, **user_kwargs)
                # Alternative way of calling:
                # user_returns = jit(cache=True)(self.user_function) \
                #                   (*user_args, **user_kwargs)
            elif numba_mode == 'generated_jit':
                user_returns = generated_jit(self.user_function,
                                             **numba_flags)(*user_args,
                                                            **user_kwargs)
            elif numba_mode == 'njit':
                numba_flags['cache'] = True  # Always force cache
                user_returns = njit(self.user_function,
                                    **numba_flags)(*user_args, **user_kwargs)
            elif numba_mode == 'vectorize':
                numba_signature = self.decorator_arguments['numba_signature']  # noqa: E501
                user_returns = vectorize(
                    numba_signature,
                    **numba_flags
                )(self.user_function)(*user_args, **user_kwargs)
            elif numba_mode == 'guvectorize':
                numba_signature = self.decorator_arguments['numba_signature']  # noqa: E501
                numba_decl = self.decorator_arguments['numba_declaration']
                user_returns = guvectorize(
                    numba_signature,
                    numba_decl,
                    **numba_flags
                )(self.user_function)(*user_args, **user_kwargs)
            elif numba_mode == 'stencil':
                user_returns = stencil(
                    **numba_flags
                )(self.user_function)(*user_args, **user_kwargs)
            elif numba_mode == 'cfunc':
                numba_signature = self.decorator_arguments['numba_signature']  # noqa: E501
                user_returns = cfunc(
                    numba_signature
                )(self.user_function).ctypes(*user_args, **user_kwargs)
            else:
                raise Exception("Unsupported numba mode.")
        else:
            try:
                # Normal task execution
                user_returns = self.user_function(*user_args, **user_kwargs)
            except COMPSsException as ce:
                compss_exception = ce
                # Check old targetDirection
                if 'targetDirection' in self.decorator_arguments:
                    target_label = 'targetDirection'
                else:
                    target_label = 'target_direction'
                compss_exception.target_direction = self.decorator_arguments[target_label]  # noqa: E501

        # Reestablish the hook if it was disabled
        if restore_hook:
            sys.setprofile(pro_f)

        return user_returns, compss_exception
Exemplo n.º 9
0
        # Add base visibilities to the output, if any
        add_coh_fn(base_vis, out)

        # Apply direction independent effects, if any
        apply_dies_fn(time_index, antenna1, antenna2, die1_jones, die2_jones,
                      tmin, out)

        return out

    return _predict_vis_fn


# inspect.getargspec doesn't work on a numba dispatcher object
# so rtd fails.
if not on_rtd():
    predict_vis = generated_jit(nopython=True, nogil=True,
                                cache=True)(predict_vis)

PREDICT_DOCS = DocstringTemplate(r"""
Multiply Jones terms together to form model visibilities according
to the following formula:

.. math::


    V_{pq} = G_{p} \left(
        B_{pq} + \sum_{s} A_{ps} X_{pqs} A_{qs}^H
        \right) G_{q}^H

where for antenna :math:`p` and :math:`q`, and source :math:`s`: