From 37f61f469bc6ea518ecb0afaf8f01103b0688724 Mon Sep 17 00:00:00 2001
From: Mahmoud Bentriou <mahmoud.bentriou@centralesupelec.fr>
Date: Tue, 17 Nov 2020 02:50:55 +0100
Subject: [PATCH] Useful brackets indexing for trajectories + test.

---
 core/observations.jl                     |  5 ++---
 tests/run_unit.jl                        |  1 +
 tests/unit/getindex_access_trajectory.jl | 11 +++++++++++
 3 files changed, 14 insertions(+), 3 deletions(-)
 create mode 100644 tests/unit/getindex_access_trajectory.jl

diff --git a/core/observations.jl b/core/observations.jl
index 99d2b55..c59d18b 100644
--- a/core/observations.jl
+++ b/core/observations.jl
@@ -35,10 +35,9 @@ function getindex(σ::AbstractTrajectory, var::String)
     end
 end
 # Get i-th state [i]
-getindex(σ::AbstractTrajectory, idx::Int) = get_state(σ, i)
+getindex(σ::AbstractTrajectory, idx::Int) = get_state(σ, idx)
 # Get i-th value of var ["I", idx]
-function getindex(σ::AbstractTrajectory, var_idx::Tuple{String,Int})
-    var, idx = var_idx[1], var_idx[2]
+function getindex(σ::AbstractTrajectory, var::String, idx::Int)
     if var  == "times"
         return σ.times[idx]
     elseif var == "transitions"
diff --git a/tests/run_unit.jl b/tests/run_unit.jl
index 9934547..ffe7a93 100644
--- a/tests/run_unit.jl
+++ b/tests/run_unit.jl
@@ -10,5 +10,6 @@ using Test
     @test include("unit/simulate_er.jl")
     @test include("unit/change_obs_var_sir.jl")
     @test include("unit/change_obs_var_sir_2.jl")
+    @test include("unit/getindex_access_trajectory.jl")
 end
 
diff --git a/tests/unit/getindex_access_trajectory.jl b/tests/unit/getindex_access_trajectory.jl
new file mode 100644
index 0000000..e5b0aaa
--- /dev/null
+++ b/tests/unit/getindex_access_trajectory.jl
@@ -0,0 +1,11 @@
+
+using MarkovProcesses
+
+load_model("SIR")
+σ = simulate(SIR)
+σ["I"]
+σ["I",2]
+σ[3]
+
+return true
+
-- 
GitLab