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 3, 2024
1 parent c0fa7e7 commit 41a51d8
Showing 1 changed file with 117 additions and 65 deletions.
182 changes: 117 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,143 @@ 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
end
ts = ts[tstart:tend]

vals = getp(sol, idx)(sol, tstart:tend)

label --> reshape(labels, 1, length(labels))
(plot_vecs...,)
# Scatterplot of points
@series begin
seriestype := :scatter
linestyle --> :dash
label --> string(hasname(idx) ? getname(idx) : idx)
ts, vals
end

@series begin
seriestype := :line
linestyle --> :dash
label := nothing

xvals = collect(Iterators.flatten(zip(ts, ts)))[2:end]
yvals = collect(Iterators.flatten(zip(vals, vals)))[1:end-1]
xvals, yvals
end
end
end

function diffeq_to_arrays(sol, plot_analytic, denseplot, plotdensity, tspan,
Expand Down

0 comments on commit 41a51d8

Please sign in to comment.