diff --git a/flux/nn_prior/nn_prior.jl b/flux/nn_prior/nn_prior.jl
index 2cd7bd18e7c5283bc4b7cd5e6bc461527ab9956d..6107bd7dc7d8e7bd636be3de86d87fbebed5378d 100644
--- a/flux/nn_prior/nn_prior.jl
+++ b/flux/nn_prior/nn_prior.jl
@@ -8,7 +8,7 @@ using OrdinaryDiffEq, LTVModels
     wdecay         = 0.0
     num_params     = 20
     stepsize       = 0.05
-    n_bootstrap    = 4
+    n_bootstrap    = 1
     num_montecarlo = 20
 end
 
@@ -22,7 +22,7 @@ if !isdefined(:AbstractEnsembleSystem)
 end
 
 sys = TwoLinkSys(N=1000, h=0.02, σ0 = 0.1)#; warn("Added noise")
-# sys = LinearSys(1,N=100, h=0.2, σ0 = 1, n = 3, ns = 3)
+# sys = LinearSys(1,N=100, h=0.2, σ0 = 1, n = 10, ns = 10)
 
 @everywhere function fit_model(opt, loss, m, x, y, u, xv, yv, uv, sys, modeltype;
     iters                 = 2000,
@@ -62,20 +62,22 @@ end
 
 @everywhere function fit_system(seed, sys, modeltype, jacprop = 0; doplot=false)
     x,y,u    = generate_data(sys, modeltype, seed)
-    iters = modeltype==System ? 3000 : 1500
+    iters = modeltype==System ? 2000 : 1000
     if jacprop > 0
-        iters = iters ÷ (jacprop+1)
+        iters = round(Int, iters / sqrt(jacprop+1))
         n,m = sys.ns,sys.n
         R1,R2 = 0.01eye(n*(n+m)), 10eye(n) # Fast
         P0 = 10000R1
         model = LTVModels.fit_model(LTVModels.KalmanModel, x,u,R1,R2,P0, extend=true)
-
-        xa = x .+ std(x,2)/10 .* randn(size(x))
-        ua = u .+ std(u,2)/10 .* randn(size(u))
-        ya = LTVModels.predict(model, xa,ua)
-        xt = [[x;u] [xa;ua]]
-        ut = [u ua]
-        yt = modeltype == VelSystem ? [y ya[3:4,:]] : [y ya]
+        xt,ut,yt = [x;u], u, y
+        for i = 1:jacprop
+            xa = x .+ std(x,2)/10 .* randn(size(x))
+            ua = u .+ std(u,2)/10 .* randn(size(u))
+            ya = LTVModels.predict(model, xa,ua)
+            xt = [xt [xa;ua]]
+            ut = [ut ua]
+            yt = modeltype ∈ [VelSystem, VelSystemD] ? [yt ya[3:4,:]] : [yt ya]
+        end
         x  = [x;u]
     else
         x        = [x;u]
@@ -91,7 +93,7 @@ end
 
         m          = modeltype(sys, num_params, activation)
         opt        = [ADAM(params(m), stepsize, decay=0.005); [expdecay(Param(p), wdecay) for p in params(m) if p isa AbstractMatrix]]
-        results    = fit_model(opt, loss(m), m, xt, yt, ut, xv, yv, uv, sys,modeltype, iters=iters, doplot=doplot, batch_size = size(yt,2))
+        results    = fit_model(opt, loss(m), m, xt, yt, ut, xv, yv, uv, sys,modeltype, iters=iters ÷ jacprop, doplot=doplot, batch_size = size(yt,2) ÷ jacprop)
         @pack results = x, u, y, xv, uv, yv
         println("Done: ", it)
         results
@@ -116,9 +118,9 @@ end
 
 # Produce figures
 # pyplot(reuse=false, show=false)
-# figs1 = fit_system(1, sys, System, 1, doplot=true);     #savefigures(figs1..., 1)
-# figs2 = fit_system(1, sys, DiffSystem, 1, doplot=true); #savefigures(figs2..., 2)
-# figs3 = fit_system(1, sys, VelSystem, doplot=true);  #savefigures(figs3..., 3)
+# figs1 = fit_system(1, sys, SystemD, 4, doplot=false);     #savefigures(figs1..., 1)
+# figs2 = fit_system(1, sys, DiffSystemD, 2, doplot=false); #savefigures(figs2..., 2)
+# figs3 = fit_system(1, sys, VelSystemD, 2, doplot=false);  #savefigures(figs3..., 3)
 # figs4 = fit_system(1, sys, AffineSystem, doplot=true);  #savefigures(figs4..., 4)
 # error()
 
@@ -162,15 +164,16 @@ end
 ##
 
 res = map(1:num_montecarlo) do it
-    r1 = @spawn fit_system(it, sys, System)
-    r2 = @spawn fit_system(it, sys, DiffSystem)
-    r3 = @spawn fit_system(it, sys, VelSystem)
-    r4 = @spawn fit_system(it, sys, System, 1)
-    r5 = @spawn fit_system(it, sys, DiffSystem, 1)
-    r6 = @spawn fit_system(it, sys, VelSystem, 1)
+    r1 = @spawn fit_system(it, sys, SystemD, 2)
+    r2 = @spawn fit_system(it, sys, DiffSystemD, 2)
+    r3 = @spawn fit_system(it, sys, VelSystemD, 2)
+    r4 = @spawn fit_system(it, sys, System, 2)
+    r5 = @spawn fit_system(it, sys, DiffSystem, 2)
+    r6 = @spawn fit_system(it, sys, VelSystem, 2)
 
     println("Done with montecarlo run $it")
     r1,r2,r3,r4,r5,r6
+    # r1,r2,r4,r5
 end
 res = [(fetch.(rs)...) for rs in res]
 
@@ -178,7 +181,8 @@ open(file->serialize(file, res), "res","w") # Save
 error()
 
 ##
-sleep(60*25)
+sleep(60*3)
+##
 using StatPlots
 res = open(file->deserialize(file), "res") # Load
 sys = res[1][1][1][:sys]
@@ -207,6 +211,7 @@ violin!((1:nr)',identity.(jacerrp), side=:right, xticks=(1:nr,labelvec), ylabel=
 # savefig2("/local/home/fredrikb/papers/nn_prior/figs/jacerr.tex")
 
 ##
+
 # TODO: ====================================================================================
 # Simulation performance
 # Prediction performance
diff --git a/flux/nn_prior/utilities.jl b/flux/nn_prior/utilities.jl
index 4744688f96904e57f9bf791acff675ccf874ba67..ac31e44a53d58f08a6f659de98531e067af6bebc 100644
--- a/flux/nn_prior/utilities.jl
+++ b/flux/nn_prior/utilities.jl
@@ -1,4 +1,5 @@
 using ForwardDiff, StatsBase, Parameters, Flux
+import StatsBase.predict
 
 abstract type AbstractSystem end
 
@@ -10,7 +11,7 @@ function System(sys, num_params, activation)
     @unpack n,ns = sys
     ny = ns
     np = num_params
-    m  = Chain(Dense(ns+n,np, activation), Dense(np, ny))
+    m  = Chain(Dense(ns+n,np, activation), Dropout(0.5), Dense(np, ny))
     System(m, sys)
 end
 (m::System)(x) = m.m(x)
@@ -24,7 +25,7 @@ function DiffSystem(sys, num_params, activation)
     @unpack n,ns = sys
     ny = ns
     np = num_params
-    m  = Chain(Dense(ns+n,np, activation), Dense(np, ny))
+    m  = Chain(Dense(ns+n,np, activation), Dropout(0.5), Dense(np, ny))
     DiffSystem(m, sys)
 end
 (m::DiffSystem)(x) = m.m(x)+x[1:m.sys.ns,:]
@@ -38,11 +39,30 @@ function VelSystem(sys, num_params, activation)
     @unpack n,ns = sys
     ny = n
     np = num_params
-    m = Chain(Dense(ns+n,np, activation), Dense(np, ny))
+    m = Chain(Dense(ns+n,np, activation), Dropout(0.5), Dense(np, ny))
     VelSystem(m, sys)
 end
 (m::VelSystem)(x) = m.m(x)
 
+struct SystemD <: AbstractSystem
+    s::System
+end
+struct DiffSystemD <: AbstractSystem
+    s::DiffSystem
+end
+struct VelSystemD <: AbstractSystem
+    s::VelSystem
+end
+for (S1,S2) in zip([:SystemD, :DiffSystemD, :VelSystemD],[:System, :DiffSystem, :VelSystem])
+    @eval $(S1)(args...) = $(S1)($(S2)(args...))
+    # @show S1
+    @eval StatsBase.predict(ms::Vector{$S1}, x, mc=10) = (mean(i->predict(getfield.(ms,:s), x, false)[1], 1:mc), std(cat(3,[predict(getfield.(ms,:s), x, false)[1] for i = 1:mc]...), 3))
+    @eval simulate(ms::Vector{$S1}, x, mc=10) = mean(i->simulate(getfield.(ms,:s), x, false), 1:mc)
+    @eval Flux.jacobian(ms::Vector{$S1}, x, mc=10) = (mean(i->jacobian(getfield.(ms,:s), x, false)[1], 1:mc), std(cat(3,[jacobian(getfield.(ms,:s), x, false)[1] for i = 1:mc]...), 3))
+    @eval (s::$(S1))(x) = s.s(x)
+end
+
+#=
 struct AffineSystem <: AbstractSystem
     m
     sys
@@ -82,21 +102,20 @@ function AffineSystem(sys, num_params, activation)
 end
 (m::AffineSystem)(x) = m.m(x)
 
+=#
+
 loss(m::AbstractSystem) = (x,y) -> sum((m(x).-y).^2)/size(x,2)
+foreach(treelike, [SystemD, DiffSystemD, VelSystemD, System, DiffSystem, VelSystem])
 
-treelike(System)
-treelike(DiffSystem)
-treelike(VelSystem)
-treelike(MatParam2)
-treelike(AffineSystem)
 const AbstractEnsembleSystem = Vector{<:AbstractSystem}
 const EnsembleSystem = Vector{System}
 const EnsembleDiffSystem = Vector{DiffSystem}
 const EnsembleVelSystem = Vector{VelSystem}
-const EnsembleAffineSystem = Vector{AffineSystem}
+# const EnsembleAffineSystem = Vector{AffineSystem}
 
 xy(::Type{<:AbstractSystem}, x, N)     = (x[:,1:N], x[:,2:N+1])
 xy(::Type{VelSystem}, x, N)  = (x[:,1:N], x[3:4,2:N+1])
+xy(::Type{VelSystemD}, x, N)  = (x[:,1:N], x[3:4,2:N+1])
 
 
 function get_minimum_loss(results, key)
@@ -104,8 +123,8 @@ function get_minimum_loss(results, key)
     return isempty(array) ? Inf : minimum(array)
 end
 
-function simulate(ms::AbstractEnsembleSystem,x)
-    Flux.testmode!.(ms)
+function simulate(ms::AbstractEnsembleSystem,x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     xsim = copy(x)
     ns = ms[1].sys.ns
     for t = 2:size(x,2)
@@ -115,8 +134,8 @@ function simulate(ms::AbstractEnsembleSystem,x)
     xsim[1:ns,:]
 end
 
-function simulate(ms::EnsembleVelSystem,x)
-    Flux.testmode!.(ms)
+function simulate(ms::EnsembleVelSystem,x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     xsim = copy(x)
     for t = 2:size(x,2)
         xsimt = map(m->m(xsim[:,t-1]), ms)
@@ -127,14 +146,14 @@ function simulate(ms::EnsembleVelSystem,x)
     xsim[1:4,:]
 end
 
-function StatsBase.predict(ms::AbstractEnsembleSystem, x)
-    Flux.testmode!.(ms)
+function StatsBase.predict(ms::AbstractEnsembleSystem, x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     y = map(m->m(x), ms)
     mean(y).data, getfield.(extrema(y), :data)
 end
 
-function StatsBase.predict(ms::EnsembleVelSystem, x)
-    Flux.testmode!.(ms)
+function StatsBase.predict(ms::EnsembleVelSystem, x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     y = map(m->m(x), ms)
     yh = mean(y).data
     h = ms[1].sys.h
@@ -143,15 +162,15 @@ function StatsBase.predict(ms::EnsembleVelSystem, x)
     yh, bounds
 end
 
-function Flux.jacobian(ms::AbstractEnsembleSystem, x)
-    Flux.testmode!.(ms)
+function Flux.jacobian(ms::AbstractEnsembleSystem, x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     jacs = [Flux.jacobian(m,x) for m in ms]
     jacmat = cat(3,jacs...)
     squeeze(mean(jacmat, 3), 3), squeeze(std(jacmat, 3), 3)
 end
 
-function Flux.jacobian(ms::EnsembleVelSystem, x)
-    Flux.testmode!.(ms)
+function Flux.jacobian(ms::EnsembleVelSystem, x, testmode=true)
+    Flux.testmode!.(ms, testmode)
     h = ms[1].sys.h
     jacs = [[[I h*I h^2/2*eye(2)];Flux.jacobian(m,x)] for m in ms] # The h²/2*I in ∇ᵤ is an approximation since there are (very small) cross terms. TODO: hard coded sample time
     jacmat = cat(3,jacs...)
@@ -164,7 +183,7 @@ function all_jacobians(results, eval=false)
     @unpack x,u,xv,uv,sys,modeltype = results[1]
     x,u = eval ? (xv,uv) : (x,u)
     N = size(xv,2)
-    J = map(1:N) do evalpoint
+    J = pmap(1:N) do evalpoint
         Jm, Js = jacobian(models(results), x[:,evalpoint])
         Jtrue = true_jacobian(sys,evalpoint, x, u)
         Jm[:], Js[:], Jtrue[:]
@@ -177,7 +196,7 @@ end
 
 function plot_jacobians(Jm, Js, Jtrue)
     N = size(Jm,1)
-    colors = [HSV(h,1,0.8) for h in linspace(0,254,24)]
+    colors = [HSV(h,1,0.8) for h in linspace(0,254,N)]
     plot(Jtrue',lab="True", l=:dash, c=colors', layout=N, legend=false)
     for i = 1:N
         plot!(Jm[i,:], ribbon=2Js[i,:], lab="Estimated", c=colors[i], subplot=i, show=false)