-
Notifications
You must be signed in to change notification settings - Fork 1
/
pvhub.py
218 lines (165 loc) · 7.38 KB
/
pvhub.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
import os
import inspect
import numpy as np
from abc import ABC
import pandas as pd
import astropy.units as u
import astropy.constants as const
from astropy.coordinates import SkyCoord
class Recon(ABC):
"""Abstract base class for the reconstruction data included in pvhub
Meant to only be accessed through the various listed reconstruction subclasses. The attributes below are inherited
by these subclasses
Attributes
----------
"""
def __init__(self, name=None, modelfile=None, extfile=None, verbose=False):
current_file = os.path.dirname(inspect.stack()[0][1])
self.data_location = os.path.normpath(current_file + "/data") + "/"
self.name = name
self.verbose = verbose
# The properties of the grid on which the reconstructions are stored.
self.dmin = -20000.0
self.dmax = 20000.0
self.nbins = 129
self.bsz = (self.dmax - self.dmin) / float(self.nbins - 1.0)
if modelfile is not None and extfile is not None:
self._load_data(modelfile, extfile)
else:
self.vmodel, self.vextmodel = None, None
def _load_data(self, modelfile, extfile):
"""Loads the relevant reconstruction model files and stores the attributes
Args
----
modelfile : string
The filepath for the file containing the reconstructed field as a function of
cartesian coordinates
extfile : string
The filepath for the file containing the external field prediction as a function of redshift
Raises
------
"""
try:
# The reconstruction model
self.vmodel = pd.read_csv(self.data_location + modelfile)
# The model beyond Vext
self.vextmodel = pd.read_csv(self.data_location + extfile, sep="\s+")
if self.verbose:
print("Loaded model " + self.name)
except:
print("ERROR: Unable to load model " + self.name)
print("Please check files " + modelfile + " and " + extfile + " are in " + self.data_location)
return
def calculate_pv(self, ra, dec, zcmb, extrapolation=True):
"""Interpolates/extrapolates the reconstruction for a set of ra, dec and CMB-frame redshifts
Args
----
ra : float, np.ndarray
A value or array of values for the Right Ascention to be queried.
dec : float, np.ndarray
A value or array of values for the Declination to be queried.
z_cmb_in: float, np.ndarray
A value or array of values for the CMB-frame redshifts to be queried.
extrapolation: bool
Whether or not to extrapolate outside of the reconstructed field based on a smooth model of the
velocity as a function of redshift. Default = True
Returns
-------
pv : np.ndarray
An array of peculiar velocities for each of the coordinates passed in
Raises
------
"""
ra, dec, zcmb = np.asarray(ra), np.asarray(dec), np.asarray(zcmb)
ccc = SkyCoord(
ra * u.degree,
dec * u.degree,
# Note that cz is *not* a distance, but treating it as such is self-consistent with 2M++
distance=const.c.value / 1000.0 * zcmb * u.km / u.s,
frame="icrs",
)
sgc = ccc.transform_to("supergalactic")
sgc.representation_type = "cartesian"
xbin = np.round(((sgc.sgx.value - self.dmin) / self.bsz), 0)
ybin = np.round(((sgc.sgy.value - self.dmin) / self.bsz), 0)
zbin = np.round(((sgc.sgz.value - self.dmin) / self.bsz), 0)
# calculate bin index even if coords outside 2M++
binindex = xbin.astype(int) * self.nbins * self.nbins + ybin.astype(int) * self.nbins + zbin.astype(int)
# set indices outside 2M++ to 0
if binindex.size > 1:
binindex[np.where((binindex < 0) | (binindex >= len(self.vmodel["vproj_2MPP"])))] = 0
else:
binindex = 0 if (binindex < 0 or binindex >= len(self.vmodel["vproj_2MPP"])) else binindex
k = np.searchsorted(self.vextmodel["z"], zcmb) # calculate bin index even if coords inside 2M++
in2MPP = (
(const.c.value / 1000.0 * zcmb < 19942) # precise redshift of 2M++ boundary
& ((self.dmin < sgc.sgx.value) & (sgc.sgx.value < self.dmax))
& ((self.dmin < sgc.sgy.value) & (sgc.sgy.value < self.dmax))
& ((self.dmin < sgc.sgz.value) & (sgc.sgz.value < self.dmax))
)
if extrapolation:
r = np.sqrt(np.sum(np.square([sgc.sgx.value, sgc.sgy.value, sgc.sgz.value]), axis=0))
vdot = (
sgc.sgx.value * self.vextmodel["Vsgx"].loc[k]
+ sgc.sgy.value * self.vextmodel["Vsgy"].loc[k]
+ sgc.sgz.value * self.vextmodel["Vsgz"].loc[k]
)
pv = np.where(
in2MPP,
self.vmodel["vproj_2MPP"].loc[binindex],
vdot / r,
)
# 2M++ is not completely spherical, so extrapolation must be extended inwards in some
# parts of the sky. The PV of the centre of the reconstruction is undefined, so we only
# use the extrapolation above a redshift securely inside 2M++ but outside the central cell
pv = np.where((np.isnan(pv)) & (zcmb > 0.01), vdot / r, pv)
pv = np.round(pv, 0)
else:
pv = np.where(in2MPP, np.round(self.vmodel["vproj_2MPP"].loc[binindex], 0), np.nan)
return pv
class TwoMPP_SDSS(Recon):
def __init__(self, verbose=False):
name = "2M++_SDSS"
model = "2MPP_SDSS.txt"
model_ext = "2MPP_SDSS_out.txt"
super().__init__(name=name, modelfile=model, extfile=model_ext, verbose=verbose)
class TwoMPP_SDSS_6dF(Recon):
def __init__(self, verbose=False):
name = "2M++_SDSS_6dF"
model = "2MPP_SDSS_6dF.txt"
model_ext = "2MPP_SDSS_6dF_out.txt"
super().__init__(name=name, modelfile=model, extfile=model_ext, verbose=verbose)
class TwoMRS_redshift(Recon):
def __init__(self, verbose=False):
name = "2MRS_redshift"
model = "2MRS_redshift.txt"
model_ext = "2MRS_redshift_out.txt"
super().__init__(name=name, modelfile=model, extfile=model_ext, verbose=verbose)
class TwoMPP_redshift(Recon):
def __init__(self, verbose=False):
name = "2M++_redshift"
model = "2MPP_redshift.txt"
model_ext = "2MPP_redshift_out.txt"
super().__init__(name=name, modelfile=model, extfile=model_ext, verbose=verbose)
def get_models():
"""Utility function to return all the reconstruction model subclasses
Args
----
None
Returns
-------
final_classes: list
A list of all the reconstruction model subclasses
"""
classes = Recon.__subclasses__()
for c in classes:
classes += c.__subclasses__()
final_classes = [c for c in classes if ABC not in c.__bases__]
return final_classes
if __name__ == "__main__":
from pvhub import *
# Get a list of all models and check they can be read in and used
# This is just a unit test. See examples.py for usage examples
print("Checking models load and run:")
for model in get_models():
model(verbose=True).calculate_pv(334.6, 40.6, 0.0029)