ci_types added to EnsembleSummary plot

hetalang · Jul 12, 2023 · 8ae3c32 · 8ae3c32
1 parent dcdca44
commit 8ae3c32
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 2 deletions.
diff --git a/cases/story_1/run.jl b/cases/story_1/run.jl
@@ -178,7 +178,11 @@ timeseries_steps_meancor(mcsim1) # Computes the correlation matrix and means at
 
 # Ensemble Summary
 ens = EnsembleSummary(mcsim1;quantiles=[0.05,0.95])
-plot(ens)
+
+plot(ens, ci_type=:SEM, title="SEM")
+plot(ens, ci_type=:std, title="std")
+plot(ens, ci_type=:variance, vars=[:a,:b], title="variance")
+plot(ens, ci_type=:quantile, vars=[:c], title="quantile")
 
 ################################## Monte-Carlo Parallel #####################
 #= 

diff --git a/src/HetaSimulator.jl b/src/HetaSimulator.jl
@@ -25,7 +25,7 @@ module HetaSimulator
   @reexport using Distributions
   using LinearAlgebra: pinv, diag
   using Distributed
-  using RecursiveArrayTools: vecvec_to_mat, VectorOfArray
+  using RecursiveArrayTools: vecvec_to_mat, VectorOfArray, vecarr_to_vectors
   using ProgressMeter
   #ProgressMeter.ijulia_behavior(:clear)
   # measurements 

diff --git a/src/plots.jl b/src/plots.jl
@@ -114,9 +114,76 @@ end
   end
 end
 
+#=
 @recipe function plot(ens::LabelledEnsembleSummary; vars=observables(ens))
   @series begin
     trajectories := [findfirst(x->x == v, observables(ens)) for v in vars]
     ens.ens
   end
 end
+=#
+
+#https://github.com/SciML/SciMLBase.jl/blob/32cae24c58f9d189a46d0068d7763a948fc5e1d7/src/ensemble/ensemble_solutions.jl#L147C1-L199C4
+@recipe function f(ens::LabelledEnsembleSummary;
+  vars=observables(ens), 
+  error_style = :ribbon, ci_type = :quantile)
+
+  trajectories = [findfirst(x->x == v, observables(ens)) for v in vars]
+  sim = ens.ens
+  if ci_type in [:SEM, :std, :variance]
+      if typeof(sim.u[1]) <: AbstractArray
+          u = vecarr_to_vectors(sim.u)
+      else
+          u = [sim.u.u]
+      end
+
+      if ci_type==:SEM 
+        val = [sqrt.(sim.v[i] / sim.num_monte) .* 1.96 for i in 1:length(sim.v)] 
+      elseif ci_type==:std 
+        val = [sqrt.(sim.v[i]) for i in 1:length(sim.v)]
+      elseif ci_type==:variance
+        val = [sim.v[i] for i in 1:length(sim.v)]
+      end
+
+      if typeof(sim.u[1]) <: AbstractArray
+          ci_low = vecarr_to_vectors(VectorOfArray(val))
+          ci_high = ci_low
+      else
+          ci_low = [val]
+          ci_high = ci_low
+      end
+  elseif ci_type == :quantile
+      if typeof(sim.med[1]) <: AbstractArray
+          u = vecarr_to_vectors(sim.med)
+      else
+          u = [sim.med.u]
+      end
+      if typeof(sim.u[1]) <: AbstractArray
+          ci_low = u - vecarr_to_vectors(sim.qlow)
+          ci_high = vecarr_to_vectors(sim.qhigh) - u
+      else
+          ci_low = [u[1] - sim.qlow.u]
+          ci_high = [sim.qhigh.u - u[1]]
+      end
+  else
+      error("ci_type choice not valid. Must be :variance or :quantile")
+  end
+
+  for i in trajectories
+      @series begin
+          legend --> false
+          linewidth --> 3
+          fillalpha --> 0.2
+          if error_style == :ribbon
+              ribbon --> (ci_low[i], ci_high[i])
+          elseif error_style == :bars
+              yerror --> (ci_low[i], ci_high[i])
+          elseif error_style == :none
+              nothing
+          else
+              error("error_style not recognized")
+          end
+          sim.t, u[i]
+      end
+  end
+end