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
