generate_light_cone.py

#!/usr/bin/env python

# Generate lightcone for DESI mocks
# from HOD catalogues generated by Shadab

import sys
import os
import configparser
import multiprocessing as mp
from astropy.io import fits
import h5py

import camb
import numpy as np
import healpy as hp
import numexpr as ne
from astropy.table import Table
from rotation_matrix import RotationMatrix

ne.set_num_threads(4)

def tp2rd(tht, phi):
    """ convert theta,phi to ra/dec """
    ra  = phi / np.pi * 180.0
    dec = -1 * (tht / np.pi * 180.0 - 90.0)
    return ra, dec


class Paths():
    def __init__(self, config_file, args, in_part_path, input_name, out_part_path, output_name):
        config     = configparser.ConfigParser()
        config.read(config_file)

        self.dir_out                 = args.dir_out
        self.dir_in                  = args.dir_in

        ### Both input and output names have empty regions
        ### that need to be completed using ".format()" 
        self.input_name              = input_name
        self.output_name             = output_name

        self.in_part_path 			 = in_part_path		
        self.out_part_path 			 = out_part_path

        if self.dir_out is None:
            self.dir_out     =  config.get('dir', 'dir_out')
        if self.dir_in is None:
            self.dir_in      =  config.get('dir', 'dir_in')

        self.shells_out_path = self.create_outpath()

        self.input_file       = self.dir_in + self.in_part_path + self.input_name		
        self.output_file      = self.shells_out_path + self.output_name

    def create_outpath(self):
        out_path = self.dir_out + "/"+ self.out_part_path
        if not os.path.exists(out_path):
            os.makedirs(out_path)
        return out_path


class LightCone():
    def __init__(self, config_file, args, parallel_BOOL=False):
        config     = configparser.ConfigParser()
        config.read(config_file)

        self.file_camb      = config.get('dir', 'file_camb')
        self.box_length     = config.getint('sim', 'box_length')
        self.shellwidth     = config.getint('sim', 'shellwidth')
        self.zmin           = config.getfloat('sim', 'zmin')
        self.zmax           = config.getfloat('sim', 'zmax')
        self.rotate         = config.getboolean('sim', 'rotate')

        self.mock_random_ic = args.mock_random_ic
        if self.mock_random_ic is None:
            self.mock_random_ic = config.get('sim', 'mock_random_ic')

        self.origin  = [0, 0, 0]
        self.clight  = 299792458.

        self.h, self.results = self.run_camb()
        file_alist     =  config.get('dir','file_alist')
        self.alist = np.loadtxt(file_alist)

        rotation_matrix_instance = RotationMatrix(config_file, args)
        self.rotation_matrix = rotation_matrix_instance.rotation_matrix
        self.parallel_BOOL = parallel_BOOL

    def run_camb(self):
        #Load all parameters from camb file
        pars = camb.read_ini(self.file_camb)
        h    = pars.h
        pars.set_for_lmax(2000, lens_potential_accuracy=3)
        pars.set_matter_power(redshifts=[0.], kmax=200.0)
        pars.NonLinearModel.set_params(halofit_version='takahashi')
        camb.set_feedback_level(level=100)
        results   = camb.get_results(pars)
        return h, results


    def compute_shellnums(self):
        shellnum_min = int(self.results.comoving_radial_distance(self.zmin) * self.h // self.shellwidth)
        shellnum_max = int(self.results.comoving_radial_distance(self.zmax) * self.h // self.shellwidth + 1)
        shellnums = list(range(shellnum_min, shellnum_max+1))
        print(f"INFO: There are {len(shellnums)} shells.")
        return shellnums


    def checkslicehit(self, chilow, chihigh, xx, yy, zz):
        """ pre-select so that we're not loading non-intersecting blocks """
        box_length = self.box_length
        origin = self.origin
        bvx = np.array([0, box_length, box_length, 0,    0,    box_length, box_length, 0])
        bvy = np.array([0,    0, box_length, box_length, 0,    0,    box_length, box_length])
        bvz = np.array([0,    0,   0,  0,    box_length, box_length, box_length, box_length])

        boo = 0
        r   = np.zeros(8)
        for i in range(0, 8):
            sx   = (bvx[i] - origin[0] + box_length * xx)
            sy   = (bvy[i] - origin[1] + box_length * yy)
            sz   = (bvz[i] - origin[2] + box_length * zz)
            r[i] = np.sqrt(sx * sx + sy * sy + sz * sz)
        if chihigh < np.min(r):
            boo = boo + 1
        if chilow > np.max(r):
            boo = boo+1

        if boo == 0:
            return True
        else:
            return False


    def convert_xyz2rdz(self, data, prefix, chilow, chiupp):
        """ Generates and saves a single lightcone shell """
        clight = self.clight
        box_length = self.box_length
        origin = self.origin

        ntiles = int(np.ceil(chiupp / box_length))
        print(prefix + "tiling [%dx%dx%d]" % (2 * ntiles, 2 * ntiles, 2 * ntiles))
        print(prefix + 'Generating map for halos in the range [%3.f - %.3f Mpc/h]' % (chilow, chiupp))

        px    = data['x']
        py    = data['y']
        pz    =	data['z']
        ngalbox = len(px)

        if self.mock_random_ic == "mock":
            vx    = data['vx']
            vy    = data['vy']
            vz    = data['vz']
            mass  = data['mass']
            id_halo   = data['id']

        elif self.mock_random_ic == "random":
            id_ = data["id"]
        elif self.mock_random_ic == "ic":
            dens = data["density"]
        else:
            print(f"ERROR!!! generate_lc, convert_xyz2rdz: You should choose between: mock, random or ic. You have chosen {self.mock_random_ic}")

        #-------------------------------------------------------------------

        totra   = np.array([])
        totdec  = np.array([])
        totz    = np.array([])
        tot_aux = np.array([])
        tot_id_halo  = np.array([])
        tot_mass= np.array([])


        [axx, axy, axz, ayx, ayy, ayz, azx, azy, azz] = self.rotation_matrix

        for xx in range(-ntiles, ntiles):
            for yy in range(-ntiles, ntiles):
                for zz in range(-ntiles, ntiles):

                    if not self.rotate:
                        slicehit = self.checkslicehit(chilow, chiupp, xx, yy, zz)             # Check if box intersects with shell
                        if not slicehit:
                            continue

                    sx_0  = ne.evaluate("px - %d + box_length * xx" % origin[0])
                    sy_0  = ne.evaluate("py - %d + box_length * yy" % origin[1])
                    sz_0  = ne.evaluate("pz - %d + box_length * zz" % origin[2])

                    if self.rotate:
                        sx = ne.evaluate("axx * sx_0 + axy * sy_0 + axz * sz_0")
                        sy = ne.evaluate("ayx * sx_0 + ayy * sy_0 + ayz * sz_0")
                        sz = ne.evaluate("azx * sx_0 + azy * sy_0 + azz * sz_0")
                    else:
                        sx = sx_0
                        sy = sy_0
                        sz = sz_0

                    r   = ne.evaluate("sqrt(sx * sx + sy * sy + sz * sz)")
#                    zi  = self.results.redshift_at_comoving_radial_distance(r / self.h) # interpolated distance from position

                    idx = np.where((r > chilow) & (r < chiupp))[0]              # only select halos that are within the shell

                    if idx.size!=0:
                        ux = sx[idx] / r[idx]
                        uy = sy[idx] / r[idx]
                        uz = sz[idx] / r[idx]
                        zp = self.results.redshift_at_comoving_radial_distance(r[idx] / self.h)
#                        zp = zi[idx]

                        if self.mock_random_ic == "mock":
                            vx_0 = vx[idx]
                            vy_0 = vy[idx]
                            vz_0 = vz[idx]

                            if self.rotate:
                                vx_1 = ne.evaluate("axx * vx_0 + axy * vy_0 + axz * vz_0")
                                vy_1 = ne.evaluate("ayx * vx_0 + ayy * vy_0 + ayz * vz_0")
                                vz_1 = ne.evaluate("azx * vx_0 + azy * vy_0 + azz * vz_0")
                            else:
                                vx_1 = vx_0
                                vy_1 = vy_0
                                vz_1 = vz_0

                            qx = vx_1 * 1000.
                            qy = vy_1 * 1000.
                            qz = vz_1 * 1000.

                            vlos    = ne.evaluate("qx * ux + qy * uy + qz * uz")
                            dz      = ne.evaluate("(vlos / clight) * (1 + zp)")
                            tot_aux = np.append(tot_aux, zp + dz)
                            tot_mass= np.append(tot_mass, mass[idx])
                            tot_id_halo  = np.append(tot_id_halo, id_halo[idx])
                        elif self.mock_random_ic == "random":
                            tot_aux = np.append(tot_aux, id_[idx])
                        elif self.mock_random_ic == "ic":
                            tot_aux = np.append(tot_aux, dens[idx])

                        tht, phi = hp.vec2ang(np.c_[ux, uy, uz])
                        ra, dec  = tp2rd(tht, phi)

                        totra   = np.append(totra , ra)
                        totdec  = np.append(totdec, dec)
                        totz    = np.append(totz  , zp)

        return totra, totdec, totz, tot_aux, ngalbox, tot_mass, tot_id_halo


    def getnearestsnap(self, zmid):
        """ get the closest snapshot """
        # zsnap  = 1/self.alist[:,1]-1.
        zsnap  = 1 / self.alist[:, 2] - 1.
        index_ = np.argmin(np.abs(zsnap - zmid))

        return int(self.alist[index_, 0]), self.alist[index_, 3]


    def obtain_data(self, infile, prefix, correct1000=False):
        try:
            if infile.endswith('.fits'):
                print('Reading fits file')
                hdul = fits.open(infile, memmap=False)
                data = hdul[1].data
                hdul.close()
            else:
                print('Reading ecsv file')
                data = Table.read(infile, format='ascii.ecsv')
            if correct1000:
                data["x"] += 1000
                data["y"] += 1000
                data["z"] += 1000
            ngalbox = len(data["x"])
            print(prefix + f"INFO: The input file is {infile} and it has {ngalbox} halos")
        except IOError:
            print(prefix + f"WARNING: Couldn't open {infile}.", file=sys.stderr)
            os._exit(1)
        return data


    def generate_shell(self, infile, subbox, prefix, chilow, chiupp, return_dict):
        ### Read Data
        data = self.obtain_data(infile, prefix)

        ### Convert XYZ to RA DEC Z
        ra0, dec0, zz0, aux0, ngalbox, mass0, id_halo0 = self.convert_xyz2rdz(data, prefix, chilow, chiupp)


        shell_subbox_dict = {"ra0": ra0, "dec0": dec0, "zz0": zz0, "aux0": aux0, "mass0": mass0, "id_halo0":id_halo0}
        return_dict[subbox] = shell_subbox_dict

        return_dict["NGAL" + str(subbox)] = ngalbox


    def generate_shells(self, path_instance, snapshot=None, redshift=None, cutsky=True, nproc=5, n_subboxes=27, cat_seed=None):
        shellnums = self.compute_shellnums()
        for shellnum in shellnums:
            chilow = self.shellwidth * (shellnum + 0)
            chiupp = self.shellwidth * (shellnum + 1)
            chimid = 0.5 * (chilow + chiupp)

            # Check whether the minimum redshift of the shell is outside the redshift range of interest 
            zlow = self.results.redshift_at_comoving_radial_distance(chilow / self.h)
            if zlow > self.zmax:
                continue

            if not cutsky:
                print("Light-cone")
                zmid = self.results.redshift_at_comoving_radial_distance(chimid / self.h)
                nearestsnap, nearestred = self.getnearestsnap(zmid)

                snapshot = nearestsnap
                redshift = "z%.3f"%(nearestred)

            out_file_name = path_instance.output_file.format(snapshot=snapshot, shellnum=str(shellnum))

            # Don't reprocess files already done
            if os.path.isfile(out_file_name):
                continue

            if self.parallel_BOOL:
                jobs = []
                manager = mp.Manager()
                return_dict = manager.dict()
                counter = 0
                for subbox in range(n_subboxes):
                    infile = path_instance.input_file.format(redshift=redshift, subbox=subbox)
                    prefix = f"[shellnum={shellnum}; subbox={subbox}]: "

                #self.generate_shell(infile, subbox, prefix, chilow, chiupp, return_dict)


                    p = mp.Process(target=self.generate_shell, args=(infile, subbox, prefix, chilow, chiupp, return_dict))
                    jobs.append(p)
                    p.start()
                    counter = counter + 1
                    if (counter == nproc) or (subbox == n_subboxes - 1):
                        for proc in jobs:
                            proc.join()
                        counter = 0
                        jobs = []
            else:
                return_dict = {}
                for subbox in range(n_subboxes):
                    infile = path_instance.input_file.format(redshift=redshift, subbox=subbox)
                    prefix = f"[shellnum={shellnum}; subbox={subbox}]: "
                    self.generate_shell(infile, subbox, prefix, chilow, chiupp, return_dict)


            ### Count the number of galaxies per shell
            n_gal_shell_all_subboxes = 0
            for subbox in range(n_subboxes):
                shell_subbox_dict = return_dict[subbox]
                n_gal_shell_all_subboxes += len(shell_subbox_dict["ra0"])

            ### Declare arrays of size
            ra0_array = np.zeros(n_gal_shell_all_subboxes)
            dec0_array = np.zeros(n_gal_shell_all_subboxes)
            zz0_array = np.zeros(n_gal_shell_all_subboxes)

            aux0_array = np.zeros(n_gal_shell_all_subboxes)
            mass0_array = np.zeros(n_gal_shell_all_subboxes)
            id_halo0_array = np.zeros(n_gal_shell_all_subboxes)

            counter_ngal = 0

            ### Fill the arrays
            index_i = 0
            index_f = 0
            for subbox in range(n_subboxes):
                counter_ngal += return_dict["NGAL" + str(subbox)]
                shell_subbox_dict = return_dict[subbox]

                index_f = index_i + len(shell_subbox_dict["ra0"])

                ra0_array[index_i: index_f]  = shell_subbox_dict["ra0"]
                dec0_array[index_i: index_f] = shell_subbox_dict["dec0"]
                zz0_array[index_i: index_f]  = shell_subbox_dict["zz0"]
                aux0_array[index_i: index_f] = shell_subbox_dict["aux0"]
                mass0_array[index_i: index_f] = shell_subbox_dict["mass0"]
                id_halo0_array[index_i: index_f] = shell_subbox_dict["id_halo0"]

                index_i = index_f

            out_file_name_tmp  = out_file_name + "_tmp"

            with h5py.File(out_file_name_tmp, 'w-') as out_file:
                out_file.create_group('galaxy')
                out_file.create_dataset('galaxy/RA',      data=ra0_array,    dtype=np.float32)
                out_file.create_dataset('galaxy/DEC',     data=dec0_array,   dtype=np.float32)
                out_file.create_dataset('galaxy/Z_COSMO', data=zz0_array,     dtype=np.float32)

                if self.mock_random_ic == "mock":
                    out_file.create_dataset('galaxy/Z_RSD',   data=aux0_array, dtype=np.float32)
                    out_file.create_dataset('galaxy/MASS',    data=mass0_array, dtype=np.float32)
                    out_file.create_dataset('galaxy/ID',      data=id_halo0_array, dtype=np.int64)
                elif self.mock_random_ic == "random":
                    out_file.create_dataset('galaxy/ID',      data=aux0_array, dtype=np.int32)
                elif self.mock_random_ic == "ic":
                    out_file.create_dataset('galaxy/ONEplusDELTA', data=1 + aux0_array, dtype=np.float32)

                out_file.attrs['NGAL']     = counter_ngal
                out_file.attrs['SHELLNUM'] = shellnum
                out_file.attrs['SNAPSHOT'] = snapshot
                out_file.attrs['CAT_SEED'] = cat_seed
                out_file.attrs['BOX_LENGTH'] = self.box_length

            os.rename(out_file_name_tmp, out_file_name)