Skip to content

Commit

Permalink
JuliaFormatter
Browse files Browse the repository at this point in the history
  • Loading branch information
@asinghvi17
asinghvi17 committed Feb 6, 2024
1 parent 0636e40 commit 57890b3
Showing 1 changed file with 131 additions and 102 deletions.
233 changes: 131 additions & 102 deletions ext/SciMLBaseMakieExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,23 +7,20 @@ import Makie.SpecApi as S

function ensure_plottrait(PT::Type, arg, desired_plottrait_type::Type)
if !(Makie.conversion_trait(PT, arg) isa desired_plottrait_type)
error(
"""
`Makie.convert_arguments` for the plot type $PT and its conversion trait $(Makie.conversion_trait(PT, arg)) was unsuccessful.
There is a recipe for the given arguments and the `$desired_plottrait_type` trait, however.
error("""
`Makie.convert_arguments` for the plot type $PT and its conversion trait $(Makie.conversion_trait(PT, arg)) was unsuccessful.
The signature that could not be converted was:
::$(string(typeof(arg)))
There is a recipe for the given arguments and the `$desired_plottrait_type` trait, however.
Makie needs to convert all plot input arguments to types that can be consumed by the backends (typically Arrays with Float32 elements).
You can define a method for `Makie.convert_arguments` (a type recipe) for these types or their supertypes to make this set of arguments convertible (See http://docs.makie.org/stable/documentation/recipes/index.html).
The signature that could not be converted was:
::$(string(typeof(arg)))
Alternatively, you can define `Makie.convert_single_argument` for single arguments which have types that are unknown to Makie but which can be converted to known types and fed back to the conversion pipeline.
"""
)
end
Makie needs to convert all plot input arguments to types that can be consumed by the backends (typically Arrays with Float32 elements).
You can define a method for `Makie.convert_arguments` (a type recipe) for these types or their supertypes to make this set of arguments convertible (See http://docs.makie.org/stable/documentation/recipes/index.html).
Alternatively, you can define `Makie.convert_single_argument` for single arguments which have types that are unknown to Makie but which can be converted to known types and fed back to the conversion pipeline.
""")
end
end

# ## `AbstractTimeseriesSolution` recipe
Expand All @@ -34,20 +31,20 @@ end
# In the recipe, we use the Makie PlotSpec API to override
Makie.plottype(sol::SciMLBase.AbstractTimeseriesSolution) = Makie.Lines

Makie.used_attributes(::Type{<: Plot}, sol::SciMLBase.AbstractTimeseriesSolution) = (:plot_analytic, :denseplot, :plotdensity, :plotat, :tspan, :tscale, :vars, :idxs)

function Makie.convert_arguments(
PT::Type{<:Plot},
sol::SciMLBase.AbstractTimeseriesSolution;
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = Makie.automatic,
plotat = nothing,
tspan = nothing,
tscale = :identity,
vars = nothing,
idxs = nothing,
)
function Makie.used_attributes(::Type{<:Plot}, sol::SciMLBase.AbstractTimeseriesSolution)
(:plot_analytic, :denseplot, :plotdensity, :plotat, :tspan, :tscale, :vars, :idxs)
end

function Makie.convert_arguments(PT::Type{<:Plot},
sol::SciMLBase.AbstractTimeseriesSolution;
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = Makie.automatic,
plotat = nothing,
tspan = nothing,
tscale = :identity,
vars = nothing,
idxs = nothing,)

# First, this recipe is specifically only for timeseries solutions.
# This means that the recipe only applies to `PointBased` plot types.
Expand All @@ -63,7 +60,7 @@ function Makie.convert_arguments(

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)
:f; force = true)
(idxs !== nothing) &&
error("Simultaneously using keywords vars and idxs is not supported. Please only use idxs.")
idxs = vars
Expand All @@ -83,19 +80,19 @@ function Makie.convert_arguments(
# level recipes
if denseplot isa Makie.Automatic
denseplot = (sol.dense ||
typeof(sol.prob) <: SciMLBase.AbstractDiscreteProblem) &&
!(typeof(sol) <: SciMLBase.AbstractRODESolution) &&
!(hasfield(typeof(sol), :interp) &&
typeof(sol.interp) <: SciMLBase.SensitivityInterpolation)
typeof(sol.prob) <: SciMLBase.AbstractDiscreteProblem) &&
!(typeof(sol) <: SciMLBase.AbstractRODESolution) &&
!(hasfield(typeof(sol), :interp) &&
typeof(sol.interp) <: SciMLBase.SensitivityInterpolation)
end

if plotdensity isa Makie.Automatic
plotdensity = min(Int(1e5),
sol.tslocation == 0 ?
(sol.prob isa SciMLBase.AbstractDiscreteProblem ?
max(1000, 100 * length(sol)) :
max(1000, 10 * length(sol))) :
1000 * sol.tslocation)
sol.tslocation == 0 ?
(sol.prob isa SciMLBase.AbstractDiscreteProblem ?
max(1000, 100 * length(sol)) :
max(1000, 10 * length(sol))) :
1000 * sol.tslocation)
end

# Originally, in the Plots recipe, `tscale`
Expand All @@ -118,13 +115,19 @@ function Makie.convert_arguments(
# TODO: broadcast across all input attributes, or figure out how to
# allow customizable colors/labels/etc if required
makie_plotspecs = if length(plot_vecs) == 2
map((x, y, label) -> PlotSpec(plot_type_sym, Point2f.(x, y); label), eachcol(plot_vecs[1]), eachcol(plot_vecs[2]), labels)
map((x, y, label) -> PlotSpec(plot_type_sym, Point2f.(x, y); label),
eachcol(plot_vecs[1]),
eachcol(plot_vecs[2]),
labels)
elseif length(plot_vecs) == 3
map((x, y, z, label) -> PlotSpec(plot_type_sym, Point3f.(x, y, z); label), eachcol(plot_vecs[1]), eachcol(plot_vecs[2]), eachcol(plot_vecs[3]), labels)
map((x, y, z, label) -> PlotSpec(plot_type_sym, Point3f.(x, y, z); label),
eachcol(plot_vecs[1]),
eachcol(plot_vecs[2]),
eachcol(plot_vecs[3]),
labels)
end

return makie_plotspecs

end

# ## Integrator recipes
Expand All @@ -133,18 +136,17 @@ end

Makie.plottype(integrator::SciMLBase.DEIntegrator) = Makie.Lines

Makie.used_attributes(::Type{<: Plot}, integrator::SciMLBase.DEIntegrator) = (:plot_analytic, :denseplot, :plotdensity, :vars, :idxs)

function Makie.convert_arguments(
PT::Type{<: Makie.Plot},
integrator::SciMLBase.DEIntegrator;
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = 10,
vars = nothing,
idxs = nothing,
)
function Makie.used_attributes(::Type{<:Plot}, integrator::SciMLBase.DEIntegrator)
(:plot_analytic, :denseplot, :plotdensity, :vars, :idxs)
end

function Makie.convert_arguments(PT::Type{<:Makie.Plot},
integrator::SciMLBase.DEIntegrator;
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = 10,
vars = nothing,
idxs = nothing,)
ensure_plottrait(PT, integrator, Makie.PointBased)

# Interpret keyword arguments
Expand All @@ -158,8 +160,8 @@ function Makie.convert_arguments(

if denseplot isa Makie.Automatic
denseplot = (integrator.opts.calck ||
integrator isa AbstractSDEIntegrator) &&
integrator.iter > 0
integrator isa AbstractSDEIntegrator) &&
integrator.iter > 0
end

# Begin deconstructing the integrator
Expand Down Expand Up @@ -195,7 +197,8 @@ function Makie.convert_arguments(
for x in int_vars
for j in 2:dims
if denseplot
if (x[j] isa Integer && x[j] == 0) || isequal(x[j],SciMLBase.getindepsym_defaultt(integrator))
if (x[j] isa Integer && x[j] == 0) ||
isequal(x[j], SciMLBase.getindepsym_defaultt(integrator))
push!(plot_vecs[j - 1], plott)
else
push!(plot_vecs[j - 1], Vector(integrator(plott; idxs = x[j])))
Expand Down Expand Up @@ -250,11 +253,15 @@ function Makie.convert_arguments(
plot_type_sym = Makie.plotsym(PT)

return if denseplot
[Makie.PlotSpec(plot_type_sym, Point2f.(plot_vecs[1][idx], plot_vecs[2][idx]); label, color = Makie.Cycled(idx)) for (idx, label) in zip(1:length(plot_vecs[1]), labels)]
[Makie.PlotSpec(plot_type_sym,
Point2f.(plot_vecs[1][idx], plot_vecs[2][idx]);
label,
color = Makie.Cycled(idx))
for (idx, label) in zip(1:length(plot_vecs[1]), labels)]
else
[S.Scatter([Point2f(plot_vecs[1][idx], plot_vecs[2][idx])]; label) for (idx, label) in zip(1:length(plot_vecs[1]), labels)]
[S.Scatter([Point2f(plot_vecs[1][idx], plot_vecs[2][idx])]; label)
for (idx, label) in zip(1:length(plot_vecs[1]), labels)]
end

end

# ## Ensemble recipes
Expand All @@ -263,37 +270,52 @@ end
Makie.plottype(sol::SciMLBase.AbstractEnsembleSolution) = Makie.Lines

# We also define the attributes that are used by the ensemble solution recipe:
Makie.used_attributes(::Type{<: Plot}, sol::SciMLBase.AbstractEnsembleSolution) = (:trajectories, :plot_analytic, :denseplot, :plotdensity, :plotat, :tspan, :tscale, :vars, :idxs)

function Makie.convert_arguments(
PT::Type{<:Lines},
sim::SciMLBase.AbstractEnsembleSolution;
trajectories = eachindex(sim),
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = Makie.automatic,
plotat = nothing,
tspan = nothing,
tscale = :identity,
vars = nothing,
idxs = nothing,
)
function Makie.used_attributes(::Type{<:Plot}, sol::SciMLBase.AbstractEnsembleSolution)
(:trajectories,
:plot_analytic,
:denseplot,
:plotdensity,
:plotat,
:tspan,
:tscale,
:vars,
:idxs)
end

function Makie.convert_arguments(PT::Type{<:Lines},
sim::SciMLBase.AbstractEnsembleSolution;
trajectories = eachindex(sim),
plot_analytic = false,
denseplot = Makie.automatic,
plotdensity = Makie.automatic,
plotat = nothing,
tspan = nothing,
tscale = :identity,
vars = nothing,
idxs = nothing,)

# First, we check if the plot type is PointBased, and if not, we throw the standard
# Makie error message for convert_arguments - just at a different place.
ensure_plottrait(PT, sim, Makie.PointBased)

@assert length(trajectories) > 0 "No trajectories to plot"
@assert length(sim.u) > 0 "No solutions to plot"
@assert length(trajectories)>0 "No trajectories to plot"
@assert length(sim.u)>0 "No solutions to plot"

plot_type_sym = Makie.plotsym(PT)

mp = [PlotSpec(plot_type_sym, sim.u[i]; plot_analytic, denseplot, plotdensity, plotat, tspan, tscale, idxs) for i in trajectories]
mp = [PlotSpec(plot_type_sym,
sim.u[i];
plot_analytic,
denseplot,
plotdensity,
plotat,
tspan,
tscale,
idxs) for i in trajectories]

# Main.Infiltrator.@infiltrate

return mp

end

# ## EnsembleSummary recipes
Expand All @@ -302,27 +324,38 @@ end

Makie.plottype(sim::SciMLBase.EnsembleSummary) = Makie.Lines

Makie.used_attributes(::Type{<: Lines}, sim::SciMLBase.EnsembleSummary) = (:trajectories, :error_style, :ci_type, :plot_analytic, :denseplot, :plotdensity, :plotat, :tspan, :tscale, :vars, :idxs)

function Makie.used_attributes(::Type{<:Lines}, sim::SciMLBase.EnsembleSummary)
(:trajectories,
:error_style,
:ci_type,
:plot_analytic,
:denseplot,
:plotdensity,
:plotat,
:tspan,
:tscale,
:vars,
:idxs)
end

# TODO: should `error_style` be Makie plot types instead? I.e. `Band`, `Errorbar`, etc
function Makie.convert_arguments(
::Type{<: Lines},
sim::SciMLBase.EnsembleSummary;
trajectories = sim.u.u[1] isa AbstractArray ? eachindex(sim.u.u[1]) :
1,
error_style = :ribbon, ci_type = :quantile,
kwargs...
)
function Makie.convert_arguments(::Type{<:Lines},
sim::SciMLBase.EnsembleSummary;
trajectories = sim.u.u[1] isa AbstractArray ? eachindex(sim.u.u[1]) :
1,
error_style = :ribbon, ci_type = :quantile,
kwargs...)
if ci_type == :SEM
if sim.u.u[1] isa AbstractArray
u = SciMLBase.vecarr_to_vectors(sim.u)
else
u = [sim.u.u]
end
if sim.u.u[1] isa AbstractArray
ci_low = SciMLBase.vecarr_to_vectors(VectorOfArray([sqrt.(sim.v.u[i] / sim.num_monte) .*
1.96 for i in 1:length(sim.v)]))
ci_low = SciMLBase.vecarr_to_vectors(VectorOfArray([sqrt.(sim.v.u[i] /
sim.num_monte) .*
1.96
for i in 1:length(sim.v)]))
ci_high = ci_low
else
ci_low = [[sqrt(sim.v.u[i] / length(sim.num_monte)) .* 1.96
Expand All @@ -349,20 +382,18 @@ function Makie.convert_arguments(
makie_plotlist = Makie.PlotSpec[]

for (count, idx) in enumerate(trajectories)
push!(
makie_plotlist,
S.Lines(sim.t, u[idx]; color = Makie.Cycled(count), label = "u[$idx]")
)
push!(makie_plotlist,
S.Lines(sim.t, u[idx]; color = Makie.Cycled(count), label = "u[$idx]"))
if error_style == :ribbon
push!(
makie_plotlist,
S.Band(sim.t, u[idx] .- ci_low[idx], u[idx] .+ ci_high[idx]; color = Makie.Cycled(count), alpha = 0.1)
)
push!(makie_plotlist,
S.Band(sim.t,
u[idx] .- ci_low[idx],
u[idx] .+ ci_high[idx];
color = Makie.Cycled(count),
alpha = 0.1))
elseif error_style == :bars
push!(
makie_plotlist,
S.Errorbars(sim.t, u[idx], ci_low[idx], ci_high[idx])
)
push!(makie_plotlist,
S.Errorbars(sim.t, u[idx], ci_low[idx], ci_high[idx]))
elseif error_style == :none
nothing
else
Expand All @@ -371,8 +402,6 @@ function Makie.convert_arguments(
end

return makie_plotlist


end

end

0 comments on commit 57890b3

Please sign in to comment.