small fixes, and formatting update

Author: frankschae

src/systems/diffeqs/sdesystem.jl

@@ -153,7 +153,7 @@ function generate_diffusion_function(sys::SDESystem, dvs = states(sys),
     return build_function(get_noiseeqs(sys),
                           map(x -> time_varying_as_func(value(x), sys), dvs),
                           map(x -> time_varying_as_func(value(x), sys), ps),
-        get_iv(sys); kwargs...)
+                          get_iv(sys); kwargs...)
 end
 
 """
@@ -205,7 +205,6 @@ function stochastic_integral_transform(sys::SDESystem, correction_factor)
               name = name, checks = false)
 end
 
-
 """
 $(TYPEDSIGNATURES)
 
@@ -259,7 +258,7 @@ simmod = solve(ensemble_probmod,EM(),dt=dt,trajectories=numtraj)
 ```
 
 """
-function Girsanov_transform(sys::SDESystem, u; θ0=1.0)
+function Girsanov_transform(sys::SDESystem, u; θ0 = 1.0)
     name = nameof(sys)
 
     # register new varible θ corresponding to 1D correction process θ(t)
@@ -269,48 +268,49 @@ function Girsanov_transform(sys::SDESystem, u; θ0=1.0)
 
     # determine the adjustable parameters `d` given `u`
     # gradient of u with respect to states
-    grad = Symbolics.gradient(u,states(sys))
+    grad = Symbolics.gradient(u, states(sys))
 
     noiseeqs = get_noiseeqs(sys)
     if typeof(noiseeqs) <: Vector
-        d = simplify.(-(noiseeqs.*grad)/u)
-        drift_correction = noiseeqs.*d
+        d = simplify.(-(noiseeqs .* grad) / u)
+        drift_correction = noiseeqs .* d
     else
-        d = simplify.(-noiseeqs*grad/u)
-        drift_correction = noiseeqs*d
+        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 = vcat([equations(sys)[i].lhs ~ equations(sys)[i].rhs - drift_correction[i]
+                 for i in eachindex(states(sys))]...)
     deqsθ = D(θ) ~ 0
-    push!(deqs,deqsθ)
+    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
+    noiseqsθ = θ * d
 
     if typeof(noiseeqs) <: Vector
         m = size(noiseeqs)
         if m == 1
-            push!(noiseeqs,noiseqsθ)
+            push!(noiseeqs, noiseqsθ)
         else
             noiseeqs = [Array(Diagonal(noiseeqs)); noiseqsθ']
         end
     else
         noiseeqs = [Array(noiseeqs); noiseqsθ']
     end
 
-    state = [states(sys);θ]
+    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)
+              defaults = Dict(θ => θ0), observed = [weight ~ θ / θ0],
+              name = name, checks = false)
 end
 
 """

test/sdesystem.jl

@@ -517,69 +517,69 @@ 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[2]^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[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
+    @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.θ, end] /
+                                                                  sol[demod.θ, 1], 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