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
|
#!/usr/bin/env python3
#
# Copyright (c) 2017 Weitna LI <weitian@aaronly.me>
# MIT License
#
"""
Average the 2D power spectrum within the EoR window (i.e., excluding the
foreground contaminated wedge) to derive the 1D spherically averaged
power spectrum.
"""
import os
import argparse
import numpy as np
import matplotlib
import matplotlib.style
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from eor_window import PS2D
# Matplotlib settings
matplotlib.style.use("ggplot")
for k, v in [("font.family", "monospace"),
("image.cmap", "jet"),
("xtick.major.size", 7.0),
("xtick.major.width", 2.0),
("xtick.minor.size", 4.0),
("xtick.minor.width", 1.5),
("ytick.major.size", 7.0),
("ytick.major.width", 2.0),
("ytick.minor.size", 4.0),
("ytick.minor.width", 1.5)]:
matplotlib.rcParams[k] = v
class PS1D:
"""
Calculate the 1D spherically averaged power spectrum from 2D PS.
"""
def __init__(self, ps2d):
self.ps2d = ps2d
self.data = ps2d.ps2d # shape: [n_k_los, n_k_perp]
self.data_err = ps2d.ps2d_err
@property
def k_perp(self):
return self.ps2d.k_perp
@property
def k_los(self):
return self.ps2d.k_los
@property
def eor_window(self):
return self.ps2d.eor_window()
def calc_ps1d(self, normalize=True):
"""
Calculate the 1D spherically averaged power spectrum by averaging
the 2D cylindrical power spectrum.
Parameters
----------
normalize : bool
Whether to normalize the 1D power spectrum to obtain the
dimensionless power spectrum, i.e.,
Δ^2(k) = (k^3 / (2*π^2)) P(k)
"""
eor_window = self.eor_window
data = self.data.copy()
data_err = self.data_err.copy()
data[~eor_window] = np.nan
data_err[~eor_window] = np.nan
k_perp = self.k_perp
k_los = self.k_los
dk_perp = k_perp[1] - k_perp[0]
dk_los = k_los[1] - k_los[0]
dk = np.sqrt(dk_perp * dk_los)
print("dk = %.6f [Mpc^-1]" % dk)
k_max = np.sqrt(k_perp[-1]**2 + k_los[-1]**2)
nk = int(k_max / dk) + 1
print("number of k points: %d" % nk)
ps1d_k = np.arange(nk) * dk
# PS1D's 3 columns: [k, ps1d, ps1d_err]
ps1d = np.zeros(shape=(nk, 3))
ps1d[:, 0] = ps1d_k
print("averaging 2D power spectrum ...")
mx, my = np.meshgrid(k_perp, k_los)
mk = np.sqrt(mx**2 + my**2)
for i, k in enumerate(ps1d_k):
ii, jj = (mk <= k).nonzero()
mk[ii, jj] = np.inf
cells = data[ii, jj]
cells = cells[np.isfinite(cells)]
if len(cells) > 0:
ps1d[i, 1] = np.mean(cells)
cells = data_err[ii, jj]
cells = cells[np.isfinite(cells)]
ps1d[i, 2] = np.sqrt(np.sum((cells/len(cells))**2))
if normalize:
coef = ps1d_k**3 / (2*np.pi**2)
ps1d[:, 1] *= coef
ps1d[:, 2] *= coef
self.ps1d_normalized = True
else:
self.ps1d_normalized = False
self.ps1d = ps1d
return ps1d
def save(self, outfile):
ps1d = self.ps1d
header = [
"EoR window:",
" FoV: %f [deg]" % self.ps2d.fov,
" e_ConvWidth: %f" % self.ps2d.e,
" k_perp_min: %f [Mpc^-1]" % self.ps2d.k_perp_min,
" k_perp_max: %f [Mpc^-1]" % self.ps2d.k_perp_max,
" k_los_min: %f [Mpc^-1]" % self.ps2d.k_los_min,
" k_los_max: %f [Mpc^-1]" % self.ps2d.k_los_max,
"",
"k: wavenumber [Mpc^-1]",
]
if self.ps1d_normalized:
header += ["ps1d: normalized power [K^2]"]
else:
header += ["ps1d: power [K^2 Mpc^3]"]
header += [
"ps1d_err: power errors",
"",
"k ps1d ps1d_err"
]
np.savetxt(outfile, ps1d, header="\n".join(header))
print("saved 1D power spectrum to file: %s" % outfile)
def plot(self, ax):
ps1d = self.ps1d
if self.ps1d_normalized:
ylabel = r"$\Delta^2(k)$ [K$^2$]"
else:
ylabel = r"$P(k)$ [K$^2$ Mpc$^3$]"
x = ps1d[:, 0]
y = ps1d[:, 1]
yerr = ps1d[:, 2]
ax.errorbar(x[1:], y[1:], yerr=yerr[1:], fmt="none")
ax.plot(x[1:], y[1:], marker="o")
ax.set(xscale="log", yscale="log",
xlabel=r"[Mpc$^{-1}$]", ylabel=ylabel,
title="1D Spherically Average Power Spectrum")
return ax
def main():
parser = argparse.ArgumentParser(
description="Calculate 1D power spectrum within the EoR window")
parser.add_argument("-C", "--clobber", dest="clobber", action="store_true",
help="overwrite the output files if already exist")
parser.add_argument("-F", "--fov", dest="fov",
type=float, required=True,
help="instrumental FoV to determine the EoR window; " +
"SKA1-Low has FoV ~ 3.12 / (nu/200MHz) [deg], i.e., " +
"~5.03 @ 124, ~3.95 @ 158, ~3.18 @ 196")
parser.add_argument("-e", "--conv-width", dest="conv_width",
type=float, default=3.0,
help="characteristic convolution width (default: 3.0)")
parser.add_argument("-p", "--k-perp-min", dest="k_perp_min", type=float,
help="minimum k wavenumber perpendicular to LoS; " +
"unit: [Mpc^-1]")
parser.add_argument("-P", "--k-perp-max", dest="k_perp_max", type=float,
help="maximum k wavenumber perpendicular to LoS")
parser.add_argument("-l", "--k-los-min", dest="k_los_min", type=float,
help="minimum k wavenumber along LoS")
parser.add_argument("-L", "--k-los-max", dest="k_los_max", type=float,
help="maximum k wavenumber along LoS")
parser.add_argument("--no-plot", dest="noplot", action="store_true",
help="do not plot and save the calculated 1D power " +
"power within the EoR window")
parser.add_argument("-i", "--infile", dest="infile", required=True,
help="2D power spectrum FITS file")
parser.add_argument("-o", "--outfile", dest="outfile", required=True,
help="output TXT file to save the PSD data")
args = parser.parse_args()
if (not args.clobber) and os.path.exists(args.outfile):
raise OSError("outfile '%s' already exists" % args.outfile)
ps2d = PS2D(args.infile, fov=args.fov, e=args.conv_width,
k_perp_min=args.k_perp_min, k_perp_max=args.k_perp_max,
k_los_min=args.k_los_min, k_los_max=args.k_los_max)
ps1d = PS1D(ps2d)
ps1d.calc_ps1d()
ps1d.save(args.outfile)
if not args.noplot:
fig = Figure(figsize=(8, 8), dpi=150)
FigureCanvas(fig)
ax = fig.add_subplot(1, 1, 1)
ps1d.plot(ax=ax)
fig.tight_layout()
plotfile = os.path.splitext(args.outfile)[0] + ".png"
fig.savefig(plotfile)
print("Plotted 1D power spectrum within EoR window: %s" % plotfile)
if __name__ == "__main__":
main()
|
