<commit_before>import pyopticon as POC
import pyopticon.samplingControl as SC
import numpy as np
import matplotlib.pyplot as plt

class doubleSlit(POC.MarkovDecisionProcess):
    def __init__(self):
        self.dt = 0.01
        self.Horizon = int(np.round(1/self.dt))
        x0 = 0
    
        A = 1.
        B = self.dt
        R = self.dt * 1e-7
        
        dynMat = np.zeros((self.Horizon,1,3))
        costMat = np.zeros((self.Horizon,3,3))
        
        for k in range(self.Horizon):
            dynMat[k] = POC.buildDynamicsMatrix(A,B)
            if k < self.Horizon-1:
                Q = 0.
            else:
                Q = 1.
            
            costMat[k] = POC.buildCostMatrix(Q,R)

        lqSys = POC.linearQuadraticSystem(dynMat,costMat,
                                          timeInvariant=False,x0=x0)

            
        # augment the cost to model the double slit.
        Slit = [(-5.,-4.),(6.,8.)]
        # Assume that the slit arises in the middle of the movement.
        SlitIndex = self.Horizon/2

            
        def augmentedCost(x,u,k):
            cost = lqSys.costStep(x,u,k)
            if k == SlitIndex:
                slitCost = 1000
                for slit in Slit:
                    if (slit[0] < x) and (slit[1] > x):
                        slitCost = 0.
                cost += slitCost
            return cost

        POC.MarkovDecisionProcess.__init__(self,x0=x0)

        self.costStep = augmentedCost
        self.step = lqSys.step

sys = doubleSlit()

T = sys.Horizon
Time = sys.dt * np.arange(T)

samplingCtrl = POC.samplingOpenLoop(SYS=sys,Horizon=T,
                                    KLWeight=1e-4,burnIn=1000,
                                    ExplorationCovariance = 2,
                                    label='Sampling')


X,U,cost = sys.simulatePolicy(samplingCtrl)

plt.figure(1)
plt.clf()
plt.plot(Time,X.squeeze())<commit_msg>got it working<commit_after>import pyopticon as POC
import pyopticon.samplingControl as SC
import numpy as np
import matplotlib.pyplot as plt

class doubleSlit(POC.MarkovDecisionProcess):
    def __init__(self):
        self.dt = 0.02
        self.Horizon = int(np.round(2/self.dt))
        x0 = 0
    
        A = 1.
        B = self.dt
        R = .5 * self.dt * 1e-1
        
        dynMat = np.zeros((self.Horizon,1,3))
        costMat = np.zeros((self.Horizon,3,3))
        
        for k in range(self.Horizon):
            dynMat[k] = POC.buildDynamicsMatrix(A,B)
            if k < self.Horizon-1:
                Q = 0.
            else:
                Q = 0.5
            
            costMat[k] = POC.buildCostMatrix(Q,R)

        lqSys = POC.linearQuadraticSystem(dynMat,costMat,
                                          timeInvariant=False,x0=x0)

        # augment the cost to model the double slit.
        Slit = [(-6.,-4.),(6.,8.)]
        # Assume that the slit arises in the middle of the movement.
        SlitIndex = self.Horizon/2

        self.Slit = Slit
        self.SlitIndex = SlitIndex

            
        def augmentedCost(x,u,k):
            cost = lqSys.costStep(x,u,k)
            if k == SlitIndex:
                slitCost = 1000
                for slit in Slit:
                    if (slit[0] < x) and (slit[1] > x):
                        slitCost = 0.
                cost += slitCost
            return cost

        POC.MarkovDecisionProcess.__init__(self,x0=x0)

        self.costStep = augmentedCost
        self.step = lqSys.step

sys = doubleSlit()

T = sys.Horizon
Time = sys.dt * np.arange(T)

slitTime = Time[sys.SlitIndex]

NumX = 10
X0 = np.linspace(-10,10,NumX)
Cost = np.zeros(NumX)

plt.figure(1)
plt.clf()
plt.plot([slitTime,slitTime],[-10,sys.Slit[0][0]],'k')
plt.plot([slitTime,slitTime],[sys.Slit[0][1],sys.Slit[1][0]],'k')
plt.plot([slitTime,slitTime],[sys.Slit[1][1],10],'k')

# Calculate appropriate noise covariance 
# dt*s = sqrt(dt) -> s = 1/sqrt(dt)

cov = 1/np.sqrt(sys.dt)

for k in range(NumX):
    print 'Initial Condition %d of %d' % (k+1,NumX)
    sys.x0 = X0[k]
    samplingCtrl = POC.samplingOpenLoop(SYS=sys,Horizon=T,
                                        KLWeight=1e-4,burnIn=1000,
                                        ExplorationCovariance = cov,
                                        label='Sampling')


    X,U,Cost[k] = sys.simulatePolicy(samplingCtrl)
    plt.plot(Time,X.squeeze())

plt.figure(2)
plt.clf()
plt.plot(X0,Cost)
{'fraction_non_alphanumeric': 0.11886537595677622, 'fraction_numerical': 0.02408824853669518, 'mean_word_length': 1.5952102803738317, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4388007', 'n_tokens_mistral': 1593, 'n_tokens_neox': 1483, 'n_words': 297}