-
Notifications
You must be signed in to change notification settings - Fork 0
/
index_test.py
220 lines (159 loc) · 6.26 KB
/
index_test.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
import numpy as np
import matplotlib.pyplot as plt
import pickle
import skyflux as sf
import skyflux.deprecated.polSims as pol
ad_hoc_labels = ["I", "Q", "U", "V"]
ad_hoc_colors = ['k', 'b', 'orange', 'g']
def probe_stokes_params(fname, rAnt1=None, rAnt2=None):
"""
out of context
"""
fcd, fs, ts, ptitle = load_wedge_sim(
fname + ".pickle", rAnt1, rAnt2
)
print("Simulation file loaded.\n")
transformed = transform_wedge(fcd, fs, ts)
print("Fourier transforms applied to simulation.\n")
wedge = collect_wedge_points(
transformed, fs, ts, rAnt1, rAnt2
)
def load_wedge_sim(fname, rAnt1, rAnt2, tGoal=None):
"""
out of context
"""
sim_file = open(fname, "rb")
meta = pickle.load(sim_file)
ptitle = meta['title']
fs = meta['frequencies']
#print(fs)
num_f = len(fs)
ts = meta['times']
num_t = len(ts)
sim = meta['picture']
fcd = {} # fourier candidate dictionary
max_counts = [0, 0, 0, 0]
for ant1 in sim.keys():
fcd[ant1] = {}
for ant2 in sim[ant1].keys():
# 0: I 1: Q 2: U 3: V
if rAnt1 is not None and rAnt1 != ant1:
continue
if rAnt2 is not None and rAnt2 != ant2:
continue
fourierc = [[], [], [], []]
for ti in range(num_t):
if tGoal is None or tGoal != np.around(ts[ti], 4):
continue
for parameter in fourierc:
parameter.append([])
for ni in range(num_f):
v = sim[ant1][ant2][ni][ti]
v_test = np.array([
np.abs(S) for S in v
])
for i in range(len(v_test)):
if v_test[i] == v_test.max():
max_counts[i] += 1
for p_idx in range(len(fourierc)):
fourierc[p_idx][ti].append(v[p_idx])
# si: Stokes index
for si in range(len(fourierc)):
parameter = fourierc[si]
parameter[ti] = np.array(parameter[ti])
# cut the next three statements
# in case you no longer want to plot this
plt.plot(fs / 1e6,
np.abs(parameter[len(parameter) - 1]),
label=str(ad_hoc_labels[si]),
color=ad_hoc_colors[si])
plt.legend(loc='upper right')
plt.title(
"Raw Visibilities (LST = " + \
str(np.around(ts[ti], 4)) + \
"). Antennae: " + \
str(ant1) + " to " + \
str(ant2)
)
plt.xlabel("Frequency [MHz]")
plt.ylabel("Brightness Magnitude [Jy]")
plt.show()
for parameter in fourierc:
parameter = np.array(parameter)
fcd[ant1][ant2] = np.array(fourierc)
print(max_counts)
return fcd, fs, ts, ptitle
def transform_wedge(original, fs, ts):
num_f = len(fs)
num_t = len(ts)
import copy
fourier_dict = copy.deepcopy(original)
window = pol.genWindow(num_f)
for ant1 in fourier_dict.keys():
for ant2 in fourier_dict[ant1].keys():
fourier = fourier_dict[ant1][ant2]
for ti in range(num_t):
for parameter in fourier:
parameter[ti] = np.fft.fft(
parameter[ti] * window
)
return fourier_dict
def collect_wedge_points(fcd, fs, ts, rAnt1, rAnt2):
"""
out of context
"""
num_t = len(ts)
num_f = len(fs)
B = fs.max() - fs.min()
visual = []
etas = pol.f2etas(fs)
k_para = np.array([
pol.k_parallel(etas[i], pol.fq2z(fs[i] / 1e9)) \
for i in range(num_f)
])
max_counts = [0, 0, 0, 0]
hmax_counts = [0, 0, 0, 0]
for ant1 in fcd.keys():
for ant2 in fcd[ant1].keys():
if rAnt1 is not None and rAnt1 != ant1:
continue
if rAnt2 is not None and rAnt2 != ant2:
continue
baselength = sf.ant.baselength(ant1, ant2)
for nu_idx in range(num_f):
nua = np.average(fs)
nu = fs[nu_idx]
z = pol.fq2z(nu / 1e9)
lambda_ = pol.C / nu
k_perp = baselength * pol.k_perp(z) / lambda_
for t_idx in range(num_t - 1):
this_instant = \
fcd[ant1][ant2][:, t_idx, nu_idx]
# [I1, Q1, U1, V1] * [I2*, Q2*, U2*, V2*]
this_test = np.array([
np.abs(S) for S in this_instant
])
for i in range(len(this_test)):
if this_test[i] == this_test.max():
max_counts[i] += 1
next_instant = \
fcd[ant1][ant2][:, t_idx + 1, nu_idx]
# [I1, Q1, U1, V1] * [I2*, Q2*, U2*, V2*]
hadamard = np.multiply(
this_instant, next_instant
)
htest = np.array([
np.abs(S) for S in hadamard
])
for i in range(len(htest)):
if htest[i] == htest.max():
hmax_counts[i] += 1
print(max_counts)
print(hmax_counts)
"""
Create a cross-sectional plot
1. Just the visibilities, with frequency as axis
2. Hadamard product of visibilities
3. Multiply by the cosmological constant
4. Divide by the normalization
"""