diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 4b637d9c..af4da6e9 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -1362,7 +1362,7 @@ This is a developer interface for plot recipe implementations. """ function add_labels!(labels, x, dims, A, strs) if ((x[2] isa Integer && x[2] == 0) || isequal(x[2], getindepsym_defaultt(A))) && - dims == 2 + dims == 2 && x[1] === DEFAULT_PLOT_FUNC push!(labels, strs[end]) elseif x[1] !== DEFAULT_PLOT_FUNC push!(labels, "f($(join(strs, ',')))") @@ -1408,11 +1408,7 @@ function diffeq_to_arrays( plott = A.t[start_idx:end_idx] end - dims = length(vars[1]) - 1 - for var in vars - @assert length(var) - 1 == dims - end - return solplot_vecs_and_labels(dims, vars, plott, A) + return solplot_vecs_and_labels(length(vars[1]) - 1, vars, plott, A) end """ @@ -1421,11 +1417,16 @@ end Build plot vectors and labels for interpreted plotting variables over the sample points `plott`. +The `dims` argument is accepted for API compatibility and ignored; the plotted +dimension is taken from each series transform's evaluated output. + This is a developer interface for plot recipe implementations. """ function solplot_vecs_and_labels(dims, vars, plott, A) plot_vecs = [] labels = String[] + # Ignore caller-supplied `dims`; recompute from evaluated series output. + dims = 0 batch_symbolic_vars = [] for x in vars for j in 2:length(x) @@ -1474,6 +1475,18 @@ function solplot_vecs_and_labels(dims, vars, plott, A) f = x[1] tmp = map(f, tmp...) + series_dims = length(tmp[1]) + if isempty(plot_vecs) + dims = series_dims + elseif series_dims != dims + throw( + ArgumentError( + "Plot idxs series must all have the same output dimension, but got $dims and $series_dims. " * + "Output dimension is the number of coordinates returned by each series transform " * + "(e.g. `(t, u)` is 2-D), not the number of input indices in the idxs tuple." + ) + ) + end tmp = tuple((getindex.(tmp, i) for i in eachindex(tmp[1]))...) for i in eachindex(tmp) if length(plot_vecs) < i diff --git a/test/Core/public_api_test.jl b/test/Core/public_api_test.jl index b41f610d..23694bda 100644 --- a/test/Core/public_api_test.jl +++ b/test/Core/public_api_test.jl @@ -51,3 +51,64 @@ end @test plot_vecs == [[0.0 0.0; 1.0 1.0], [1.0 2.0; 3.0 4.0]] @test plot_labels == ["u[1]", "u[2]"] end + +# #360: check plot dims from transform output, not idxs input arity. +@testset "plot idxs with mixed input arity but matching output dims (#360)" begin + t = [0.0, 0.25, 0.5, 0.75, 1.0] + u = [[1.0, 2.0, 10.0 + tt, 20.0 + 2tt] for tt in t] + A = DiffEqArray(u, t) + + adder(tt, a, b) = (tt, a + b) + adder3(tt, a, b) = (tt, a, b) + g(tt::Float64, x::Float64) = (tt, x) + dom(tt, x) = (tt, sqrt(x - 2)) + + function plot_sparse(idxs) + vars = interpret_vars(idxs, A) + return diffeq_to_arrays(A, false, 100, nothing, vars, :identity, nothing) + end + + u3 = [uu[3] for uu in u] + u3pu4 = [uu[3] + uu[4] for uu in u] + + for idxs in ([3, (adder, 0, 3, 4)], [(0, 3), (adder, 0, 3, 4)]) + plot_vecs, labels = plot_sparse(idxs) + @test length(plot_vecs) == 2 + @test size(plot_vecs[1], 2) == 2 + @test plot_vecs[1][:, 1] ≈ t + @test plot_vecs[2][:, 1] ≈ u3 + @test plot_vecs[1][:, 2] ≈ t + @test plot_vecs[2][:, 2] ≈ u3pu4 + @test labels == ["u[3]", "f(t,u[3],u[4])"] + end + + _, labels01 = plot_sparse((0, 1)) + @test labels01 == ["u[1]"] + + plot_vecs_g, labels_g = plot_sparse([(g, 0, 3)]) + @test plot_vecs_g[1][:, 1] ≈ t + @test plot_vecs_g[2][:, 1] ≈ u3 + @test labels_g == ["f(t,u[3])"] + + plot_vecs_dom, labels_dom = plot_sparse([(dom, 0, 3)]) + @test plot_vecs_dom[1][:, 1] ≈ t + @test plot_vecs_dom[2][:, 1] ≈ sqrt.(u3 .- 2) + @test labels_dom == ["f(t,u[3])"] + + # Exported 4-arg form: `dims` is accepted and ignored. + vars4 = interpret_vars([3, (adder, 0, 3, 4)], A) + plot_vecs4, labels4 = solplot_vecs_and_labels(99, vars4, t, A) + @test length(plot_vecs4) == 2 + @test plot_vecs4[2][:, 1] ≈ u3 + @test plot_vecs4[2][:, 2] ≈ u3pu4 + @test labels4 == ["u[3]", "f(t,u[3],u[4])"] + + err = try + plot_sparse([3, (adder3, 0, 3, 4)]) + nothing + catch e + e + end + @test err isa ArgumentError + @test occursin("output dimension", sprint(showerror, err)) +end