def create_image_array(kernel,
xmin,
xmax,
ymin,
ymax,
max_iter,
base_accuracy,
splits=None,
*args):
if abs(xmax - xmin) > abs(ymax - ymin):
ny = base_accuracy
nx = int((base_accuracy * abs(xmax - xmin) / abs(ymax - ymin)))
else:
nx = base_accuracy
ny = int(base_accuracy * abs(ymax - ymin) / abs(xmax - xmin))
xstride = abs(xmax - xmin) / nx
ystride = abs(ymax - ymin) / ny
topleft = nb.complex128(xmin + 1j * ymax)
image_array = np.zeros((ny, nx, 3), dtype=np.uint8)
if splits:
run_kernel_split(kernel, image_array, topleft, xstride, ystride,
max_iter, splits, *args)
else:
run_kernel(kernel, image_array, topleft, xstride, ystride, max_iter,
*args)
return image_array
def random_image(kernel, xmi, xma, ymi, yma, log_scale_min, log_scale_max,
log_lambda_min, log_lambda_max, log_julia_min, log_julia_max):
scale = 10**rnd.uniform(log_scale_min, log_scale_max)
xc = rnd.uniform(xmi, xma)
yc = rnd.uniform(ymi, yma)
xmin = xc - scale / 2
ymax = yc + scale / 2
nx = 1000
ny = 1000
xstride = scale / nx
ystride = scale / ny
topleft = nb.complex128(xmin + 1j * ymax)
rot = rnd.uniform(0, 360)
max_iter = 1000
lambd = 10**rnd.uniform(log_lambda_min, log_lambda_max)
julia = lambd * 10 * rnd.uniform(log_julia_min, log_julia_max) * cexp(
1j * rnd.uniform(0, 360))
image_array = np.zeros((ny, nx), dtype=np.uint16)
run_kernel(kernel, image_array, topleft, xstride, ystride, max_iter, lambd,
julia, rot)
print(xc, yc, scale, lambd, abs(julia))
return image_array
def mandelbrot(image_array, topleft, xstride, ystride, max_iter):
y, x = cuda.grid(2)
if x < image_array.shape[1] and y < image_array.shape[0]:
c = complex128(topleft + x * xstride - 1j * y * ystride)
z = 0
i = 0
while i < max_iter and z.real * z.real + z.imag * z.imag < 4:
z = z * z + c
i += 1
get_log_color_rgb(image_array, x, y, i, max_iter)
def exp_m(image_array, topleft, xstride, ystride, max_iter):
y, x = cuda.grid(2)
if x < image_array.shape[1] and y < image_array.shape[0]:
c = complex128(topleft + x * xstride - 1j * y * ystride)
z = c
i = 0
while i < max_iter and not isinf(z):
z = exp(z) + c
i += 1
get_log_color_rgb(image_array, x, y, i, max_iter)
def lambert(image_array, topleft, xstride, ystride, max_iter):
y, x = cuda.grid(2)
if x < image_array.shape[1] and y < image_array.shape[0]:
c = complex128(topleft + x * xstride - complex128(1j) * y * ystride)
c = exp(c * exp(-c))
z = c
o = complex128(0.0)
for i in range(max_iter):
z = power(c, z)
if isinf(z):
get_log_color_rgb(image_array, x, y, i, max_iter)
return
if is_close(z, o):
get_log_color_b(image_array, x, y, i, max_iter)
return
if i % 3 == 0:
o = z
def rational2_2(image_array, topleft, xstride, ystride, lambd, julia, rot):
y, x = cuda.grid(2)
if x < image_array.shape[1] and y < image_array.shape[0]:
z = complex128(topleft + x * xstride - 1j * y * ystride) * exp(
1j * rot)
i = 0
while i < 1000 and z.real * z.real + z.imag * z.imag < 4:
z = z * z - lambd / (z * z) + julia
i += 1
k = real_log(float64(i)) / real_log(999.0)
norm = z.real * z.real + z.imag * z.imag
if norm > 4:
l = uint16(real_log(z.real * z.real + z.imag * z.imag - 4) * 200)
# if l>255: l = 255
image_array[y, x] = uint16(255 * k) * 256 - l
def test_types_to_ftylist():
import numba
types_to_ftylist = quaternionic.utilities.convert_numpy_ufunc_type_to_numba_ftylist
types = '?bhilqpBHILQPfdgF->D'
ftylist = numba.complex128(
numba.boolean,
numba.byte,
numba.short,
numba.intc,
numba.int_,
numba.longlong,
numba.intp,
numba.char,
numba.ushort,
numba.uintc,
numba.uint,
numba.ulonglong,
numba.uintp,
numba.float32,
numba.float_,
numba.double,
numba.complex64,
)
assert types_to_ftylist([types]) == [ftylist]
import numba as nb
import numpy as np
from numba import jit, prange
@nb.vectorize([nb.complex64(nb.float32), nb.complex128(nb.float64)])
def complex_exponential(x):
return np.cos(x) + 1.j * np.sin(x)
@nb.vectorize([nb.float32(nb.complex64), nb.float64(nb.complex128)])
def abs2(x):
return x.real**2 + x.imag**2
@jit(nopython=True, nogil=True, parallel=True)
def interpolate_radial_functions(array, disc_indices, positions, v, r, dvdr,
sampling):
n = r.shape[0]
dt = np.log(r[-1] / r[0]) / (n - 1)
for i in range(positions.shape[0]):
for j in prange(disc_indices.shape[0]):
k = int(round(positions[i, 0] / sampling[0]) + disc_indices[j, 0])
l = int(round(positions[i, 1] / sampling[1]) + disc_indices[j, 1])
if ((k < array.shape[0]) & (l < array.shape[1]) & (k >= 0) &
(l >= 0)):
r_interp = np.sqrt((k * sampling[0] - positions[i, 0])**2 +
(l * sampling[1] - positions[i, 1])**2)
idx = int(np.floor(np.log(r_interp / r[0] + 1e-7) / dt))
err += abs(img[cx + x, cy + y] - img[cx - x, cy - y])
out[xs + d, ys + d] = err / cn
return out.argmin()
@_numba.vectorize(
[_numba.float64(_numba.complex128),
_numba.float32(_numba.complex64)],
target="parallel",
)
def abs2(x):
return x.real * x.real + x.imag * x.imag
@_numba.vectorize(
[_numba.complex128(_numba.complex128),
_numba.complex64(_numba.complex64)],
target="parallel",
)
def abs2c(x):
return x.real * x.real + x.imag * x.imag + 0j
def fill(data, invalid=None):
"""
fill invalid values by closest valid value. invalid: mask of invalid values, default: np.isnan(data)
"""
if invalid is None:
invalid = _np.isnan(data)
ind = _snd.distance_transform_edt(invalid,
return_distances=False,
mat_dim_list,
num_reps,
dtype=np.csingle)
#%%
'''
NUMBA TESTS
'''
#%% Function declarations
import cmath
from numba import vectorize, complex128, complex64
@vectorize(
[complex64(complex64, complex64),
complex128(complex128, complex128)],
target='parallel')
def nb_add(x, y):
return x + y
@vectorize(
[complex64(complex64, complex64),
complex128(complex128, complex128)],
target='parallel')
def nb_subtract(x, y):
return x + y
@vectorize(
[complex64(complex64, complex64),
j + 1 + lmax - m)] * f[m:(lmax + 1), 2, 1] + a[2, j:(
j + 1 + lmax - m)] * f[m:(lmax + 1), 2, 2]
j += 1 + lmax - m
return z
@vectorize([float64(float64, float64)], nopython=True, target='parallel')
def numba_array_addition(x, y):
"""
Use numba to do array addition (x + y) where both x and y have same shape
"""
return x + y
@vectorize([complex128(complex128, complex128)],
nopython=True,
target='parallel')
def numba_array_addition_harmonic(x, y):
"""
Use numba to do array addition (x + y) where both x and y have same shape
"""
return x + y
@vectorize([float64(float64, float64)], nopython=True, target='parallel')
def numba_array_subtraction(x, y):
"""
Use numba to do array subtraction (x - y) where both x and y have same shape
"""
return x - y
def __init__(self, re: float32, im: float32):
self.element = complex128(re, im)
def power(z, x):
return abs(z) ** x * exp(phase(z) * x * complex128(1j))
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Aug 6 13:16:55 2018
@author: dietz
"""
from cmath import exp
from math import sin, cos
import numba
@numba.njit(numba.complex128(numba.int64, numba.int64, numba.float64,
numba.float64),
nogil=True)
def sph_harm_hard(l, m, theta, phi):
"""hard coded spherical harmonics extracted from sympy"""
if l == 0 and m == 0:
return 0.282094791773878
elif l == 1 and m == -1:
return 0.345494149471335 * exp(-1j * phi) * sin(theta)
elif l == 1 and m == 0:
return 0.48860251190292 * cos(theta)
elif l == 1 and m == 1:
return -0.345494149471335 * exp(1j * phi) * sin(theta)
elif l == 2 and m == -2:
return -0.38627420202319 * exp(-2 * 1j * phi) * cos(
theta)**2 + 0.38627420202319 * exp(-2 * 1j * phi)
elif l == 2 and m == -1:
return 0.772548404046379 * exp(-1j * phi) * sin(theta) * cos(theta)
"""This module contains functions to compute the layer mediated particle
coupling coefficients."""
import numpy as np
import scipy.special
from numba import complex128, int64, jit
import smuthi.fields
import smuthi.fields.transformations as trf
import smuthi.layers as lay
import smuthi.utility.math as sf
@jit(complex128(complex128[:], complex128[:]),
nopython=True,
cache=True,
nogil=True)
def numba_trapz(y, x):
out = 0.0 + 0.0j
#TODO implement some (optional) advanced summation?
#e.g. https://github.com/nschloe/accupy/blob/master/accupy/sums.py
#or better Sum2 from https://doi.org/10.1137/030601818 (Algorithm 4.4)
#Note, that this may need to have exact summation for x and y, and exact product.
for i in range(len(y) - 2):
out += (x[i + 1] - x[i]) * (y[i + 1] + y[i]) / 2.0
return out
@jit((complex128[:], complex128[:, :, :], complex128[:, :, :, :],
complex128[:, :, :], int64),
nopython=True,
cache=True,
"""Variants of standard maths functions written for speed."""
import numpy as np
import numba
@numba.vectorize(
[numba.complex64(numba.float32),
numba.complex128(numba.float64)])
def expj2pi(x):
"""Equivalent to ``expj(2 * np.pi * x)`` where `x` is real.
x is reduced to a small value before multiplication, which
improves precision at a small cost in performance.
"""
y = 2 * np.pi * (x - np.rint(x))
return complex(np.cos(y), np.sin(y))
def nansum(x, *args, **kwargs):
"""Like np.nansum, provided `x` is a floating-point type."""
return np.sum(x, *args, where=~np.isnan(x), **kwargs)
Parameters:
val (uint32): Input unsigned int.
Returns:
uint32: Output unsigned int.
"""
val = (val & 0x55555555) + ((val >> 1) & 0x55555555)
val = (val & 0x33333333) + ((val >> 2) & 0x33333333)
val = (val & 0x0f0f0f0f) + ((val >> 4) & 0x0f0f0f0f)
val = (val & 0x00ff00ff) + ((val >> 8) & 0x00ff00ff)
val = (val & 0x0000ffff) + ((val >> 16) & 0x0000ffff)
return val
@jit(complex128(complex128[:], int32[:], int32[:], complex128[:]),
nopython=True,
cache=True)
def expect_psi_csr(data, ind, ptr, vec):
"""Computes the expectation value of a complex matrix in CSR sparse format.
Note that this routine returns a complex type regardless of whether
the matrix is Hermitian or not. If so, take the real part.
Parameters:
data (ndarray): A complex128 array of data.
ind (ndarray): A int32 array of indices.
ptr (ndarray): A int32 array of indptrs.
vec (ndarray): A complex128 array for the statevector.
Returns:
half_window = (window_size - 1) // 2
# Precompute coefficients
b = np.mat([[k**i for i in order_range]
for k in range(-half_window, half_window + 1)])
m = np.linalg.pinv(b).A[deriv] * rate**deriv * factorial(deriv)
# Pad the signal at the extremes with values taken from the signal itself
firstvals = y[0] - np.abs(y[1:half_window + 1][::-1] - y[0])
lastvals = y[-1] + np.abs(y[-half_window - 1:-1][::-1] - y[-1])
y = np.concatenate((firstvals, y, lastvals))
return np.convolve(m[::-1], y, mode='valid')
@jit(complex128(int32, complex128, int32, int32))
def xcov(wid, spectra_full, overlap, average):
"""
Calculation of the array covariance matrix from the array data vectors
stored in the spectra matrix (should be n_traces x n_windows x n_freq),
over one set of averaged windows.
"""
n_traces, n_windows, n_frequencies = spectra_full.shape
beg = overlap * wid
end = beg + average
spectra = copy.deepcopy(spectra_full[:, beg:end, :])
X = spectra[:, None, 0, :] * np.conj(spectra[:, 0, :])
for swid in range(1, average - 1):
X += np.conj(spectra[:, None, swid, :]) * spectra[:, swid, :]
return X
from numba import jit, complex128
import numpy as np
@jit(complex128(complex128, complex128), nopython=True)
def F(z, c):
return z * z + c
