# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# +
import numpy as np

list_SD36 = ['1.038550e+04','1.091930e+04','1.111270e+04','1.135200e+04','1.199560e+04'
            ,'1.201480e+04','1.253880e+04','1.286780e+04','1.321270e+04','1.351050e+04'
            ,'1.405970e+04','1.439510e+04','1.478630e+04','1.512690e+04','1.526000e+04'
            ,'1.568450e+04','1.614020e+04','1.619980e+04','1.654070e+04','1.691070e+04'
            ,'1.741890e+04','1.768180e+04','1.797620e+04','1.833130e+04','1.892010e+04'
            ,'1.912230e+04','1.963680e+04','1.978310e+04','7.623220e+03','8.133860e+03'
            ,'8.335080e+03','8.585430e+03','9.214770e+03','9.533420e+03','9.747990e+03'
            ,'9.931470e+03']


list_COV36 = ['1.038550e+04_ng1.988930e+01_nl3.422830e+01','1.091930e+04_ng2.087710e+01_nl3.616160e+01',
              '1.111270e+04_ng3.288210e+01_nl6.051840e+01','1.135200e+04_ng1.918510e+01_nl3.285780e+01',
              '1.199560e+04_ng1.265530e+01_nl2.050290e+01','1.201480e+04_ng1.000950e+01_nl1.571620e+01',
              '1.253880e+04_ng1.632340e+01_nl2.736000e+01','1.286780e+04_ng1.330660e+01_nl2.170300e+01',
              '1.321270e+04_ng2.489560e+01_nl4.414800e+01','1.351050e+04_ng2.252650e+01_nl3.941680e+01',
              '1.405970e+04_ng2.359790e+01_nl4.154870e+01','1.439510e+04_ng1.197870e+01_nl1.926470e+01',
              '1.478630e+04_ng3.339750e+01_nl6.159480e+01','1.512690e+04_ng1.038020e+01_nl1.637770e+01',
              '1.526000e+04_ng2.583270e+01_nl4.603640e+01','1.568450e+04_ng1.234240e+01_nl1.992920e+01',
              '1.614020e+04_ng2.703690e+01_nl4.847670e+01','1.619980e+04_ng1.863040e+01_nl3.178300e+01',
              '1.654070e+04_ng1.074150e+01_nl1.702530e+01','1.691070e+04_ng2.200120e+01_nl3.837660e+01',
              '1.741890e+04_ng9.072180e+00_nl1.405870e+01','1.768180e+04_ng3.536430e+01_nl6.572260e+01',
              '1.797620e+04_ng1.424180e+01_nl2.344000e+01','1.833130e+04_ng1.574630e+01_nl2.626600e+01',
              '1.892010e+04_ng1.517110e+01_nl2.518120e+01','1.912230e+04_ng1.743810e+01_nl2.948740e+01',
              '1.963680e+04_ng3.804120e+01_nl7.139000e+01','1.978310e+04_ng9.587190e+00_nl1.496670e+01',
              '7.623220e+03_ng1.399100e+01_nl2.297260e+01','8.133860e+03_ng3.687280e+01_nl6.890940e+01',
              '8.335080e+03_ng2.959040e+01_nl5.369840e+01','8.585430e+03_ng2.848750e+01_nl5.143540e+01',
              '9.214770e+03_ng2.692600e+01_nl4.825130e+01','9.533420e+03_ng3.146080e+01_nl5.756200e+01',
              '9.747990e+03_ng1.111050e+01_nl1.768990e+01','9.931470e+03_ng1.706900e+01_nl2.878090e+01']


dir_datav = "datav/"
dir_bary = "datav/"
dir_COV = "/users/timeifler/Dropbox/cosmolike_store/LSST_emu/npcov/"

dir_LPC = "baryons/"

class gen_LPCi():
    def __init__(self,id_SD):
        '''
            id_SD = 0~35, int
        '''
        self.id_SD = id_SD
        self.list_bary_scenario = ['_EAGLE','_Illustris','_TNG100','_MB2','_OWLS_AGN','_HzAGN']
        self.datav_dmo = self.load_datav(dir_datav,bary_scenario="",isCutted=1)
        print self.id_SD
        self.load_COV()
        print self.id_SD
        print self.datav_dmo.shape, self.COV_cut.shape
        
        self.cal_invL()
        self.Ndata = len(self.datav_dmo)
        self.Nscenario = len(self.list_bary_scenario)
        
        self.build_Ratio()
        self.build_Delta()
        self.build_DeltaChy()
        self.SVD(in_Delta=self.DeltaChy)
        self.write_LPC()
        
    
    def load_datav(self,dir_path,bary_scenario,isCutted=0):
        fname_datav = dir_path+'3x2pt'+bary_scenario+'_Y10_area'+list_SD36[self.id_SD]
        print fname_datav
        datav = np.genfromtxt(fname_datav,skip_header=0,dtype="int,double",usecols=[0,1],names=["ind","obs"])["obs"]
        
        self.Ndata_full  = len(datav)
        self.takeout_ID  = np.where(datav > 1e-20)[0]
        self.zero_ID     = np.where(datav < 1e-20)[0]
        
        if isCutted==1:
            return np.array(datav)[self.takeout_ID]
        else:
            return np.array(datav)
        
            
    def load_COV(self):
        
        self.fname_COV = dir_COV+'npcov_LSST_Y10_area'+list_COV36[self.id_SD]+'.npy'
        print self.fname_COV
        self.COV_full = np.load(self.fname_COV)
        self.COV_cut  = self.COV_full[self.takeout_ID][:,self.takeout_ID]
    
    
    def cal_invL(self):
        self.L    = np.linalg.cholesky(self.COV_cut)
        self.invL =  np.linalg.inv(self.L)
        
    def build_Ratio(self):
        self.Ratio     = np.zeros((self.Ndata,self.Nscenario))
        
        for j in range(self.Nscenario):
        
            datav_bary = self.load_datav(dir_bary,bary_scenario=self.list_bary_scenario[j],isCutted=1)
            self.Ratio.T[j] = datav_bary/self.datav_dmo
            
    def build_Delta(self):
        DeltaT = self.Ratio.T*self.datav_dmo-self.datav_dmo
        self.Delta = DeltaT.T
    
    def build_DeltaChy(self):
        self.DeltaChy = np.dot(self.invL,self.Delta)
        
    def SVD(self,in_Delta):
        # self.U stores PC modes
        # PC1 = self.U.T[0]  ; PC2 = self.U.T[1]
        self.U, self.Sdig, VT = np.linalg.svd(in_Delta,full_matrices=True)
    
    def gen_datav_cosmolike_format(self,datav_cut):
        
        datav_full = np.zeros(self.Ndata_full)
        datav_full[self.takeout_ID] = datav_cut
        
        return datav_full

    
    def write_LPC(self):
        
        self.fname_out = dir_LPC+'LPC_area'+list_SD36[self.id_SD]
        
        self.LPC1 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[0]))
        self.LPC2 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[1]))
        self.LPC3 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[2]))
        self.LPC4 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[3]))
        self.LPC5 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[4]))
        self.LPC6 = self.gen_datav_cosmolike_format(np.dot(self.L,self.U.T[5]))
        
        col_names='LPC1 LPC2 LPC3 LPC4 LPC5 LPC6'
        
        LPC_Matrix=np.column_stack((self.LPC1,self.LPC2,self.LPC3,self.LPC4,self.LPC5,self.LPC6))

        np.savetxt(self.fname_out,LPC_Matrix, fmt='%.8e %.8e %.8e %.8e %.8e %.8e')

       
# -

for x in range(36):
    gen_LPCi(x)



StarcoderdataJupyterScriptsDedupFiltered