Skip to content

Commit 14a0488

Browse files
calumgordoncalumgordon
committed
works barely
1 parent 63d3cbc commit 14a0488

15 files changed

+694
-26
lines changed

‎lya_2pt/.ipynb_checkpoints/correlation-checkpoint.py

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
import sys
2+
import numpy as np
3+
from numba import njit
4+
5+
import lya_2pt.global_data as globals
6+
from lya_2pt.tracer_utils import get_angle
7+
from lya_2pt.utils import gen_gamma
8+
9+
10+
import pdb
11+
class ForkedPdb(pdb.Pdb):
12+
"""A Pdb subclass that may be used
13+
from a forked multiprocessing child
14+
15+
"""
16+
def interaction(self, *args, **kwargs):
17+
_stdin = sys.stdin
18+
try:
19+
sys.stdin = open('/dev/stdin')
20+
pdb.Pdb.interaction(self, *args, **kwargs)
21+
finally:
22+
sys.stdin = _stdin
23+
24+
25+
def compute_xi(healpix_id):
26+
"""Compute correlation function
27+
28+
Parameters
29+
----------
30+
tracers1 : array of lya_2pt.tracer.Tracer
31+
First set of tracers
32+
tracers2 : array of lya_2pt.tracer.Tracer
33+
Second set of tracers
34+
config : ConfigParser
35+
Internal configuration object containing the settings section
36+
auto_flag : bool, optional
37+
Flag for auto-correlation, by default False
38+
39+
Returns
40+
-------
41+
(array, array, array, array, array, array)
42+
correlation function, sum of weights in each bin, line-of-sight separation grid,
43+
transverse separation grid, redshift grid, number of pixel pairs in each bin
44+
"""
45+
hp_neighs = [other_hp for other_hp in globals.healpix_neighbours[healpix_id]
46+
if other_hp in globals.tracers2]
47+
48+
hp_neighs += [healpix_id]
49+
50+
total_size = int(globals.num_bins_rp * globals.num_bins_rt)
51+
52+
xi_grid = np.zeros(total_size)
53+
weights_grid = np.zeros(total_size)
54+
rp_grid = np.zeros(total_size)
55+
rt_grid = np.zeros(total_size)
56+
z_grid = np.zeros(total_size)
57+
num_pairs_grid = np.zeros(total_size, dtype=np.int32)
58+
59+
#init gamma auto and cross arrays
60+
gamma_grid = np.zeros(total_size)
61+
delta_gamma_grid = np.zeros(total_size)
62+
# delta_gamma_p_grid = np.zeros(total_size)
63+
64+
sigma_v = globals.gamma_z_error
65+
66+
if not globals.cont_polynomial is None:
67+
# #load polynomials
68+
los_ids = globals.cont_polynomial['los_ids']
69+
zall = globals.cont_polynomial['ztrue']
70+
71+
# p_ids = globals.cont_polynomial['los_ids']
72+
# bqs_1 = globals.cont_polynomial['bqs']
73+
# aqs_1 = globals.cont_polynomial['aqs']
74+
# bqs_2 = globals.cont_polynomial['bqs_500']
75+
# aqs_2 = globals.cont_polynomial['aqs_500']
76+
77+
for tracer1 in globals.tracers1[healpix_id]:
78+
with globals.lock:
79+
xicounter = round(globals.counter.value * 100. / globals.num_tracers, 2)
80+
if (globals.counter.value % 1000 == 0):
81+
print(("computing xi: {}%").format(xicounter))
82+
sys.stdout.flush()
83+
globals.counter.value += 1
84+
85+
86+
87+
#calculate rest frame waves and gamma function of tracer 1
88+
lambda_rest_1 = 1215.67 * (1 + tracer1.z) / (1 + tracer1.z_qso)
89+
gamma_1 = gen_gamma(lambda_rest_1,sigma_v)
90+
91+
potential_neighbours = [tracer2 for hp in hp_neighs for tracer2 in globals.tracers2[hp]]
92+
# for hp in hp_neighs:
93+
# if hp not in globals.tracers2:
94+
# continue
95+
# else:
96+
# potential_neighbours.append(globals.tracers2[hp])
97+
# potential_neighbours = np.concatenate(potential_neighbours)
98+
#ForkedPdb().set_trace()
99+
100+
wlam_1 = (lambda_rest_1 > 1000) & (lambda_rest_1 < 1300)
101+
102+
if not globals.cont_polynomial is None:
103+
104+
#where the forests are outside of the 1200A upper limit
105+
id_match_1 = los_ids == tracer1.los_id
106+
ztrue_1 = zall[id_match_1]
107+
if len(ztrue_1)==0:
108+
continue
109+
lrest_true_1 = 1215.67 * (1 + tracer1.z) / (1 + ztrue_1)
110+
wlam_1 = (lrest_true_1 > 1040) & (lrest_true_1 < 1200)
111+
112+
113+
neighbours = tracer1.get_neighbours(
114+
potential_neighbours, globals.auto_flag,
115+
globals.z_min, globals.z_max,
116+
globals.rp_max, globals.rt_max
117+
)
118+
119+
# #get aq,bq value
120+
# id_match = p_ids == tracer1.los_id
121+
# bq_1 = bqs_1[id_match]
122+
123+
# if bq_1.size==0:
124+
# gamma_p1 = np.zeros_like(lambda_rest_1).astype(float)
125+
# continue
126+
# else:
127+
# aq_1 = aqs_1[id_match]
128+
# aq_2 = aqs_2[id_match]
129+
# bq_2 = bqs_2[id_match]
130+
# L = (tracer1.log_lambda - tracer1.log_lambda.min()) / (tracer1.log_lambda.max() - tracer1.log_lambda.min())
131+
# gamma_p1 = (aq_2 + bq_2 * L) / (aq_1 + ( bq_1 * L ))
132+
#calculate rest frame waves and gamma function of tracer 2
133+
134+
for tracer2 in neighbours:
135+
136+
lambda_rest_2 = 1215.67 * (1 + tracer2.z) / (1 + tracer2.z_qso)
137+
gamma_2 = gen_gamma(lambda_rest_2,sigma_v)
138+
139+
wlam_2 = (lambda_rest_2> 1000) & (lambda_rest_2 < 1300)
140+
141+
if not globals.cont_polynomial is None:
142+
#where the forests are outside of the 1200A upper limit
143+
id_match_2 = los_ids == tracer2.los_id
144+
ztrue_2 = zall[id_match_2]
145+
if len(ztrue_2)==0:
146+
continue
147+
lrest_true_2 = 1215.67 * (1 + tracer2.z) / (1 + ztrue_2)
148+
wlam_2 = (lrest_true_2> 1040) & (lrest_true_2 < 1200)
149+
150+
151+
152+
# #get aq,bq value
153+
# id_match = p_ids == tracer2.los_id
154+
# bq_1 = bqs_1[id_match]
155+
156+
# if bq_1.size==0:
157+
# gamma_p2 = np.zeros_like(lambda_rest_2).astype(float)
158+
# continue
159+
# else:
160+
# aq_1 = aqs_1[id_match]
161+
# aq_2 = aqs_2[id_match]
162+
# bq_2 = bqs_2[id_match]
163+
# L = (tracer2.log_lambda - tracer2.log_lambda.min()) / (tracer2.log_lambda.max() - tracer2.log_lambda.min())
164+
# gamma_p2 = (aq_2 + bq_2 * L) / (aq_1 + ( bq_1 * L ))
165+
166+
167+
angle = get_angle(
168+
tracer1.x_cart, tracer1.y_cart, tracer1.z_cart, tracer1.ra, tracer1.dec,
169+
tracer2.x_cart, tracer2.y_cart, tracer2.z_cart, tracer2.ra, tracer2.dec
170+
)
171+
172+
compute_xi_pair(
173+
tracer1.deltas[wlam_1], tracer1.weights[wlam_1], tracer1.z[wlam_1], tracer1.dist_c[wlam_1], tracer1.dist_m[wlam_1],
174+
tracer2.deltas[wlam_2], tracer2.weights[wlam_2], tracer2.z[wlam_2], tracer2.dist_c[wlam_2], tracer2.dist_m[wlam_2],
175+
angle, xi_grid, weights_grid, rp_grid, rt_grid, z_grid, num_pairs_grid,
176+
gamma_grid, delta_gamma_grid, gamma_1[wlam_1], gamma_2[wlam_2]
177+
)
178+
179+
# Normalize correlation and average coordinate grids
180+
w = weights_grid > 0
181+
xi_grid[w] /= weights_grid[w]
182+
rp_grid[w] /= weights_grid[w]
183+
rt_grid[w] /= weights_grid[w]
184+
z_grid[w] /= weights_grid[w]
185+
186+
#normalise gamma cross and auto terms
187+
gamma_grid[w] /= weights_grid[w]
188+
delta_gamma_grid[w] /= weights_grid[w]
189+
#delta_gamma_p_grid[w] /= weights_grid[w]
190+
191+
return healpix_id, (xi_grid, weights_grid, rp_grid, rt_grid, z_grid, num_pairs_grid, gamma_grid, delta_gamma_grid)
192+
193+
194+
@njit
195+
def compute_xi_pair(
196+
deltas1, weights1, z1, dist_c1, dist_m1,
197+
deltas2, weights2, z2, dist_c2, dist_m2, angle,
198+
xi_grid, weights_grid, rp_grid, rt_grid, z_grid, num_pairs_grid,
199+
gamma_grid, delta_gamma_grid, gamma_1, gamma_2
200+
):
201+
sin_angle = np.sin(angle / 2)
202+
cos_angle = np.cos(angle / 2)
203+
204+
for i in range(deltas1.size):
205+
if weights1[i] == 0:
206+
continue
207+
for j in range(deltas2.size):
208+
if weights2[j] == 0:
209+
continue
210+
211+
# Comoving separation between the two pixels
212+
rp = (dist_c1[i] - dist_c2[j]) * cos_angle
213+
rt = (dist_m1[i] + dist_m2[j]) * sin_angle
214+
if globals.auto_flag:
215+
rp = np.abs(rp)
216+
217+
# Skip if pixel pair is too far apart
218+
if (rp < globals.rp_min) or (rp >= globals.rp_max) or (rt >= globals.rt_max):
219+
continue
220+
221+
# Compute bin in the correlation function to asign the pixel pair to
222+
bins_rp = np.floor((rp - globals.rp_min) / (globals.rp_max - globals.rp_min)
223+
* globals.num_bins_rp)
224+
bins_rt = np.floor(rt / globals.rt_max * globals.num_bins_rt)
225+
bins = int(bins_rt + globals.num_bins_rt * bins_rp)
226+
227+
# Compute and write correlation and associated quantities
228+
weight12 = weights1[i] * weights2[j]
229+
xi_grid[bins] += deltas1[i] * deltas2[j] * weight12
230+
weights_grid[bins] += weight12
231+
rp_grid[bins] += rp * weight12
232+
rt_grid[bins] += rt * weight12
233+
z_grid[bins] += (z1[i] + z2[j]) / 2 * weight12
234+
num_pairs_grid[bins] += 1
235+
236+
#<delta gamma>
237+
delta_gamma_grid[bins] += deltas1[i] * gamma_2[j] * weight12
238+
#<delta gamma 2>
239+
#delta_gamma_p_grid[bins] += ( (((1+deltas1[i]-gamma_1[i])/gamma_p1[i]) - 1) * (((1+deltas2[j]-gamma_2[j])/gamma_p2#[j]) - 1) * weight12 )
240+
#<gamma gamma>
241+
gamma_grid[bins] += gamma_1[i] * gamma_2[j] * weight12

‎lya_2pt/.ipynb_checkpoints/global_data-checkpoint.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
2+
3+
counter = None
4+
lock = None
5+
6+
tracers1 = None
7+
tracers2 = None
8+
num_tracers = None
9+
healpix_neighbours = None
10+
11+
z_min = None
12+
z_max = None
13+
auto_flag = None
14+
15+
rp_min = None
16+
rp_max = None
17+
rt_max = None
18+
num_bins_rp = None
19+
num_bins_rt = None
20+
num_bins_rp_model = None
21+
num_bins_rt_model = None
22+
23+
get_old_distortion = None
24+
rejection_fraction = None
25+
26+
gamma_z_error = 0
27+
measured_gamma_interp = None
28+
cont_polynomial = None

0 commit comments

Comments
 (0)