Use idxs instead of vars for plotting (#252)

nathanaelbosch · Oct 13, 2023 · 70b6d57 · 70b6d57
1 parent 616b04f
commit 70b6d57
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/docs/src/tutorials/dynamical_odes.md b/docs/src/tutorials/dynamical_odes.md
@@ -58,7 +58,7 @@ u0, du0 = [0.0, 0.1], [0.5, 0.0]
 tspan = (0.0, 100.0)
 prob = ODEProblem(Hénon_Heiles, [du0; u0], tspan)
 sol = solve(prob, EK1());
-plot(sol, vars=(3, 4)) # where `vars=(3,4)` is used to plot x agains y
+plot(sol, idxs=(3, 4)) # where `idxs=(3,4)` is used to plot x agains y
 ```
 
 ### Solving the second-order ODE directly
@@ -77,7 +77,7 @@ function Hénon_Heiles2(ddu, du, u, p, t)
 end
 prob2 = SecondOrderODEProblem(Hénon_Heiles2, du0, u0, tspan)
 sol2 = solve(prob2, EK1());
-plot(sol2, vars=(3, 4))
+plot(sol2, idxs=(3, 4))
 ```
 
 ### Benchmark: Solving second order ODEs is _faster_

diff --git a/src/solution_plotting.jl b/src/solution_plotting.jl
@@ -3,12 +3,24 @@
 ########################################################################################
 @recipe function f(
     sol::AbstractProbODESolution;
-    vars=nothing,
+    idxs=nothing,
     denseplot=sol.dense,
     plotdensity=1000,
     tspan=nothing,
     ribbon_width=1.96,
+    vars=nothing,
 )
+    if vars !== nothing
+        Base.depwarn(
+            "To maintain consistency with solution indexing, keyword argument vars will be removed in a future version. Please use keyword argument idxs instead.",
+            :f; force=true)
+        (idxs !== nothing) &&
+            error(
+                "Simultaneously using keywords vars and idxs is not supported. Please only use idxs.",
+            )
+        idxs = vars
+    end
+
     tstart, tend = isnothing(tspan) ? (sol.t[1], sol.t[end]) : tspan
     times = denseplot ? range(tstart, tend, length=plotdensity) : sol.t
     sol_rvs = denseplot ? sol(times).u : sol.pu
@@ -19,16 +31,16 @@
     values = stack(mean.(sol_rvs))
     stds = stack(std.(sol_rvs))
 
-    if isnothing(vars)
+    if isnothing(idxs)
         ribbon --> ribbon_width * stds
         xguide --> "t"
         yguide --> "u(t)"
         label --> hcat(["u$(i)(t)" for i in 1:length(sol.u[1])]...)
         return times, values
     else
-        int_vars = interpret_vars(vars, sol, getsyms(sol))
+        int_vars = interpret_vars(idxs, sol, getsyms(sol))
         varsizes = unique(length.(int_vars))
-        @assert length(varsizes) == 1 "`vars` argument has an errors"
+        @assert length(varsizes) == 1 "`idxs` argument has an errors"
         ndims = varsizes[1] - 1  # First argument is not about dimensions
         if ndims == 2
             _x = []

diff --git a/test/solution.jl b/test/solution.jl
@@ -105,8 +105,8 @@ using ODEProblemLibrary: prob_ode_lotkavolterra
     @testset "Plotting" begin
         @test_nowarn plot(sol)
         @test_nowarn plot(sol, denseplot=false)
-        @test_nowarn plot(sol, vars=(1, 2))
-        @test_nowarn plot(sol, vars=(1, 1, 2))
+        @test_nowarn plot(sol, idxs=(1, 2))
+        @test_nowarn plot(sol, idxs=(1, 1, 2))
         @test_nowarn plot(sol, tspan=prob.tspan)
     end