def csd(*args): args_array = np.array(args) args_k_re = args_array + step_re args_k_im = args_array + step_im return 2 * (func_re(*to_fixed_tuple(args_k_re, arg_len)) - func_im(*to_fixed_tuple(args_k_im, arg_len)) ).real / im_step_size**2
def fin_diff(*args): diff = 0. arg_array = np.array(args) for coeff, step in zip(coeff_prod, step_prod): args_k = arg_array + np.array(step) diff += coeff * func(*to_fixed_tuple(args_k, arg_len)) return diff / step_size ** total_deriv_order
def kernel(lhs, rhs): location = cuda.grid(1) if location >= n_elements: # bake n_elements into the kernel, better than passing it in # as another argument. return # [0, :] is the to-index (into `lhs`) # [1, :] is the from-index (into `rhs`) idx = cuda.local.array(shape=(2, ndim), dtype=types.int64) for i in range(ndim - 1, -1, -1): idx[0, i] = location % lhs.shape[i] idx[1, i] = (location % lhs.shape[i]) * (rhs.shape[i] > 1) location //= lhs.shape[i] lhs[to_fixed_tuple(idx[0], ndim)] = rhs[to_fixed_tuple(idx[1], ndim)]
def unravel_index(index, shape): sizes = np.zeros(len(shape), dtype=np.int64) result = np.zeros(len(shape), dtype=np.int64) sizes[-1] = 1 for i in range(len(shape) - 2, -1, -1): sizes[i] = sizes[i + 1] * shape[i + 1] remainder = index for i in range(len(shape)): result[i] = remainder // sizes[i] remainder %= sizes[i] return to_fixed_tuple(result, len(shape))
def kernel(lhs, rhs): location = cuda.grid(1) n_elements = 1 for i in range(lhs.ndim): n_elements *= lhs.shape[i] if location >= n_elements: # bake n_elements into the kernel, better than passing it in # as another argument. return # [0, :] is the to-index (into `lhs`) # [1, :] is the from-index (into `rhs`) idx = cuda.local.array( shape=(2, ndim), dtype=types.int64) for i in range(ndim - 1, -1, -1): idx[0, i] = location % lhs.shape[i] idx[1, i] = (location % lhs.shape[i]) * (rhs.shape[i] > 1) location //= lhs.shape[i] lhs[to_fixed_tuple(idx[0], ndim)] = rhs[to_fixed_tuple(idx[1], ndim)]
def tuple_with_length(array, length): return to_fixed_tuple(array, length)
def foo(array): a = to_fixed_tuple(array, length=1) b = to_fixed_tuple(array, 2) c = to_fixed_tuple(array, const) d = to_fixed_tuple(array, 0) return a, b, c, d
def tuple_with_length(array, length): return to_fixed_tuple(array, length)
def foo(array): a = to_fixed_tuple(array, length=1) b = to_fixed_tuple(array, 2) c = to_fixed_tuple(array, const) d = to_fixed_tuple(array, 0) return a, b, c, d
def csd(*args): args_k = np.array(args) + step * 1j * arg_select return func_im( *to_fixed_tuple(args_k, arg_len)).imag / step