Put a DiffEq interface on PDE solve and structure the tests

Project.toml

@@ -7,6 +7,7 @@ version = "0.0.1"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]

src/NeuralNetDiffEq.jl

@@ -5,16 +5,10 @@ using Reexport, Statistics
 using Flux
 
 abstract type NeuralNetDiffEqAlgorithm <: DiffEqBase.AbstractODEAlgorithm end
-struct nnode{C,O} <: NeuralNetDiffEqAlgorithm
-    chain::C
-    opt::O
-end
-nnode(chain;opt=Flux.ADAM(0.1)) = nnode(chain,opt)
-export nnode
 
+include("ode_solve.jl")
+include("pde_solve.jl")
 
-include("solve.jl")
-include("PdeSolve.jl")
-
+export NNODE, TerminalPDEProblem, NNPDEHan
 
 end # module

src/solve.jl → src/ode_solve.jl

@@ -1,3 +1,9 @@
+struct NNODE{C,O} <: NeuralNetDiffEqAlgorithm
+    chain::C
+    opt::O
+end
+NNODE(chain;opt=Flux.ADAM(0.1)) = NNODE(chain,opt)
+
 function DiffEqBase.solve(
     prob::DiffEqBase.AbstractODEProblem,
     alg::NeuralNetDiffEqAlgorithm,

src/pde_solve.jl

@@ -0,0 +1,86 @@
+struct TerminalPDEProblem{G,F,Mu,Sigma,X,T,P} <: DiffEqBase.DEProblem
+    g::G
+    f::F
+    μ::Mu
+    σ::Sigma
+    X0::X
+    tspan::Tuple{T,T}
+    p::P
+    TerminalPDEProblem(g,f,μ,σ,X0,tspan,p=nothing) = new{typeof(g),typeof(f),
+                                                         typeof(μ),typeof(σ),
+                                                         typeof(X0),eltype(tspan),
+                                                         typeof(p)}(
+                                                         g,f,μ,σ,X0,tspan,p)
+end
+
+Base.summary(prob::TerminalPDEProblem) = string(nameof(typeof(prob)))
+
+function Base.show(io::IO, A::TerminalPDEProblem)
+  println(io,summary(A))
+  print(io,"timespan: ")
+  show(io,A.tspan)
+end
+
+struct NNPDEHan{C1,C2,O} <: NeuralNetDiffEqAlgorithm
+    u0::C1
+    σᵀ∇u::C2
+    opt::O
+end
+NNPDEHan(u0,σᵀ∇u;opt=Flux.ADAM(0.1)) = NNPDEHan(u0,σᵀ∇u,opt)
+
+function DiffEqBase.solve(
+    prob::TerminalPDEProblem,
+    alg::NNPDEHan;
+    timeseries_errors = true,
+    save_everystep=true,
+    adaptive=false,
+    abstol = 1f-6,
+    verbose = false,
+    maxiters = 100,
+    dt,
+    trajectories)
+
+    X0 = prob.X0
+    ts = prob.tspan[1]:dt:prob.tspan[2]
+    d  = length(X0)
+    g,f,μ,σ,p = prob.g,prob.f,prob.μ,prob.σ,prob.p
+
+    data = Iterators.repeated((), maxiters)
+
+
+    #hidden layer
+    opt = alg.opt
+    u0 = alg.u0
+    σᵀ∇u = alg.σᵀ∇u
+    ps = Flux.params(u0, σᵀ∇u...)
+
+    function sol()
+        map(1:trajectories) do j
+            u = u0(X0)[1]
+            X = X0
+            for i in 1:length(ts)-1
+                t = ts[i]
+                _σᵀ∇u = σᵀ∇u[i](X)
+                dW = sqrt(dt)*randn(d)
+                u = u - f(X, u, _σᵀ∇u, p, t)*dt + _σᵀ∇u'*dW
+                X  = X .+ μ(X,p,t)*dt .+ σ(X,p,t)*dW
+            end
+            X,u
+        end
+    end
+
+    function loss()
+        mean(sum(abs2,g(X) - u) for (X,u) in sol())
+    end
+
+
+    cb = function ()
+        l = loss()
+        verbose && println("Current loss is: $l")
+        l < abstol && Flux.stop()
+    end
+
+    Flux.train!(loss, ps, data, opt; cb = cb)
+
+    u0(X0)[1].data
+end #pde_solve

test/NNODE_tests.jl

@@ -0,0 +1,58 @@
+using Test, Flux, NeuralNetDiffEq
+using DiffEqDevTools
+
+# Run a solve on scalars
+linear = (u,p,t) -> cos(2pi*t)
+tspan = (0.0f0, 1.0f0)
+u0 = 0.0f0
+prob = ODEProblem(linear, u0 ,tspan)
+chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
+opt = Flux.ADAM(0.1, (0.9, 0.95))
+sol = solve(prob, NeuralNetDiffEq.NNODE(chain,opt), dt=1/20f0, verbose = true,
+            abstol=1e-10, maxiters = 200)
+
+# Run a solve on vectors
+linear = (u,p,t) -> [cos(2pi*t)]
+tspan = (0.0f0, 1.0f0)
+u0 = [0.0f0]
+prob = ODEProblem(linear, u0 ,tspan)
+chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
+opt = Flux.ADAM(0.1, (0.9, 0.95))
+sol = solve(prob, NeuralNetDiffEq.NNODE(chain,opt), dt=1/20f0, abstol=1e-10,
+            verbose = true, maxiters=200)
+
+#Example 1
+linear = (u,p,t) -> @. t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u*(t + ((1+3*(t^2))/(1+t+t^3)))
+linear_analytic = (u0,p,t) -> [exp(-(t^2)/2)/(1+t+t^3) + t^2]
+prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),[1f0],(0.0f0,1.0f0))
+chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
+opt = Flux.ADAM(0.1, (0.9, 0.95))
+sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/5f0)
+err = sol.errors[:l2]
+sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/20f0)
+sol.errors[:l2]/err < 0.5
+
+#=
+dts = 1f0 ./ 2f0 .^ (6:-1:2)
+sim = test_convergence(dts, prob, NeuralNetDiffEq.NNODE(chain, opt))
+@test abs(sim.𝒪est[:l2]) < 0.1
+@test minimum(sim.errors[:l2]) < 0.5
+=#
+
+#Example 2
+linear = (u,p,t) -> -u/5 + exp(-t/5).*cos(t)
+linear_analytic = (u0,p,t) ->  exp(-t/5)*(u0 + sin(t))
+prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),0.0f0,(0.0f0,1.0f0))
+chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
+opt = Flux.ADAM(0.1, (0.9, 0.95))
+sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/5f0)
+err = sol.errors[:l2]
+sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/20f0)
+sol.errors[:l2]/err < 0.5
+
+#=
+dts = 1f0 ./ 2f0 .^ (6:-1:2)
+sim = test_convergence(dts, prob, NeuralNetDiffEq.NNODE(chain, opt))
+@test abs(sim.𝒪est[:l2]) < 0.5
+@test minimum(sim.errors[:l2]) < 0.1
+=#