Figure 8.10:

Open-loop simulation for \eqref {eq:vdp} using collocation.

Code for Figure 8.10

Text of the GNU GPL.

main.py

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
# [makes] collocation_integrator.mat
# Simulates a system using Gauss-Legendre collocation schemes of order 2, 4 and 6
# Joel Andersson, UW Madison 2017
import numpy as np
import casadi
from numpy.polynomial.polynomial import polyval as _polyval, polyder, polyint
from scipy.io import savemat


def _polymul(a, b):
    """Convolution of two coefficient sequences (like MATLAB conv)."""
    return np.convolve(a, b)


def _eval(coeff, x):
    """Evaluate polynomial with MATLAB's highest-degree-first coefficient order."""
    # Convert from highest-degree-first (MATLAB) to lowest-degree-first (numpy).
    return _polyval(x, coeff[::-1])


# Declare model variables
x1 = casadi.SX.sym('x1')
x2 = casadi.SX.sym('x2')
u = casadi.SX.sym('u')

# Model equations
x1_dot = (1 - x2**2) * x1 - x2 + u
x2_dot = x1

# Objective function (integral term)
quad = x1**2 + x2**2 + u**2

# Dimensions
nx = 2
np_ = 1

# Test value
x_test = np.array([0.0, 1.0])
u_test = 0.5

# End time
T = 10.0

# Number of integrator steps
N = 20

# Time step
dt = T / N

# Time grid for plotting
tgrid = np.linspace(0, T, N + 1)

# Continuous-time dynamics
f = casadi.Function('f', [casadi.vertcat(x1, x2), u],
                    [casadi.vertcat(x1_dot, x2_dot), quad],
                    ['x', 'p'], ['ode', 'quad'])

results = {}  # keyed by d

# Simulate with collocation, order 1, 2 and 3
for d in [1, 2, 3]:
    if d == 1:
        tau_root = np.array([0.0, 0.5])
    elif d == 2:
        tau_root = np.array([0.0, 0.5 - np.sqrt(3) / 6, 0.5 + np.sqrt(3) / 6])
    elif d == 3:
        tau_root = np.array([0.0, 0.5 - np.sqrt(15) / 10, 0.5,
                             0.5 + np.sqrt(15) / 10])
    else:
        raise ValueError('order not supported')

    # Degree of interpolating polynomial
    d = len(tau_root) - 1

    C = np.zeros((d + 1, d + 1))
    D = np.zeros(d + 1)
    B = np.zeros(d + 1)

    # Construct polynomial basis (Lagrange, highest-degree-first coefficients)
    for j in range(d + 1):
        coeff = np.array([1.0])
        for r in range(d + 1):
            if r != j:
                coeff = _polymul(coeff, np.array([1.0, -tau_root[r]]))
                coeff = coeff / (tau_root[j] - tau_root[r])
        D[j] = _eval(coeff, 1.0)
        pder = np.polyder(coeff)  # numpy handles highest-degree-first
        for r in range(d + 1):
            C[j, r] = np.polyval(pder, tau_root[r])
        pint = np.polyint(coeff)
        B[j] = np.polyval(pint, 1.0)

    # Start with an empty nonlinear system of equations
    w = []
    w0 = []
    rhs = []

    # x0, p are parameters of the Newton solver
    X0 = casadi.MX.sym('X0', nx)
    P = casadi.MX.sym('P', np_)

    # State vector at collocation points are unknowns
    Xj = []
    for j in range(d):
        Xj.append(casadi.MX.sym('X_{}'.format(j + 1), nx))
        w.append(Xj[-1])
        w0.append(np.zeros(nx))

    # State derivatives at collocation points
    Xj_dot = []
    for j in range(d):
        d_expr = C[0, j + 1] * X0
        for r in range(d):
            d_expr = d_expr + C[r + 1, j + 1] * Xj[r]
        Xj_dot.append(d_expr)

    # Evaluate f at collocation points
    ode_j = []
    quad_j = []
    for j in range(d):
        oj, qj = f(Xj[j], P)
        ode_j.append(oj)
        quad_j.append(qj)

    # Collocation residuals
    for j in range(d):
        rhs.append(dt * ode_j[j] - Xj_dot[j])

    # State at the end of the interval
    Xf = D[0] * X0
    for j in range(d):
        Xf = Xf + D[j + 1] * Xj[j]

    # Quadrature
    Qf = 0
    for j in range(d):
        Qf = Qf + B[j + 1] * quad_j[j] * dt

    w_vec = casadi.vertcat(*w)
    w0_vec = casadi.vertcat(*w0)
    rhs_vec = casadi.vertcat(*rhs)

    rfp = casadi.Function('rfp', [w_vec, X0, P], [rhs_vec, Xf, Qf],
                          ['w0', 'x0', 'p'], ['w', 'xf', 'qf'])
    solver = casadi.rootfinder('solver', 'newton', rfp)

    # Simulate
    x_sim = x_test.reshape(nx, 1)
    q_sim = np.array([[0.0]])
    w_sim = np.asarray(w0_vec).flatten()
    for k in range(N):
        sol = solver(x0=x_sim[:, -1], p=u_test, w0=w_sim)
        x_sim = np.hstack((x_sim, np.asarray(sol['xf'])))
        q_sim = np.hstack((q_sim, q_sim[:, -1:] + float(sol['qf'])))
        w_sim = np.asarray(sol['w']).flatten()

    results[d] = (x_sim, q_sim)

x_gl2, q_gl2 = results[1]
x_gl4, q_gl4 = results[2]
x_gl6, q_gl6 = results[3]

# Compare with CVODES with high accuracy
prob = dict(x=casadi.vertcat(x1, x2), p=u,
            ode=casadi.vertcat(x1_dot, x2_dot), quad=quad)
opts = dict(grid=tgrid.tolist(), output_t0=True, reltol=1e-13, abstol=1e-13)
integ = casadi.integrator('integ', 'cvodes', prob, opts)
sol = integ(x0=x_test, p=u_test)
x_cvodes = np.array(sol['xf'])
q_cvodes = np.array(sol['qf'])

savemat('collocation_integrator.mat',
        {'tgrid': tgrid, 'x_gl2': x_gl2, 'x_gl4': x_gl4, 'x_gl6': x_gl6,
         'x_cvodes': x_cvodes, 'q_gl2': q_gl2, 'q_gl4': q_gl4, 'q_gl6': q_gl6,
         'q_cvodes': q_cvodes})