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# -*- coding: utf-8 -*-
25from casadi import *
26import numpy as NP
27import matplotlib.pyplot as plt
28
29# Degree of interpolating polynomial
30d = 3
31
32# Choose collocation points
33tau_root = [0] + collocation_points(d, "radau")
34
35# Coefficients of the collocation equation
36C = NP.zeros((d+1,d+1))
37
38# Coefficients of the continuity equation
39D = NP.zeros(d+1)
40
41# Coefficients of the quadrature function
42F = NP.zeros(d+1)
43
44# Construct polynomial basis
45for j in range(d+1):
46  # Construct Lagrange polynomials to get the polynomial basis at the collocation point
47  p = NP.poly1d([1])
48  for r in range(d+1):
49    if r != j:
50      p *= NP.poly1d([1, -tau_root[r]]) / (tau_root[j]-tau_root[r])
51
52  # Evaluate the polynomial at the final time to get the coefficients of the continuity equation
53  D[j] = p(1.0)
54
55  # Evaluate the time derivative of the polynomial at all collocation points to get the coefficients of the continuity equation
56  pder = NP.polyder(p)
57  for r in range(d+1):
58    C[j,r] = pder(tau_root[r])
59
60  # Evaluate the integral of the polynomial to get the coefficients of the quadrature function
61  pint = NP.polyint(p)
62  F[j] = pint(1.0)
63
64# Control discretization
65nk = 20
66
67# End time
68tf = 10.0
69
70# Size of the finite elements
71h = tf/nk
72
73# All collocation time points
74T = NP.zeros((nk,d+1))
75for k in range(nk):
76  for j in range(d+1):
77    T[k,j] = h*(k + tau_root[j])
78
79# Declare variables (use scalar graph)
80t  = SX.sym("t")    # time
81u  = SX.sym("u")    # control
82x  = SX.sym("x",2)  # state
83
84# ODE rhs function and quadratures
85xdot = vertcat((1 - x[1]*x[1])*x[0] - x[1] + u, \
86               x[0])
87qdot = x[0]*x[0] + x[1]*x[1] + u*u
88f = Function('f', [t,x,u],[xdot, qdot])
89
90# Control bounds
91u_min = -0.75
92u_max = 1.0
93u_init = 0.0
94
95u_lb = NP.array([u_min])
96u_ub = NP.array([u_max])
97u_init = NP.array([u_init])
98
99# State bounds and initial guess
100x_min =  [-inf, -inf]
101x_max =  [ inf,  inf]
102xi_min = [ 0.0,  1.0]
103xi_max = [ 0.0,  1.0]
104xf_min = [ 0.0,  0.0]
105xf_max = [ 0.0,  0.0]
106x_init = [ 0.0,  0.0]
107
108# Dimensions
109nx = 2
110nu = 1
111
112# Total number of variables
113NX = nk*(d+1)*nx      # Collocated states
114NU = nk*nu              # Parametrized controls
115NXF = nx                # Final state
116NV = NX+NU+NXF
117
118# NLP variable vector
119V = MX.sym("V",NV)
120
121# All variables with bounds and initial guess
122vars_lb = NP.zeros(NV)
123vars_ub = NP.zeros(NV)
124vars_init = NP.zeros(NV)
125offset = 0
126
127# Get collocated states and parametrized control
128X = NP.resize(NP.array([],dtype=MX),(nk+1,d+1))
129U = NP.resize(NP.array([],dtype=MX),nk)
130for k in range(nk):
131  # Collocated states
132  for j in range(d+1):
133    # Get the expression for the state vector
134    X[k,j] = V[offset:offset+nx]
135
136    # Add the initial condition
137    vars_init[offset:offset+nx] = x_init
138
139    # Add bounds
140    if k==0 and j==0:
141      vars_lb[offset:offset+nx] = xi_min
142      vars_ub[offset:offset+nx] = xi_max
143    else:
144      vars_lb[offset:offset+nx] = x_min
145      vars_ub[offset:offset+nx] = x_max
146    offset += nx
147
148  # Parametrized controls
149  U[k] = V[offset:offset+nu]
150  vars_lb[offset:offset+nu] = u_min
151  vars_ub[offset:offset+nu] = u_max
152  vars_init[offset:offset+nu] = u_init
153  offset += nu
154
155# State at end time
156X[nk,0] = V[offset:offset+nx]
157vars_lb[offset:offset+nx] = xf_min
158vars_ub[offset:offset+nx] = xf_max
159vars_init[offset:offset+nx] = x_init
160offset += nx
161
162# Constraint function for the NLP
163g = []
164lbg = []
165ubg = []
166
167# Objective function
168J = 0
169
170# For all finite elements
171for k in range(nk):
172
173  # For all collocation points
174  for j in range(1,d+1):
175
176    # Get an expression for the state derivative at the collocation point
177    xp_jk = 0
178    for r in range (d+1):
179      xp_jk += C[r,j]*X[k,r]
180
181    # Add collocation equations to the NLP
182    fk,qk = f(T[k,j], X[k,j], U[k])
183    g.append(h*fk - xp_jk)
184    lbg.append(NP.zeros(nx)) # equality constraints
185    ubg.append(NP.zeros(nx)) # equality constraints
186
187    # Add contribution to objective
188    J += F[j]*qk*h
189
190  # Get an expression for the state at the end of the finite element
191  xf_k = 0
192  for r in range(d+1):
193    xf_k += D[r]*X[k,r]
194
195  # Add continuity equation to NLP
196  g.append(X[k+1,0] - xf_k)
197  lbg.append(NP.zeros(nx))
198  ubg.append(NP.zeros(nx))
199
200# Concatenate constraints
201g = vertcat(*g)
202
203# NLP
204nlp = {'x':V, 'f':J, 'g':g}
205
206## ----
207## SOLVE THE NLP
208## ----
209
210# Set options
211opts = {}
212opts["expand"] = True
213#opts["ipopt.max_iter"] = 4
214opts["ipopt.linear_solver"] = 'ma27'
215
216# Allocate an NLP solver
217solver = nlpsol("solver", "ipopt", nlp, opts)
218arg = {}
219
220# Initial condition
221arg["x0"] = vars_init
222
223# Bounds on x
224arg["lbx"] = vars_lb
225arg["ubx"] = vars_ub
226
227# Bounds on g
228arg["lbg"] = NP.concatenate(lbg)
229arg["ubg"] = NP.concatenate(ubg)
230
231# Solve the problem
232res = solver(**arg)
233
234# Print the optimal cost
235print("optimal cost: ", float(res["f"]))
236
237# Retrieve the solution
238v_opt = NP.array(res["x"])
239
240# Get values at the beginning of each finite element
241x0_opt = v_opt[0::(d+1)*nx+nu]
242x1_opt = v_opt[1::(d+1)*nx+nu]
243u_opt = v_opt[(d+1)*nx::(d+1)*nx+nu]
244tgrid = NP.linspace(0,tf,nk+1)
245tgrid_u = NP.linspace(0,tf,nk)
246
247# Plot the results
248plt.figure(1)
249plt.clf()
250plt.plot(tgrid,x0_opt,'--')
251plt.plot(tgrid,x1_opt,'-.')
252plt.step(tgrid_u,u_opt,'-')
253plt.title("Van der Pol optimization")
254plt.xlabel('time')
255plt.legend(['x[0] trajectory','x[1] trajectory','u trajectory'])
256plt.grid()
257plt.show()
258
259