Skip to content

Commit

Permalink
feat: add discrete variables support to solution plot recipe
Browse files Browse the repository at this point in the history
  • Loading branch information
@AayushSabharwal
AayushSabharwal committed Jul 7, 2024
1 parent 3c6b617 commit 34db82f
Showing 1 changed file with 112 additions and 65 deletions.
177 changes: 112 additions & 65 deletions src/solutions/solution_interface.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 34db82f

Please sign in to comment.