Add trajectories to Kolmogorov Solver and fix bugs.

src/NeuralNetDiffEq.jl

@@ -53,9 +53,9 @@ function Base.show(io::IO, A::KolmogorovPDEProblem)
   print(io,"xspan: ")
   show(io,A.xspan)
   println(io , "μ")
-  show(io , A.μ)
+  show(io , A.f)
   println(io,"Sigma")
-  show(io , A.sigma)
+  show(io , A.g)
 end
 
 include("ode_solve.jl")

src/kolmogorov_solve.jl

@@ -8,10 +8,11 @@ NNKolmogorov(chain  ; opt=Flux.ADAM(0.1) , sdealg = EM() , ensemblealg = Ensembl
 
 function DiffEqBase.solve(
     prob::Union{KolmogorovPDEProblem,SDEProblem},
-    alg::NeuralNetDiffEqAlgorithm;
+    alg::NNKolmogorov;
     abstol = 1f-6,
     verbose = false,
     maxiters = 300,
+    trajectories = 1000,
     save_everystep = false,
     dt,
     kwargs...
@@ -32,7 +33,7 @@ function DiffEqBase.solve(
         phi = prob.phi
     end
     ts = tspan[1]:dt:tspan[2]
-    xs = xspan[1]:0.001:xspan[2]
+    xs = xspan[1]:0.0001:xspan[2]
     N = size(ts)
     T = tspan[2]
 
@@ -42,15 +43,15 @@ function DiffEqBase.solve(
     sdealg = alg.sdealg
     ensemblealg = alg.ensemblealg
     ps     = Flux.params(chain)
-    xi     = rand(xs ,d ,N[1])
+    xi     = rand(xs ,d ,trajectories)
     #Finding Solution to the SDE having initial condition xi. Y = Phi(S(X , T))
     sdeproblem = SDEProblem(μ,sigma,xi,tspan,noise_rate_prototype = noise_rate_prototype)
     function prob_func(prob,i,repeat)
       SDEProblem(prob.f , prob.g , xi[: , i] , prob.tspan ,noise_rate_prototype = prob.noise_rate_prototype)
     end
     output_func(sol,i) = (sol[end],false)
     ensembleprob = EnsembleProblem(sdeproblem , prob_func = prob_func , output_func = output_func)
-    sim = solve(ensembleprob, sdealg, ensemblealg , dt=0.01,trajectories=N[1],adaptive=false)
+    sim = solve(ensembleprob, sdealg, ensemblealg , dt=dt, trajectories=trajectories,adaptive=false)
     x_sde  = reshape([],d,0)
     # sol = solve(sdeproblem, sdealg ,dt=0.01 , save_everystep=false , kwargs...)
     # x_sde = sol[end]

src/stopping_solve.jl

@@ -8,7 +8,7 @@ NNStopping(chain  ; opt=Flux.ADAM(0.1) , sdealg = EM() , ensemblealg = EnsembleT
 
 function DiffEqBase.solve(
     prob::SDEProblem,
-    alg::NeuralNetDiffEqAlgorithm;
+    alg::NNStopping;
     abstol = 1f-6,
     verbose = false,
     maxiters = 300,
@@ -66,7 +66,7 @@ function DiffEqBase.solve(
         return 10000 - reward
     end
     dataset = Iterators.repeated(() , maxiters)
-    
+
     cb = function ()
         l = loss()
         un = []

test/NNKolmogorov_tests.jl

@@ -1,32 +1,41 @@
-using Test, Flux, NeuralNetDiffEq , StochasticDiffEq
-using DiffEqDevTools
+using Test, Flux, StochasticDiffEq
+println("Kolmogorov Tests")
+using DiffEqDevTools , NeuralNetDiffEq
 using Distributions
 
 
 #Using SDEProblem for the Algorithm.
 # For a diract delta take u0 = Normal(0 , sigma) where sigma --> 0
-u0 = Normal(0 , 1)
-xspan = (-10.0 , 10.0)
+u0 = Normal(1.00 , 1.00)
+xspan = (-2.0 , 6.0)
 tspan = (0.0 , 1.0)
-g(u , p , t) = 1
-f(u , p , t) = 0
+g(u , p , t) = 2.00
+f(u , p , t) = -2.00
 d = 1
 sdealg = EM()
 prob = SDEProblem(f , g , u0 , (0.0 , 1.0) ; xspan = xspan , d = d)
-opt = Flux.ADAM(0.0001)
-m = Chain(Dense(1, 256, elu) ,Dense(256 , 512 , elu), Dense(512 , 1))
-sol = solve(prob, NeuralNetDiffEq.NNKolmogorov(m,opt , sdealg), verbose = true, dt = 0.001,
-            abstol=1e-6, maxiters = 1600)
+opt = Flux.ADAM(0.01)
+m = Chain(Dense(1, 5, elu),Dense(5, 5, elu) , Dense(5 , 5 , elu) , Dense(5 , 1))
+ensemblealg = EnsembleThreads()
+sol = solve(prob, NNKolmogorov(m,opt , sdealg,ensemblealg) , verbose = true, dt = 0.01,
+            abstol=1e-10, trajectories = 100000 ,  maxiters = 500)
 # using Plots
+#
+# xs = -2:0.00001:6
 # x_val = collect(xs)
 # x_val= reshape(x_val , 1 , size(x_val)[1])
 # y_val = m(x_val)
-# y_val = reshape(y_val , 20001 , 1)
+# y_val = reshape(y_val , 800001 , 1)
 # x_val = collect(xs)
-# plot(x_val , y_val)
-# plot!(x_val , analytical(x_val))
+# plot(x_val , y_val,linewidth=3,title="Solution to the linear ODE with a thick line",
+#      xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!")
+# # plot(x_val , y_val)
+# plot!(x_val , analytical(x_val),linewidth=3,title="Solution to the linear ODE with a thick line",
+#      xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!")
+#
+
 ## The solution is obtained taking the Fourier Transform.
-analytical(xi) = pdf.(Normal(0 , 1.414) , xi)
+analytical(xi) = pdf.(Normal(3 , sqrt(1.0 + 5.00)) , xi)
 ##Validation
 x_1 = rand(xs , 1 , 1000)
 err_l2 = Flux.mse(analytical(x_1) , m(x_1))
@@ -52,9 +61,9 @@ d2 = 1
 sdealg2 = EM()
 prob2 = KolmogorovPDEProblem(f2 , g2, phi , xspan2 , tspan2, d2)
 opt2 = Flux.ADAM(0.01)
-m2 = Chain(Dense(1, 512, elu) , Dense(512 , 1024 , elu),Dense(1024 , 512 , elu), Dense(512 , 1))
-sol = solve(prob2, NeuralNetDiffEq.NNKolmogorov(m2,opt2 , sdealg2), verbose = true, dt = 0.001,
-            abstol=1e-6, maxiters = 270)
+m2 = Chain(Dense(1, 16, elu) , Dense(16 , 32 , elu),Dense(32 , 16 , elu), Dense(16 , 1))
+sol = solve(prob2, NeuralNetDiffEq.NNKolmogorov(m2,opt2 , sdealg2, ensemblealg), verbose = true, dt = 0.01,
+            trajectories = 1000 , abstol=1e-6, maxiters = 300)
 
 
 function analytical2(xi)
@@ -67,7 +76,7 @@ function analytical2(xi)
     y = reshape(y , size(xi)[1] , size(xi)[2] )
     return y
 end
-xs2 = xspan2[1]:0.001:xspan2[2]
+xs2 = -5.00:0.01:5.00
 x_val2 = rand(xs2 , d2 , 50)
 errorl2 = Flux.mse(analytical2(x_val2) , m2(x_val2))
 println("error_l2 = ", errorl2, "\n")
@@ -93,8 +102,8 @@ xspan3 = (-10.0 , 10.0)
 tspan3 = (0.0 , 1.0)
 d3 = 2
 prob = SDEProblem(f_noise , g_noise , uo3 , (0.0 , 1.0) ; xspan = xspan3 , d = d3 , noise_rate_prototype=zeros(2,4))
-opt = Flux.ADAM(0.0001)
-m3 = Chain(Dense(d3, 256, elu) ,Dense(256 , 512 , elu), Dense(512 , 1))
+opt = Flux.ADAM(0.01)
+m3 = Chain(Dense(d3, 32, elu) ,Dense(32 , 64 , elu), Dense(64 , 1))
 sol3 = solve(prob, NeuralNetDiffEq.NNKolmogorov(m3,opt , sdealg3 , EnsembleThreads()), verbose = true, dt = 0.001,
-            abstol=1e-6, maxiters = 1000)
+            abstol=1e-6, trajectories = 10000,maxiters = 200)
 println("Non-Diagonal test working.")