<commit_before>import numpy as np
from scipy.linalg import cholesky
from numpy.random import randn
from bovy_mcmc.elliptical_slice import elliptical_slice as eslice
import Controller as ctrl

#### Basic Helper Functions

def initializeFlatOpenLoop(SYS, initialPolicy, Horizon=1):
    if initialPolicy is None:
        U = np.zeros(Horizon * SYS.NumInputs)
    else:
        UFull = SYS.simulatePolicy(initialPolicy)[1]
        U = UFull[:Horizon].flatten()

    return U


#### Slice Sampling Optimizer ####

def sliceSample(sampleObj,X,burnIn=1,resetObject=False):
    """
    sampleObj - an object to be sampled. Must have the following attributes
    sampleObj.priorChol - cholesky factorization of Gaussian prior
    sampleObj.loglikelihood - a log-likelihood function 
    sampleObj.displaySampleInfo - a function that prints relevant info to screen

    if resetObject is True then sampleObj.reset must also be defined
    X - an initial sample
    burnIn - length of burn-in period
    resetObject - flag to reset sampleObj after each sample
    """

    lenW = sampleObj.priorChol.shape[0]    
    logLik = sampleObj.loglikelihood(X)
    bestLik = -np.inf

    eps = 1e-3
        
    samp = 0
    likSequence = [logLik]
    
    while samp < burnIn:
        samp += 1
        W = np.dot(sampleObj.priorChol,randn(lenW))
        if resetObject:
            # resets data associated with sample object
            # This is needed for some stochastic variants
            sampleObj.reset()
            # when the sample object has been reset
            # Older values of the log-likelihood function may not
            # be relevant
            X,newLogLik = eslice(X,W,sampleObj.loglikelihood)
        else:
            X,newLogLik = eslice(X,W,sampleObj.loglikelihood,cur_lnpdf=logLik)
        if newLogLik > bestLik:
            bestLik = newLogLik
            bestX = X
                
        likChange = newLogLik - logLik
        logLik = newLogLik

        likSequence.append(logLik)

        if (burnIn == np.inf) and (likChange <= 0):
            break

        sampleObj.displaySampleInfo(logLik,bestLik,samp,burnIn)
            
    return X, bestX, likSequence




def trajectoryNoiseMatrix(Cov,Horizon):
    if isinstance(Cov,np.ndarray):
        cholSig = cholesky(Cov,lower=True)
        NoiseGain = np.kron(np.eye(Horizon),cholSig)
    else:
        NoiseGain = np.sqrt(Cov) * np.eye(Horizon)
    return NoiseGain

def sliceOptimizationDisplay(logLik,bestLik,samp,burnIn,KLWeight):
    cost = -logLik * KLWeight
    bestCost = -bestLik * KLWeight
        
    if burnIn < np.inf:
        if samp % 10 == 0:
            print 'run %d of %d, Cost: %g, Best Cost: %g' % \
                (samp,burnIn,cost,bestCost)
    else:
        if samp % 10 == 0:
            print 'run %d, Cost: %g, Cost Change %g' % \
                (samp,cost,costChange)


def nullDisplay(logLik,bestLik,samp,burnIn,KLWeight):
    pass

#### Slice Sampling Controllers ####


class samplingOpenLoop(ctrl.flatOpenLoopPolicy):
    def __init__(self,
                 SYS = None,
                 KLWeight=1,
                 burnIn=0,
                 ExplorationCovariance=1.,
                 initialPolicy = None,
                 displayFun = sliceOptimizationDisplay,
                 *args, **kwargs):

        self.displayFun = displayFun
        self.SYS = SYS
        self.KLWeight = KLWeight

        ctrl.flatOpenLoopPolicy.__init__(self,NumInputs=SYS.NumInputs,
                                    *args,**kwargs)


        self.priorChol = trajectoryNoiseMatrix(ExplorationCovariance,
                                               self.Horizon)

        lenW = self.priorChol.shape[0]


        self.U = initializeFlatOpenLoop(SYS, initialPolicy, self.Horizon)


        self.burnIn = burnIn
        self.updatePolicy()

    def updatePolicy(self):
        Ucur, self.U, likSequence = sliceSample(self,self.U,self.burnIn)
        self.costSequence = -np.array(likSequence) * self.KLWeight
        
    def loglikelihood(self,U):
        self.U = U
        cost = self.SYS.simulatePolicy(self)[2]
        return -cost / self.KLWeight

    def displaySampleInfo(self,logLik,bestLik,samp,burnIn):
        self.displayFun(logLik,bestLik,samp,burnIn,self.KLWeight)

class samplingStochasticAffine(ctrl.flatVaryingAffine):
    def __init__(self,SYS=None,Horizon=1,
                 NumSamples = 1,
                 KLWeight=1,burnIn=0,
                 ExplorationCovariance=1.,
                 initialPolicy = None,
                 *args, **kwargs):

        self.NumSamples = NumSamples<commit_msg>removed some extraneous code for intiializing policies<commit_after>import numpy as np
from scipy.linalg import cholesky
from numpy.random import randn
from bovy_mcmc.elliptical_slice import elliptical_slice as eslice
import Controller as ctrl

#### Basic Helper Functions

def initializeFlatOpenLoop(SYS, initialPolicy, Horizon=1):
    if initialPolicy is None:
        U = np.zeros(Horizon * SYS.NumInputs)
    else:
        UFull = SYS.simulatePolicy(initialPolicy)[1]
        U = UFull[:Horizon].flatten()

    return U


#### Slice Sampling Optimizer ####

def sliceSample(sampleObj,X,burnIn=1,resetObject=False):
    """
    sampleObj - an object to be sampled. Must have the following attributes
    sampleObj.priorChol - cholesky factorization of Gaussian prior
    sampleObj.loglikelihood - a log-likelihood function 
    sampleObj.displaySampleInfo - a function that prints relevant info to screen

    if resetObject is True then sampleObj.reset must also be defined
    X - an initial sample
    burnIn - length of burn-in period
    resetObject - flag to reset sampleObj after each sample
    """

    lenW = sampleObj.priorChol.shape[0]    
    logLik = sampleObj.loglikelihood(X)
    bestLik = -np.inf

    eps = 1e-3
        
    samp = 0
    likSequence = [logLik]
    
    while samp < burnIn:
        samp += 1
        W = np.dot(sampleObj.priorChol,randn(lenW))
        if resetObject:
            # resets data associated with sample object
            # This is needed for some stochastic variants
            sampleObj.reset()
            # when the sample object has been reset
            # Older values of the log-likelihood function may not
            # be relevant
            X,newLogLik = eslice(X,W,sampleObj.loglikelihood)
        else:
            X,newLogLik = eslice(X,W,sampleObj.loglikelihood,cur_lnpdf=logLik)
        if newLogLik > bestLik:
            bestLik = newLogLik
            bestX = X
                
        likChange = newLogLik - logLik
        logLik = newLogLik

        likSequence.append(logLik)

        if (burnIn == np.inf) and (likChange <= 0):
            break

        sampleObj.displaySampleInfo(logLik,bestLik,samp,burnIn)
            
    return X, bestX, likSequence




def trajectoryNoiseMatrix(Cov,Horizon):
    if isinstance(Cov,np.ndarray):
        cholSig = cholesky(Cov,lower=True)
        NoiseGain = np.kron(np.eye(Horizon),cholSig)
    else:
        NoiseGain = np.sqrt(Cov) * np.eye(Horizon)
    return NoiseGain

def sliceOptimizationDisplay(logLik,bestLik,samp,burnIn,KLWeight):
    cost = -logLik * KLWeight
    bestCost = -bestLik * KLWeight
        
    if burnIn < np.inf:
        if samp % 10 == 0:
            print 'run %d of %d, Cost: %g, Best Cost: %g' % \
                (samp,burnIn,cost,bestCost)
    else:
        if samp % 10 == 0:
            print 'run %d, Cost: %g, Cost Change %g' % \
                (samp,cost,costChange)


def nullDisplay(logLik,bestLik,samp,burnIn,KLWeight):
    pass

#### Slice Sampling Controllers ####


class samplingOpenLoop(ctrl.flatOpenLoopPolicy):
    def __init__(self,
                 SYS = None,
                 KLWeight=1,
                 burnIn=0,
                 ExplorationCovariance=1.,
                 initialPolicy = None,
                 displayFun = sliceOptimizationDisplay,
                 *args, **kwargs):

        self.displayFun = displayFun
        self.SYS = SYS
        self.KLWeight = KLWeight

        ctrl.flatOpenLoopPolicy.__init__(self,NumInputs=SYS.NumInputs,
                                    *args,**kwargs)


        self.priorChol = trajectoryNoiseMatrix(ExplorationCovariance,
                                               self.Horizon)


        self.U = initializeFlatOpenLoop(SYS, initialPolicy, self.Horizon)


        self.burnIn = burnIn
        self.updatePolicy()

    def updatePolicy(self):
        Ucur, self.U, likSequence = sliceSample(self,self.U,self.burnIn)
        self.costSequence = -np.array(likSequence) * self.KLWeight
        
    def loglikelihood(self,U):
        self.U = U
        cost = self.SYS.simulatePolicy(self)[2]
        return -cost / self.KLWeight

    def displaySampleInfo(self,logLik,bestLik,samp,burnIn):
        self.displayFun(logLik,bestLik,samp,burnIn,self.KLWeight)

class samplingStochasticAffine(ctrl.flatVaryingAffine):
    def __init__(self,SYS=None,Horizon=1,
                 NumSamples = 1,
                 KLWeight=1,burnIn=0,
                 ExplorationCovariance=1.,
                 initialPolicy = None,
                 *args, **kwargs):

        self.NumSamples = NumSamples
{'fraction_non_alphanumeric': 0.08775355244603536, 'fraction_numerical': 0.005315110098709187, 'mean_word_length': 2.1651218674905595, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5660751', 'n_tokens_mistral': 2930, 'n_tokens_neox': 2772, 'n_words': 637}