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#
24import casadi as ca
25import numpy as np
26import matplotlib.pyplot as plt
27
28# Degree of interpolating polynomial
29d = 3
30
31# Get collocation points
32tau_root = np.append(0, ca.collocation_points(d, 'legendre'))
33
34# Coefficients of the collocation equation
35C = np.zeros((d+1,d+1))
36
37# Coefficients of the continuity equation
38D = np.zeros(d+1)
39
40# Coefficients of the quadrature function
41B = np.zeros(d+1)
42
43# Construct polynomial basis
44for j in range(d+1):
45    # Construct Lagrange polynomials to get the polynomial basis at the collocation point
46    p = np.poly1d([1])
47    for r in range(d+1):
48        if r != j:
49            p *= np.poly1d([1, -tau_root[r]]) / (tau_root[j]-tau_root[r])
50
51    # Evaluate the polynomial at the final time to get the coefficients of the continuity equation
52    D[j] = p(1.0)
53
54    # Evaluate the time derivative of the polynomial at all collocation points to get the coefficients of the continuity equation
55    pder = np.polyder(p)
56    for r in range(d+1):
57        C[j,r] = pder(tau_root[r])
58
59    # Evaluate the integral of the polynomial to get the coefficients of the quadrature function
60    pint = np.polyint(p)
61    B[j] = pint(1.0)
62
63# Time horizon
64T = 10.
65
66# Declare model variables
67x1 = ca.SX.sym('x1')
68x2 = ca.SX.sym('x2')
69x = ca.vertcat(x1, x2)
70u = ca.SX.sym('u')
71
72# Model equations
73xdot = ca.vertcat((1-x2**2)*x1 - x2 + u, x1)
74
75# Objective term
76L = x1**2 + x2**2 + u**2
77
78# Continuous time dynamics
79f = ca.Function('f', [x, u], [xdot, L], ['x', 'u'], ['xdot', 'L'])
80
81# Control discretization
82N = 20 # number of control intervals
83h = T/N
84
85# Start with an empty NLP
86w=[]
87w0 = []
88lbw = []
89ubw = []
90J = 0
91g=[]
92lbg = []
93ubg = []
94
95# For plotting x and u given w
96x_plot = []
97u_plot = []
98
99# "Lift" initial conditions
100Xk = ca.MX.sym('X0', 2)
101w.append(Xk)
102lbw.append([0, 1])
103ubw.append([0, 1])
104w0.append([0, 1])
105x_plot.append(Xk)
106
107# Formulate the NLP
108for k in range(N):
109    # New NLP variable for the control
110    Uk = ca.MX.sym('U_' + str(k))
111    w.append(Uk)
112    lbw.append([-1])
113    ubw.append([1])
114    w0.append([0])
115    u_plot.append(Uk)
116
117    # State at collocation points
118    Xc = []
119    for j in range(d):
120        Xkj = ca.MX.sym('X_'+str(k)+'_'+str(j), 2)
121        Xc.append(Xkj)
122        w.append(Xkj)
123        lbw.append([-0.25, -np.inf])
124        ubw.append([np.inf,  np.inf])
125        w0.append([0, 0])
126
127    # Loop over collocation points
128    Xk_end = D[0]*Xk
129    for j in range(1,d+1):
130       # Expression for the state derivative at the collocation point
131       xp = C[0,j]*Xk
132       for r in range(d): xp = xp + C[r+1,j]*Xc[r]
133
134       # Append collocation equations
135       fj, qj = f(Xc[j-1],Uk)
136       g.append(h*fj - xp)
137       lbg.append([0, 0])
138       ubg.append([0, 0])
139
140       # Add contribution to the end state
141       Xk_end = Xk_end + D[j]*Xc[j-1];
142
143       # Add contribution to quadrature function
144       J = J + B[j]*qj*h
145
146    # New NLP variable for state at end of interval
147    Xk = ca.MX.sym('X_' + str(k+1), 2)
148    w.append(Xk)
149    lbw.append([-0.25, -np.inf])
150    ubw.append([np.inf,  np.inf])
151    w0.append([0, 0])
152    x_plot.append(Xk)
153
154    # Add equality constraint
155    g.append(Xk_end-Xk)
156    lbg.append([0, 0])
157    ubg.append([0, 0])
158
159# Concatenate vectors
160w = ca.vertcat(*w)
161g = ca.vertcat(*g)
162x_plot = ca.horzcat(*x_plot)
163u_plot = ca.horzcat(*u_plot)
164w0 = np.concatenate(w0)
165lbw = np.concatenate(lbw)
166ubw = np.concatenate(ubw)
167lbg = np.concatenate(lbg)
168ubg = np.concatenate(ubg)
169
170# Create an NLP solver
171prob = {'f': J, 'x': w, 'g': g}
172solver = ca.nlpsol('solver', 'ipopt', prob);
173
174# Function to get x and u trajectories from w
175trajectories = ca.Function('trajectories', [w], [x_plot, u_plot], ['w'], ['x', 'u'])
176
177# Solve the NLP
178sol = solver(x0=w0, lbx=lbw, ubx=ubw, lbg=lbg, ubg=ubg)
179x_opt, u_opt = trajectories(sol['x'])
180x_opt = x_opt.full() # to numpy array
181u_opt = u_opt.full() # to numpy array
182
183# Plot the result
184tgrid = np.linspace(0, T, N+1)
185plt.figure(1)
186plt.clf()
187plt.plot(tgrid, x_opt[0], '--')
188plt.plot(tgrid, x_opt[1], '-')
189plt.step(tgrid, np.append(np.nan, u_opt[0]), '-.')
190plt.xlabel('t')
191plt.legend(['x1','x2','u'])
192plt.grid()
193plt.show()
194