-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathwigner3j.py
62 lines (54 loc) · 1.98 KB
/
wigner3j.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import numpy as np
import pywigxjpf as wig
import multiprocessing as mp
def compute_3j_symbol_one_ell(l1, ell_sum_max):
'''
ARGUMENTS
---------
l1: int, value of first ell for Wigner 3j symbol
ell_sum_max: int, maximum l2,l3 for which to calculate Wigner 3j symbols
RETURNS
-------
wig3j_one_ell: 2D numpy array, indexed as wig3j_one_ell[l2,l3], contains all 3j symbols for l1
'''
wig.wig_table_init(ell_sum_max*2,9)
wig.wig_temp_init(ell_sum_max*2)
tj0 = lambda l1,l2,l3: wig.wig3jj(2*l1,2*l2,2*l3,0,0,0)
wig3j_one_ell = np.zeros((ell_sum_max+1, ell_sum_max+1), dtype=np.float32)
for l2 in range(ell_sum_max+1):
for l3 in range(ell_sum_max+1):
wig3j_one_ell[l2,l3] = tj0(l1,l2,l3)
return wig3j_one_ell
def compute_3j(ell_sum_max):
'''
Computes Wigner 3j symbols with multiprocessing (parallelization)
ARGUMENTS
---------
ell_sum_max: int, maximum ell for which to calculate Wigner 3j symbols
RETURNS
-------
wig3j: 3D numpy array, contains wigner 3j symbols, indexed as wig3j[l1,l2,l3]
'''
pool = mp.Pool(16)
wig3j = pool.starmap(compute_3j_symbol_one_ell, [(l1, ell_sum_max) for l1 in range(ell_sum_max+1)])
pool.close()
return np.array(wig3j)
def compute_3j_no_multiprocessing(ell_sum_max):
'''
Computes Wigner 3j symbols without parallelization
ARGUMENTS
---------
ell_sum_max: int, maximum l2,l3 for which to calculate Wigner 3j symbols
RETURNS
-------
wig3j: 3D numpy array, contains wigner 3j symbols, indexed as wig3j[l1,l2,l3]
'''
wig.wig_table_init(ell_sum_max*2,9)
wig.wig_temp_init(ell_sum_max*2)
tj0 = lambda l1,l2,l3: wig.wig3jj(2*l1,2*l2,2*l3,0,0,0)
wig3j = np.zeros((ell_sum_max+1, ell_sum_max+1, ell_sum_max+1), dtype=np.float32)
for l1 in range(ell_sum_max+1):
for l2 in range(ell_sum_max+1):
for l3 in range(ell_sum_max+1):
wig3j[l1,l2,l3] = tj0(l1,l2,l3)
return wig3j