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