Skip to content

Commit be42045

Browse files
author
Arnaud De-Mattia
committed
pip and angular weights in all pair counters, with fallback kernel
1 parent 0180645 commit be42045

35 files changed

+3499
-460
lines changed

‎Corrfunc/mocks/DDrppi_mocks.py

+20-4
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
2424
zbin_refine_factor=1, max_cells_per_dim=100,
2525
copy_particles=True, enable_min_sep_opt=True,
2626
c_api_timer=False, isa='fastest',
27-
weight_type=None, bin_type='custom'):
27+
weight_type=None, bin_type='custom',
28+
pair_weights=None, sep_pair_weights=None):
2829
"""
2930
Calculate the 2-D pair-counts corresponding to the projected correlation
3031
function, :math:`\\xi(r_p, \pi)`. Pairs which are separated by less
@@ -236,6 +237,12 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
236237
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
237238
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
238239
240+
pair_weights : array-like, optional. Default: None.
241+
Array of pair weights.
242+
243+
sep_pair_weights : array-like, optional. Default: None.
244+
Array of separations corresponding to ``pair_weights``.
245+
239246
Returns
240247
--------
241248
@@ -361,17 +368,26 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
361368
weights1, weights2 = process_weights(weights1, weights2, RA1, RA2, weight_type, autocorr)
362369

363370
# Ensure all input arrays are native endian
364-
RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2 = [
371+
RA1, DEC1, CZ1, RA2, DEC2, CZ2 = [
365372
convert_to_native_endian(arr, warn=True) for arr in
366-
[RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2]]
373+
[RA1, DEC1, CZ1, RA2, DEC2, CZ2]]
367374

368375
fix_ra_dec(RA1, DEC1)
369376
if autocorr == 0:
370377
fix_ra_dec(RA2, DEC2)
371378

379+
if weight_type is not None:
380+
weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
381+
weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
382+
383+
if pair_weights is not None:
384+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
385+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
386+
372387
# Passing None parameters breaks the parsing code, so avoid this
373388
kwargs = {}
374-
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2']:
389+
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2',
390+
'pair_weights', 'sep_pair_weights']:
375391
v = locals()[k]
376392
if v is not None:
377393
kwargs[k] = v

‎Corrfunc/mocks/DDsmu_mocks.py

+20-4
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ def DDsmu_mocks(autocorr, cosmology, nthreads, mu_max, nmu_bins, binfile,
2323
zbin_refine_factor=1, max_cells_per_dim=100,
2424
copy_particles=True, enable_min_sep_opt=True,
2525
c_api_timer=False, isa='fastest',
26-
weight_type=None, bin_type='custom'):
26+
weight_type=None, bin_type='custom',
27+
pair_weights=None, sep_pair_weights=None):
2728
"""
2829
Calculate the 2-D pair-counts corresponding to the projected correlation
2930
function, :math:`\\xi(s, \mu)`. The pairs are counted in bins of
@@ -238,6 +239,12 @@ def DDsmu_mocks(autocorr, cosmology, nthreads, mu_max, nmu_bins, binfile,
238239
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
239240
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
240241
242+
pair_weights : array-like, optional. Default: None.
243+
Array of pair weights.
244+
245+
sep_pair_weights : array-like, optional. Default: None.
246+
Array of separations corresponding to ``pair_weights``.
247+
241248
Returns
242249
--------
243250
@@ -295,18 +302,27 @@ def DDsmu_mocks(autocorr, cosmology, nthreads, mu_max, nmu_bins, binfile,
295302
weights1, weights2 = process_weights(weights1, weights2, RA1, RA2, weight_type, autocorr)
296303

297304
# Ensure all input arrays are native endian
298-
RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2 = [
305+
RA1, DEC1, CZ1, RA2, DEC2, CZ2 = [
299306
convert_to_native_endian(arr, warn=True) for arr in
300-
[RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2]]
307+
[RA1, DEC1, CZ1, RA2, DEC2, CZ2]]
301308

302309

303310
fix_ra_dec(RA1, DEC1)
304311
if autocorr == 0:
305312
fix_ra_dec(RA2, DEC2)
306313

314+
if weight_type is not None:
315+
weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
316+
weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
317+
318+
if pair_weights is not None:
319+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
320+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
321+
307322
# Passing None parameters breaks the parsing code, so avoid this
308323
kwargs = {}
309-
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2']:
324+
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2',
325+
'pair_weights', 'sep_pair_weights']:
310326
v = locals()[k]
311327
if v is not None:
312328
kwargs[k] = v

‎Corrfunc/mocks/DDtheta_mocks.py

+20-4
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@ def DDtheta_mocks(autocorr, nthreads, binfile,
2424
dec_refine_factor=2, max_cells_per_dim=100,
2525
copy_particles=True, enable_min_sep_opt=True,
2626
c_api_timer=False, isa='fastest',
27-
weight_type=None, bin_type='custom'):
27+
weight_type=None, bin_type='custom',
28+
pair_weights=None, sep_pair_weights=None):
2829
"""
2930
Function to compute the angular correlation function for points on
3031
the sky (i.e., mock catalogs or observed galaxies).
@@ -215,6 +216,12 @@ def DDtheta_mocks(autocorr, nthreads, binfile,
215216
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
216217
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
217218
219+
pair_weights : array-like, optional. Default: None.
220+
Array of pair weights.
221+
222+
sep_pair_weights : array-like, optional. Default: None.
223+
Array of separations corresponding to ``pair_weights``.
224+
218225
Returns
219226
--------
220227
@@ -312,20 +319,29 @@ def DDtheta_mocks(autocorr, nthreads, binfile,
312319
RA1, RA2, weight_type, autocorr)
313320

314321
# Ensure all input arrays are native endian
315-
RA1, DEC1, weights1, RA2, DEC2, weights2 = [
322+
RA1, DEC1, RA2, DEC2 = [
316323
convert_to_native_endian(arr, warn=True) for arr in
317-
[RA1, DEC1, weights1, RA2, DEC2, weights2]]
324+
[RA1, DEC1, RA2, DEC2]]
318325

319326
fix_ra_dec(RA1, DEC1)
320327
if autocorr == 0:
321328
fix_ra_dec(RA2, DEC2)
322329

330+
if weight_type is not None:
331+
weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
332+
weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
333+
334+
if pair_weights is not None:
335+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
336+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
337+
323338
if link_in_ra:
324339
link_in_dec = True
325340

326341
# Passing None parameters breaks the parsing code, so avoid this
327342
kwargs = {}
328-
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2']:
343+
for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2',
344+
'pair_weights', 'sep_pair_weights']:
329345
v = locals()[k]
330346
if v is not None:
331347
kwargs[k] = v

‎Corrfunc/theory/DDrppi.py

+19-4
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ def DDrppi(autocorr, nthreads, pimax, binfile, X1, Y1, Z1, weights1=None,
2121
zbin_refine_factor=1, max_cells_per_dim=100,
2222
copy_particles=True, enable_min_sep_opt=True,
2323
c_api_timer=False, isa='fastest',
24-
weight_type=None, bin_type='custom'):
24+
weight_type=None, bin_type='custom',
25+
pair_weights=None, sep_pair_weights=None):
2526
"""
2627
Calculate the 3-D pair-counts corresponding to the real-space correlation
2728
function, :math:`\\xi(r_p, \pi)` or :math:`\\wp(r_p)`. Pairs which are
@@ -174,6 +175,12 @@ def DDrppi(autocorr, nthreads, pimax, binfile, X1, Y1, Z1, weights1=None,
174175
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
175176
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
176177
178+
pair_weights : array-like, optional. Default: None.
179+
Array of pair weights.
180+
181+
sep_pair_weights : array-like, optional. Default: None.
182+
Array of separations corresponding to ``pair_weights``.
183+
177184
Returns
178185
--------
179186
@@ -290,14 +297,22 @@ def DDrppi(autocorr, nthreads, pimax, binfile, X1, Y1, Z1, weights1=None,
290297
weights1, weights2 = process_weights(weights1, weights2, X1, X2, weight_type, autocorr)
291298

292299
# Ensure all input arrays are native endian
293-
X1, Y1, Z1, weights1, X2, Y2, Z2, weights2 = [
300+
X1, Y1, Z1, X2, Y2, Z2 = [
294301
convert_to_native_endian(arr, warn=True) for arr in
295-
[X1, Y1, Z1, weights1, X2, Y2, Z2, weights2]]
302+
[X1, Y1, Z1, X2, Y2, Z2]]
303+
if weight_type is not None:
304+
weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
305+
weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
306+
307+
if pair_weights is not None:
308+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
309+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
296310

297311
# Passing None parameters breaks the parsing code, so avoid this
298312
kwargs = {}
299313
for k in ['weights1', 'weights2', 'weight_type',
300-
'X2', 'Y2', 'Z2', 'boxsize']:
314+
'X2', 'Y2', 'Z2', 'boxsize',
315+
'pair_weights', 'sep_pair_weights']:
301316
v = locals()[k]
302317
if v is not None:
303318
kwargs[k] = v

‎Corrfunc/theory/DDsmu.py

+19-4
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ def DDsmu(autocorr, nthreads, binfile, mu_max, nmu_bins,
2222
zbin_refine_factor=1, max_cells_per_dim=100,
2323
copy_particles=True, enable_min_sep_opt=True,
2424
c_api_timer=False, isa='fastest',
25-
weight_type=None, bin_type='custom'):
25+
weight_type=None, bin_type='custom',
26+
pair_weights=None, sep_pair_weights=None):
2627
"""
2728
Calculate the 2-D pair-counts corresponding to the redshift-space
2829
correlation function, :math:`\\xi(s, \mu)` Pairs which are separated
@@ -182,6 +183,12 @@ def DDsmu(autocorr, nthreads, binfile, mu_max, nmu_bins,
182183
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
183184
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
184185
186+
pair_weights : array-like, optional. Default: None.
187+
Array of pair weights.
188+
189+
sep_pair_weights : array-like, optional. Default: None.
190+
Array of separations corresponding to ``pair_weights``.
191+
185192
Returns
186193
--------
187194
results : A python list
@@ -305,14 +312,22 @@ def DDsmu(autocorr, nthreads, binfile, mu_max, nmu_bins,
305312
weights1, weights2 = process_weights(weights1, weights2, X1, X2, weight_type, autocorr)
306313

307314
# Ensure all input arrays are native endian
308-
X1, Y1, Z1, weights1, X2, Y2, Z2, weights2 = [
315+
X1, Y1, Z1, X2, Y2, Z2 = [
309316
convert_to_native_endian(arr, warn=True) for arr in
310-
[X1, Y1, Z1, weights1, X2, Y2, Z2, weights2]]
317+
[X1, Y1, Z1, X2, Y2, Z2]]
318+
if weight_type is not None:
319+
weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
320+
weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
321+
322+
if pair_weights is not None:
323+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
324+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
311325

312326
# Passing None parameters breaks the parsing code, so avoid this
313327
kwargs = {}
314328
for k in ['weights1', 'weights2', 'weight_type',
315-
'X2', 'Y2', 'Z2', 'boxsize']:
329+
'X2', 'Y2', 'Z2', 'boxsize',
330+
'pair_weights', 'sep_pair_weights']:
316331
v = locals()[k]
317332
if v is not None:
318333
kwargs[k] = v

‎Corrfunc/theory/wp.py

+17-4
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,8 @@ def wp(boxsize, pimax, nthreads, binfile, X, Y, Z,
285285
xbin_refine_factor=2, ybin_refine_factor=2,
286286
zbin_refine_factor=1, max_cells_per_dim=100,
287287
copy_particles=True, enable_min_sep_opt=True,
288-
c_api_timer=False, c_cell_timer=False, isa='fastest', bin_type='custom'):
288+
c_api_timer=False, c_cell_timer=False, isa='fastest', bin_type='custom',
289+
pair_weights=None, sep_pair_weights=None):
289290
"""
290291
Function to compute the projected correlation function in a
291292
periodic cosmological box. Pairs which are separated by less
@@ -423,6 +424,12 @@ def wp(boxsize, pimax, nthreads, binfile, X, Y, Z,
423424
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
424425
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
425426
427+
pair_weights : array-like, optional. Default: None.
428+
Array of pair weights.
429+
430+
sep_pair_weights : array-like, optional. Default: None.
431+
Array of separations corresponding to ``pair_weights``.
432+
426433
Returns
427434
--------
428435
@@ -504,12 +511,18 @@ def wp(boxsize, pimax, nthreads, binfile, X, Y, Z,
504511
weights, _ = process_weights(weights, None, X, None, weight_type, autocorr=True)
505512

506513
# Ensure all input arrays are native endian
507-
X, Y, Z, weights = [convert_to_native_endian(arr, warn=True)
508-
for arr in [X, Y, Z, weights]]
514+
X, Y, Z = [convert_to_native_endian(arr, warn=True)
515+
for arr in [X, Y, Z]]
516+
if weight_type is not None:
517+
weights = [convert_to_native_endian(arr, warn=True) for arr in weights]
518+
519+
if pair_weights is not None:
520+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
521+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
509522

510523
# Passing None parameters breaks the parsing code, so avoid this
511524
kwargs = {}
512-
for k in ['weights', 'weight_type']:
525+
for k in ['weights', 'weight_type', 'pair_weights', 'sep_pair_weights']:
513526
v = locals()[k]
514527
if v is not None:
515528
kwargs[k] = v

‎Corrfunc/theory/xi.py

+17-4
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ def xi(boxsize, nthreads, binfile, X, Y, Z,
2020
xbin_refine_factor=2, ybin_refine_factor=2,
2121
zbin_refine_factor=1, max_cells_per_dim=100,
2222
copy_particles=True, enable_min_sep_opt=True,
23-
c_api_timer=False, isa='fastest', bin_type='custom'):
23+
c_api_timer=False, isa='fastest', bin_type='custom',
24+
pair_weights=None, sep_pair_weights=None):
2425
"""
2526
Function to compute the projected correlation function in a
2627
periodic cosmological box. Pairs which are separated by less
@@ -142,6 +143,12 @@ def xi(boxsize, nthreads, binfile, X, Y, Z,
142143
``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
143144
of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
144145
146+
pair_weights : array-like, optional. Default: None.
147+
Array of pair weights.
148+
149+
sep_pair_weights : array-like, optional. Default: None.
150+
Array of separations corresponding to ``pair_weights``.
151+
145152
Returns
146153
--------
147154
@@ -214,12 +221,18 @@ def xi(boxsize, nthreads, binfile, X, Y, Z,
214221
weights, _ = process_weights(weights, None, X, None, weight_type, autocorr=True)
215222

216223
# Ensure all input arrays are native endian
217-
X, Y, Z, weights = [convert_to_native_endian(arr, warn=True)
218-
for arr in [X, Y, Z, weights]]
224+
X, Y, Z = [convert_to_native_endian(arr, warn=True)
225+
for arr in [X, Y, Z]]
226+
if weight_type is not None:
227+
weights = [convert_to_native_endian(arr, warn=True) for arr in weights]
228+
229+
if pair_weights is not None:
230+
pair_weights = convert_to_native_endian(pair_weights, warn=True)
231+
sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
219232

220233
# Passing None parameters breaks the parsing code, so avoid this
221234
kwargs = {}
222-
for k in ['weights', 'weight_type']:
235+
for k in ['weights', 'weight_type', 'pair_weights', 'sep_pair_weights']:
223236
v = locals()[k]
224237
if v is not None:
225238
kwargs[k] = v

‎common.mk

+1-1
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ CLINK ?=
1919
## Set the python command (supply the full path to python you want to
2020
## use, if different from directly calling `python` on the shell,
2121
## as can be the case if python is set via an alias)
22-
PYTHON:=/local/home/adematti/anaconda3/envs/cosmopipe-dev/bin/python
22+
PYTHON:=python
2323

2424
## Important note -> if you directly call /some/path/to/python
2525
## then the previous two variables will be updated to point

‎mocks/DDrppi_mocks/countpairs_rp_pi_mocks_impl.c.src

+2-1
Original file line numberDiff line numberDiff line change
@@ -637,7 +637,7 @@ int countpairs_mocks_DOUBLE(const int64_t ND1, DOUBLE *ra1, DOUBLE *dec1, DOUBLE
637637
this_cell_pair->min_dx, this_cell_pair->min_dy, this_cell_pair->min_dz,
638638
this_cell_pair->closest_x1, this_cell_pair->closest_y1, this_cell_pair->closest_z1,
639639
this_rpavg, npairs,
640-
this_weightavg, extra->weight_method, options->bin_type);
640+
this_weightavg, extra->weight_method, extra->pair_weight, options->bin_type);
641641

642642
/* This actually causes a race condition under OpenMP - but mostly
643643
I care that an error occurred - rather than the exact value of
@@ -771,6 +771,7 @@ int countpairs_mocks_DOUBLE(const int64_t ND1, DOUBLE *ra1, DOUBLE *dec1, DOUBLE
771771
// Keep in mind this is an autocorrelation (i.e. only one particle set to consider)
772772
weight_func_t_DOUBLE weight_func = get_weight_func_by_method_DOUBLE(extra->weight_method);
773773
pair_struct_DOUBLE pair = {.num_weights = extra->weights0.num_weights,
774+
.num_integer_weights = extra->weights0.num_integer_weights,
774775
.dx.d=0., .dy.d=0., .dz.d=0., // always 0 separation
775776
.parx.d=0., .pary.d=0., .parz.d=0.};
776777
for(int64_t j = 0; j < ND1; j++){

‎mocks/DDrppi_mocks/countpairs_rp_pi_mocks_impl.h.src

+2-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ extern "C" {
3030
const DOUBLE min_xdiff, const DOUBLE min_ydiff, const DOUBLE min_zdiff,
3131
const DOUBLE closest_icell_xpos, const DOUBLE closest_icell_ypos, const DOUBLE closest_icell_zpos,
3232
DOUBLE *src_rpavg, uint64_t *src_npairs,
33-
DOUBLE *src_weightavg, const weight_method_t weight_method, const bin_type_t bin_type);
33+
DOUBLE *src_weightavg, const weight_method_t weight_method, const pair_weight_struct pair_weight,
34+
const bin_type_t bin_type);
3435

3536
extern countpairs_mocks_func_ptr_DOUBLE countpairs_rp_pi_mocks_driver_DOUBLE(const struct config_options *options) __attribute__((warn_unused_result));
3637

0 commit comments

Comments
 (0)