feat: add discrete variables support to solution plot recipe

Author: AayushSabharwal

src/solutions/solution_interface.jl

@@ -176,91 +176,138 @@ DEFAULT_PLOT_FUNC(x, y, z) = (x, y, z) # For v0.5.2 bug
     end
 
     idxs = idxs === nothing ? (1:length(sol.u[1])) : idxs
-
+    disc_idxs = []
+    cont_idxs = []
+    for idx in idxs
+        tsidxs = get_all_timeseries_indexes(sol, idx)
+        if ContinuousTimeseries() in tsidxs
+            push!(cont_idxs, idx)
+        else
+            push!(disc_idxs, (idx, only(tsidxs)))
+        end
+    end
+    idxs = identity.(cont_idxs)
     if !(idxs isa Union{Tuple, AbstractArray})
         vars = interpret_vars([idxs], sol)
     else
         vars = interpret_vars(idxs, sol)
     end
-
-    tscale = get(plotattributes, :xscale, :identity)
-    plot_vecs, labels = diffeq_to_arrays(sol, plot_analytic, denseplot,
-        plotdensity, tspan, vars, tscale, plotat)
-
     tdir = sign(sol.t[end] - sol.t[1])
     xflip --> tdir < 0
     seriestype --> :path
 
-    # Special case labels when idxs = (:x,:y,:z) or (:x) or [:x,:y] ...
-    if idxs isa Tuple && vars[1][1] === DEFAULT_PLOT_FUNC
-        val = hasname(vars[1][2]) ? String(getname(vars[1][2])) : vars[1][2]
-        if val isa Integer
-            if val == 0
-                val = "t"
-            else
-                val = "u[$val]"
+    @series begin
+        if isempty(idxs)
+            label --> nothing
+            ([], [])
+        else
+            tscale = get(plotattributes, :xscale, :identity)
+            plot_vecs, labels = diffeq_to_arrays(sol, plot_analytic, denseplot,
+                plotdensity, tspan, vars, tscale, plotat)
+
+
+            # Special case labels when idxs = (:x,:y,:z) or (:x) or [:x,:y] ...
+            if idxs isa Tuple && vars[1][1] === DEFAULT_PLOT_FUNC
+                val = hasname(vars[1][2]) ? String(getname(vars[1][2])) : vars[1][2]
+                if val isa Integer
+                    if val == 0
+                        val = "t"
+                    else
+                        val = "u[$val]"
+                    end
+                end
+                xguide --> val
+                val = hasname(vars[1][3]) ? String(getname(vars[1][3])) : vars[1][3]
+                if val isa Integer
+                    if val == 0
+                        val = "t"
+                    else
+                        val = "u[$val]"
+                    end
+                end
+                yguide --> val
+                if length(idxs) > 2
+                    val = hasname(vars[1][4]) ? String(getname(vars[1][4])) : vars[1][4]
+                    if val isa Integer
+                        if val == 0
+                            val = "t"
+                        else
+                            val = "u[$val]"
+                        end
+                    end
+                    zguide --> val
+                end
             end
-        end
-        xguide --> val
-        val = hasname(vars[1][3]) ? String(getname(vars[1][3])) : vars[1][3]
-        if val isa Integer
-            if val == 0
-                val = "t"
-            else
-                val = "u[$val]"
+
+            if (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 1))) &&
+                getindex.(vars, 1) == zeros(length(vars))) ||
+               (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 2))) &&
+                getindex.(vars, 2) == zeros(length(vars))) ||
+               all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 1)) ||
+               all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 2))
+                xguide --> "$(getindepsym_defaultt(sol))"
             end
-        end
-        yguide --> val
-        if length(idxs) > 2
-            val = hasname(vars[1][4]) ? String(getname(vars[1][4])) : vars[1][4]
-            if val isa Integer
-                if val == 0
-                    val = "t"
+            if length(vars[1]) >= 3 &&
+               ((!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 3))) &&
+                 getindex.(vars, 3) == zeros(length(vars))) ||
+                all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 3)))
+                yguide --> "$(getindepsym_defaultt(sol))"
+            end
+            if length(vars[1]) >= 4 &&
+               ((!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 4))) &&
+                 getindex.(vars, 4) == zeros(length(vars))) ||
+                all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 4)))
+                zguide --> "$(getindepsym_defaultt(sol))"
+            end
+
+            if (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 2))) &&
+                getindex.(vars, 2) == zeros(length(vars))) ||
+               all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 2))
+                if tspan === nothing
+                    if tdir > 0
+                        xlims --> (sol.t[1], sol.t[end])
+                    else
+                        xlims --> (sol.t[end], sol.t[1])
+                    end
                 else
-                    val = "u[$val]"
+                    xlims --> (tspan[1], tspan[end])
                 end
             end
-            zguide --> val
-        end
-    end
 
-    if (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 1))) &&
-        getindex.(vars, 1) == zeros(length(vars))) ||
-       (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 2))) &&
-        getindex.(vars, 2) == zeros(length(vars))) ||
-       all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 1)) ||
-       all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 2))
-        xguide --> "$(getindepsym_defaultt(sol))"
-    end
-    if length(vars[1]) >= 3 &&
-       ((!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 3))) &&
-         getindex.(vars, 3) == zeros(length(vars))) ||
-        all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 3)))
-        yguide --> "$(getindepsym_defaultt(sol))"
-    end
-    if length(vars[1]) >= 4 &&
-       ((!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 4))) &&
-         getindex.(vars, 4) == zeros(length(vars))) ||
-        all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 4)))
-        zguide --> "$(getindepsym_defaultt(sol))"
+            label --> reshape(labels, 1, length(labels))
+            (plot_vecs...,)
+        end
     end
-
-    if (!any(!isequal(NotSymbolic()), symbolic_type.(getindex.(vars, 2))) &&
-        getindex.(vars, 2) == zeros(length(vars))) ||
-       all(t -> Symbol(t) == getindepsym_defaultt(sol), getindex.(vars, 2))
-        if tspan === nothing
-            if tdir > 0
-                xlims --> (sol.t[1], sol.t[end])
-            else
-                xlims --> (sol.t[end], sol.t[1])
+    for (idx, tsidx) in disc_idxs
+        partition = sol.discretes[tsidx]
+        ts = current_time(partition)
+        if tspan !== nothing
+            tstart = searchsortedfirst(ts, tspan[1])
+            tend = searchsortedlast(ts, tspan[2])
+            if tstart == lastindex(ts) + 1 || tend == firstindex(ts) - 1
+                continue
             end
         else
-            xlims --> (tspan[1], tspan[end])
+            tstart = firstindex(ts)
+            tend = lastindex(ts)
+        end
+        ts = ts[tstart:tend]
+
+        vals = getp(sol, idx)(sol, tstart:tend)
+        # Scatterplot of points
+        @series begin
+            seriestype := :line
+            linestyle --> :dash
+            markershape --> :o
+            markersize --> repeat([2, 0], length(ts)-1)
+            markeralpha --> repeat([1, 0], length(ts)-1)
+            label --> string(hasname(idx) ? getname(idx) : idx)
+
+            x = vec([ts[1:end-1]'; ts[2:end]'])
+            y = repeat(vals, inner=2)[1:end-1]
+            x, y
         end
     end
-
-    label --> reshape(labels, 1, length(labels))
-    (plot_vecs...,)
 end
 
 function diffeq_to_arrays(sol, plot_analytic, denseplot, plotdensity, tspan,