Solve a SDE using Neural network

I am trying to solve SDE using Neural networks recently. I have an idea to solve a SDE in this form:

d X_{t}=f\left(X_{t}, t,r,\mu_t \right) d t+g\left(X_{t}, t,r,\mu_t \right) d W_{t}

But I am not sure how to implement it in DiffEqFlux package.
I have read an example by Chris Rackauckas like below


But it seems it only contains information about X_t and t , how can I add r,\mu(known) into the Neural Network?

Note that this method doesn’t solve the SDEs, it trains neural networks to be components of the SDEs. If you want to solve the SDE via neural networks you have to do something fairly different.

That said, if what you’re looking to do is add known models to the models you’re estimating, that’s what’s done in the universal differential equation tutorials, so look at:

https://diffeqflux.sciml.ai/dev/examples/LV-Univ/

as a very simple example, or for more complex examples, take a look at GitHub - ChrisRackauckas/universal_differential_equations: Repository for the Universal Differential Equations for Scientific Machine Learning paper, describing a computational basis for high performance SciML which is the repository for [2001.04385] Universal Differential Equations for Scientific Machine Learning

You could make the extension as simple as making the f and g neural networks with two more inputs, and then indeed you just define that SDE and train it.

1 Like

Thanks a lot!

I tried to replicate the experiment in your first link. But it seems the Status turns out to be failure. Is this algorithm an unstable one? Or should I set some random seed or a fixed initial point to ensure I can reach the optimal point. Is it common that you will encounter failure Status while training the model?

Updates: I rerun the code and get a success status. I guess it may because of the change of initial point.

Status is always a failure if it hits maxiters, which is somewhat misleading. So it’s working, just the printout is bad. Using the Optim output types doesn’t make sense, and working to correct that.

Hi, I try to train a SDE with neural network as part of it. But it turns out something wrong with the function sciml_train. I follow the example provided by you which solved an ODE in :
https://diffeqflux.sciml.ai/dev/examples/LV-Univ/

The code are as below




image
image

The error message is given as below

I am not familiar with Julia Language and I’m a novice in DiffEqFlux package. But this is such a great tool that can help us solve Differential Equation in a new way. Could you please take a look at the code and help me figure out why there is something wrong with my code? I would much appreciate it.

I can’t run any of this code because it’s all screenshots. Did you try the neural SDE tutorial? This works on latest versions.

Sorry for only screenshots. I had read the tutorial you provided but it seems it doesn’t fit what I need. In the example, there are only neural network as part of the equation so you can solve it by the function NeuralDSDE . Below is the code which I try to implement:

using Flux, DiffEqFlux, StochasticDiffEq, Plots, DiffEqBase.EnsembleAnalysis, Statistics,PyCall

sde_data = 1400 .+ 10*randn(50)
u_new = [1400.7,0.00054,0.0115] #the initial condotion
t  =  range(0.0,49.0,length=50); 
fixed = FastChain(FastDense(3, 50, tanh),FastDense(50, 50, relu),FastDense(50, 1))
wiener = FastChain(FastDense(3, 50, relu),FastDense(50, 1))
p_model_1 = initial_params(fixed)
n_weights_1 = length(p_model_1)
p_model_2 = initial_params(wiener)
n_weights_2 = length(p_model_2)
p_all = [p_model_1;p_model_2]

function dudt1(du, u, p, t)
    # Destructure the parameters
    model_weights = p[1:n_weights_1]
    # The neural network outputs a control taken by the system
    # The system then produces an output
    s, r,sigma = u

    # Dynamics of the control and system
    ds = fixed(u, model_weights)[1] 
    dr = 0  #收益率的变动
    dsigma = 0  # 波动率的变动

    # Update in place
    du[1] = ds
    du[2] = dr
    du[3] = dsigma
end

#random part
function dudt2(du, u, p, t)
    # Destructure the parameters
    model_weights = p[n_weights_1+1:end]
    # The neural network outputs a control taken by the system
    # The system then produces an output
    s, r,sigma = u

    # Dynamics of the control and system
    ds = wiener(u, model_weights)[1] #股票价格的变动
    dr = 0  #收益率的变动
    dsigma = 0  # 波动率的变动

    # Update in place
    du[1] = ds
    du[2] = dr
    du[3] = dsigma
end

prob_univ = SDEProblem(dudt1,dudt2,u_new,(0.0f0,49.0f0),p_all)
ensemble_nprob_univ = EnsembleProblem(prob_univ)
ensemble_nsol_univ = solve(ensemble_nprob_univ,SOSRI(),trajectories = 100, saveat = t)
ensemble_nsum_univ = EnsembleSummary(ensemble_nsol_univ)

function predict_univ(p_all)
    prob_univ = SDEProblem(dudt1,dudt2,u_new,(0.0f0,49.0f0),p_all)
    ensemble_nprob_univ = EnsembleProblem(prob_univ)
    ensemble_nsol_univ = solve(ensemble_nprob_univ,SOSRI(),trajectories = 100, saveat = t)
    ensemble_nsum_univ = EnsembleSummary(ensemble_nsol_univ)
    return(Array(ensemble_nsum_univ[1,:]))
end

loss_univ(p_all) = sum(abs2, predict_univ(p_all) - sde_data) ##sde_data is the data I try to fit
l = loss_univ(p_all)

callback = function (p_all,l) #callback function to observe training
  global iter
  if iter%50 == 0
      sample = predict_univ(p_all)
      # loss against current data
      display(sqrt(sum(abs2,sde_data .- sample)/length(sde_data)))
      # plot current prediction against data
      pl = scatter(t,sde_data,label="data")
      scatter!(pl,t,sample,label="prediction")
      display(plot(pl))
      println(l)
  end
  iter += 1
  return false
end

using DiffEqFlux, Flux, Optim, OrdinaryDiffEq, Plots

result_univ = DiffEqFlux.sciml_train(loss_univ ,p_all,BFGS(initial_stepnorm = 0.01);cb=callback,maxiters=300)

And the error message is below(I tried data = Iterators.repeated((),300 ), didn’t work)

Thank you again for your help!

What I tried to do is to fit a differential equation like below:

dX_t = f(X_t,r,u,t)dt + g(X_t,r,u,t)dt
dr/dt = 0 \\ du/dt=0

Where f,g are neural network and X_t,r,u as input. r,u are constant as part of the neural network. Do you have any efficent way that I can fit this equation by a given data. Since I don’t know how to add the piror imformation r,u into the neural network, I can just do it in this stupid way.

using Flux, DiffEqFlux, StochasticDiffEq, Plots, DiffEqBase.EnsembleAnalysis, Statistics, DiffEqSensitivity, RecursiveArrayTools

sde_data = 1400 .+ 10*randn(50)
u_new = Float32[1400.7,0.00054,0.0115] #the initial condotion
t  =  range(0.0,49.0,length=50);
fixed = FastChain(FastDense(3, 50, tanh),FastDense(50, 50, relu),FastDense(50, 1))
wiener = FastChain(FastDense(3, 50, relu),FastDense(50, 1))
p_model_1 = initial_params(fixed)
n_weights_1 = length(p_model_1)
p_model_2 = initial_params(wiener)
n_weights_2 = length(p_model_2)
p_all = [p_model_1;p_model_2]

function dudt1(u, p, t)
    # Destructure the parameters
    model_weights = p[1:n_weights_1]
    # The neural network outputs a control taken by the system
    # The system then produces an output
    s, r,sigma = u

    # Dynamics of the control and system
    ds = fixed(u, model_weights)[1]
    dr = 0  #收益率的变动
    dsigma = 0  # 波动率的变动

    [ds,dr,dsigma]
end

#random part
function dudt2(u, p, t)
    # Destructure the parameters
    model_weights = p[n_weights_1+1:end]
    # The neural network outputs a control taken by the system
    # The system then produces an output
    s, r,sigma = u

    # Dynamics of the control and system
    ds = wiener(u, model_weights)[1] #股票价格的变动
    dr = 0  #收益率的变动
    dsigma = 0  # 波动率的变动

    [ds,dr,dsigma]
end

prob_univ = SDEProblem{false}(dudt1,dudt2,u_new,(0.0f0,49.0f0),p_all)
ensemble_nprob_univ = EnsembleProblem(prob_univ)
ensemble_nsol_univ = solve(ensemble_nprob_univ,SOSRI(),trajectories = 100, saveat = t)
ensemble_nsum_univ = EnsembleSummary(ensemble_nsol_univ)

mean(ensemble_nsol_univ,dims=3)

function predict_univ(p_all)
    prob_univ = SDEProblem{false}(dudt1,dudt2,u_new,(0.0f0,49.0f0),p_all)
    ensemble_nsol_univ = [Array(solve(prob_univ,SOSRI(),dt = 0.1, saveat = t, sensealg=TrackerAdjoint())) for i in 1:100]
    ensemble_nsum_univ = sum(ensemble_nsol_univ)/length(ensemble_nsol_univ)
    return(Array(ensemble_nsum_univ[1,:]))
end

loss_univ(p_all) = sum(abs2, predict_univ(p_all) - sde_data) ##sde_data is the data I try to fit
l = loss_univ(p_all)

iter = 0
callback = function (p_all,l) #callback function to observe training
  global iter
  println(l)
  if iter%50 == 0
      sample = predict_univ(p_all)
      # loss against current data
      display(sqrt(sum(abs2,sde_data .- sample)/length(sde_data)))
      # plot current prediction against data
      pl = scatter(t,sde_data,label="data")
      scatter!(pl,t,sample,label="prediction")
      display(plot(pl))
      println(l)
  end
  iter += 1
  return false
end

result_univ = DiffEqFlux.sciml_train(loss_univ ,p_all,ADAM(0.01);cb=callback,maxiters=300)

This updates your example to follow what’s done in the tutorial and works. It’s not as elegant as it could be, but I am currently working on fixing that so spending more time on this would run counter to that effort so I’ll leave it here and when the tutorial updates in about a month it should be about 10x faster and require less workarounds.

So grateful for fast reply. But as I copy your code and run it in my Julia Pro IDE, it comes up with a error message which I cannot find an answer in the Internet, which is as below:


Need an adjoint for constructor RODESolution{Float32,2,Array{Array{Float32,1},1},Nothing,Nothing,Array{Float32,1},DiffEqNoiseProcess.NoiseProcess{Float32,2,Float32,Array{Float32,1},Array{Float32,1},Array{Array{Float32,1},1},typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE),false,DataStructures.Stack{Tuple{Float32,Array{Float32,1},Array{Float32,1}}},ResettableStacks.ResettableStack{Tuple{Float32,Array{Float32,1},Array{Float32,1}},false},DiffEqNoiseProcess.RSWM{:RSwM3,Float64},RandomNumbers.Xorshifts.Xoroshiro128Plus},SDEProblem{Array{Float32,1},Tuple{Float32,Float32},false,Array{Float32,1},Nothing,SDEFunction{false,typeof(dudt1),typeof(dudt2),LinearAlgebra.UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(dudt2),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},SOSRI,StochasticDiffEq.LinearInterpolationData{Array{Array{Float32,1},1},Array{Float32,1}},DiffEqBase.DEStats}. Gradient is of type Array{Float64,2}

I wonder is this because I didn’t add any package?

That shouldn’t show up in the version I posted. I just copy-pasted and ran it again in a clean REPL and it was fine, so check your package versions.