Testing module with Corrfunc for halotools

test_wp_and_xi.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)
import numpy as np
from Corrfunc._countpairs import countpairs_wp as wp
from Corrfunc._countpairs import countpairs_xi as xi


def string_replace(string, old, new):
    from sys import version_info
    if version_info[0] < 3:
        return string.decode("utf-8").replace(old, new)
    else:
        return string.replace(old, new)


def main():
    data_files = ["sample1_position_array.npy", "sample2_position_array.npy"]
    first_three = [np.array([21.92387009, 24.82151985, 24.42320061]),
                   np.array([8.32623005, 23.28944016, 16.8574295])]
    last_three = [np.array([223.17524719, 238.28190613, 115.92366028]),
                  np.array([229.25822449,  233.86260986,  129.75932312])]

    rbins = np.load("rp_bins_array.npy")
    with open('binfile', 'w') as f:
        for i in range(len(rbins)-1):
            f.write("{0} {1}\n".format(rbins[i], rbins[i+1]))

    nthreads = 4
    boxsize = 250.0
    pimax = 40.0
            
    for first, last, f in zip(first_three, last_three, data_files):
        arr = np.load(f)
        assert np.allclose(arr[0, :], first)
        assert np.allclose(arr[-1, :], last)
        x1 = arr[:, 0]
        y1 = arr[:, 1]
        z1 = arr[:, 2]

        my_wp = wp(boxsize, pimax, nthreads, 'binfile', x1, y1, z1)
        # my_wp is a list of lists with form [[rmin, rmax, rpavg, wp, npairs]]
        # for each bin in the 'binfile'. item[-2] selects wp, and
        # item[-1] selects npairs.
        wp_array = np.array([[item[-2], np.int64(item[-1])] for item in my_wp])
        print("wp_array = \n{0}".format(wp_array))

        # create a consistent numpy saved filename from the input data file
        npy_save_fname = string_replace(f, "array", "array_wp")
        np.save(npy_save_fname, wp_array)
        
        my_xi = xi(boxsize, nthreads, 'binfile', x1, y1, z1)
        # my_xi is a list of lists with form [[rmin, rmax, ravg, xi, npairs]]
        # for each bin in the binfile. item[-2] selects xi, and
        # item[-1] selects npairs.
        xi_array = np.array([[item[-2], np.int64(item[-1])] for item in my_xi])
        print("xi_array = \n{0}".format(xi_array))

        # create a consistent numpy saved filename from the input data file
        npy_save_fname = string_replace(f, "array", "array_xi")
        np.save(npy_save_fname, xi_array)


if __name__ == '__main__':
    main()