Python numba.vectorize() Examples

def performance_example():
    # Simple function, already supports vectorization
    f_vec = sampling_function(
        lambda x: x ** 2, domain=odl.IntervalProd(0, 1)
    )

    # Vectorized with NumPy's poor man's vectorization function
    f_novec = np.vectorize(lambda x: x ** 2)

    # We test both versions with 10000 evaluation points. The natively
    # vectorized version should be much faster than the one using
    # numpy.vectorize.
    points = np.linspace(0, 1, 10000)

    print('Vectorized runtime:     {:5f}'
          ''.format(timeit.timeit(lambda: f_vec(points), number=100)))
    print('Non-vectorized runtime: {:5f}'
          ''.format(timeit.timeit(lambda: f_novec(points), number=100))) 

def decorator_apply(func, batcher=None, cache=None, vectorize=None):
	def wrapper_func(*args, **kwargs):
		return Apply(func, args=args[1:], kwargs= kwargs, batcher= batcher, cache= cache,
					vectorize= vectorize).transform(args[0])
	return wrapper_func 

def batch_transform(args):
	f= args[1]
	f_args= args[2]
	f_kwargs= args[3]
	if args[5] is not None:
		from numba import vectorize
		return vectorize(args[5], fastmath=True)(f)(*zip(*args[0]))
	if args[4] is not None:
		from functools import lru_cache
		f= lru_cache(maxsize=args[4])(f)
	#Applying per DataFrame row is very slow, use ApplyBatch instead
	if isinstance(args[0], pd.DataFrame):  return args[0].apply(lambda x: f(x, *f_args, **f_kwargs), axis=1)
	return [f(row, *f_args, **f_kwargs) for row in args[0]] 

def __init__(self, function, batcher=None, args=[], kwargs={}, cache=None, vectorize=None):
		if batcher is None:   self.batcher= wordbatch.batcher.Batcher()
		else:  self.batcher= batcher
		self.function= function
		self.args= [args]
		self.kwargs= [kwargs]
		self.cache = [cache]
		self.vectorize = [vectorize] 

def transform(self, data, input_split=False, merge_output=True, minibatch_size=None, batcher=None):
		if batcher is None:  batcher = self.batcher
		return batcher.process_batches(batch_transform, data,
		                               [self.function] + self.args + self.kwargs + self.cache + self.vectorize,
		                               input_split=input_split, merge_output=merge_output,
		                               minibatch_size= minibatch_size)

# import wordbatch.batcher as batcher
# b= batcher.Batcher(minibatch_size=2)#, method="serial")
# import numpy as np
# a= Apply(np.power, b, [2],{})
# print(a.transform([1, 2, 3, 4])) 

def transform_lists_to_arrays(module, to_change, __funcs, vec=False, cache_blacklist=set([])):
    if vec:
        conv_fun = numba.vectorize
        extra_args = extra_args_vec
    else:
        conv_fun = numba.njit
        extra_args = extra_args_std
    for s in to_change:
        func = s.split('.')[-1]
        mod = '.'.join(s.split('.')[:-1])
        fake_mod = __funcs[mod]

        try:
            real_mod = getattr(module, mod)
        except:
            real_mod = module
            for s in mod.split('.'):
                real_mod = getattr(real_mod, s)

        orig_func = getattr(real_mod, func)
        source = inspect.getsource(orig_func)
        source = remove_for_numba(source) # do before anything else
        source = return_value_numpy(source)
        source = re.sub(list_mult_expr, numpy_not_list_expr, source)
#        if 'longitude_obliquity_nutation' in s:
#        print(source)
        numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
        new_func = fake_mod.__dict__[func]
        do_cache = caching and func not in cache_blacklist
        obj = conv_fun(cache=do_cache, **extra_args)(new_func)
        __funcs[func] = obj
        fake_mod.__dict__[func] = obj
        obj.__doc__ = ''


#set_signatures = {'Clamond': [numba.float64(numba.float64, numba.float64, numba.boolean),
#                              numba.float64(numba.float64, numba.float64, numba.optional(numba.boolean))
#                              ]
#                    } 

def transform_complete_ht(replaced, __funcs, __all__, normal, vec=False):
    replaced, NUMERICS_SUBMOD = fluids.numba.create_numerics(replaced, vec=vec)
    
    __funcs.update(fluids.numba.__dict__.copy())
    fluids.numba.transform_module(normal, __funcs, replaced, vec=vec)

    if vec:
        conv_fun = numba.vectorize
    else:
        conv_fun = numba.jit
    
    to_change_AvailableMethods = []
    to_change_full_output = []
    
    to_change = {k: 'AvailableMethods' for k in to_change_AvailableMethods}
    to_change.update({k: 'full_output' for k in to_change_full_output})
    to_change['hx.Ntubes_Phadkeb'] = 'square_C1s is None'
    to_change['boiling_nucleic.Gorenflo'] = 'h0 is None: # NUMBA: DELETE'

    for s, bad_branch in to_change.items():
        mod, func = s.split('.')
        source = inspect.getsource(getattr(getattr(normal, mod), func))
        fake_mod = __funcs[mod]
        source = fluids.numba.remove_branch(source, bad_branch)
        fluids.numba.numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
        new_func = fake_mod.__dict__[func]
        obj = conv_fun(cache=caching)(new_func)
        __funcs[func] = obj
        globals()[func] = obj
        obj.__doc__ = ''
    
    to_change = ['air_cooler.Ft_aircooler']
    fluids.numba.transform_lists_to_arrays(normal, to_change, __funcs)

        
    __funcs['hx']._load_coeffs_Phadkeb() 

def ray_trace(x, y, poly):
    """
    Determines for some set of x and y coordinates, which of those coordinates is within `poly`. Ray trace is \
    generally called as an internal function, see :func:`.poly_clip`

    :param x: A 1D numpy array of x coordinates.
    :param y: A 1D numpy array of y coordinates.
    :param poly: The coordinates of a polygon as a numpy array (i.e. from geo_json['coordinates']
    :return: A 1D boolean numpy array, true values are those points that are within `poly`.
    """

    @vectorize([bool_(float64, float64)])
    def ray(x, y):
        # where xy is a coordinate
        n = len(poly)
        inside = False
        p2x = 0.0
        p2y = 0.0
        xints = 0.0
        p1x, p1y = poly[0]
        for i in range(n + 1):
            p2x, p2y = poly[i % n]
            if y > min(p1y, p2y):
                if y <= max(p1y, p2y):
                    if x <= max(p1x, p2x):
                        if p1y != p2y:
                            xints = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
                        if p1x == p2x or x <= xints:
                            inside = not inside
            p1x, p1y = p2x, p2y
        return inside

    return ray(x, y)