Add a vars parameter to plot phase diagrams

.travis.yml

@@ -5,7 +5,9 @@ os:
   - osx
 julia:
   - 0.5
-  - nightly
+matrix:
+  allow_failures:
+    - julia: nightly
 notifications:
   email: false
 # uncomment the following lines to override the default test script

appveyor.yml

@@ -1,10 +1,9 @@
 environment:
   matrix:
-  - JULIAVERSION: "julialang/bin/winnt/x86/0.5/julia-0.5-latest-win32.exe"
-  - JULIAVERSION: "julialang/bin/winnt/x64/0.5/julia-0.5-latest-win64.exe"
-  - JULIAVERSION: "julianightlies/bin/winnt/x86/julia-latest-win32.exe"
-  - JULIAVERSION: "julianightlies/bin/winnt/x64/julia-latest-win64.exe"
-
+    - JULIAVERSION: "julialang/bin/winnt/x64/0.5/julia-0.5-latest-win64.exe"
+matrix:
+  allow_failures:
+    - JULIAVERSION: "julianightlies/bin/winnt/x64/julia-latest-win64.exe"
 branches:
   only:
     - master

src/solutions/solution_interface.jl

@@ -47,47 +47,116 @@ function show(io::IO,sol::DESolution)
 end
 =#
 
-@recipe function f(sol::AbstractODESolution;plot_analytic=false,denseplot=true,plotdensity=100)
-  plotseries = Vector{Any}(0)
-  if typeof(sol) <: AbstractSDESolution; denseplot=false; end
+@recipe function f(sol::AbstractODESolution;
+                   plot_analytic=false,denseplot=true,plotdensity=100,vars=nothing)
 
-  if denseplot && sol.dense # Generate the points from the plot from dense function
+  if typeof(sol) <: AbstractSDESolution
+    denseplot = false
+  end
+
+  if vars == nothing
+    # Default: plot all timeseries
+    if typeof(sol[1]) <: AbstractArray
+      vars = collect((0, i) for i in eachindex(sol[1]))
+    else
+      vars = [(0, 1)]
+    end
+  end
+  if typeof(vars) <: Integer
+    vars = [(0, vars)]
+  end
+  if typeof(vars) <: AbstractArray
+    # If list given, its elements should be tuples, or we assume x = time
+    vars = [if typeof(x) <: Tuple; x else (0, x) end for x in vars]
+  end
+  if typeof(vars) <: Tuple
+    # If tuple given...
+    if typeof(vars[1]) <: AbstractArray
+      if typeof(vars[2]) <: AbstractArray
+        # If both axes are lists we zip (will fail if different lengths)
+        vars = collect(zip(vars[1], vars[2]))
+      else
+        # Just the x axis is a list
+        vars = [(x, vars[2]) for x in vars[1]]
+      end
+    else
+      if typeof(vars[2]) <: AbstractArray
+        # Just the y axis is a list
+        vars = [(vars[1], y) for y in vars[2]]
+      else
+        # Both axes are numbers
+        vars = [vars]
+      end
+    end
+  end
+
+  # Here `vars` should be a list of tuples (x, y).
+  assert(typeof(vars) <: AbstractArray)
+  assert(eltype(vars) <: Tuple)
+
+  if denseplot && sol.dense
+    # Generate the points from the plot from dense function
     plott = collect(Ranges.linspace(sol.t[1],sol.t[end],plotdensity))
     plot_timeseries = sol(plott)
     if plot_analytic
       plot_analytic_timeseries = [sol.prob.analytic(t,sol.prob.u0) for t in plott]
     end
-  else # Plot for not dense output use the timeseries itself
+  else
+    # Plot for sparse output: use the timeseries itself
+    plott = sol.t
     plot_timeseries = sol.u
     if plot_analytic
       plot_analytic_timeseries = sol.u_analytic
     end
-    plott = sol.t
   end
 
-  # Make component-wise plots
-  if typeof(sol[1]) <:AbstractArray
-    for i in eachindex(sol[1])
+  function u_n(timeseries::AbstractArray, n::Int)
+    # Returns the nth variable from the timeseries, t if n == 0
+    if n == 0
+      plott
+    elseif n == 1 && !(typeof(sol[1]) <: AbstractArray)
+      timeseries
+    else
       tmp = Vector{eltype(sol[1])}(length(plot_timeseries))
       for j in 1:length(plot_timeseries)
-        tmp[j] = plot_timeseries[j][i]
+        tmp[j] = plot_timeseries[j][n]
       end
-      push!(plotseries,tmp)
+      tmp
     end
-  else
-    push!(plotseries,plot_timeseries)
   end
+
+  plotx = Vector{Any}(0)
+  ploty = Vector{Any}(0)
+  labels = Array{String, 2}(1, length(vars)*(1+plot_analytic))
+  for (i, (x, y)) in enumerate(vars)
+    push!(plotx, u_n(plot_timeseries, x))
+    push!(ploty, u_n(plot_timeseries, y))
+    if y == 0
+      ly = "t"
+    else
+      ly = "u$y"
+    end
+    if x == 0
+      labels[i] = "$ly(t)"
+    else
+      labels[i] = "(u$x, $ly)"
+    end
+  end
+
   if plot_analytic
-    if typeof(sol[1]) <: AbstractArray
-      for i in eachindex(sol[1])
-        tmp = Vector{eltype(sol[1])}(length(plot_timeseries))
-        for j in 1:length(plot_timeseries)
-          tmp[j] = plot_analytic_timeseries[j][i]
-        end
-        push!(plotseries,tmp)
+    for (i, (x, y)) in enumerate(vars)
+      push!(plotx, u_n(plot_analytic_timeseries, x))
+      push!(ploty, u_n(plot_analytic_timeseries, y))
+      if y == 0
+        ly = "t"
+      else
+        ly = "au$y"
+      end
+      if x == 0
+        labels[i] = "$ly(t)"
+      else
+        labels[i] = "(au$x, $ly)"
       end
-    else
-      push!(plotseries,plot_analytic_timeseries)
     end
   end
 
@@ -97,5 +166,6 @@ end
   #ytickfont --> font(11)
   #legendfont --> font(11)
   #guidefont  --> font(11)
-  plott, plotseries
+  label --> labels
+  plotx, ploty
 end