forked from SgmAstro/DESI
-
Notifications
You must be signed in to change notification settings - Fork 1
/
smith_kcorr.py
248 lines (186 loc) · 9.62 KB
/
smith_kcorr.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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
import os
import numpy as np
import matplotlib.pyplot as plt
import os
from scipy.interpolate import interp1d
from pkg_resources import resource_filename
from tmr_kcorr import tmr_kcorr
raw_dir = os.environ['CODE_ROOT'] + '/data/'
class GAMA_KCorrection(object):
def __init__(self, band, kind="linear"):
"""
Colour-dependent polynomial fit to the GAMA K-correction (Fig. 13 of Smith+17),
used to convert between SDSS r-band Petrosian apparent magnitudes, and rest
frame absolute manigutues at z_ref = 0.1
Args:
k_corr_file: file of polynomial coefficients for each colour bin
z0: reference redshift. Default value is z0=0.1
kind: type of interpolation between colour bins,
e.g. "linear", "cubic". Default is "linear"
"""
k_corr_file = raw_dir + '/ajs_kcorr_{}band_z01.dat'.format(band.lower())
# read file of parameters of polynomial fit to k-correction
# polynomial k-correction is of the form
# A*(z-z0)^4 + B*(z-z0)^3 + C*(z-z0)^2 + D*(z-z0) + E
col_min, col_max, A, B, C, D, E, col_med = \
np.loadtxt(k_corr_file, unpack=True)
self.z0 = 0.1 # reference redshift
self.nbins = len(col_min) # number of colour bins in file
self.colour_min = np.min(col_med)
self.colour_max = np.max(col_med)
self.colour_med = col_med
# functions for interpolating polynomial coefficients in rest-frame color.
self.__A_interpolator = self.__initialize_parameter_interpolator(A, col_med, kind=kind)
self.__B_interpolator = self.__initialize_parameter_interpolator(B, col_med, kind=kind)
self.__C_interpolator = self.__initialize_parameter_interpolator(C, col_med, kind=kind)
self.__D_interpolator = self.__initialize_parameter_interpolator(D, col_med, kind=kind)
self.__E = E[0]
# Linear extrapolation for z > 0.5
self.__X_interpolator = lambda x: None
self.__Y_interpolator = lambda x: None
self.__X_interpolator, self.__Y_interpolator = self.__initialize_line_interpolators()
def __initialize_parameter_interpolator(self, parameter, median_colour, kind="linear"):
# returns function for interpolating polynomial coefficients, as a function of colour
return interp1d(median_colour, parameter, kind=kind, fill_value="extrapolate")
def __initialize_line_interpolators(self):
# linear coefficients for z>0.5
X = np.zeros(self.nbins)
Y = np.zeros(self.nbins)
# find X, Y at each colour
redshift = np.array([0.48,0.5])
arr_ones = np.ones(len(redshift))
for i in range(self.nbins):
k = self.k(redshift, arr_ones*self.colour_med[i])
X[i] = (k[1]-k[0]) / (redshift[1]-redshift[0])
Y[i] = k[0] - X[i]*redshift[0]
X_interpolator = interp1d(self.colour_med, X, kind='linear', fill_value="extrapolate")
Y_interpolator = interp1d(self.colour_med, Y, kind='linear', fill_value="extrapolate")
return X_interpolator, Y_interpolator
def __A(self, colour):
# coefficient of the z**4 term
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__A_interpolator(colour_clipped)
def __B(self, colour):
# coefficient of the z**3 term
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__B_interpolator(colour_clipped)
def __C(self, colour):
# coefficient of the z**2 term
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__C_interpolator(colour_clipped)
def __D(self, colour):
# coefficient of the z**1 term
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__D_interpolator(colour_clipped)
def __X(self, colour):
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__X_interpolator(colour_clipped)
def __Y(self, colour):
colour_clipped = np.clip(colour, self.colour_min, self.colour_max)
return self.__Y_interpolator(colour_clipped)
def k(self, redshift, restframe_colour, median=False):
"""
Polynomial fit to the GAMA K-correction for z<0.5
The K-correction is extrapolated linearly for z>0.5
Args:
redshift: array of redshifts
colour: array of ^0.1(g-r) colour
Returns:
array of K-corrections
"""
K = np.zeros(len(redshift))
idx = redshift <= 0.5
if median:
restframe_colour = np.copy(restframe_colour)
# Fig. 13 of https://arxiv.org/pdf/1701.06581.pdf
restframe_colour = 0.603 * np.ones_like(restframe_colour)
K[idx] = self.__A(restframe_colour[idx])*(redshift[idx]-self.z0)**4 + \
self.__B(restframe_colour[idx])*(redshift[idx]-self.z0)**3 + \
self.__C(restframe_colour[idx])*(redshift[idx]-self.z0)**2 + \
self.__D(restframe_colour[idx])*(redshift[idx]-self.z0) + self.__E
idx = redshift > 0.5
K[idx] = self.__X(restframe_colour[idx])*redshift[idx] + self.__Y(restframe_colour[idx])
return K
def k_nonnative_zref(self, refz, redshift, restframe_colour, median=False):
refzs = refz * np.ones_like(redshift)
return self.k(redshift, restframe_colour, median=median) - self.k(refzs, restframe_colour, median=median) - 2.5 * np.log10(1. + refz)
def rest_gmr_index(self, rest_gmr, kcoeff=False):
bins = np.array([-100., 0.18, 0.35, 0.52, 0.69, 0.86, 1.03, 100.])
idx = np.digitize(rest_gmr, bins=bins)
'''
if kcoeff==True:
for i in enumerate(rest_gmr):
ddict = {i:{col_med, A[0], B[0], C[0], D[0]}}
'''
return idx
class GAMA_KCorrection_color():
def __init__(self):
self.kRcorr = GAMA_KCorrection(band='R')
self.kGcorr = GAMA_KCorrection(band='G')
def obs_gmr(self, rest_gmr):
return rest_gmr + self.kRcorr.k(z, rest_gmr) - self.kGcorr.k(z, rest_gmr)
def rest_gmr_nonnative(self, native_rest_gmr):
refzs = np.zeros_like(native_rest_gmr)
return native_rest_gmr + self.kGcorr.k(refzs, native_rest_gmr) - self.kRcorr.k(refzs, native_rest_gmr)
def test_plots(axes):
kcorr_r = GAMA_KCorrection(band='R')
kcorr_g = GAMA_KCorrection(band='G')
z = np.arange(-0.01,0.601,0.01)
cols = 0.130634, 0.298124, 0.443336, 0.603434, 0.784644, 0.933226, 1.06731
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
# make r-band k-correction plot
for i, c in enumerate(cols):
col = np.ones(len(z)) * c
k = kcorr_r.k(z, col)
axes[0].plot(z, k, label=r"$^{0.1}(g-r)_\mathrm{med}=%.3f$"%c, c=colors[i])
axes[0].set_xlabel(r"$z$")
axes[0].set_ylabel(r"$^{0.1}K_r(z)$")
axes[0].set_xlim(0,0.6)
axes[0].set_ylim(-0.6,1)
axes[0].legend(loc="upper left").draw_frame(False)
# make g-band k-correction plot
for i, c in enumerate(cols):
col = np.ones(len(z)) * c
k = kcorr_g.k(z, col)
axes[1].plot(z, k, label=r"$^{0.1}(g-r)_\mathrm{med}=%.3f$"%c, c=colors[i])
axes[1].set_xlabel(r"$z$")
axes[1].set_ylabel(r"$^{0.1}K_g(z)$")
axes[1].set_xlim(-0.01,0.6)
axes[1].set_ylim(-0.4,1.4)
axes[1].legend(loc="upper left").draw_frame(False)
def test_nonnative_plots(kE, zref, axes=None):
kcorr_tmr = tmr_kcorr()
kcorr_r = GAMA_KCorrection(band='R')
kcorr_g = GAMA_KCorrection(band='G')
z = np.arange(0.01, 0.601, 0.01)
cols = [0.130634, 0.298124, 0.443336, 0.603434, 0.784644, 0.933226, 1.06731]
# TMR test
# cols = [0.158, 0.298, 0.419, 0.553, 0.708, 0.796, 0.960]
if axes == None:
fig, axes = plt.subplots(1, 2, figsize=(20,10))
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
for idx in np.unique(kE['REST_GMR_0P1_INDEX']):
isin = (kE['REST_GMR_0P1_INDEX'] == idx)
axes[0].scatter(kE['ZSURV'][isin], kE['KCORR_R0P0'][isin], s=0.25)
axes[1].scatter(kE['ZSURV'][isin], kE['KCORR_G0P0'][isin], s=0.25)
# make r-band k-correction plot
for i, c in enumerate(cols):
col = np.ones(len(z)) * c
k = kcorr_r.k_nonnative_zref(zref, z, col)
axes[0].plot(z, k, label=r"$^{0.0}(g-r)_\mathrm{med}=%.3f$"%c, c=colors[i], alpha=1.)
k = kcorr_tmr.ref_eval(c, z)
'''
for i, c in enumerate(cols):
k = kcorr_tmr.ref_eval(c, z)
isin = (z <= 0.5)
axes[0].plot(z[isin], k[isin], '--', c=colors[i], alpha=0.75)
'''
# make g-band k-correction plot
for i, c in enumerate(cols):
col = np.ones(len(z)) * c
k = kcorr_g.k_nonnative_zref(zref, z, col)
axes[1].plot(z, k, label=r"$^{0.0}(g-r)_\mathrm{med}=%.3f$"%c, c=colors[i])
axes[0].set_xlim(0, 0.3)
axes[1].set_xlim(0, 0.3)
axes[0].set_ylim(-0.2, 1.2)
axes[1].set_ylim(-0.2, 1.2)