Merge branch 'SciML:master' into func-affect

Author: dodoplus

Project.toml

@@ -91,11 +91,12 @@ OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]
+test = ["AmplNLWriter", "BenchmarkTools", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]

docs/src/basics/FAQ.md

@@ -41,5 +41,58 @@ ERROR: TypeError: non-boolean (Num) used in boolean context
 
 then it's likely you are trying to trace through a function which cannot be
 directly represented in Julia symbols. The techniques to handle this problem,
-such as `@register_symbolic`, are described in detail 
+such as `@register_symbolic`, are described in detail
 [in the Symbolics.jl documentation](https://symbolics.juliasymbolics.org/dev/manual/faq/#Transforming-my-function-to-a-symbolic-equation-has-failed.-What-do-I-do?-1).
+
+## Using ModelingToolkit with Optimization / Automatic Differentiation
+
+If you are using ModelingToolkit inside of a loss function and are having issues with
+mixing MTK with automatic differentiation, getting performance, etc... don't! Instead, use
+MTK outside of the loss function to generate the code, and then use the generated code
+inside of the loss function.
+
+For example, let's say you were building ODEProblems in the loss function like:
+
+```julia
+function loss(p)
+    prob = ODEProblem(sys, [], [p1 => p[1], p2 => p[2]])
+    sol = solve(prob, Tsit5())
+    sum(abs2,sol)
+end
+```
+
+Since `ODEProblem` on a MTK `sys` will have to generate code, this will be slower than
+caching the generated code, and will required automatic differentiation to go through the
+code generation process itself. All of this is unnecessary. Instead, generate the problem
+once outside of the loss function, and remake the prob inside of the loss function:
+
+```julia
+prob = ODEProblem(sys, [], [p1 => p[1], p2 => p[2]])
+function loss(p)
+    remake(prob,p = ...)
+    sol = solve(prob, Tsit5())
+    sum(abs2,sol)
+end
+```
+
+Now, one has to be careful with `remake` to ensure that the parameters are in the right
+order. One can use the previously mentioned indexing functionality to generate index
+maps for reordering `p` like:
+
+```julia
+p = @parameters x y z
+idxs = ModelingToolkit.varmap_to_vars([p[1] => 1, p[2] => 2, p[3] => 3], p)
+p[idxs]
+```
+
+Using this, the fixed index map can be used in the loss function. This would look like:
+
+```julia
+prob = ODEProblem(sys, [], [p1 => p[1], p2 => p[2]])
+idxs = Int.(ModelingToolkit.varmap_to_vars([p1 => 1, p2 => 2], p))
+function loss(p)
+    remake(prob,p = p[idxs])
+    sol = solve(prob, Tsit5())
+    sum(abs2,sol)
+end
+```

docs/src/systems/SDESystem.md

@@ -24,6 +24,7 @@ sde = SDESystem(ode, noiseeqs)
 ```@docs
 structural_simplify
 alias_elimination
+Girsanov_transform
 ```
 
 ## Analyses

src/systems/diffeqs/sdesystem.jl

@@ -205,6 +205,114 @@ function stochastic_integral_transform(sys::SDESystem, correction_factor)
               name = name, checks = false)
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Measure transformation method that allows for a reduction in the variance of an estimator `Exp(g(X_t))`.
+Input:  Original SDE system and symbolic function `u(t,x)` with scalar output that
+        defines the adjustable parameters `d` in the Girsanov transformation. Optional: initial
+        condition for `θ0`.
+Output: Modified SDESystem with additional component `θ_t` and initial value `θ0`, as well as
+        the weight `θ_t/θ0` as observed equation, such that the estimator `Exp(g(X_t)θ_t/θ0)`
+        has a smaller variance.
+
+Reference:
+Kloeden, P. E., Platen, E., & Schurz, H. (2012). Numerical solution of SDE through computer
+experiments. Springer Science & Business Media.
+
+# Example
+
+```julia
+using ModelingToolkit
+
+@parameters α β
+@variables t x(t) y(t) z(t)
+D = Differential(t)
+
+eqs = [D(x) ~ α*x]
+noiseeqs = [β*x]
+
+@named de = SDESystem(eqs,noiseeqs,t,[x],[α,β])
+
+# define u (user choice)
+u = x
+θ0 = 0.1
+g(x) = x[1]^2
+demod = ModelingToolkit.Girsanov_transform(de, u; θ0=0.1)
+
+u0modmap = [
+    x => x0
+]
+
+parammap = [
+    α => 1.5,
+    β => 1.0
+]
+
+probmod = SDEProblem(demod,u0modmap,(0.0,1.0),parammap)
+ensemble_probmod = EnsembleProblem(probmod;
+          output_func = (sol,i) -> (g(sol[x,end])*sol[demod.weight,end],false),
+          )
+
+simmod = solve(ensemble_probmod,EM(),dt=dt,trajectories=numtraj)
+```
+
+"""
+function Girsanov_transform(sys::SDESystem, u; θ0 = 1.0)
+    name = nameof(sys)
+
+    # register new varible θ corresponding to 1D correction process θ(t)
+    t = get_iv(sys)
+    D = Differential(t)
+    @variables θ(t), weight(t)
+
+    # determine the adjustable parameters `d` given `u`
+    # gradient of u with respect to states
+    grad = Symbolics.gradient(u, states(sys))
+
+    noiseeqs = get_noiseeqs(sys)
+    if typeof(noiseeqs) <: Vector
+        d = simplify.(-(noiseeqs .* grad) / u)
+        drift_correction = noiseeqs .* d
+    else
+        d = simplify.(-noiseeqs * grad / u)
+        drift_correction = noiseeqs * d
+    end
+
+    # transformation adds additional state θ: newX = (X,θ)
+    # drift function for state is modified
+    # θ has zero drift
+    deqs = vcat([equations(sys)[i].lhs ~ equations(sys)[i].rhs - drift_correction[i]
+                 for i in eachindex(states(sys))]...)
+    deqsθ = D(θ) ~ 0
+    push!(deqs, deqsθ)
+
+    # diffusion matrix is of size d x m (d states, m noise), with diagonal noise represented as a d-dimensional vector
+    # for diagonal noise processes with m>1, the noise process will become non-diagonal; extra state component but no new noise process.
+    # new diffusion matrix is of size d+1 x M
+    # diffusion for state is unchanged
+
+    noiseqsθ = θ * d
+
+    if typeof(noiseeqs) <: Vector
+        m = size(noiseeqs)
+        if m == 1
+            push!(noiseeqs, noiseqsθ)
+        else
+            noiseeqs = [Array(Diagonal(noiseeqs)); noiseqsθ']
+        end
+    else
+        noiseeqs = [Array(noiseeqs); noiseqsθ']
+    end
+
+    state = [states(sys); θ]
+
+    # return modified SDE System
+    SDESystem(deqs, noiseeqs, get_iv(sys), state, parameters(sys);
+              defaults = Dict(θ => θ0), observed = [weight ~ θ / θ0],
+              name = name, checks = false)
+end
+
 """
 ```julia
 function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.ps;

test/sdesystem.jl

@@ -1,6 +1,7 @@
 using ModelingToolkit, StaticArrays, LinearAlgebra
 using StochasticDiffEq, OrdinaryDiffEq, SparseArrays
 using Random, Test
+using Statistics
 
 # Define some variables
 @parameters t σ ρ β
@@ -513,3 +514,71 @@ noiseeqs = [0.1 * x]
     @test observed(ode) == [weight ~ x * 10]
     @test solode[weight] == 10 * solode[x]
 end
+
+@testset "Measure Transformation for variance reduction" begin
+    @parameters α β
+    @variables t x(t) y(t) z(t)
+    D = Differential(t)
+
+    # Evaluate Exp [(X_T)^2]
+    # SDE: X_t = x + \int_0^t α X_z dz + \int_0^t b X_z dW_z
+    eqs = [D(x) ~ α * x]
+    noiseeqs = [β * x]
+
+    @named de = SDESystem(eqs, noiseeqs, t, [x], [α, β])
+
+    g(x) = x[1]^2
+    dt = 1 // 2^(7)
+    x0 = 0.1
+
+    ## Standard approach
+    # EM with 1`000 trajectories for stepsize 2^-7
+    u0map = [
+        x => x0,
+    ]
+
+    parammap = [
+        α => 1.5,
+        β => 1.0,
+    ]
+
+    prob = SDEProblem(de, u0map, (0.0, 1.0), parammap)
+
+    function prob_func(prob, i, repeat)
+        remake(prob, seed = seeds[i])
+    end
+    numtraj = Int(1e3)
+    seed = 100
+    Random.seed!(seed)
+    seeds = rand(UInt, numtraj)
+
+    ensemble_prob = EnsembleProblem(prob;
+                                    output_func = (sol, i) -> (g(sol[end]), false),
+                                    prob_func = prob_func)
+
+    sim = solve(ensemble_prob, EM(), dt = dt, trajectories = numtraj)
+    μ = mean(sim)
+    σ = std(sim) / sqrt(numtraj)
+
+    ## Variance reduction method
+    u = x
+    demod = ModelingToolkit.Girsanov_transform(de, u; θ0 = 0.1)
+
+    probmod = SDEProblem(demod, u0map, (0.0, 1.0), parammap)
+
+    ensemble_probmod = EnsembleProblem(probmod;
+                                       output_func = (sol, i) -> (g(sol[x, end]) *
+                                                                  sol[demod.weight, end],
+                                                                  false),
+                                       prob_func = prob_func)
+
+    simmod = solve(ensemble_probmod, EM(), dt = dt, trajectories = numtraj)
+    μmod = mean(simmod)
+    σmod = std(simmod) / sqrt(numtraj)
+
+    display("μ = $(round(μ, digits=2)) ± $(round(σ, digits=2))")
+    display("μmod = $(round(μmod, digits=2)) ± $(round(σmod, digits=2))")
+
+    @test μ≈μmod atol=2σ
+    @test σ > σmod
+end