<commit_before>'''multifield.py - Classes for multi field cosmological models

Author: <NAME>
For license and copyright information see LICENSE.txt which was distributed with this file.

'''

import numpy as np
from scipy import interpolate

import cosmomodels as c

class MultiFieldModels(c.CosmologicalModel):
    '''
    Parent class for all multifield models. 
    '''


    def __init__(self, *args, **kwargs):
        """Call superclass init method."""
        super(MultiFieldModels, self).__init__(*args, **kwargs)
        
        #Set the number of fields using keyword argument, defaults to 1.
        self.nfields = kwargs.get("nfields", 1)
        
        #Set field indices. These can be used to select only certain parts of
        #the y variable, e.g. y[self.bg_ix] is the array of background values.
        self.H_ix = self.nfields*2
        self.bg_ix = slice(0,self.nfields*2)
        self.pert_ix = slice(self.nfields*2+1, None)
        
    def findH(self, U, y):
        """Return value of Hubble variable, H at y for given potential.
        
        y - array_like: variable array at one time step. So y[1] should be phidot
                        for the first field and y[n+1] for each subsequent field. 
        """
        #Get phidots from y, should be second variable for each field.
        phidots = y[self.bg_ix][1::2]
        
        
        #Expression for H
        H = np.sqrt(U/(3.0-0.5*(np.sum(phidots**2))))
        return H
    
    def potentials(self, y, pot_params=None):
        """Return value of potential at y, along with first and second derivs."""
        pass
    
    def findinflend(self):
        """Find the efold time where inflation ends,
            i.e. the hubble flow parameter epsilon >1.
            Returns tuple of endefold and endindex (in tresult)."""
        
        self.epsilon = self.getepsilon()
        if not any(self.epsilon>1):
            raise c.ModelError("Inflation did not end during specified number of efoldings. Increase tend and try again!")
        endindex = np.where(self.epsilon>=1)[0][0]
        
        #Interpolate results to find more accurate endpoint
        tck = interpolate.splrep(self.tresult[:endindex], self.epsilon[:endindex])
        t2 = np.linspace(self.tresult[endindex-1], self.tresult[endindex], 100)
        y2 = interpolate.splev(t2, tck)
        endindex2 = np.where(y2>1)[0][0]
        #Return efold of more accurate endpoint
        endefold = t2[endindex2]
        
        return endefold, endindex
    
    def getepsilon(self):
        """Return an array of epsilon = -\dot{H}/H values for each timestep."""
        #Find Hdot
        if len(self.yresult.shape) == 3:
            phidots = self.yresult[:,self.bg_ix,0][1::2]
        else:
            phidots = self.yresult[:,self.bg_ix][1::2]
        epsilon = - 0.5*np.sum(phidots**2)
        return epsilon
        
class MultiFieldBackground(MultiFieldModels):
    """Basic model with background equations for multiple fields
        Array of dependent variables y is given by:
        
       y[0] - \phi_a : Background inflaton
       y[1] - d\phi_a/d\n : First deriv of \phi_a
       ...
       y[self.nfields*2] - H: Hubble parameter
    """
        
    def __init__(self,  *args, **kwargs):
        """Initialize variables and call superclass"""
        
        super(MultiFieldBackground, self).__init__(*args, **kwargs)
        
        #Set initial H value if None
        if np.all(self.ystart[self.H_ix] == 0.0):
            U = self.potentials(self.ystart, self.pot_params)[0]
            self.ystart[self.H_ix] = self.findH(U, self.ystart)
    
    def derivs(self, y, t, **kwargs):
        """Basic background equations of motion.
            dydx[0] = dy[0]/dn etc"""
        #get potential from function
        U, dUdphi, d2Udphi2 = self.potentials(y, self.pot_params)[0:3]       
        
        #Set derivatives
        dydx = np.zeros_like(y)
        
        bg_indices = np.arange(0,self.H_ix,2)
        firstderiv_indices = np.arange(1,self.H_ix,2)
        
        #d\phi_0/dn = y_1
        dydx[bg_indices] = y[firstderiv_indices] 
        
        #dphi^prime/dn
        dydx[firstderiv_indices] = -(U*y[firstderiv_indices] + dUdphi)/(y[self.H_ix]**2)
        
        #dH/dn
        dydx[self.H_ix] = -0.5*(np.sum(y[firstderiv_indices]**2))*y[self.H_ix]

        return dydx<commit_msg>Sum only over the fields not time as well.<commit_after>'''multifield.py - Classes for multi field cosmological models

Author: <NAME>
For license and copyright information see LICENSE.txt which was distributed with this file.

'''

import numpy as np
from scipy import interpolate

import cosmomodels as c

class MultiFieldModels(c.CosmologicalModel):
    '''
    Parent class for all multifield models. 
    '''


    def __init__(self, *args, **kwargs):
        """Call superclass init method."""
        super(MultiFieldModels, self).__init__(*args, **kwargs)
        
        #Set the number of fields using keyword argument, defaults to 1.
        self.nfields = kwargs.get("nfields", 1)
        
        #Set field indices. These can be used to select only certain parts of
        #the y variable, e.g. y[self.bg_ix] is the array of background values.
        self.H_ix = self.nfields*2
        self.bg_ix = slice(0,self.nfields*2)
        self.phis_ix = slice(0,self.nfields*2,2)
        self.phidots_ix = slice(1,self.nfields*2,2)
        self.pert_ix = slice(self.nfields*2+1, None)
        
    def findH(self, U, y):
        """Return value of Hubble variable, H at y for given potential.
        
        y - array_like: variable array at one time step. So y[1] should be phidot
                        for the first field and y[n+1] for each subsequent field. 
        """
        #Get phidots from y, should be second variable for each field.
        phidots = y[self.phidots_ix]
        
        
        #Expression for H
        H = np.sqrt(U/(3.0-0.5*(np.sum(phidots**2))))
        return H
    
    def potentials(self, y, pot_params=None):
        """Return value of potential at y, along with first and second derivs."""
        pass
    
    def findinflend(self):
        """Find the efold time where inflation ends,
            i.e. the hubble flow parameter epsilon >1.
            Returns tuple of endefold and endindex (in tresult)."""
        
        self.epsilon = self.getepsilon()
        if not any(self.epsilon>1):
            raise c.ModelError("Inflation did not end during specified number of efoldings. Increase tend and try again!")
        endindex = np.where(self.epsilon>=1)[0][0]
        
        #Interpolate results to find more accurate endpoint
        tck = interpolate.splrep(self.tresult[:endindex], self.epsilon[:endindex])
        t2 = np.linspace(self.tresult[endindex-1], self.tresult[endindex], 100)
        y2 = interpolate.splev(t2, tck)
        endindex2 = np.where(y2>1)[0][0]
        #Return efold of more accurate endpoint
        endefold = t2[endindex2]
        
        return endefold, endindex
    
    def getepsilon(self):
        """Return an array of epsilon = -\dot{H}/H values for each timestep."""
        #Find Hdot
        if len(self.yresult.shape) == 3:
            phidots = self.yresult[:,self.phidots_ix,0]
        else:
            phidots = self.yresult[:,self.phidots_ix]
        #Make sure to do sum across only phidot axis (1 in this case)
        epsilon = - 0.5*np.sum(phidots**2, axis=1)
        return epsilon
        
class MultiFieldBackground(MultiFieldModels):
    """Basic model with background equations for multiple fields
        Array of dependent variables y is given by:
        
       y[0] - \phi_a : Background inflaton
       y[1] - d\phi_a/d\n : First deriv of \phi_a
       ...
       y[self.nfields*2] - H: Hubble parameter
    """
        
    def __init__(self,  *args, **kwargs):
        """Initialize variables and call superclass"""
        
        super(MultiFieldBackground, self).__init__(*args, **kwargs)
        
        #Set initial H value if None
        if np.all(self.ystart[self.H_ix] == 0.0):
            U = self.potentials(self.ystart, self.pot_params)[0]
            self.ystart[self.H_ix] = self.findH(U, self.ystart)
    
    def derivs(self, y, t, **kwargs):
        """Basic background equations of motion.
            dydx[0] = dy[0]/dn etc"""
        #get potential from function
        U, dUdphi, d2Udphi2 = self.potentials(y, self.pot_params)[0:3]       
        
        #Set derivatives
        dydx = np.zeros_like(y)
        
        bg_indices = np.arange(0,self.H_ix,2)
        firstderiv_indices = np.arange(1,self.H_ix,2)
        
        #d\phi_0/dn = y_1
        dydx[bg_indices] = y[firstderiv_indices] 
        
        #dphi^prime/dn
        dydx[firstderiv_indices] = -(U*y[firstderiv_indices] + dUdphi)/(y[self.H_ix]**2)
        
        #dH/dn
        dydx[self.H_ix] = -0.5*(np.sum(y[firstderiv_indices]**2))*y[self.H_ix]

        return dydx
{'fraction_non_alphanumeric': 0.1116672227783339, 'fraction_numerical': 0.014903792681570459, 'mean_word_length': 2.1729004940014116, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '900300', 'n_tokens_mistral': 2908, 'n_tokens_neox': 2828, 'n_words': 779}