I want to convert this piece of code in order to make it compatible with Numba. The only sort method that Numba support is sorted() but not with the key arg. I have to manualy sort without other lib imports or maybe just some numpy. Someone could give me an efficient way to do this sort ? Thanks
JavaScript
x
9
1
import random
2
3
n = 1000
4
index = list(range(n))
5
keys = list(range(n))
6
random.shuffle(keys)
7
8
index.sort(key=lambda x: keys[x])) <= HOW TO CONVERT THIS ?
9
Edit :
JavaScript
1
38
38
1
import numpy as np
2
from numba import jit
3
4
@jit(nopython=True)
5
def fourier_fit_extra(data, harmonic, extra=0):
6
size = len(data)
7
x = np.arange(0, size, 1)
8
m = np.ones((x.shape[0], 2))
9
m[:, 1] = x
10
scale = np.empty((2,))
11
for n in range(0, 2):
12
norm = np.linalg.norm(m[:, n])
13
scale[n] = norm
14
m[:, n] /= norm
15
lsf = (np.linalg.lstsq(m, data, rcond=-1)[0] / scale)[::-1]
16
lsd = data - lsf[0] * x
17
size_lsd = len(lsd)
18
four = np.zeros(size_lsd, dtype=np.complex128)
19
for i in range(size_lsd):
20
sum_f = 0
21
for n in range(size_lsd):
22
sum_f += lsd[n] * np.exp(-2j * np.pi * i * n * (1 / size_lsd))
23
four[i] = sum_f
24
freq = np.empty(size)
25
mi = (size - 1) // 2 + 1
26
freq[:mi] = np.arange(0, mi)
27
freq[mi:] = np.arange(-(size // 2), 0)
28
freq *= 1.0 / size
29
lx = np.arange(0, size + extra)
30
out = np.zeros(lx.shape)
31
32
# IT'S USED TO SORT FOURIER REALS
33
index = [v for _, v in sorted([(np.absolute(four[v]), v) for v in list(range(size))])][::-1]
34
35
for i in index[:1 + harmonic * 2]:
36
out += (abs(four[i]) / size) * np.cos(2 * np.pi * freq[i] * lx + np.angle(four[i]))
37
return out + lsf[0] * lx
38
Advertisement
Answer
For this particular kind of input, you can achieve the sorting with:
JavaScript
1
3
1
for value in index[:]:
2
index[keys[value]] = value
3
If the keys are not a permutation from a range(n)
(like in your question), then create temporary tuples, call sorted
and then extract the value again from the tuples:
JavaScript
1
4
1
result = [value for _, value in sorted(
2
[(keys[value], value) for value in index]
3
)]
4