text:
'''
code in python3
Needs to be in the same directory as SONG  (/.../song/)
The python part of song call songy is written in python2. Make the following change:
    In /song/python/songy.py line 170:
         range(0,N*(N+1)/2) --> list(range(0,int(N*(N+1)/2)))
'''
####################################################################################### import
import numpy as np
import os
import sys
from numba import jit,prange
song_path='/home/thomas/song/'                  # set your one path to song
gev_path ='/home/thomas/song/gevolution-1.2/'   # set your one path to gevolution
sys.path.insert(0, song_path+'python')          # path to python module of song
import songy as s
import h5py
from subprocess import call
import matplotlib.pyplot as plt

#######################################################################################
####################################################################################### Parameters
h=0.67556
omega_b=0.022032/h**2
omega_cdm=0.12038/h**2
omega_k=0
H0= 100*h/3/10**5
z_song=[100,101] # second one for time derivative (velocity)
z=z_song[0]
H=100*np.sqrt((omega_cdm+omega_b)*(1+z))/3/10**5 #song provide a slightly different H
fnl=0
A_s = 2.215e-9
n_s = 0.9619
k_pivot = 0.05 # 1/Mpc

#######################################################################################
####################################################################################### SONG wrapper
def run_song(kmax,kmin,N,opt):
    ''' Call this function in song repository:
        It will create the ini and pre files from the global parameters and run song
        '''

    ini_file=r"""output = delta_cdm_bk
T_cmb = 2.7255
N_eff = 3.046
reio_parametrization = reio_none
tau_reio = 0.0952
k_pivot = 0.05
A_s = {}
n_s = {}
YHe = 0.2477055
gauge = newtonian
output_single_precision = yes
output_class_perturbations = yes
background_verbose = 1
thermodynamics_verbose = 1
primordial_verbose = 1
spectra_verbose = 1
nonlinear_verbose = 1
lensing_verbose = 1
output_verbose = 1
perturbations_verbose = 1
perturbations2_verbose = 2
transfer_verbose = 1
transfer2_verbose = 1
bessels_verbose = 1
bessels2_verbose = 1
bispectra_verbose = 1
fisher_verbose = 1
format = camb
write parameters = yes
h={}
omega_b={}
omega_cdm={}
Omega_k={}
primordial_local_fnl_phi={}
z_out={}"""

    pre_file=r"""sources2_k3_sampling = {}
k3_size = {}
k_min_tau0 = 0.05
k_max_tau0_over_l_max = 2.8
k_step_sub =0.1
k_logstep_super = 1.2
k_step_super = 0.025
k_step_transition = 0.2
quadsources_time_interpolation = cubic
sources_time_interpolation = linear
sources_k3_interpolation = cubic #linear
tau_start_evolution_song = 0
start_small_k_at_tau_c_over_tau_h_song = 0.001
start_large_k_at_tau_h_over_tau_k_song = 0.04
sources2_k_sampling = {}
k_min_custom_song = {}
k_max_custom_song = {}
k_size_custom_song = {}"""

    k_modulus_max=2*kmax

    ini="./m{}matter.ini".format(N)
    pre="./m{}matter.pre".format(N)

    file = open(ini, "w")
    if len(z_song)>1: file.write(ini_file.format(A_s,n_s,h,omega_b*h**2,omega_cdm*h**2,omega_k,fnl,str(z_song[0])+','+str(z_song[1])))
    else: file.write(ini_file.format(A_s,n_s,h,omega_b*h**2,omega_cdm*h**2,omega_k,fnl,z))
    file.close()

    file = open(pre, "w")
    file.write(pre_file.format(opt,N,opt,kmin,k_modulus_max,int(N*2)))
    file.close()

    os.system("./song "+ini+' '+pre)

    os.system("cp "+song_path+'output/sources_song_z000.dat '+song_path+"output/sources_song_z000_N{}.dat".format(N))
    os.system("cp "+song_path+'output/sources_song_z001.dat '+song_path+"output/sources_song_z001_N{}.dat".format(N))

def song_output(kmax,kmin,N,opt,filename='sources_song_z000.dat'):
    ''' Once song has run, this function load the output by using
        songy (see song/python/songy.py) routine FixedTauFile.
        It return the needed output:
            -song.get_source(b'delta_cdm') = song :second order kernel mulptiply by two transfer functions i.e.
                                                K(k1,k2,k3)*T_delta(k1)*T_delta(k2) in the expression
                                             int_k1_k2 (K(k1,k2,k3) T_delta(k1) T_delta(k2) zeta(k1) zeta(k2))
            -song.tau conformal time corresponding to the redshift. It is needed to get the velocity potential (dK/dtau)
            -song.k1, song.k2, song.k3: grie of mode
            -song.flatidx: Song output shape is weird ! see song/python/songy.py
            -dk12,dk3: step of the grid
        '''

    if not os.path.isfile(song_path+'output/'+filename):
        print(song_path+'output/sources_song_z000.dat not found')
        print('Run song ...')
        run_song(kmax,kmin,N,opt)
    song=s.FixedTauFile(song_path+'output/'+filename)

    if len(song.k1)!=2*N:
        print('The output '+song_path+'output/ found does not not has the right shape')
        print('Run song ...')
        run_song(kmax,kmin,N,opt)
        song=s.FixedTauFile(song_path+'output/'+filename)

    dk12=song.k1[1]-song.k1[0]
    dk3=np.diff(song.k3)[:,0]
    return np.array(song.get_source(b'delta_cdm')),song.tau,song.k1,np.array(song.k2,dtype=object),np.array(song.k3),song.flatidx,dk12,dk3

def song_main(kmax,kmin,N,opt):
    '''Main function for SONG '''

    s1="sources_song_z000_N{}.dat".format(N)
    s2="sources_song_z001_N{}.dat".format(N)
    source,tau,k1,k2,k3,flatidx,dk12,dk3=song_output(kmax,kmin,N,opt,s1)
    source1,tau1,_,_,_,_,_,_=song_output(kmax,kmin,N,opt,s2)
    return source,k1/h,k2/h,k3/h,flatidx,dk12/h,dk3/h,source1,tau1-tau

#######################################################################################
####################################################################################### first order transfert fct
def trans():
    '''
       This function returns the Primordial power spectrum, the transfer functions of delta_cdm and phi, and
       derivative of the last transfer function.

        -Primordial power spectrum: Primordial = A_s(k/k_0)**(ns-1) / (k**3/(2*np.pi**2)).
        -delta_cdm transfer function: tr_delta_cdm(k,z)*zeta(k)=delta_cdm(k,z)
        -potential transfer function: tr_phi(z,k)*zeta(k)=phi(z,k)
        '''
    import pylab as plt

    song=s.FixedTauFile(song_path+"output/sources_song_z000_N{}.dat".format(N))
    k=song.first_order_sources['k']/h

    Primordial= A_s*(k/(k_pivot/h))**(n_s-1)/k**3*2*np.pi**2

    tr_delta_cdm=song.first_order_sources[b'delta_cdm']
    tr_delta_b  =song.first_order_sources[b'delta_b']
    tr_phi= (omega_b/(omega_b+omega_cdm)*tr_delta_b + omega_cdm/(omega_b+omega_cdm)*tr_delta_cdm)*(-3*H**2/2)/(k**2+3*H**2)

    dk=np.diff(np.append(k,k[-1]*2-k[-2]))
    dT=np.diff(np.append(tr_phi,tr_phi[-1]*2-tr_phi[-2]))

    return np.array([k,tr_delta_cdm]),np.array([k,tr_phi]),np.array([k,dT/dk])

#######################################################################################
####################################################################################### First order stocastic potential
def zeta_realisation(k_grid):
    '''
        Generate the linear curvature perturbation field (N//2+1,N,N) at redshift z in half of Fourier space.
        The reality condition ensure the other half.
        The computation is in 3 steps:
            -compute the modulus of k in the grid (k)
            -interpolate transfer function and primordial power spectrum tr=T(k) and P=P(k)
            -randomly draw the real/imaginary part of the primordial curvature zeta following a Gaussian PDF with std=sqrt(P(k)/2)
        '''
    def random (k):
        P= A_s*(k/(k_pivot/h))**(n_s-1)/k**3*2*np.pi**2

        zeta_ini_Re=np.random.normal(0,N**3*np.sqrt(P/2*dk**3/(2*np.pi)**3),k.shape) # https://nms.kcl.ac.uk/eugene.lim/AdvCos/lecture2.pdf
        zeta_ini_Im=np.random.normal(0,N**3*np.sqrt(P/2*dk**3/(2*np.pi)**3),k.shape)

        # equivalent :
        #rho =  np.random.normal(0,N**3*np.sqrt(P*dk**3/(2*np.pi)**3),k.shape)
        #phase = np.random.uniform(0,2*np.pi,k.shape)
        #zeta_ini_Re=rho*np.cos(phase)
        #zeta_ini_Im=rho*np.sin(phase)
        return  zeta_ini_Re+zeta_ini_Im*1j

    zeta = np.zeros((N//2+1,N,N),dtype=np.complex64)

    # fill density ignoring the plan z=0 (zeta[0])
    k=np.sqrt(k_grid[0][N//2+1:]**2+k_grid[1]**2+k_grid[2]**2)
    zeta[1:,:,:]= random (k)
    # fill half of the plan z=0 ignoring the line (z=0,y=N//2)
    k=np.sqrt(k_grid[0][N//2+1:]**2+k_grid[1]**2)[:,:,0]
    zeta[0,N//2+1:]=random(k)
    # fill half of the line (z=0,y=N//2)
    k=k_grid[0][:,0,0][N//2+1:]
    zeta[0,N//2,N//2+1:]=random(k)

    # Even N in real space give a N+1 FFT grid with symmetries !
    zeta[1:-1,-1,1:-1]=zeta[1:-1,0,1:-1]  #z&x Plan
    zeta[1:-1,1:-1,-1]=zeta[1:-1,1:-1,0]  #z&y Plan

    # Zmax plan Surfaces
    zeta[-1,1:N//2,1:N//2]   =np.conjugate(zeta[-1,N//2+1:-1,N//2+1:-1][::-1,::-1])
    zeta[-1,N//2+1:-1,1:N//2]=np.conjugate(zeta[-1,1:N//2,N//2+1:-1][::-1,::-1])

    # Zmax plan lines X constant and Y constant
    zeta[-1,N//2,1:N//2]=np.conjugate(zeta[-1,N//2,N//2+1:-1][::-1])
    zeta[-1,1:N//2,N//2]=np.conjugate(zeta[-1,N//2+1:-1,N//2][::-1])

    r=zeta[:-1,-1,0] # All edges (x=0,y=0),(x=0,y=-1),(x=-1,y=0) and (x=-1,y=-1) are equal
    zeta[:-1,-1,-1],zeta[:-1,0,0],zeta[:-1,0,-1]=r,r,r

    r=zeta[-1,0,1:N//2]  # Zmax edges sym with Y constant
    zeta[-1,-1,1:N//2],zeta[-1,0,N//2+1:-1],zeta[-1,-1,N//2+1:-1]=r,np.conjugate(r[::-1]),np.conjugate(r[::-1])

    r=zeta[-1,1:N//2,0]# Zmax edges sym with X constant
    zeta[-1,1:N//2,-1],zeta[-1,N//2+1:-1,0],zeta[-1,N//2+1:-1,-1]=r,np.conjugate(r[::-1]),np.conjugate(r[::-1])

    r=zeta[-1,0,0].real    # Zmax plan corners all equal and real
    zeta[-1,0,0],zeta[-1,-1,0],zeta[-1,-1,-1],zeta[-1,0,-1]=r,r,r,r

    r=zeta[-1,N//2,0].real # Zmax plan: middle point of edges
    zeta[-1,N//2,0],zeta[-1,N//2,-1]=r,r
    r=zeta[-1,0,N//2].real
    zeta[-1,0,N//2],zeta[-1,-1,N//2]=r,r

    # Zmax middle point real
    zeta[-1,N//2,N//2]=zeta[-1,N//2,N//2].real

    # z=0 Plan

    zeta[0,N//2,-1]=zeta[0,N//2,-1].real

    zeta[0,-1,N//2]=zeta[0,-1,N//2].real

    zeta[0,N//2:-1,-1]=zeta[0,N//2:-1,0]
    zeta[0,-1,N//2+1:-1]=np.conjugate(zeta[0,-1,1:N//2][::-1])

    r=zeta[0,-1,0].real
    zeta[0,-1,0],zeta[0,-1,-1]=r,r

    zeta[0,:N//2]     =np.conjugate(zeta[0,N//2+1:][::-1,::-1])
    zeta[0,N//2,:N//2]=np.conjugate(zeta[0,N//2,N//2+1:][::-1])
    return zeta

#######################################################################################
####################################################################################### From initial potential to displacement field
def order1(k):
    '''Compute the first order quantities X1/delta1 at a given k.
       X being: potential phi1 (==psi1), displacement field xi1, velocity v1
       See equation (36) of the note.
        '''
    phi1=-3*H**2 / 2 / (3*H**2+k**2)
    xi1= (1-3*phi1)/k**2
    v1=-2/3/H*phi1
    return phi1,xi1,v1

def song2xi(song,k1,k2,k,flatidx):
    '''Compute the second order displacement field xi2:
       input:
            -song is the output of SONG: the second order density kernel
            -k1,k2,k are the SONG output mode modulus k1,k2,k3
            -flatidx is the output of SONG flatidx
       See note for the detailed computation.

            -k1dk2 is scalar product of k1 and k2
            -phi2p is the time derivative of the second order potential phi
        '''
    xi2  =np.copy(song)
    phi2 =np.copy(song)
    phi2p=np.copy(song)
    chi2 =np.copy(song)
    for ind1,kk1 in enumerate(k1):
        for ind2,kk2 in enumerate(k2[ind1]):

            kk3 = k3[flatidx[ind1,ind2]]

            phi1_k1,xi1_k1,v1_k1 = order1(kk1)
            phi1_k2,xi1_k2,v1_k2 = order1(kk2)
            d2 = song[flatidx[ind1,ind2]]

            k1dk2=(kk3**2-kk1**2-kk2**2)/2

            chi2[flatidx[ind1,ind2]]=3./2/kk3**4*( 2*(kk1**2+kk2**2)*k1dk2/3+k1dk2**2/3+kk1**2*kk2**2 )\
                 *(3*H**2*v1_k1*v1_k2+2*phi1_k1*phi1_k2)

            phi2p[flatidx[ind1,ind2]]=(-3*H**2*k1dk2*v1_k1*v1_k2-2*kk3**2*chi2[flatidx[ind1,ind2]]\
                                            +k1dk2*phi1_k1*phi1_k2)/21/H

            phi2[flatidx[ind1,ind2]] = (3*H**2/2/(3*H**2+kk3**2))*(-d2+2*chi2[flatidx[ind1,ind2]]+ \
                    (2+k1dk2/3/H**2-2./3/H**2*(kk1**2+kk2**2))*phi1_k1*phi1_k2\
                    -2*phi2p[flatidx[ind1,ind2]]/H )

            xi2[flatidx[ind1,ind2]]  = 1/kk3**2*(d2-3*phi2[flatidx[ind1,ind2]]-3./2*\
                    (kk2**2*phi1_k1*xi1_k2+kk1**2*phi1_k2*xi1_k1)\
                    +1./2*k1dk2*v1_k1*v1_k2-9./2*phi1_k1*phi1_k2+1./2*(kk1**2*kk2**2-k1dk2**2)*xi1_k1*xi1_k2)

    return xi2,phi2,phi2p,chi2

def zeta2fields(field,zeta,k_grid,tr_delta_cdm=0,tr_phi=0):
    '''Compute the whole first order stocastic field from the first order density
       computed in zeta_realisation().
       field can be 'phi','psi','xi','v','chi'.
        '''
    k=(np.sqrt(k_grid[0][N//2:]**2+k_grid[1]**2+k_grid[2]**2))

    tr_d=np.interp(k,tr_delta_cdm[0],tr_delta_cdm[1])
    tr_p=np.interp(k,tr_phi[0],tr_phi[1])

    if field=='delta':
        return zeta*tr_d
    elif field=='xi':
        xi=zeta*(tr_d-3*tr_p)/k**2
        xi[np.where(np.isnan(xi))]=0+0*1j
        return xi
    elif field=='v':
        return -zeta*2*tr_p/3/H
    elif field=='chi':
        return 0
    elif field=='phi':
        return tr_p*zeta
    else:
        return 'input field not understood'

def plot_staff():
    ## Phi transfer function comparision with CLASS
    #plt.figure()
    #claSS=np.loadtxt('../class_public/output/myclasstk.dat')
    #plt.loglog(claSS[:,0],np.abs(claSS[:,3]))
    #plt.loglog(k,np.abs(tr_delta_cdm))
    #plt.show()

    # Power spectrum comparison with Gevolution

    #plt.figure()
    #plt.loglog(k,Primordial*(omega_b/(omega_b+omega_cdm)*tr_delta_b + omega_cdm/(omega_b+omega_cdm)*tr_delta_cdm)**2)
    #claSS=np.loadtxt('../class_public/output/myclasspk.dat')
    #plt.loglog(claSS[:,0],np.abs(claSS[:,1]))
    ##gev=np.loadtxt('gevolution-1.2/output/lcdm_pk000_phi.dat')
    ##plt.loglog(gev[:,0],gev[:,1])
    #plt.show()

    #################################################################Pure class
    #claSS=np.loadtxt('../class_public/output/myclasstk.dat')
    #k=claSS[:,0]
    #Primordial= A_s*(k/(k_pivot/h))**(n_s-1)/k**3*2*np.pi**2
    #tr_delta_cdm=claSS[:,3]
    #tr_phi=claSS[:,7]

    fig, axs = plt.subplots(2)

    gev=np.loadtxt('gevolution-1.2/output/Pk_basic000_phi.dat')
    axs[0].loglog(gev[:,0],gev[:,1],label='basic')
    gev=np.loadtxt('gevolution-1.2/output/Pk_SONG000_phi.dat')
    axs[0].loglog(gev[:,0],gev[:,1],label='SONG')
    axs[0].loglog(k,Primordial*k**3/2/np.pi**2 *tr_phi**2,label='Analytic')
    axs[0].legend()
    axs[0].set_xlabel('k')
    axs[0].set_ylabel('Power spectra')
    #plt.show()

    P=np.interp(gev[:,0],k,Primordial*k**3/2/np.pi**2 *tr_phi**2)
    gev=np.loadtxt('gevolution-1.2/output/Pk_basic000_phi.dat')
    axs[1].loglog(gev[:,0],np.abs(gev[:,1]-P)/P,label='basic')
    gev=np.loadtxt('gevolution-1.2/output/Pk_SONG000_phi.dat')
    axs[1].loglog(gev[:,0],np.abs(gev[:,1]-P)/P,label='SONG')
    axs[1].legend()
    axs[1].set_xlabel('k')
    axs[1].set_ylabel('errors')
    plt.show()

def we(kz,ky,kx,di,p):
    """debuging staff"""
    import pylab as plt
    from matplotlib.colors import LogNorm
    TT  = di[kz[0]:kz[1],ky[0]:ky[1],kx[0]:kx[1]].ravel()
    vu=np.zeros_like(di)
    ind=0
    for i in TT:
        if i in di:
            w=np.where(i==di)
            vu[w]=i

    #for i in range(N//2+1):
    for i in [p]:
        plt.figure()
        plt.imshow(np.abs(vu[i].real))#,norm=LogNorm())
    #plt.imshow(vu)
    plt.show()
    return vu

#######################################################################################
####################################################################################### Kernels
Nopython=True
@jit(nopython=Nopython,fastmath=True)
def squeezed_term(k,trans,dTdk):
    '''
        Interpolate the squeezed term in the VR+ analytic approximation.
        This function is used by Kernel() for the term gamma. To get the VR approximation,
        this function has to return 0
        '''
    T=np.interp(k,trans[0],trans[1])
    dT=np.interp(k,dTdk[0],dTdk[1])
    return k/T*dT

@jit(nopython=Nopython,fastmath=True)
def Kernel_analytic(k1,k2,k3,trans,dTdk):
    '''
        Analytic kernel VR+ approximation following equation 5.1 of 1602.05933.
        The term dividedby is implicit in 1602.05933 from equation 3.7 (VR definition)
        '''
    Hk3=H/k3
    alpha=2./7+59./14.*Hk3**2+45./2*Hk3**4
    beta =1.-0.5*Hk3**2+54.*Hk3**4
    gamma=-3./2*Hk3**2+9./2.*Hk3**4
    gamma=gamma-5./4*(Hk3**2+3*Hk3**4)*squeezed_term(k3,trans,dTdk)
    k1k2=(k3**2-k1**2-k2**2)/2/k1/k2
    trk1=np.interp(k1,trans[0],trans[1])
    trk2=np.interp(k2,trans[0],trans[1])
    return 4*trk1*trk2*k1**2*k2**2*(beta-alpha+beta/2*k1k2*(k2/k1+k1/k2)+alpha*k1k2**2+gamma*(k1/k2-k2/k1)**2)/9/H**4

@jit(nopython=Nopython,fastmath=True)
def Kernel_song(kk1,kk2,kk3,K,flatidx):
    out=np.zeros((len(kk1)))
    count=0
    for ind in range(len(kk1)):
        if kk1[ind]==0 or kk2[ind]==0:
            continue
        pkk1,pkk2 = np.int(np.around((kk1[ind]-dk)/dk12)) , np.int(np.around((kk2[ind]-dk)/dk12))
        if kk3>k3[flatidx[max(pkk1,pkk2),min(pkk1,pkk2)]][-1]:
            pkk3=-1
        else:
            pkk3=np.int(np.around((kk3-k3[flatidx[max(pkk1,pkk2),min(pkk1,pkk2)]][0])/dk3[flatidx[max(pkk1,pkk2),min(pkk1,pkk2)]]))
        out[ind]=K[flatidx[max(pkk1,pkk2),min(pkk1,pkk2)]][pkk3]
    return out

#######################################################################################
####################################################################################### intergration
@jit(nopython=Nopython,fastmath=True,parallel=False)
def I2_TT_LL(Int,N,xidx,yidx,zidx,k_modulus,k_max,zeta,kern,kern_arg1,kern_arg2):
    '''
        Compute LL integral with cut-off k_max. Here k_max is an index defined in
        function integre(): k_max=int(np.where(klin==k_max)[0])-N//2.
        If k_max=N//2, this function return the raw integral.
        input:
            -Int: integral result array(N//2+1,N,N)
            -N: size of grid
            -xidx,yidx,zidx: index of the k-grid where the integral is computed. As for zeta computed in zeta_realisation,
                             this splitting allow this integral to be computed in different part (Volume+plan+line).
                             More details in function integre().
            -k_modulus: (half)whole grid modulus
            -k_max: cut-off index
            -zeta: linear density
            -kern,kern_arg1,kern_arg2: see integre() intput description

        Computation is split in two parts. For zeta(k1)*zeta(k-k1):

            -if 0<=k1_z<=k_z
                zeta(k1)=zeta_pos (pos=positive because then k_z-k1_z>=0).
                In this particular case: zeta(k-k1)=zeta_pos[::-1]
            -if k1_z>k_z
                zeta(k1)  =zeta_neg_k
                zeta(k-k1)=np.conjugate(zeta_neg_kmk) (kmk for k minus k1)
                    (Thanks to reality and because we only have in memory density for k_z>=0)

            a third part would be considered:
                -if k1_z<0:
                    let us call zeta_d_k and zeta_d_kmk the corresponding density
                    We can show that zeta_d_k=zeta_neg_kmk and zeta_neg_k=zeta_d_kmk

            Taking into account reality condition, we finally get:
            Int[kz,ky,kx]=(np.nansum(             zeta_pos     *zeta_pos[::-1]*kern_pos)\
                        +2*np.nansum(np.conjugate(zeta_neg_kmk)*zeta_neg_k    *kern_neg))*dk**3

        '''
    for iikx in prange(len(xidx)):
        kx=xidx[iikx]
        for iiky in prange(len(yidx)):
            ky=yidx[iiky]
            for iikz in prange(len(zidx)):
                kz=zidx[iikz]

                k_norm=k_modulus[kz,ky,kx]

                kz_min_lim =max(0,kz-k_max)
                kz_max_lim1=min(kz+1,k_max+1)

                kx_min_lim=max(N//2-k_max,kx-k_max)
                kx_max_lim=min(kx+k_max+1,N//2+k_max+1)

                ky_min_lim=max(N//2-k_max,ky-k_max)
                ky_max_lim=min(ky+k_max+1,N//2+k_max+1)

                zeta_pos=np.copy(zeta[kz_min_lim:kz_max_lim1,\
                                      ky_min_lim:ky_max_lim,\
                                      kx_min_lim:kx_max_lim])

                kk_pos_k=k_modulus[kz_min_lim:kz_max_lim1, ky_min_lim:ky_max_lim, kx_min_lim:kx_max_lim]

                #if kz_min_lim<=kz<kz_max_lim1 and ky_min_lim<=ky<ky_max_lim and kx_min_lim<=kx<kx_max_lim:
                #    zeta_pos [kz-kz_min_lim,ky-ky_min_lim,kx-kx_min_lim]=0

                kern_pos=kern(kk_pos_k.ravel(),kk_pos_k[::-1,::-1,::-1].ravel(),k_norm,kern_arg1,kern_arg2)
                I_pos=np.nansum(zeta_pos.ravel()*zeta_pos[::-1,::-1,::-1].ravel()*kern_pos)

                # if kz>k_max, the zeta_neg part vanishes
                if k_max-kz+1>0:
                    kz_max_lim2=k_max+1
                    zeta_neg_k=zeta[kz_max_lim1:kz_max_lim2,\
                                             ky_min_lim:ky_max_lim,\
                                             kx_min_lim:kx_max_lim].ravel()

                    kk_neg_k   =k_modulus   [kz_max_lim1:kz_max_lim2,\
                                             ky_min_lim:ky_max_lim,\
                                             kx_min_lim:kx_max_lim].ravel()

                    zeta_neg_kmk=zeta[1:k_max-kz+1,N-ky_max_lim:N-ky_min_lim,N-kx_max_lim:N-kx_min_lim].ravel()

                    kk_neg_kmk   =k_modulus   [1:k_max-kz+1,N-ky_max_lim:N-ky_min_lim,N-kx_max_lim:N-kx_min_lim].ravel()

                    kern_neg=kern(kk_neg_k,kk_neg_kmk\
                                         ,k_norm,kern_arg1,kern_arg2)

                    Int[kz,ky,kx]=I_pos+2*np.nansum(np.conjugate(zeta_neg_kmk)\
                                             *zeta_neg_k*kern_neg)
                else:
                    Int[kz,ky,kx]=I_pos

    return Int

@jit(nopython=Nopython,fastmath=True,parallel=False)
def compute(zeta,cc,kz,ky,kx,kmkz,kmky,kmkx,k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=True):

    zeta_k   = zeta[kz[0]:kz[1]    ,ky[0]:ky[1]    ,kx[0]:kx[1]].ravel()
    if reverse:
        kernel=kern(k_modulus[kz[0]:kz[1], ky[0]:ky[1],kx[0]:kx[1]].ravel()\
                     ,k_modulus[kmkz[0]:kmkz[1],kmky[0]:kmky[1],kmkx[0]:kmkx[1]][::-1,::-1,::-1].ravel()\
                     ,k_norm,kern_arg1,kern_arg2)
        zeta_kmk = zeta[kmkz[0]:kmkz[1],kmky[0]:kmky[1],kmkx[0]:kmkx[1]]\
                             [::-1,::-1,::-1].ravel()
    else:
        kernel=kern(k_modulus[kz[0]:kz[1], ky[0]:ky[1],kx[0]:kx[1]].ravel()\
                     ,k_modulus[kmkz[0]:kmkz[1],kmky[0]:kmky[1],kmkx[0]:kmkx[1]].ravel()\
                     ,k_norm,kern_arg1,kern_arg2)
        zeta_kmk = zeta[kmkz[0]:kmkz[1],kmky[0]:kmky[1],kmkx[0]:kmkx[1]].ravel()

    if cc==1:
        zeta_kmk=np.conjugate(zeta_kmk)
    elif cc==2:
        zeta_k=np.conjugate(zeta_k)
    elif cc==3:
        zeta_kmk=np.conjugate(zeta_kmk)
        zeta_k=np.conjugate(zeta_k)
    return np.nansum(zeta_k*zeta_kmk*kernel)

@jit(nopython=Nopython,fastmath=True,parallel=False)
def I2_SL(Int,N,xidx,yidx,zidx,k_modulus,k_max,zeta,kern,kern_arg1,kern_arg2):
    '''
        Compute SL integral, see function I2_TT_LL() for description of inputs

        Computation is split in three parts. For zeta(k1)*zeta(k-k1):

            -if 0<=k1_z<=k_z
                zeta(k1)=zeta_pos (pos=positive because then k_z-k1_z>=0).
                In this particular case: zeta(k-k1)=zeta_pos[::-1]
            -if k1_z>k_z
                zeta(k1)  =zeta_neg_k
                zeta(k-k1)=np.conjugate(zeta_neg_kmk) (kmk for k minus k1)
                    (Thanks to reality and because we only have in memory density for k_z>=0)
            -if k1_z<0:
                zeta(k1)  =np.conjugate(zeta_d_k )
                    (Thanks to reality and because we only have in memory density for k_z>=0)
                zeta(k-k1)=zeta_d_kmk

            Taking into account reality condition, we finally get:

            Int[kz,ky,kx]=(np.nansum(             zeta_pos_kmk[::-1]*zeta_pos_k            *kern_bis_pos)\
                          +np.nansum(np.conjugate(zeta_neg_kmk)     *zeta_neg_k            *kern_bis_neg)\
                          +np.nansum(             zeta_d_kmk        *np.conjugate(zeta_d_k)*kern_bis_d))*dk**3
        '''
    Nd2=N//2
    for iikx in prange(len(xidx)):
        print(iikx)
        kx=xidx[iikx]
        for iiky in prange(len(yidx)):
            ky=yidx[iiky]
            for iikz in prange(len(zidx)):
                kz=zidx[iikz]

                k_norm=k_modulus[kz,ky,kx]

                kz_min_lim=max(0,kz-k_max)
                kz_max_lim1=min(kz+k_max+1,Nd2+1)
                kz_max_lim2=kz+1

                kx_min_lim=max(0,kx-k_max)
                kx_max_lim=min(kx+k_max+1,N)
                ky_min_lim=max(0,ky-k_max)
                ky_max_lim=min(ky+k_max+1,N)

                kym1=-ky_min_lim+ky
                kym2= ky_max_lim-ky-1
                kxm1=-kx_min_lim+kx
                kxm2= kx_max_lim-kx-1

                ###########################################################################################################################
                Xpos,Xneg,Ypos,Yneg,Xd,Yd,XYpos,XYneg,XYd,Z,ZX,ZY,ZXY=0,0,0,0,0,0,0,0,0,0,0,0,0
                if kx-k_max<0 or kx+k_max>=N:

                    if kx-k_max<0:
                        kx_min,kx_max=N-(k_max-kx)-1,N-1
                        kmkx_min,kmkx_max=Nd2+kx+1,Nd2+k_max+1
                        kmkx_min_neg,kmkx_max_neg=Nd2-k_max,Nd2-kx

                    if kx+k_max>=N:
                        kx_min,kx_max=1,kx+k_max-N+2
                        kmkx_min,kmkx_max=Nd2-k_max,kx-Nd2
                        kmkx_min_neg,kmkx_max_neg=Nd2-kx+N,Nd2+k_max+1

                    Xpos=compute(zeta,0,[kz_min_lim,kz_max_lim2],[ky_min_lim,ky_max_lim],
                                   [kx_min,kx_max],[0,kz_max_lim2-kz_min_lim],[Nd2-kym2,Nd2+kym1+1],
                                   [kmkx_min,kmkx_max],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    Xneg=compute(zeta,1,[kz_max_lim2,kz_max_lim1],[ky_min_lim,ky_max_lim],
                                   [kx_min,kx_max],[1,-kz_max_lim2+kz_max_lim1+1],[Nd2-kym1,Nd2+kym2+1],
                                   [kmkx_min_neg,kmkx_max_neg],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                    if kz-k_max<0:
                        if kx<Nd2:
                            kx_min_kzneg,kx_max_kzneg=1,k_max-kx+1
                        else:
                            kx_min_kzneg,kx_max_kzneg=N-(kx+k_max-N)-2,N-1

                        Xd=compute(zeta,2,[1,k_max-kz+1],[N-ky_max_lim,N-ky_min_lim],
                                   [kx_min_kzneg,kx_max_kzneg],[kz_max_lim2-kz_min_lim,k_max+1],[Nd2-kym2,Nd2+kym1+1],
                                   [kmkx_min,kmkx_max],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                ###########################################################################################################################"
                if ky-k_max<0 or ky+k_max>=N:

                    if ky-k_max<0:
                        ky_min,ky_max=N-(k_max-ky)-1,N-1
                        kmky_min,kmky_max=Nd2+ky+1,Nd2+k_max+1
                        kmky_min_neg,kmky_max_neg=Nd2-k_max,Nd2-ky

                    if ky+k_max>=N:
                        ky_min,ky_max=1,ky+k_max-N+2
                        kmky_min,kmky_max=Nd2-k_max,ky-N//2
                        kmky_min_neg,kmky_max_neg=Nd2-ky+N,Nd2+k_max+1

                    Ypos=compute(zeta,0,[kz_min_lim,kz_max_lim2],[ky_min,ky_max],
                                   [kx_min_lim,kx_max_lim],[0,kz_max_lim2-kz_min_lim],[kmky_min,kmky_max],
                                   [Nd2-kxm2,Nd2+kxm1+1],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    Yneg=compute(zeta,1,[kz_max_lim2,kz_max_lim1],[ky_min,ky_max],
                                   [kx_min_lim,kx_max_lim],[1,-kz_max_lim2+kz_max_lim1+1],[kmky_min_neg,kmky_max_neg],
                                   [Nd2-kxm1,Nd2+kxm2+1],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                    if kz-k_max<0:
                        if ky<Nd2:
                            ky_min_kzneg,ky_max_kzneg=1,k_max-ky+1
                        else:
                            ky_min_kzneg,ky_max_kzneg=N-(ky+k_max-N)-2,N-1

                        Yd=compute(zeta,2,[1,k_max-kz+1],[ky_min_kzneg,ky_max_kzneg],
                                   [N-kx_max_lim,N-kx_min_lim],[kz_max_lim2-kz_min_lim,k_max+1],[kmky_min,kmky_max],
                                   [Nd2-kxm2,Nd2+kxm1+1],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                #########################################################################################################################
                if ky-k_max<0:
                    ky_min_corner,ky_max_corner=1,k_max-ky+1
                    kmky_min_corner,kmky_max_corner=Nd2-kym2,Nd2-kym2+k_max-ky
                if ky+k_max>=N:
                    ky_min_corner,ky_max_corner=N-(k_max+ky-N)-2,N-1
                    kmky_min_corner,kmky_max_corner=Nd2+kym1-(ky+k_max-N),Nd2+kym1+1
                if kx-k_max<0:
                    kx_min_corner,kx_max_corner=1,k_max-kx+1
                    kmkx_min_corner,kmkx_max_corner=Nd2-kxm2,Nd2-kxm2+k_max-kx
                if kx+k_max>=N:
                    kx_min_corner,kx_max_corner=N-(k_max+kx-N)-2,N-1
                    kmkx_min_corner,kmkx_max_corner=Nd2+kxm1-(kx+k_max-N),Nd2+kxm1+1

                if (kx-k_max<0 or kx+k_max>=N) and (ky-k_max<0 or ky+k_max>=N):
                    # x&y side pos
                    XYpos=compute(zeta,0,[kz_min_lim,kz_max_lim2],[N-ky_max_corner,N-ky_min_corner],
                                 [N-kx_max_corner,N-kx_min_corner],[0,kz_max_lim2-kz_min_lim],[kmky_min,kmky_max],
                                 [kmkx_min,kmkx_max],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    if kz<Nd2+1:
                        # x&y side neg
                        XYneg=compute(zeta,1,[kz_max_lim2,kz_max_lim1],[ky_min,ky_max],
                                 [kx_min,kx_max],[1,-kz_max_lim2+kz_max_lim1+1],[N-kmky_max,N-kmky_min],
                                 [N-kmkx_max,N-kmkx_min],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                    if kz-k_max<0:
                        # x&y side d
                        XYd=compute(zeta,2,[1,k_max-kz+1],[N-ky_max,N-ky_min],
                                 [N-kx_max,N-kx_min],[kz_max_lim2-kz_min_lim,k_max+1],[kmky_min,kmky_max],
                                 [kmkx_min,kmkx_max],k_modulus,kern,k_norm,kern_arg1,kern_arg2,reverse=False)

                if kz+k_max>Nd2:
                    ky_min,ky_max=N-(ky_max_lim),N-(ky_min_lim)
                    kx_min,kx_max=N-(kx_max_lim),N-(kx_min_lim)
                    kz_min,kz_max=N-kz-k_max-1,N//2

                    # pure z>N//2+1
                    Z=compute(zeta,3,[kz_min,kz_max],[ky_min,ky_max],
                                   [kx_min,kx_max],[1,kz+k_max-Nd2+1],[Nd2-kym1,Nd2+kym2+1],
                                   [Nd2-kxm1,Nd2+kxm2+1],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    if (kx-k_max<0 or kx+k_max>=N):
                        ZX=compute(zeta,3,[kz_min,kz_max],[ky_min,ky_max],
                                   [kx_min_corner,kx_max_corner],[1,kz+k_max-Nd2+1],[N-(Nd2+kym1+1),N-(Nd2-kym2)],
                                   [kmkx_min_corner,kmkx_max_corner],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    if (ky-k_max<0 or ky+k_max>=N):
                        ZY=compute(zeta,3,[kz_min,kz_max],[ky_min_corner,ky_max_corner],
                                   [kx_min,kx_max],[1,kz+k_max-N//2+1],[kmky_min_corner,kmky_max_corner],
                                   [N-(Nd2+kxm1+1),N-(Nd2-kxm2)],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                    if (kx-k_max<0 or kx+k_max>=N) and (ky-k_max<0 or ky+k_max>=N):
                        ZXY=compute(zeta,3,[kz_min,kz_max],[ky_min_corner,ky_max_corner],
                                   [kx_min_corner,kx_max_corner],[1,kz+k_max-Nd2+1],[kmky_min_corner,kmky_max_corner],
                                   [kmkx_min_corner,kmkx_max_corner],k_modulus,kern,k_norm,kern_arg1,kern_arg2)

                I_out=Xpos+Xneg+Ypos+Yneg+Xd+Yd+XYpos+XYneg+XYd+Z+ZX+ZY+ZXY

                zeta_pos_k=np.copy(zeta[kz_min_lim:kz_max_lim2,\
                                                 ky_min_lim:ky_max_lim,
                                                 kx_min_lim:kx_max_lim]).ravel()
                kk_pos_k=k_modulus[kz_min_lim:kz_max_lim2,\
                                   ky_min_lim:ky_max_lim,\
                                   kx_min_lim:kx_max_lim].ravel()

                zeta_neg_k=np.copy(zeta[kz_max_lim2:kz_max_lim1,ky_min_lim:ky_max_lim,kx_min_lim:kx_max_lim]).ravel()
                kk_neg_k  =k_modulus[kz_max_lim2:kz_max_lim1,ky_min_lim:ky_max_lim,kx_min_lim:kx_max_lim].ravel()

                kzlower= k_max+kz+1-kz_max_lim1

                zeta_pos_kmk =np.copy(zeta[:kz_max_lim2-kz_min_lim,\
                                                  Nd2-kym2:Nd2+kym1+1,\
                                                  Nd2-kxm2:Nd2+kxm1+1][::-1,::-1,::-1]).ravel()
                kk_pos_kmk    =        k_modulus   [:kz_max_lim2-kz_min_lim,\
                                                  Nd2-kym2:Nd2+kym1+1,\
                                                  Nd2-kxm2:Nd2+kxm1+1][::-1,::-1,::-1].ravel()
                zeta_neg_kmk =np.copy(zeta[1:k_max+1-kzlower,\
                                                  Nd2-kym1:Nd2+kym2+1,
                                                  Nd2-kxm1:Nd2+kxm2+1]).ravel()
                kk_neg_kmk    =        k_modulus   [1:k_max+1-kzlower,\
                                                  Nd2-kym1:Nd2+kym2+1,\
                                                  Nd2-kxm1:Nd2+kxm2+1].ravel()

                if kz-k_max<k_max+1:

                    if Nd2-2*k_max<=kx<Nd2+2*k_max+1 and Nd2-2*k_max<=ky<Nd2+2*k_max+1:

                        kxupper=min(Nd2+k_max-(N-kx_max_lim)+1,kx_max_lim-kx_min_lim)
                        kxlower=max(Nd2-k_max-(N-kx_max_lim),0)
                        kyupper=min(Nd2+k_max-(N-ky_max_lim)+1,ky_max_lim-ky_min_lim)
                        kylower=max(Nd2-k_max-(N-ky_max_lim),0)
                        kzupper=min(2*k_max-kz+1,k_max+1)

                        zeta_pos_k.reshape((kz_max_lim2-kz_min_lim,\
                                                ky_max_lim-ky_min_lim,\
                                                kx_max_lim-kx_min_lim))[:kzupper,kylower:kyupper,kxlower:kxupper]=0
                        if kz<k_max+1:
                            zeta_neg_k.reshape((kz_max_lim1-kz_max_lim2,ky_max_lim-ky_min_lim,kx_max_lim-kx_min_lim))\
                                            [:k_max-kz,kylower:kyupper,kxlower:kxupper]=0

                kern_bis_pos=kern(kk_pos_k,kk_pos_kmk,k_norm,kern_arg1,kern_arg2)

                if kz==Nd2:
                    Int[kz,ky,kx]=(np.nansum(zeta_pos_kmk*zeta_pos_k*kern_bis_pos))+I_out#*dk**3
                elif kz-k_max>=0:
                    kern_bis_neg=kern(kk_neg_k,kk_neg_kmk      ,k_norm,kern_arg1,kern_arg2)
                    Int[kz,ky,kx]=(np.nansum(zeta_pos_kmk*zeta_pos_k*kern_bis_pos)\
                       +np.nansum(np.conjugate(zeta_neg_kmk)*zeta_neg_k*kern_bis_neg))+I_out
                else:
                    kern_bis_neg=kern(kk_neg_k,kk_neg_kmk      ,k_norm,kern_arg1,kern_arg2)

                    zeta_d_k  =np.copy(zeta[1:k_max-kz+1,N-ky_max_lim:N-ky_min_lim,N-kx_max_lim:N-kx_min_lim].ravel())
                    zeta_d_kmk =np.copy(zeta[1:k_max-kz+1,Nd2-kym2:Nd2+kym1+1,Nd2-kxm2:Nd2+kxm1+1]).ravel()
                    kk_d_k  =k_modulus[1:k_max-kz+1,N-ky_max_lim:N-ky_min_lim,N-kx_max_lim:N-kx_min_lim].ravel()
                    kk_d_kmk=k_modulus[1:k_max-kz+1,Nd2-kym2:Nd2+kym1+1,Nd2-kxm2:Nd2+kxm1+1].ravel()
                    kern_bis_d=kern(kk_d_k,kk_d_kmk,k_norm,kern_arg1,kern_arg2)

                    if Nd2-2*k_max<=kx<Nd2+2*k_max+1 and Nd2-2*k_max<=ky<Nd2+2*k_max+1:
                        kxupper=min(Nd2+k_max-(N-kx_max_lim)+1,kx_max_lim-kx_min_lim)
                        kxlower=max(Nd2-k_max-(N-kx_max_lim),0)
                        kyupper=min(Nd2+k_max-(N-ky_max_lim)+1,ky_max_lim-ky_min_lim)
                        kylower=max(Nd2-k_max-(N-ky_max_lim),0)
                        kzupper=min(2*k_max-kz+1,k_max+1)

                        zeta_d_k.reshape((k_max-kz,ky_max_lim-ky_min_lim,kx_max_lim-kx_min_lim))\
                                               [:kzupper,kylower:kyupper,kxlower:kxupper]=0

                    Int[kz,ky,kx]=(np.nansum(zeta_pos_kmk*zeta_pos_k*kern_bis_pos)\
                              +np.nansum(np.conjugate(zeta_neg_kmk)*zeta_neg_k*kern_bis_neg)\
                              +np.nansum(zeta_d_kmk*np.conjugate(zeta_d_k)*kern_bis_d))+I_out
    return Int

def integre(which,klin,N,k_max,klin_grid,zeta,kern,kern_arg1,kern_arg2):
    '''
        inputs:
        -which: whether it is 'TT','LL', 'SL' or 'SS' to be computed. The 'SS' is not implemented yet
        -klin: one dimensional list of k coordinate to be considered
        -N: size of the grid
        -k_max: cut-off in h/Mpc^-1
        -klin_grid: sparse grid of kx,ky,kz
        -zeta: Primordial curvature
        -kern: which kernel function to use. Either Kernel_song or Kernel_analytic
                if Kernel_song:
                    kern_arg1 has to be the SONG output song and kern_arg2 has to be flatidx
                if Kernel_analytic:
                    kern_arg1 and kern_arg2 have top be the two last output of function trans():
                                tr_pi=potential transfert function and dTdk its derivativ wrt k.
                                It will be interpolated in squeezed_term()

        Computation is split in three part in the space like in zeta_realisation()
        in order to compute only independent quantity regarding reality conditions

        -kz>=1
        -kz=0,ky>N//2
        -kz=0,ky=N//2,kx>N//2
        '''
    Int=np.zeros(zeta.shape,dtype=np.complex64)
    k_modulus=(np.sqrt(klin_grid[0][N//2:]**2+klin_grid[1]**2+klin_grid[2]**2))
    k_max=int(np.where(klin==k_max)[0])-N//2
    print('k_max={}'.format(k_max))
    klin=np.arange(len(klin))-N//2

    if which=='SL':
        print('integre SL ...')
        for opt in range(3):
            if opt==0:
                xidx=np.arange(N)
                yidx=np.arange(N)
                zidx=np.arange(1,N//2+1)
                #zidx=np.arange(N)
            elif opt==1:
                xidx=np.arange(N)
                yidx=np.arange(N//2+1,N)
                zidx=np.arange(1)
            elif opt==2:
                xidx=np.arange(N//2+1,N)
                yidx=np.array([N//2])
                zidx=np.arange(1)

            I2_SL(Int,N,xidx,yidx,zidx,k_modulus,k_max,zeta,kern,kern_arg1,kern_arg2)

    else:
        if which=='TT' or which=='LL':
            print('integre LL or TT ...')
            upper=k_max*2+1

            xidx1=np.where(np.abs(klin)<=upper)[0]
            yidx1=np.where(np.abs(klin)<=upper)[0]
            zidx1=np.where(np.logical_and(klin[N//2:]<=upper,klin[N//2:]>10**-15))[0]

            xidx2=np.where(np.abs(klin)<=upper)[0]
            yidx2=np.where(np.logical_and(np.abs(klin)<=upper,klin>0))[0]
            zidx2=np.arange(1)

            xidx3=np.where(np.logical_and(np.abs(klin)<=upper,klin>0))[0]
            yidx3=np.array([N//2])
            zidx3=np.arange(1)
        elif which=='SS':
            print('not yet')
            yidx1=np.arange(N)
            zidx1=np.arange(1,N//2+1)

            xidx2=np.arange(N)
            yidx2=np.arange(N//2+1,N)
            zidx2=np.array([0])

            xidx3=np.arange(N//2+1,N)
            yidx3=np.array([N//2])
            zidx3=np.array([0])
            return 0

        for opt in range(3):
            if opt==0:
                xidx=xidx1
                yidx=yidx1
                zidx=zidx1
            elif opt==1:
                xidx=xidx2
                yidx=yidx2
                zidx=zidx2
            elif opt==2:
                xidx=xidx3
                yidx=yidx3
                zidx=zidx3

            I2_TT_LL(Int,N,xidx,yidx,zidx,k_modulus,k_max,zeta,kern,kern_arg1,kern_arg2)
    return Int*dk**3/2**3/np.pi**3

#######################################################################################
####################################################################################### mode grid
def k_distrib(k_min,N,klbd,absolute=True):
    ''' Inputs:
            -k_min: Minimum mode to be consider. Setting k_min automatically set the step dk=k_min
             because in order for k-k1 to be always on the grid k1, we need to include 0 and to have a
             constant step dk.
            -N size of the grid. In order to include 0. If it is not odd, we set N+=1
                                            (the final fft return the right even N grid)
            -klbd: k_lambda:
                if absolute==True:
                    the function will return the closest in the grid
                else:
                    klbd is considered as being a ratio, return kL=k[N//2:][int(klbd*N//2)]
        output:
        klin_concat,kmax,N,dk,klambda
            -k: list of k coordinate
            -kmax: largest mode to be considered
            -N like input
            -k_min in float32
            -kL: actual k_lambda

        '''
    if N%2==0:
        print('N has to be odd to include 0: N+=1')
        N+=1
    k=np.linspace(-(N//2)*k_min,N//2*k_min,N,dtype=np.float32)
    if absolute:
        idxL=np.where(np.abs(klbd-k[N//2:])==np.min(np.abs(klbd-k[N//2:])))[0]
        kL=k[N//2:][idxL][0]
    else:
        kL=k[N//2:][int(klbd*N//2)]
    print('klambda='+str(kL))
    return k,np.float32(N//2*k_min),N,np.float32(k_min),kL

def fft(field):
    '''This function performs the inverse Fourier transform. It uses the numpy function irfftn.
       The input array has first to be re-organized.
       In this code, the array filed is organized this way field=(z=0:Nyquist,y=-Nyquist:0:Nyquist,x=-Nyquist:0:Nyquist)
       which means shape(field)=(N//2+1,N,N) (Reminder: in the code, N is always odd while N_input is even, N=N_input+1).
       The python modules takes as an input an array organized as follow:
            field=(x=0:Nyquist-1:-1:-Nyquist, y=0:Nyquist-1:-1:-Nyquist, z=0:Nyquist) which means shape(field)=(N//2+1,N-1,N-1)
            Note that -Nyquist=+Nyquist since N_input is even.
        '''
    new_field=np.zeros((N//2+1,N-1,N-1),dtype=np.complex64)

    field[0,N//2,:N//2]=np.conjugate(field[0,N//2,N//2+1:][::-1])
    field[0,:N//2,:]   =np.conjugate(field[0,N//2+1:,:][::-1,::-1])

    new_field[:,N//2+1:,N//2+1:]=field[:,1:N//2,1:N//2]
    new_field[:,:N//2+1,:N//2+1]=field[:,N//2:,N//2:]

    new_field[:,:N//2+1,N//2+1:]=field[:,N//2:,1:N//2]
    new_field[:,N//2+1:,:N//2+1]=field[:,1:N//2,N//2:]

    return np.fft.irfftn(new_field.transpose(),(N-1,N-1,N-1))

#######################################################################################
####################################################################################### Gevolution place

def settings(kmin,setting_name='msetting.ini',ICg='SONG',disp_file='displacement.h5',vel_file='velocitypotential.h5',pot_file='phi.h5'):
    '''create a setting file with the global parameters used in the programme'''
    setx=r"""
template file = sc1_crystal.dat
Tk file = class_tk.dat
baryon treatment = ignore
seed = 42
correct displacement = yes
k-domain = cube
Courant factor      = 48.0
T_cmb       = 2.7255
N_ur        = 3.046
time step limit     = 0.04
gravity theory      = GR
vector method       = parabolic
output path         = output/
generic file base   = lcdm
Pk outputs          = phi
Pk bins             = 1024
output              = mPk
gauge               = Newtonian
P_k_ini type        = analytic_Pk
P_k_max_h/Mpc       = 192
z_pk                = 100
root                = class_
background_verbose  = 1
spectra_verbose     = 1
output_verbose      = 1

Pk redshifts        = {}
Pk file base        = {}
IC generator      = {}
displacement file = {}
velocity file     = {}
metric file       = {}
tiling factor =     {}
k_pivot =           {}
A_s =               {}
n_s =               {}
h           =       {}
omega_b     =       {}
omega_cdm   =       {}
initial redshift    = {}
boxsize             = {}
Ngrid               = {}"""

    file = open(gev_path+setting_name, "w")
    file.write(setx.format(z,'Pk_'+ICg,ICg,disp_file,vel_file,pot_file,N//4,k_pivot,A_s,n_s,h,omega_b*h**2,omega_cdm*h**2\
                            ,z,2*np.pi/kmin,N-1))
    file.close()

##########################################################################################################################
##########################################################################################################################
########################################################################################################################## Main

N=128                                                 # Number of mode to be considered (on laptop N=10 is quick enough)
print('N='+str(N))                                   #
kmin=np.float32(0.01)                                # The non-zero smallest mode
dk=kmin                                              # which has to be equal to the step of the grid
klin_concat,kmax,N,dk,klambda=k_distrib(kmin,N,0.1) # Generate the list of mode coordinate (assume h/Mpc)

                                                     ##########################
song,k1,k2,k3,flatidx,dk12,dk3,song1,d_eta=song_main(kmax*h,kmin*h,N,'lin')  # Get song outputs
                                                                             #
transcdm,transphi,dTdk=trans()                                               # Get transfer functions and primordial power spectrum

k_grid_lin=np.array(np.meshgrid(klin_concat,klin_concat,klin_concat\
        ,sparse=True,indexing='ij'),dtype=object)                            # mode Grid with 'ij' indexing

#TT =integre('TT',klin_concat,N,kmax   ,k_grid_lin,phi,transphi,dTdk)                 # total integrale
#LL =integre('LL',klin_concat,N,klambda,k_grid_lin,phi,Kernel_analytic,transphi,dTdk) # LL analytic approx int
#SL =integre('SL',klin_concat,N,klambda,k_grid_lin,phi,Kernel_analytic,transphi,dTdk) # SL analytic approx int

#####debug
#k_modulus=(np.sqrt(k_grid_lin[0][N//2:]**2+k_grid_lin[1]**2+k_grid_lin[2]**2))
#xi2,phi2,_,_=song2xi(song,k1,k2,k3,flatidx)
#G=np.ones(np.shape(k_modulus))
#SL=integre('SL',klin_concat,N,klambda,k_grid_lin,G,Kernel_song,xi2,flatidx)
#LL=integre('LL',klin_concat,N,klambda,k_grid_lin,G,Kernel_song,xi2,flatidx)
#####

#################################################################################### displacement field computation
print('delta:')                                                                    #
for realisation in range(5):
    zeta=zeta_realisation(k_grid_lin)                                              # first order stocastic density

    delta_LL=integre('LL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,song,flatidx)  # Integration
    delta_SL=integre('SL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,song,flatidx)  #
    delta1 = zeta2fields('delta',zeta,k_grid_lin,transcdm,transphi)                        # first order displacement field
                                                                                #
    hf = h5py.File(gev_path+'density1_{}.h5'.format(realisation), 'w')                     # Save in h5 format
    hf.create_dataset('data', data=fft(delta1))                                            #
    hf.close()                                                                             #
                                                                                #
    hf = h5py.File(gev_path+'density2_{}.h5'.format(realisation), 'w')                     # Save in h5 format
    hf.create_dataset('data', data=fft(delta_LL+2*delta_SL))                               #
    hf.close()

##################################################################################### displacement field computation
#print('displacement:')                                                             #
##xi2,phi2,_,_=song2xi(song,k1,k2,k3,flatidx)                                        # Compute the second order kernels xi2 and phi2
##xi_LL=integre('LL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,xi2,flatidx)  # Integration
##xi_SL=integre('SL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,xi2,flatidx)  #
#xi1    = fft(zeta2fields('xi',zeta,k_grid_lin,transcdm,transphi))                       # first order displacement field
#                                                                                #
#hf = h5py.File(gev_path+'displacement.h5', 'w')                                    # Save in h5 format
#hf.create_dataset('data', data=xi1)                                                #
#hf.close()                                                                         #
#
#print('potential:')                                                                #
##phi_LL=integre('LL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,phi2,flatidx)# Integration
##phi_SL=integre('SL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,phi2,flatidx)#
#                                                                                #
#phi1 = fft(zeta2fields('phi',zeta,k_grid_lin,transcdm,transphi))                   # first order displacement field
#hf = h5py.File(gev_path+'phi.h5', 'w')                                             # Save in h5 format
#hf.create_dataset('data', data=phi1)                                               #
#hf.close()                                                                         #
#
################################################################################# scalar velocity computation
#print('Scalar velocity:')                                                       #
##xi2_1,_,_,_=song2xi(song1,k1,k2,k3,flatidx)                                    # Compute the second order kernels velocity v2
##v2 = (xi2_1-xi2)/d_eta                                                         #
##v_LL=integre('LL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,v2,flatidx)# Integration
##v_SL=integre('SL',klin_concat,N,klambda,k_grid_lin,zeta,Kernel_song,v2,flatidx)#
#                                                                                #
#v1  =fft(zeta2fields('v',zeta,k_grid_lin,transcdm,transphi))                    # first order velocity
#hf = h5py.File(gev_path+'velocitypotential.h5', 'w')                            # Save in h5 format
#hf.create_dataset('data', data=v1)                                              #
#hf.close()                                                                      #
#
###################################################################################### Gevolution
#settings(kmin,'msetting.ini')
#settings(kmin,'gsetting.ini',ICg='basic',disp_file='',vel_file='',pot_file='')

reason: import numpy,from numba


text:
# -*- coding: utf-8 -*-

# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""These are some tool function."""
import logging
import numpy as np
from .ops import constant_values, compute_funcs

def get_upstreams(dag, node):
    """Return the upstream nodes of this node."""
    nodes = dag.nodes
    res = []
    for key, value in nodes.items():
        if node in value:
            res.append(key)
    return res

def dag2compute(dag, input_data):
    """Calculate the output according to the dag."""
    wait_cal_list = []
    has_cal_list = []
    res = {}
    all_nodes = dag.topological_sort()
    for node in all_nodes:
        values = {"status": 0, "data": None}
        res[node] = values

    logging.debug("The calculation nodes is :{}.".format(dag.nodes))
    for node in all_nodes:
        if node.split('-')[0] == 'in':
            cal_res = input_data
            res[node]['data'] = cal_res
            has_cal_list.append(node)
        elif node.startswith('const'):
            constant_type = node.split('-')[0]
            cal_res = constant_values[constant_type]
            res[node]['data'] = cal_res
            has_cal_list.append(node)
        elif node == 'out':
            upstream_node = get_upstreams(dag, node)[0]
            cal_res = res[upstream_node]['data']
            res[node]['data'] = cal_res
            has_cal_list.append(node)
        else:
            upstream_nodes = get_upstreams(dag, node)
            upstream_done = True
            inputs = []
            for upstream_node in upstream_nodes:
                if res[upstream_node]['data'] is None:
                    upstream_done = False
                    wait_cal_list.append(dag.next_nodes(node=upstream_node))
                else:
                    inputs.append(res[upstream_node]['data'])

            if upstream_done:
                node_type = node.split('-')[0]
                if len(inputs) == 1:
                    if node_type in ['rec'] and not isinstance(inputs[0], np.ndarray) and inputs[0] == 0:
                        logging.debug("To avoid zero div, y will be added a smller number.")
                        inputs[0] = inputs[0] + 1e-15
                    cal_res = compute_funcs[node_type](inputs[0])
                elif len(inputs) == 2:
                    if node_type in ['div'] and not isinstance(inputs[1], np.ndarray) and inputs[1] == 0:
                        logging.debug("To avoid zero div, y will be added a smller number.")
                        inputs[1] = inputs[1] + 1e-15
                    cal_res = compute_funcs[node_type](inputs[0], inputs[1])
                else:
                    raise ValueError("The op {} only support one or two inputs, but got {}.".format(node, len(inputs)))
                res[node]['data'] = cal_res
                has_cal_list.append(node)

    return res['out']['data']

def cal_mish(x):
    """Calculate mish."""
    return x * np.tanh(np.log(np.exp(x) + 1))

def cal_gelu(x):
    """Calculate gelu."""
    w = np.sqrt(2 / np.pi)
    return 0.5 * x * (1 + np.tanh(w * (x + 0.044715 * x * x * x)))

def cal_tanh(x):
    """Calculate tanh."""
    return np.tanh(x)

def cal_sqrt(x):
    """Calculate sqrt."""
    return np.sqrt(x)

def cal_softplus(x):
    """Calculate softplus."""
    return np.log(np.exp(x) + 1)

def cal_error_threshold(a, b):
    """Calculate error threshold."""
    threshold = np.abs(a - b) / (np.abs(b) + 1)
    return threshold.max()

def is_close(arr1, arr2):
    """Check if two array is close or not."""
    return isinstance(arr1, np.ndarray) and len(arr1) == len(arr2) and np.allclose(arr1, arr2, atol=1, rtol=1)

def is_close_l2(arr1, arr2):
    """Check if two array is close or not."""
    return len(arr1) == len(arr2) and np.allclose(arr1, arr2, atol=1e-4, rtol=1e-4)

def cal_fitness(output, real_value):
    """Calculate the relative error of two array."""
    diff = output - real_value
    r_error = np.abs(diff) / np.abs(real_value)
    mean_error = np.mean(r_error)
    return mean_error

ext: py

max_stars_repo_name: zjzh/vega

max_issues_repo_name: zjzh/vega

max_forks_repo_name: zjzh/vega

reason: import numpy


text:
"""
Functions that take in command line arguments and build the model.
"""
import json
import numpy as np
from holoprot.models.affinity_pred import AffinityPred
from holoprot.models.prot_classifier import ProtClassifier
from holoprot.feat import (AMINO_ACIDS, SECONDARY_STRUCTS, ATOM_FDIM,
        BOND_FDIM, CONTACT_FDIM, SURFACE_NODE_FDIM, SURFACE_EDGE_FDIM,
        PATCH_NODE_FDIM, PATCH_EDGE_FDIM)
from holoprot.utils.metrics import EVAL_METRICS, DATASET_METRICS, METRICS

MODEL_CLASSES = {
    'pdbbind': AffinityPred,
    'reacbio': AffinityPred,
    'scope': ProtClassifier,
    'enzyme': ProtClassifier
}

def build_mpn_config(config):
    mpn_config = {}
    if config['prot_mode'] == 'backbone':
        mpn_config['node_fdim'] = len(AMINO_ACIDS) + len(SECONDARY_STRUCTS) + 2
        mpn_config['edge_fdim'] = CONTACT_FDIM
        mpn_config['hsize'] = config['bhsize']
        mpn_config['depth'] = config['bdepth']
        mpn_config['dropout_p'] = config['dropout_mpn']
        mpn_config['activation'] = config.get("activation", 'relu') #default
        mpn_config['jk_pool'] = config.get("jk_pool", None) #default

    elif config['prot_mode'] == 'surface':
        mpn_config['node_fdim'] = SURFACE_NODE_FDIM
        mpn_config['edge_fdim'] = SURFACE_EDGE_FDIM + CONTACT_FDIM
        mpn_config['hsize'] = config['shsize']
        mpn_config['depth'] = config['sdepth']
        mpn_config['dropout_p'] = config['dropout_mpn']
        mpn_config['activation'] = config.get("activation", "relu")
        mpn_config['jk_pool'] = config.get("jk_pool", None) #default

    else:
        bconfig, sconfig = {}, {}
        bconfig['node_fdim'] = len(AMINO_ACIDS) + len(SECONDARY_STRUCTS) + 2
        bconfig['edge_fdim'] = CONTACT_FDIM
        bconfig['hsize'] = config['bhsize']
        bconfig['depth'] = config['bdepth']
        bconfig['dropout_p'] = config['dropout_mpn']
        bconfig['activation'] = config.get('activation', 'relu')
        bconfig['jk_pool'] = config.get("jk_pool", None) #default

        sconfig['hsize'] = config['shsize']
        sconfig['depth'] = config['sdepth']
        sconfig['dropout_p'] = config['dropout_mpn']
        sconfig['activation'] = config.get('activation', 'relu')
        sconfig['jk_pool'] = config.get("jk_pool", None) #default

        if config['prot_mode'] == 'surface2backbone' or config['prot_mode'] == 'backbone2surface':
            sconfig['node_fdim'] = SURFACE_NODE_FDIM
            sconfig['edge_fdim'] = SURFACE_EDGE_FDIM + CONTACT_FDIM
        elif config['prot_mode'] == 'backbone2patch' or config['prot_mode'] == 'patch2backbone':
            sconfig['node_fdim'] = PATCH_NODE_FDIM
            sconfig['edge_fdim'] = PATCH_EDGE_FDIM

        if config.get("use_mpn_in_patch", False):
            sp_config = {}
            sp_config['node_fdim'] = SURFACE_NODE_FDIM
            sp_config['edge_fdim'] = SURFACE_EDGE_FDIM + CONTACT_FDIM
            sp_config['hsize'] = config['shsize'] # Use the same hsize as patch
            sp_config['depth'] = config['sdepth'] # Use the same depth as patch
            sp_config['activation'] = config.get('activation', 'relu')
            sp_config['dropout_p'] = config['dropout_mpn']
            sp_config['jk_pool'] = config.get("jk_pool", None) # default
            mpn_config['sp_config'] = sp_config

        mpn_config['bconfig'] = bconfig
        mpn_config['sconfig'] = sconfig
        mpn_config['dropout_mlp'] = config['dropout_mlp'] # Check this part

    mpn_config['graph_pool'] = config.get("graph_pool", "sum_pool")
    return mpn_config

def affinity_pred_config(loaded_config):
    model_config = {}
    config = {}
    toggles = {}
    mpn_config = build_mpn_config(loaded_config)

    config['mpn_config'] = mpn_config
    config['atom_fdim'] = ATOM_FDIM
    config['bond_fdim'] = BOND_FDIM
    config['lig_hsize'] = loaded_config['lig_hsize']
    if isinstance(loaded_config['hsize'], int):
        config['hsize'] = [loaded_config['hsize']]
    else:
        config['hsize'] = loaded_config['hsize']
    config['encoder'] = loaded_config['encoder']
    config['lig_depth'] = loaded_config['lig_depth']
    config['dropout_mlp'] = loaded_config['dropout_mlp']
    config['dropout_mpn'] = loaded_config['dropout_mpn']
    config['prot_mode'] = loaded_config['prot_mode']
    config['activation'] = loaded_config.get("activation", "relu")
    config['jk_pool'] = loaded_config.get("jk_pool", None)
    config['graph_pool'] = loaded_config.get("graph_pool", "sum_pool")

    toggles['ligand_only'] = loaded_config.get('ligand_only', False)
    toggles['use_attn'] = loaded_config.get('use_attn', False)
    toggles['use_mpn_in_patch'] = loaded_config.get('use_mpn_in_patch', False)

    if config['prot_mode'] in ['backbone', 'surface']:
        toggles['use_attn'] = False

    model_config['config'] = config
    model_config['toggles'] = toggles
    return model_config

def prot_class_config(loaded_config):
    model_config = {}
    config = {}
    toggles = {}

    mpn_config = build_mpn_config(loaded_config)
    config['mpn_config'] = mpn_config
    config['n_classes'] = loaded_config['n_classes']
    if isinstance(loaded_config['hsize'], int):
        config['hsize'] = [loaded_config['hsize']]
    else:
        config['hsize'] = loaded_config['hsize']
    config['encoder'] = loaded_config['encoder']
    config['dropout_mlp'] = loaded_config['dropout_mlp']
    config['prot_mode'] = loaded_config['prot_mode']
    config['activation'] = loaded_config.get("activation", "relu")
    config['jk_pool'] = loaded_config.get("jk_pool", None)

    toggles['use_attn'] = loaded_config.get('use_attn', False)
    toggles['use_mpn_in_patch'] = loaded_config.get('use_mpn_in_patch', False)

    if config['prot_mode'] in ['backbone', 'surface']:
        toggles['use_attn'] = False
    model_config['config'] = config
    model_config['toggles'] = toggles
    return model_config

CONFIG_FNS = {
    'pdbbind': affinity_pred_config,
    'reacbio': affinity_pred_config,
    'scope': prot_class_config,
    'enzyme': prot_class_config
    }

def compute_function_weights(config):
    raw_dir = f"{config['data_dir']}/raw/enzyme/metadata"
    labels_to_idx_file = f"{raw_dir}/labels_to_idx.json"
    labels_file = f"{raw_dir}/function_labels.json"
    with open(labels_to_idx_file, "r") as f:
        labels_to_idx = json.load(f)

    with open(labels_file, "r") as f:
        function_labels = json.load(f)

    weights = np.full((len(labels_to_idx)), 0, dtype=np.int32)
    for pdb_id in function_labels:
        weights[labels_to_idx[function_labels[pdb_id]]] += 1
    weights = weights.astype(np.float32)/(float(len(function_labels))*0.5)

    print("Min occurence: ", np.amin(weights))
    print("Max occurence: ", np.amax(weights))

    weights_log = 1.0 / np.log(1.2 + weights)
    return weights_log

def build_model(loaded_config, device='cpu'):
    config_fn = CONFIG_FNS.get(loaded_config['dataset'])
    dataset_metrics = DATASET_METRICS.get(loaded_config['dataset'])
    if not isinstance(dataset_metrics, list):
        dataset_metrics = [dataset_metrics]
    model_config = config_fn(loaded_config)

    metrics = {}
    for metric in dataset_metrics:
        metrics[metric], _, _ = METRICS.get(metric)

    if loaded_config['dataset'] == 'enzyme':
        weights = compute_function_weights(loaded_config)
        model_config['class_weights'] = weights

    model_config['metrics'] = metrics
    model_class = MODEL_CLASSES.get(loaded_config['dataset'])
    model = model_class(**model_config, device=device)
    return model

max_stars_repo_name: vsomnath/holoprot

max_issues_repo_name: vsomnath/holoprot

max_forks_repo_name: vsomnath/holoprot


text:
from rllab.sampler import parallel_sampler
# from rllab.sampler.base import BaseSampler
from sandbox.finetuning.sampler.hier_base import BaseSampler
import numpy as np
from rllab.misc import special
from rllab.misc import tensor_utils
from rllab.algos import util
import rllab.misc.logger as logger

class HierBatchSampler(BaseSampler):
    '''
    this class truncates the paths at a multiple of self.period
    '''
    def __init__(self, algo, period=None):
        """
        :type algo: BatchPolopt
        """
        self.algo = algo
        self.period = period

    def start_worker(self):
        parallel_sampler.populate_task(self.algo.env, self.algo.policy, scope=self.algo.scope)

    def shutdown_worker(self):
        parallel_sampler.terminate_task(scope=self.algo.scope)

    def obtain_samples(self, itr):
        cur_params = self.algo.policy.get_param_values()
        raw_paths = parallel_sampler.sample_paths(
            policy_params=cur_params,
            max_samples=self.algo.batch_size,
            max_path_length=self.algo.max_path_length,
            scope=self.algo.scope,
        )
        if self.period is None:  # hippo random p
            paths = raw_paths
        else:
            #todo: this will break for environments where the rollout terminates after goal is reached
            paths = []
            for path in raw_paths:
                new_length = (len(path['rewards']) // self.period) * self.period
                for key in path.keys():
                    if isinstance(path[key], dict):
                        for key2 in path[key].keys():
                            path[key][key2] = path[key][key2][:new_length]
                    else:
                        path[key] = path[key][:new_length]
                if len(path['rewards']) > 0:
                    paths.append(path)

                # num_padding = self.period - (len(path['rewards']) % self.period)
                # for key in path.keys():
                #     if isinstance(path[key], dict):
                #         for key2 in path[key].keys():
                #             path[key][key2].
            # paths = raw_paths

        if self.algo.whole_paths:
            return paths
        else:
            paths_truncated = parallel_sampler.truncate_paths(paths, self.algo.batch_size)
            return paths_truncated

    def process_samples(self, itr, paths):
        samples_data = super().process_samples(itr, paths)
        if hasattr(self.algo, "use_skill_dependent_baseline") and self.algo.use_skill_dependent_baseline:
            skill_samples_data = self.process_samples_skill_dependent(itr, paths)
            samples_data['skill_advantages'] = skill_samples_data['advantages']
        return samples_data

    # almost all copied over from hier_base
    def process_samples_skill_dependent(self, itr, paths):
        # need to generate the correct observations using the outer product
        new_paths = []
        for i in range(len(paths)):
            latents = paths[i]['agent_infos']['latents']
            observations = paths[i]['observations']
            # insert the time_remaining
            time_remaining = paths[i]['agent_infos']['time_remaining'].reshape(len(observations), 1)
            extended_obs = np.concatenate([observations, time_remaining], axis=1)
            # new_observations = np.matmul(observations[:, :, np.newaxis], latents[:, np.newaxis, :]).reshape(observations.shape[0], -1)
            new_observations = np.matmul(extended_obs[:, :, np.newaxis], latents[:, np.newaxis, :]).reshape(extended_obs.shape[0], -1)
            new_observations = np.concatenate([new_observations, extended_obs, latents], axis=1)
            new_paths.append(dict(observations=new_observations, rewards=paths[i]['rewards'], returns=paths[i]['returns']))
        paths = new_paths

        baselines = []
        returns = []

        if hasattr(self.algo.skill_dependent_baseline, "predict_n"):
            all_path_baselines = self.algo.skill_dependent_baseline.predict_n(paths)
        else:
            all_path_baselines = [self.algo.skill_dependent_baseline.predict(path) for path in paths]

        for idx, path in enumerate(paths):
            path_baselines = np.append(all_path_baselines[idx], 0)
            deltas = path["rewards"] + \
                     self.algo.discount * path_baselines[1:] - \
                     path_baselines[:-1]
            path["advantages"] = special.discount_cumsum(
                deltas, self.algo.discount * self.algo.gae_lambda)
            path["returns"] = special.discount_cumsum(path["rewards"], self.algo.discount)
            baselines.append(path_baselines[:-1])
            returns.append(path["returns"])

        ev = special.explained_variance_1d(
            np.concatenate(baselines),
            np.concatenate(returns)
        )

        if not self.algo.policy.recurrent:
            advantages = tensor_utils.concat_tensor_list([path["advantages"] for path in paths])

            if self.algo.center_adv:
                advantages = util.center_advantages(advantages)

            if self.algo.positive_adv:
                advantages = util.shift_advantages_to_positive(advantages)

            samples_data = dict(
                advantages=advantages,
            )
        else:
            max_path_length = max([len(path["advantages"]) for path in paths])

            if self.algo.center_adv:
                raw_adv = np.concatenate([path["advantages"] for path in paths])
                adv_mean = np.mean(raw_adv)
                adv_std = np.std(raw_adv) + 1e-8
                adv = [(path["advantages"] - adv_mean) / adv_std for path in paths]
            else:
                adv = [path["advantages"] for path in paths]

            adv = np.asarray([tensor_utils.pad_tensor(a, max_path_length) for a in adv])

            samples_data = dict(
                advantages=adv,
            )

        logger.log("fitting skill-depdendent baseline...")
        if hasattr(self.algo.skill_dependent_baseline, 'fit_with_samples'):
            self.algo.skill_dependent_baseline.fit_with_samples(paths, samples_data)
        else:
            self.algo.skill_dependent_baseline.fit(paths)
        logger.log("fitted skill-dependent baseline")

        logger.record_tabular('SkillBaselineExplainedVariance', ev)
        return samples_data

max_issues_repo_name: WeiChengTseng/rllab-finetuning

max_forks_repo_name: WeiChengTseng/rllab-finetuning


text:
# src/train.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import config
import models
import argparse
import create_folds

import pandas as pd
import numpy as np
import pickle
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import MinMaxScaler

import keras.backend as K
from keras.preprocessing import timeseries_dataset_from_array
from keras.models import Sequential
from keras.layers import Dense, GRU, LSTM

from numpy.random import seed
seed(1)
from tensorflow.random import set_seed
set_seed(1)

def run(Error, test = 1):
    if test == 1:
        Producto_no = 0
        data, Fechas = create_folds.datos_agro()
        for Producto in data.Producto.unique(): #data.Producto.unique()[0:test]
            Producto_no += 1
            print(f'Producto número: {Producto_no} , Total de productos: {len(data.Producto.unique())} , avance: {round(Producto_no/len(data.Producto.unique()))*100}%')
            print(f'Corriendo el producto: {Producto}')
            # Filtrando por producto
            temp = data.query("".join(["'", Producto,"'", "== Producto"])).groupby('Fecha').agg(np.sum).Pedido
            temp = temp.reindex(Fechas, fill_value = 0)
            pickle.dump(temp, open(os.path.join(config.TEMP_FILE, 'temp.dat'), 'wb'))
            pickle.dump(Producto, open(os.path.join(config.TEMP_FILE, 'Producto.dat'), 'wb'))
            pickle.dump(Error, open(os.path.join(config.TEMP_FILE, 'Error.dat'), 'wb'))

            models.run_study(f'Study_Producto_{Producto}')
            #break
    else:
        Producto_no = 0
        data, Fechas = create_folds.datos_agro()
        for Producto in data.Producto.unique(): #data.Producto.unique()[0:test]
            Producto_no += 1
            print(f'Producto número: {Producto_no} , Total de productos: {len(data.Producto.unique())} , avance: {round(Producto_no/len(data.Producto.unique()), 2)*100}%')
            print(f'Corriendo el producto: {Producto}')
            # Filtrando por producto
            temp = data.query("".join(["'", Producto,"'", "== Producto"])).groupby('Fecha').agg(np.sum).Pedido
            temp = temp.reindex(Fechas, fill_value = 0)

            # Creando train y test
            size = int(len(temp)*0.9)
            train_Y, test_Y = temp.values.reshape(-1, 1)[0:size], temp.values.reshape(-1, 1)[size:len(temp)]
            #print(train_Y, test_Y)

            # Realizando transformación MaxMinScaler
            scaler = MinMaxScaler()
            scaler.fit(train_Y, y=None)
            scaled_train_data = scaler.transform(train_Y)
            scaled_test_data = scaler.transform(test_Y)

            # Construyendo el generador
            n_input = train_Y.shape[0]; n_features= 1
            generator = timeseries_dataset_from_array(scaled_train_data, scaled_train_data, sequence_length=n_input, sampling_rate = n_features, batch_size=1, shuffle=False)

            # Creando modelo
            lstm_model = Sequential()
            lstm_model.add(GRU(200, input_shape=(n_input, n_features)))
            #lstm_model.add(LSTM(units=50, return_sequences = True))
            lstm_model.add(Dense(1))
            lstm_model.compile(optimizer='adam', loss='mean_squared_error')
            #lstm_model.summary()

            # Entrenando modelo
            lstm_model.fit(generator, epochs=20, verbose = 0)

            # Viendo perdida en el modelo
            losses_lstm = lstm_model.history.history['loss']
            plt.figure(figsize=(20,5))
            plt.xticks(np.arange(0,21,1))
            plt.title(f'Loss Product {Producto}')
            plt.plot(range(len(losses_lstm)),losses_lstm)
            plt.savefig(os.path.join(config.PLOT_LOSS_OUTPUT, f'loss_product_{Producto}.png'));

            # realizando roling forecast
            lstm_predictions_scaled = list()
            batch = scaled_train_data[-n_input:]
            current_batch = batch.reshape((1, n_input, n_features))

            for i in range(len(test_Y)):
                lstm_pred = K.round(lstm_model.predict(current_batch)[0])
                lstm_predictions_scaled.append(lstm_pred)
                current_batch = np.append(current_batch[:,1:,:],[[lstm_pred]],axis=1)

            # Transformando del scaler
            lstm_predictions = scaler.inverse_transform(lstm_predictions_scaled)
            rolling_error_LSTM = mean_squared_error(test_Y, lstm_predictions) ** 0.5
            print(f'    SI rolling forecast -> Test RMSE: {rolling_error_LSTM}')

            # Grafincando el rollback
            RollBack=pd.concat([pd.DataFrame({'TEST':np.concatenate(test_Y)}), pd.DataFrame({'LSTM':np.concatenate(lstm_predictions, axis=0)})], axis=1)
            # agregando fechas al rollback
            RollBack = pd.concat([RollBack,pd.DataFrame({'Time':temp.index[size:]})],axis=1)

            # Colocando en index
            RollBack.set_index('Time', inplace=True)
            RollBack.head()
            RollBack.plot(figsize=(20,5), linewidth=2, fontsize=10)
            plt.title(f'Forecast Product {Producto}')
            plt.xlabel('time', fontsize=15)
            plt.savefig(os.path.join(config.PLOT_FORECAST_OUTPUT, f'forecast_product_{Producto}.png'));

            # Agregando valores al Error de todos los productos
            error_LSTM = mean_squared_error(RollBack.TEST, RollBack.LSTM)
            print(f'    NO rolling forecast -> Test RMSE: {error_LSTM}')
            Error = Error.append({'Producto': Producto, 'Rolling Forecast RMSE': rolling_error_LSTM, 'RMSE': error_LSTM}, ignore_index=True)
            Error.to_csv(os.path.join(config.ERROR_OUTPUT, 'Errores.csv'))

            # Guardando modelo
            lstm_model.save(os.path.join(config.MODEL_OUTPUT_RNN, f'lstm__model_{Producto}.h5'))
            break

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
    "--test",
    type=int
    )
    args = parser.parse_args()

    Error = pd.DataFrame(columns = ['Producto', 'Rolling Forecast RMSE', 'RMSE'])
    run(Error, test = args.test)

reason: import numpy,from numpy


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Independent Component Analysis
#
# Hyvärinen, Aapo, and Erkki Oja.
# "Independent component analysis: algorithms and applications."
# Neural networks 13.4-5 (2000): 411-430.
#
#
#     @article{hyvarinen2000independent,
#       title={Independent component analysis: algorithms and applications},
#       author={Hyv{\"a}rinen, Aapo and Oja, Erkki},
#       journal={Neural networks},
#       volume={13},
#       number={4-5},
#       pages={411--430},
#       year={2000},
#       publisher={Elsevier}
#     }

# +
import nussl
import matplotlib.pyplot as plt
import time
import warnings
import numpy as np

warnings.filterwarnings("ignore")
start_time = time.time()

def visualize_and_embed(sources):
    plt.figure(figsize=(10, 6))
    plt.subplot(211)
    nussl.utils.visualize_sources_as_masks(sources,
        y_axis='mel', db_cutoff=-40, alpha_amount=2.0)
    plt.subplot(212)
    nussl.utils.visualize_sources_as_waveform(
        sources, show_legend=False)
    plt.show()
    nussl.play_utils.multitrack(sources)

musdb = nussl.datasets.MUSDB18(download=True)
i = 39
# -

# Setting up a signal that ICA can work on

# +
item = musdb[i]
sources = [item['sources']['other'], item['sources']['vocals']]

a = nussl.mixing.pan_audio_signal(sources[0], -35)
a_delays = [np.random.randint(1, 20) for _ in range(a.num_channels)]
a = nussl.mixing.delay_audio_signal(a, a_delays)

b = nussl.mixing.pan_audio_signal(sources[1], 35)
b_delays = [np.random.randint(1, 20) for _ in range(b.num_channels)]
b = nussl.mixing.delay_audio_signal(b, b_delays)

mix = a + b

audio_signals = [
    mix.make_audio_signal_from_channel(ch)
    for ch in range(mix.num_channels)
]
# -

# Now running ICA

# +
separator = nussl.separation.factorization.ICA(audio_signals)
estimates = separator()

estimates = {
    f'Source {i}': e for i, e in enumerate(estimates)
}

visualize_and_embed(estimates)
# -

end_time = time.time()
time_taken = end_time - start_time
print(f'Time taken: {time_taken:.4f} seconds')

max_issues_repo_name: titocaco/nussl

max_forks_repo_name: titocaco/nussl


text:
import numpy as np
import pandas as pd

from powersimdata.input.grid import Grid

class MockInputData:
    """
    MockInputData is a mock of powersimdata.input.input_data.InputData
    that generates random profiles.

    Exactly 3 of {`start_time`, `end_time`, `periods`, `freq`} must be specified. See <https://pandas.pydata.org/docs/reference/api/pandas.date_range.html>.

    :param powersimdata.input.grid.Grid grid: instance of Grid object.
    :param str start_time: when profiles begin.
    :param str end_time: when profiles end.
    :param int periods: number of times in profile.
    :param str freq: frequency of times in profile.
    :param int random_seed: used to initialize the random generator.
    :raises ValueError: raised if `field_name` specified in `get_data()` is not specified by this mock
    :return: (*powersimdata.tests.mock_input_data.MockInputData*)
    """

    _RESOURCES = {
        "wind": {"wind", "wind_offshore"},
        "solar": {"solar"},
        "hydro": {"hydro"},
    }

    def __init__(
        self,
        grid: Grid,
        start_time="2016-01-01 00:00:00",
        end_time=None,
        periods=24,
        freq="H",
        random_seed=6669,
    ):
        self._grid = grid
        self._start_time = start_time
        self._end_time = end_time
        self._periods = periods
        self._freq = freq
        self._random = np.random.default_rng(seed=random_seed)

        self._profiles = {
            "demand": self._get_demand(),
            **{
                resource: self._get_resource_profile(resource)
                for resource in self._RESOURCES.keys()
            },
        }
        self._profiles.update(self._get_demand_flexibility())

    def get_data(self, scenario_info, field_name):
        """Returns fake profile data.

        :param dict scenario_info: not used.
        :param str field_name: Can be any of *'demand'*, *'hydro'*, *'solar'*, *'wind'*,
            *'demand_flexibility_up'*, or *'demand_flexibility_dn'*.
        :return: (*pandas.DataFrame*) -- fake profile data
        """
        profile = self._profiles.get(field_name)

        if profile is None:
            raise ValueError(f"No profile specified for {field_name}!")

        return profile

    def _get_demand(self):
        """Returns fake demand data.

        :return: (*pandas.DataFrame*) -- fake demand data
        """
        zone_ids = set(self._grid.plant["zone_id"])
        fake_demand_profile = self._create_fake_profile(zone_ids)
        return fake_demand_profile

    def _get_demand_flexibility(self):
        """Returns fake flexible demand data.

        :return: (*dict*) -- dictionary of fake flexibile demand data
        """
        zone_ids = set(self._grid.plant["zone_id"])
        demand_flexibility_profile_types = [
            "demand_flexibility_up",
            "demand_flexibility_dn",
        ]
        fake_demand_flexibility_profiles = {
            p: self._create_fake_profile(zone_ids, demand_flexibility=True)
            for p in demand_flexibility_profile_types
        }
        return fake_demand_flexibility_profiles

    def _get_resource_profile(self, resource_type):
        """Returns fake data for given resource_type.

        :param str resource_type: Can be any of *'hydro'*, *'solar'*, or *'wind'*.
        :return: (*pandas.DataFrame*) -- fake data for resource
        """
        plant_ids = self._get_plant_ids_for_type(resource_type)
        fake_resource_profile = self._create_fake_profile(plant_ids)
        return fake_resource_profile

    def _get_plant_ids_for_type(self, resource_type):
        """Retrieves plant_ids for plants of `resource_type` from the grid.

        :param str resource_type: Can be any of *'hydro'*, *'solar'*, or *'wind'*.
        :return: (*list*) -- list of plant_ids
        """
        resources = self._RESOURCES[resource_type]
        plant_ids = list(self._grid.plant[lambda ds: ds.type.isin(resources)].index)
        return plant_ids

    def _create_fake_profile(self, columns, demand_flexibility=False):
        """Generates a fake profile.

        :param list columns: columns for the DataFrame
        :param bool demand_flexibility: indicates whether the fake profile being created
            is a demand flexibility profile.
        :return: (*pandas.DataFrame*) -- a fake profile
        """
        times = pd.date_range(
            start=self._start_time,
            end=self._end_time,
            periods=self._periods,
            freq=self._freq,
        )
        index = pd.Index(times, name="UTC Time")
        data = self._random.uniform(low=0, high=1, size=(len(times), len(columns)))
        fake_profile = pd.DataFrame(data=data, index=index, columns=columns)
        if demand_flexibility:
            fake_profile.columns = [f"zone.{z}" for z in columns]
        return fake_profile

max_stars_repo_name: lanesmith/PowerSimData

max_issues_repo_name: lanesmith/PowerSimData

max_forks_repo_name: lanesmith/PowerSimData


text:
import numpy as np
from utils import CROP_ENCODING

RANDOM_SEED = 16
IMG_SIZE = 366

AUTHORS = 'Papoutsis I., Sykas D., Zografakis D., Sdraka M.'
DATASET_VERSION = '21.03'
DATASET_URL = 'https://www.sen4agrinet.space.noa.gr/'

LICENSES = [
    {
        'url': 'https://creativecommons.org/licenses/by-sa/4.0/',
        'id': 1,
        'name': 'Attribution-ShareAlike 4.0 International (CC BY-SA 4.0)'
    },
]

# Band names and their resolutions
BANDS = {
    'B02': 10, 'B03': 10, 'B04': 10, 'B08': 10,
    'B05': 20, 'B07': 20, 'B06': 20, 'B8A': 20, 'B11': 20, 'B12': 20,
    'B01': 60, 'B09': 60, 'B10': 60
}

# Maps arbitrary number of classes to a discrete range of numbers starting from 0
LINEAR_ENCODER = {val: i + 1 for i, val in enumerate(sorted(list(CROP_ENCODING.values())))}
LINEAR_ENCODER[0] = 0

# Extract patches based on this band
REFERENCE_BAND = 'B02'

# Output DTYPE for patches (bands), range is 0-10k, use smaller int to reduce size required
BAND_OUT_DTYPE = np.uint16

# Output DTYPE for patches (labels/parcels), use bigger ints cause parcels have greater range
LABEL_OUT_DTYPE = np.uint32

# Divider for normalizing tiff data to [0-1] range
NORMALIZATION_DIV = 10000