Merge 1358783 into ea0d9f5

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.2"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ProximalOperators = "a725b495-10eb-56fe-b38b-717eba820537"
@@ -14,6 +15,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Compat = "2.2, 3.0"
+DSP = "0.6"
 ModelingToolkit = "1.1.3"
 ProximalOperators = "0.10"
 QuadGK = "2.3.1"

examples/SInDy_Examples2.jl

@@ -0,0 +1,100 @@
+using DataDrivenDiffEq
+using ModelingToolkit
+using OrdinaryDiffEq
+
+using LinearAlgebra
+using Plots
+gr()
+
+# Create a test problem
+function lorenz(u,p,t)
+    x, y, z = u
+
+    ẋ = 10.0*(y - x)
+    ẏ = x*(28.0-z) - y
+    ż = x*y - (8/3)*z
+    return [ẋ, ẏ, ż]
+end
+
+u0 = [1.0;0.0;0.0]
+tspan = (0.0,100.0)
+prob = ODEProblem(lorenz,u0,tspan)
+sol = solve(prob, Tsit5(), saveat = 0.005)
+
+plot(sol,vars=(1,2,3))
+
+# Differential data from equations
+X = Array(sol)
+DX = similar(X)
+for (i, xi) in enumerate(eachcol(X))
+    DX[:,i] = lorenz(xi, [], 0.0)
+end
+
+# Estimate differential data from state variables via a Savitzky-Golay filter
+# Test on a single variable
+windowSize, polyOrder = 9, 4
+DX1_sg = savitzky_golay(X[1,:], windowSize, polyOrder, deriv=1, dt=0.005)
+
+# Exclude borders, where the estimation is less accurate
+halfWindow = Int(ceil((windowSize+1)/2))
+DX1_sg = DX1_sg[halfWindow+1:end-halfWindow]
+DX = DX[:,halfWindow+1:end-halfWindow]
+
+# Check if the estimated derivatives are approximate to the "ground truth"
+isapprox(DX1_sg, DX[1,:], rtol=1e-2)
+
+plot(DX1_sg, label = "Estimated with Savitzky-Golay filter")
+plot!(DX[1,:],label="Ground truth")
+
+
+# Now let's estimate the derivatives for all variables and use them infer the equations
+DX_sg = similar(X)
+for i =1:size(X,1)
+    DX_sg[i,:] = savitzky_golay(X[i,:], windowSize, polyOrder, deriv=1, dt=0.005)
+end
+
+# Exclude borders, where the estimation is less accurate
+DX_sg = DX_sg[:,halfWindow+1:end-halfWindow]
+X_cropped = X[:,halfWindow+1:end-halfWindow]
+
+
+# Create a basis
+@variables u[1:3]
+
+# Lots of polynomials
+polys = [u[1]^0]
+for i ∈ 0:3
+    for j ∈ 0:3
+        for  k ∈ 0:3
+            push!(polys, u[1]^i * u[2]^j * u[3]^k)
+        end
+    end
+end
+
+# And some other stuff
+h = [1u[1];1u[2]; cos(u[1]); sin(u[1]); u[1]*u[2]; u[1]*sin(u[2]); u[2]*cos(u[2]); polys...]
+basis = Basis(h, u)
+
+# Get the reduced basis via the sparse regression
+opt = STRRidge(0.1)
+Ψ = SInDy(X_cropped, DX_sg, basis, maxiter = 100, opt = opt)
+print(Ψ)
+
+
+# Now let's try adding some noise
+using Random
+seed = MersenneTwister(3)
+X_noisy = X + 0.01*randn(seed,size(X))
+
+DX_sg = similar(X)
+X_smoothed = similar(X)
+for i =1:size(X,1)
+    DX_sg[i,:] = savitzky_golay(X_noisy[i,:], windowSize, polyOrder, deriv=1, dt=dt)
+    X_smoothed[i,:] = savitzky_golay(X_noisy[i,:], windowSize, polyOrder, deriv=0, dt=dt)
+end
+
+DX_sg = DX_sg[:,halfWindow+1:end-halfWindow]
+X_smoothed = X_smoothed[:,halfWindow+1:end-halfWindow]
+
+Ψ = SInDy(X_smoothed, DX_sg, basis, maxiter = 100, opt = opt)
+print(Ψ)

src/DataDrivenDiffEq.jl

@@ -2,7 +2,7 @@ module DataDrivenDiffEq
 
 using LinearAlgebra
 using ModelingToolkit
-using QuadGK, Statistics
+using QuadGK, Statistics, DSP
 using Compat
 
 abstract type abstractBasis end;
@@ -45,6 +45,6 @@ export ISInDy
 
 include("./utils.jl")
 export AIC, AICC, BIC
-export hankel, optimal_shrinkage, optimal_shrinkage!
+export hankel, optimal_shrinkage, optimal_shrinkage!, savitzky_golay
 
 end # module

src/utils.jl

@@ -100,3 +100,35 @@ function optimal_shrinkage!(X::AbstractArray{T, 2}) where T <: Number
     X .= U*Diagonal(S)*V'
     return
 end
+
+function savitzky_golay(x::Vector, windowSize::Integer, polyOrder::Integer; deriv::Integer=0, dt::Real=1.0)
+	# Polynomial smoothing with the Savitzky Golay filters
+	# Adapted from: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
+	# More information: https://pdfs.semanticscholar.org/066b/7534921b308925f6616480b4d2d2557943d1.pdf
+	# Requires LinearAlgebra and DSP modules loaded.
+
+	# Some error checking
+	@assert isodd(windowSize) "Window size must be an odd integer."
+	@assert polyOrder < windowSize "Polynomial order must be less than window size."
+
+	# Form the design matrix A
+	halfWindow = Int(ceil((windowSize - 1)/2))
+	A = zeros(windowSize, polyOrder+1)
+	for order = 0:polyOrder
+		A[:, order+1] = (-halfWindow:halfWindow).^(order)
+	end
+
+	# Compute the required column of the inverse of A'*A
+	# and calculate filter coefficients
+	ei = zeros(polyOrder+1)
+	ei[deriv+1] = 1.0
+	inv_col = (A'*A) \ ei
+	filterCoeffs = A*inv_col * factorial(deriv) ./(dt^deriv);
+
+	# Pad the signal with the endpoints and convolve with filter
+	paddedX = [x[1]*ones(halfWindow); x; x[end]*ones(halfWindow)]
+	y = conv(filterCoeffs[end:-1:1], paddedX)
+
+	# Return the valid midsection
+	return y[2*halfWindow+1:end-2*halfWindow]
+end

test/runtests.jl

@@ -274,4 +274,34 @@ end
     @test BIC(k, X, Y) == -2*log(sum(abs2, X -Y)) + k*log(size(X)[2])
     @test AICC(k, X, Y, likelyhood = (X,Y)->sum(abs, X-Y)) == AIC(k, X, Y, likelyhood = (X,Y)->sum(abs, X-Y))+ 2*(k+1)*(k+2)/(size(X)[2]-k-2)
 
+    function lorenz(u,p,t)
+        x, y, z = u
+        ẋ = 10.0*(y - x)
+        ẏ = x*(28.0-z) - y
+        ż = x*y - (8/3)*z
+        return [ẋ, ẏ, ż]
+    end
+    u0 = [1.0;0.0;0.0]
+    tspan = (0.0,50.0)
+    dt = 0.005
+    prob = ODEProblem(lorenz,u0,tspan)
+    sol = solve(prob, Tsit5(), saveat = dt)
+
+    X = Array(sol)
+    DX = similar(X)
+    for (i, xi) in enumerate(eachcol(X))
+        DX[:,i] = lorenz(xi, [], 0.0)
+    end
+
+    windowSize, polyOrder = 9, 4
+    halfWindow = Int(ceil((windowSize+1)/2))
+    DX_sg = similar(X)
+    for i =1:size(X,1)
+        DX_sg[i,:] = savitzky_golay(X[i,:], windowSize, polyOrder, deriv=1, dt=dt)
+    end
+
+    DX_sg = DX_sg[:,halfWindow+1:end-halfWindow]
+    DX = DX[:,halfWindow+1:end-halfWindow]
+
+    @test isapprox(DX_sg, DX, rtol=1e-2)
 end