Customising the legend display order in Plots.jl

[jmair]

I am using Plots.jl with the pgfplotsx() backend and I want to have a somewhat custom legend, which has two columns and shows the marker for the expected curve next to actual samples on that curve. A MWE is given below:

using Plots
using LaTeXStrings
pgfplotsx()

x_actual = LinRange(0, 2*pi, 1000)
x_obs = LinRange(0, 2*pi, 10)

# Plot underlying relationship
plot(x_actual, (A->sin.(x_actual).*sqrt(A)).(1:3); 
    labels=[L"A^2=1" L"A^2=2" L"A^2=3"],
    markershape=:none, 
    linestyle=:dash
)
# Plot sampled points with markers
plot!(x_obs, (A->sin.(x_obs).*sqrt(A)).(1:3); 
    labels=[L"A^2=1" L"A^2=2" L"A^2=3"],
    markershapes=:auto, 
    linealpha=0
)

plot!(; legend_column=2, legend=:topright, grid=nothing, legend_foreground_color=nothing)

This gives:

The main issue is that I want the A^2=1 labels to be next to one another. I do not want to have to individually plot each separate line in a different order to get the right legend ordering, as this will reduce code reuse. Preferably, I would like the plot to look something like below:

I understand that this is similar to a solution here - Plot legend with multiple columns - #7 by rafael.guerra, but I am wondering if there is a better way to do this.

[jmair]

I have found that I can edit the order directly, but the solution is not very elegant.

using Plots
using LaTeXStrings
pgfplotsx()

x_actual = LinRange(0, 2*pi, 1000)
x_obs = LinRange(0, 2*pi, 10)

n_lines = 3
# Plot underlying relationship
plt = plot(x_actual, (A->sin.(x_actual).*sqrt(A)).(1:n_lines); 
    labels=[L"" L"" L""],
    markershape=:none, 
    linestyle=:dash
)
# Plot sampled points with markers
plot!(x_obs, (A->sin.(x_obs).*sqrt(A)).(1:n_lines); 
    labels=[L"A^2=1" L"A^2=2" L"A^2=3"],
    markershapes=:auto, 
    linealpha=0
)

# Manually edit the order of the series
n_rows = n_lines
n_cols = 2
plt.series_list = [plt.series_list[1 + (i ÷ n_cols % n_rows) + (i % n_cols) * n_rows] for i in 0:length(plt.series_list)-1]
# Edit attributes that are todo with the indices
for i in 1:length(plt.series_list)
    plt.series_list[i].plotattributes[:series_index] = i
    plt.series_list[i].plotattributes[:series_plotindex] = i
end
# Update the subplot series list, as it still points at the original object
plt.subplots[1].series_list = plt.series_list

plot!(; legend_column=2, legend=:topright, grid=nothing, legend_foreground_color=nothing)

The result is below: