1#
2#     This file is part of CasADi.
3#
4#     CasADi -- A symbolic framework for dynamic optimization.
5#     Copyright (C) 2010-2014 Joel Andersson, Joris Gillis, Moritz Diehl,
6#                             K.U. Leuven. All rights reserved.
7#     Copyright (C) 2011-2014 Greg Horn
8#
9#     CasADi is free software; you can redistribute it and/or
10#     modify it under the terms of the GNU Lesser General Public
11#     License as published by the Free Software Foundation; either
12#     version 3 of the License, or (at your option) any later version.
13#
14#     CasADi is distributed in the hope that it will be useful,
15#     but WITHOUT ANY WARRANTY; without even the implied warranty of
16#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17#     Lesser General Public License for more details.
18#
19#     You should have received a copy of the GNU Lesser General Public
20#     License along with CasADi; if not, write to the Free Software
21#     Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
22#
23#
24
25from pylab import *
26
27# End time
28T = 10.
29
30# Number of control intervals
31N = 20
32
33# Number of Runge-Kutta 4 steps per interval and step size
34NK = 20
35DT = T/(N*NK)
36
37# Number of discrete control values
38NU = 101
39
40# Number of discrete state values
41NX = 101
42
43# System dynamics, can be called with matricex
44def f(x1,x2,u):
45  x1_dot = (1 - x2*x2)*x1 - x2 + u
46  x2_dot = x1
47  q_dot  = x1*x1 + x2*x2 + u*u
48  return (x1_dot, x2_dot, q_dot)
49
50# Control enumeration
51U  = linspace(-1,1,NU)
52
53# State space enumeration
54x1 = linspace(-1,1,NX)
55x2 = linspace(-1,1,NX)
56X1,X2 = meshgrid(x1,x2)
57
58# For each control action and state, precalculate next state and stage cost
59stage_J = []
60next_x1 = []
61next_x2 = []
62for u in U:
63  # Take number of integration steps
64  X1_k = copy(X1)
65  X2_k = copy(X2)
66  Q_k = zeros(X1.shape)
67  for k in range(NK):
68    # RK4 integration for x1, x2 and q
69    k1_x1, k1_x2, k1_q = f(X1_k,                X2_k,                u)
70    k2_x1, k2_x2, k2_q = f(X1_k + DT/2 * k1_x1, X2_k + DT/2 * k1_x2, u)
71    k3_x1, k3_x2, k3_q = f(X1_k + DT/2 * k2_x1, X2_k + DT/2 * k2_x2, u)
72    k4_x1, k4_x2, k4_q = f(X1_k + DT   * k3_x1, X2_k + DT   * k3_x2, u)
73    X1_k += DT/6*(k1_x1 + 2*k2_x1 + 2*k3_x1 + k4_x1)
74    X2_k += DT/6*(k1_x2 + 2*k2_x2 + 2*k3_x2 + k4_x2)
75    Q_k  += DT/6*(k1_q  + 2*k2_q  + 2*k3_q  + k4_q )
76
77  # Find out which state comes next (index)
78  X1_k = matrix.round((X1_k+1)/2*(NX-1)).astype(int)
79  X2_k = matrix.round((X2_k+1)/2*(NX-1)).astype(int)
80
81  # Infinite cost if state gets out-of-bounds
82  I = X1_k  <  0; Q_k[I]=inf; X1_k[I]=0
83  I = X2_k  <  0; Q_k[I]=inf; X2_k[I]=0
84  I = X1_k >= NX; Q_k[I]=inf; X1_k[I]=0
85  I = X2_k >= NX; Q_k[I]=inf; X2_k[I]=0
86
87  # Save the stage cost and next state
88  next_x1.append(X1_k)
89  next_x2.append(X2_k)
90  stage_J.append(Q_k)
91
92# Calculate cost-to-go (no end cost) and optimal control
93J = zeros(X1.shape)
94U_opt = []
95for k in reversed(list(range(N))):
96  # Cost to go for the previous step, optimal control action
97  J_prev = inf*ones(X1.shape)
98  u_prev = -ones(X1.shape,dtype=int)
99
100  # Test all control actions
101  for uind in range(NU):
102    J_prev_test = J[next_x2[uind],next_x1[uind]]+stage_J[uind]
103    better = J_prev_test<J_prev
104    u_prev[better] = uind
105    J_prev[better] = J_prev_test[better]
106
107  # Update cost-to-go and save optimal control
108  J = J_prev
109  U_opt.append(u_prev)
110
111# Reorder U_opt by stage
112U_opt.reverse()
113
114# Find optimal control starting at x1=0, x2=1
115i1 = NX//2
116i2 = NX-1
117u_opt = []
118x1_opt = [x1[i1]]
119x2_opt = [x2[i2]]
120cost = 0
121for k in range(N):
122  # Get the optimal control and go to next step
123  u_ind = U_opt[k][i2,i1]
124  cost += stage_J[u_ind][i2,i1]
125  i1, i2 = next_x1[u_ind][i2,i1], next_x2[u_ind][i2,i1]
126
127  # Save the trajectories
128  u_opt.append(U[u_ind])
129  x1_opt.append(x1[i1])
130  x2_opt.append(x2[i2])
131
132# Optimal cost
133print("Minimal cost: ", cost)
134assert abs(cost-J[NX-1,NX//2])<1e-8 # Consistency check
135
136# Plot
137figure(1)
138clf()
139
140# Plot optimal cost-to-go
141subplot(121)
142contourf(X1,X2,J)
143colorbar()
144xlabel('x1')
145ylabel('x2')
146title('Cost-to-go')
147
148subplot(122)
149plot(linspace(0,T,N+1),x1_opt,'--')
150plot(linspace(0,T,N+1),x2_opt,'-.')
151step(linspace(0,T,N),u_opt,'-')
152plt.title("Dynamic programming solution")
153plt.xlabel('time')
154plt.legend(['x1 trajectory','x2 trajectory','u trajectory'])
155grid(True)
156show()
157