---

Raj Dandekar, 20th May 2020

---

Sparse Identification of Dynamical Systems with Julia

This is a Julia implementation of the examples presented in "Discovering governing equations from data by sparse identification of nonlinear dynamical systems", Brunton et al, PNAS, 2016. Three examples are presented. (a) Lorenz attractor. (b) Vortex shedding behind a cylinder. (c) Hopf bifurcation. For details regarding the physics of each example, refer to the above paper.

using DifferentialEquations, Plots, LinearAlgebra, MAT, LaTeXStrings

Helper function 1: Collect basis

Through this function, we accumulate polynomial basis upto degree 4. It takes as its input: (a) the data matrix: x; (b) the number of variables: nvar and (c) the maximum polynomial order: npoly. Each column of the resulting matrix will be for a particular function (X, XY, X^{2}Y etc..) for all times considered.

function collect_basis(x :: AbstractArray, nvar, npoly)
    basis = ones(size(x,1), 1)
    basis = hcat(basis, x)

    if npoly >= 2
        for i in 1:nvar
            for j in i:nvar
                basis = hcat(basis, (x[:, i] .* x[:, j]) )
            end
        end
    end


    if npoly >= 3
        for i in 1:nvar
            for j in i:nvar
                for k in j:nvar
                basis = hcat(basis, (x[:,i] .* x[:, j] .* x[:, k]) )
            end
            end
        end
    end


    if npoly >= 4
        for i in 1:nvar
            for j in i:nvar
                for k in j:nvar
                    for l in k:nvar
                basis =  hcat(basis, (x[:,i].* x[:, j] .* x[:,k] .* x[:,l]))
            end
            end
        end
    end
end


    return basis
end

collect_basis (generic function with 1 method)

Helper function 2: Collect annotated table

This function takes as its arguments: (a) the string of variables: s (for eg: ('x' 'y' 'z')); (b) the sparse matrix: ϵ; (c) the number of variables: nvar and (d) the maximum polynomal order: npoly. It adds an addtional column to the sparse matrix ϵ, with the string names. Through this, we can easily identify the variables which correspond to non-zero entries in the sparse matrix and hence, need to be retained for constructing the ODE's.

function collect_table(s, epsilon, nvar, npoly)
    table = []
    push!(table, "")

    push!(table, "1")

    for i in 1:nvar
        push!(table, s[i])
    end

    if npoly >= 2
        for i in 1:nvar
            for j in i:nvar
                push!(table, string(s[i], s[j]))
            end
        end
    end


    if npoly >= 3
        for i in 1:nvar
            for j in i:nvar
                for k in j:nvar
                  push!(table, string(s[i], s[j], s[k]))
            end
            end
        end
    end


    if npoly >= 4
        for i in 1:nvar
            for j in i:nvar
                for k in j:nvar
                    for l in k:nvar
                      push!(table, string(s[i], s[j], s[k], s[l]))

            end
            end
        end
    end
end

header = []

for i in 1:length(s)
    push!(header, string(s[i], dot))
end

header = permutedims(header)

table = hcat(table, vcat(header, epsilon))

    return table
end

collect_table (generic function with 1 method)

Helper function 3: Generate sparse matrix using Sequential Thresholding Algorithm

This function takes as its input: (a) the basis matrix returned from the collect_basis function: θ; (b) the derivative data: dx; (c) the thresholding parameter: λ; (d) number of variables: nvars and (e) maximum number of iterations: maxiter. It returns a sparse matrix ϵ.

function sparse_matrix(Theta :: AbstractArray, dx :: AbstractArray, lambda, nvars, maxiter)
   Epsilon = Theta\dx

   Epsilon_c = similar(Epsilon)
   Epsilon_c .= Epsilon

    for i in 1:maxiter
        index_remove = abs.(Epsilon) .<= lambda
        Epsilon[index_remove] .= zero(eltype(Theta))
        for j in 1:nvars
            index_retain = @. !index_remove[:,j]
            Epsilon[index_retain, j] = Theta[:, index_retain]\dx[:,j]
        end
    end
    return Epsilon
end

sparse_matrix (generic function with 1 method)

Example 1. The Lorenz Attractor

Extracting and storing data from the original system

function lorenz!(u,p,t)
    σ,ρ,β = p
    du = zeros(3, 1)
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
    return [du[1], du[2], du[3]]
end

u0 = [-8.0,8.0,27]

p = (10,28,8/3)

tspan = (0.0,100.0)
dt = 0.001
t = range(tspan[1],tspan[2],length=100001)


prob = ODEProblem(lorenz!, u0, tspan, p)
sol = solve(prob, Tsit5(), saveat = dt)

X = Array(sol);

DX = similar(X);

for (i, xi) in enumerate(eachcol(X))
    DX[:,i] = lorenz!(xi, p, 0.0)
end

X = X';

DX = DX';

gr()
plot(sol,vars=(1,2,3), lw=1, label = "True Lorenz Attractor")

Basis collection and sparsity functions

Theta = collect_basis(X, 3,4)

print("Basis collection done")

Basis collection done

Apply sparsity algorithm to generate sparse matrix

lambda = 0.025

Xi = sparse_matrix(Theta,DX,lambda,3, 20)

print("Sparse Matrix Generation done")

Sparse Matrix Generation done

Generate annotated table

s = ['x' 'y' 'z']
T = collect_table(s, Xi, 3, 4)

36×4 Array{Any,2}:
 ""         "xdot"    "ydot"    "zdot"
 "1"       0.0       0.0       0.0    
 'x'     -10.0      28.0       0.0    
 'y'      10.0      -1.0       0.0    
 'z'       0.0       0.0      -2.66667
 "xx"      0.0       0.0       0.0    
 "xy"      0.0       0.0       1.0    
 "xz"      0.0      -1.0       0.0    
 "yy"      0.0       0.0       0.0    
 "yz"      0.0       0.0       0.0    
 ⋮                                    
 "xyyy"    0.0       0.0       0.0    
 "xyyz"    0.0       0.0       0.0    
 "xyzz"    0.0       0.0       0.0    
 "xzzz"    0.0       0.0       0.0    
 "yyyy"    0.0       0.0       0.0    
 "yyyz"    0.0       0.0       0.0    
 "yyzz"    0.0       0.0       0.0    
 "yzzz"    0.0       0.0       0.0    
 "zzzz"    0.0       0.0       0.0

Construct equations based on variables retained from annotated table and plot recovered solution

σ1, σ2 = Xi[:,1][findall(x -> x !=0, Xi[:,1])]
ρ1, ρ2, ρ3 = Xi[:,2][findall(x -> x !=0, Xi[:,2])]
β1, β2 = Xi[:,3][findall(x -> x !=0, Xi[:,3])]



function lorenz_recover!(u,p,t)
    du = zeros(3, 1)
    du[1] =  σ2*u[2] + σ1*u[1]
    du[2] =  ρ1*u[1] + ρ2*u[1]*u[3] + ρ3*u[2]
    du[3] = β2*u[1]*u[2] + β1*u[3]
    return [du[1], du[2], du[3]]
end

u0 = [-8.0,8.0,27]

tspan = (0.0,100.0)
dt = 0.001
t = range(tspan[1],tspan[2],length=100001)


prob = ODEProblem(lorenz_recover!, u0, tspan, [])

sol_recover = solve(prob, Tsit5(), saveat = dt)

X_recover= Array(sol_recover)

X_recover = X_recover'

gr()

Plots.GRBackend()

Generate animation

plt = plot3d(
    1,
    xlim = (-30, 30),
    ylim = (-30, 30),
    zlim = (0, 60),
    xlabel = "u1",
    ylabel = "u2",
    zlabel = "u3",
    title = "Lorenz Attractor",
    marker = 2,
    label =   "SINDY Recovered solution"
)


@gif for i=1:5: 20000
    push!(plt, X_recover[:,1][i], X_recover[:,2][i], X_recover[:,3][i])
end every 10

Plot original solution on top of recovered solution to compare

pyplot()

using Measures
plot(X[:,1][1:10:end], X[:,2][1:10:end], X[:,3][1:10:end], label = "True Lorenz Attractor")

plot!(X_recover[:,1][1:10:end], X_recover[:,2][1:10:end], X_recover[:,3][1:10:end], label = "Recovered Lorenz Attractor", right_margin = 10mm, xlabel = "u1", ylabel = "u2", zlabel = "u3")

Example 2. Vortex Shedding behind a cylinder.

Data is obtained from Brunton et al, PNAS, 2016. Extracting and storing data from the original system

vars = matread("C:/Users/Raj/Desktop/SINDy/sparsedynamics/sparsedynamics/data/PODcoefficients.mat")

x = [vars["alpha"][1:5000, 1:2] vars["alphaS"][1:5000, 1]]

dx = zeros(size(x, 1)-5, size(x, 2))
dt = 0.02
for i in 3: size(x, 1) -3
    for j in 1:size(x, 2)
        dx[i-2, j] = (1/(12*dt))*( - x[i+2, j] + 8*x[i+1, j] -8*x[i-1, j] + x[i-2,j])
    end
end

vars = matread("C:/Users/Raj/Desktop/SINDy/sparsedynamics/sparsedynamics/data/PODcoefficients_run1.mat")

x1 = [vars["alpha"][1:3000, 1:2] vars["alphaS"][1:3000, 1]]

dx1 = zeros(size(x1, 1)-5, size(x1, 2))

for i in 3: size(x1, 1) -3
    for j in 1:size(x1, 2)
        dx1[i-2, j] = (1/(12*dt))*( - x1[i+2, j] + 8*x1[i+1, j] -8*x1[i-1, j] + x1[i-2,j])
    end
end

x = [x[3:end-3, :]; x1[3:end-3, :]]

dx = [dx; dx1]

print("Collection of data done")

Collection of data done

Basis collection and sparsity functions

Theta = collect_basis(x, 3,4)

print("Basis collection done")

Basis collection done

Apply sparsity algorithm to generate sparse matrix

lambda = 0.00002

Xi = sparse_matrix(Theta,dx,lambda, 3, 50)

print("Sparse Matrix Generation done")

Sparse Matrix Generation done

Generate annotated table

s = ['x' 'y' 'z']
T = collect_table(s, Xi, 3, 4)

36×4 Array{Any,2}:
 ""        "xdot"        "ydot"         "zdot"    
 "1"     -0.122493     -0.0569325    -21.9004     
 'x'     -0.00917506    1.03466       -0.000909618
 'y'     -1.02243       0.00467241     5.32038e-5 
 'z'     -0.000922031  -0.000449316   -0.311718   
 "xx"     0.0           0.0            0.00113813 
 "xy"     0.0           0.0            0.000222186
 "xz"     0.000212607   0.00218728     0.0        
 "yy"     0.0           0.0            0.000914273
 "yz"    -0.00193115   -0.0017535      0.0        
 ⋮                                                
 "xyyy"   0.0           0.0            0.0        
 "xyyz"   0.0           0.0            0.0        
 "xyzz"   0.0           0.0            0.0        
 "xzzz"   0.0           0.0            0.0        
 "yyyy"   0.0           0.0            0.0        
 "yyyz"   0.0           0.0            0.0        
 "yyzz"   0.0           0.0            0.0        
 "yzzz"   0.0           0.0            0.0        
 "zzzz"   0.0           0.0            0.0

Construct equations based on variables retained from annotated table and plot recovered solution The solution looks very similar to that identified through CFD simulations (refer Brunton et al, PNAS, 2016)

σ = Xi[:,1][findall(x -> x !=0, Xi[:,1])]
β = Xi[:,2][findall(x -> x !=0, Xi[:,2])]
γ = Xi[:,3][findall(x -> x !=0, Xi[:,3])]


function cylinder_recover!(u,p,t)
    du = zeros(3, 1)
    du[1] =  σ[1] + σ[2]*u[1] + σ[3]*u[2] + σ[4]*u[3] + σ[5]*u[1]*u[3] + σ[6]*u[2]*u[3]
    du[2] =  β[1] + β[2]*u[1] + β[3]*u[2] + β[4]*u[3] + β[5]*u[1]*u[3] + β[6]*u[2]*u[3]
    du[3] =  γ[1] + γ[2]*u[1] + γ[3]*u[2] + γ[4]*u[3] + γ[5]*u[1]*u[1] + γ[6]*u[1]*u[2] + γ[7]*u[2]*u[2] + γ[8]*u[3]*u[3]
    return [du[1], du[2], du[3]]
end

u0 = x[1,:]

tspan = (0.0,500.0)
dt = 0.02

t = range(tspan[1],tspan[2],length= 5001)

prob = ODEProblem(cylinder_recover!, u0, tspan, [])
sol_recover = solve(prob, Tsit5(), saveat = dt)

X_recover = Array(sol_recover)

X_recover = X_recover'

gr()

plt = plot3d(
    1,
    xlim = (-150, 150),
    xlabel = "x",
    ylabel = "y",
    zlabel = "z",
    ylim = (-150, 150),
    zlim = (-150, 0),
    title = "Vortex Shedding past a cylinder",
    marker = 2,
    label =   "SINDY Recovered solution"
)


@gif for i=500: 3: 15000
    push!(plt, X_recover[:,1][i], X_recover[:,2][i], X_recover[:,3][i])
end every 10

Example 3. Hopf Bifurcation

Data is obtained from Brunton et al, PNAS, 2016. Extracting and storing data from the original system

vars = matread("C:/Users/Raj/Desktop/SINDy/sparsedynamics/sparsedynamics/Hopf_SINDY.mat")

x = vars["x"]
dx = vars["dx"]

print("Collection of data done")

Collection of data done

Basis collection and sparsity functions

Theta = collect_basis(x, 3,4)
print("Basis collection done")

Basis collection done

Apply sparsity algorithm to generate sparse matrix

lambda = 0.85

Xi = sparse_matrix(Theta,dx,lambda, 3, 20)

print("Sparse Matrix Generation done")

Sparse Matrix Generation done

Generate annotated table

s = ['x' 'y' 'u']
T = collect_table(s, Xi, 3, 4)

36×4 Array{Any,2}:
 ""        "xdot"    "ydot"    "udot"
 "1"      0.0       0.0       0.0    
 'x'      0.0       0.991598  0.0    
 'y'     -0.991902  0.0       0.0    
 'u'      0.0       0.0       0.0    
 "xx"     0.0       0.0       0.0    
 "xy"     0.0       0.0       0.0    
 "xu"     0.92845   0.0       0.0    
 "yy"     0.0       0.0       0.0    
 "yu"     0.0       0.927828  0.0    
 ⋮                                   
 "xyyy"   0.0       0.0       0.0    
 "xyyu"   0.0       0.0       0.0    
 "xyuu"   0.0       0.0       0.0    
 "xuuu"   0.0       0.0       0.0    
 "yyyy"   0.0       0.0       0.0    
 "yyyu"   0.0       0.0       0.0    
 "yyuu"   0.0       0.0       0.0    
 "yuuu"   0.0       0.0       0.0    
 "uuuu"   0.0       0.0       0.0

Construct equations based on variables retained from annotated table and plot recovered solution

σ = Xi[:,1][findall(x -> x !=0, Xi[:,1])]
β = Xi[:,2][findall(x -> x !=0, Xi[:,2])]
γ = Xi[:,3][findall(x -> x !=0, Xi[:,3])]

function hopf_recover!(u,p,t)
    du = zeros(3, 1)
    du[1] =  σ[1]*u[2] + σ[2]*u[1]*u[3] + σ[3]*u[1]*u[1]*u[1] + σ[4]*u[1]*u[2]*u[2]
    du[2] =  β[1]*u[1] + β[2]*u[2]*u[3] + β[3]*u[1]*u[1]*u[2] + β[4]*u[2]*u[2]*u[2]
    du[3] =  0
    return [du[1], du[2], du[3]]
end

μ = [-0.15, -0.05, 0.05, 0.15, 0.25, 0.35, 0.45, 0.55]
tspan = (0.0025,75.0)
dt = 0.0025

u0 = [2, 0, μ[1]]
prob = ODEProblem(hopf_recover!, u0, tspan, [])
sol_recover = solve(prob, Tsit5(), saveat = dt)
X_recover = Array(sol_recover)
X_recover = X_recover'

for i = 2:length(μ)
    global X_recover
    u0 = [2, 0, μ[i]]
    prob = ODEProblem(hopf_recover!, u0, tspan, [])
    sol_recover = solve(prob, Tsit5(), saveat = dt)
    X = Array(sol_recover)
    X = X'
    X_recover = vcat(X_recover, X)
end

gr()

plt = plot3d(
    1,
    xlim = (-0.2, 0.6),
    ylim = (-0.6, 0.6),
    zlim = (-0.7, 0.7),
    xlabel = "Bifurcation parameter",
    ylabel = "x",
    zlabel = "y",
    title = " Hopf Bifurcation Cone Identification",
    marker = 2,
    label =   "SINDY Recovered solution",

)

@gif for i=1:30: 240000
    push!(plt, X_recover[:,3][i], X_recover[:,1][i], X_recover[:,2][i])
end every 10

Plot cone of bifurcation

datasize = 100
Rvspan = (0, sqrt(.6))
Tvspan = (0, 2*π)
Rv = range(Rvspan[1],Rvspan[2],length=datasize)
Tv = range(Tvspan[1],Tvspan[2],length=datasize)

x = zeros(1, length(Rv)*length(Tv))
y = zeros(1, length(Rv)*length(Tv))

count = 1
for i = 1:length(Tv)
    for j = 1:length(Rv)
        global count
        x[count] = Rv[j]*cos(Tv[i])
        y[count] = Rv[j]*sin(Tv[i])
        count+=1
    end
end

f(x,y) = x.^2 + y.^2

g = f(x, y)

plot!(g', x', y', label = "")

---