Support publication plot with replications of different lengths (#788)

src/visualization/publication_plot.jl

@@ -12,19 +12,25 @@ function publication_data(
 )
     max_stages = maximum(length.(dataset))
     output_array = fill(NaN, length(quantiles), max_stages)
-    for stage in 1:max_stages
-        stage_data = stage_function.([data[stage] for data in dataset])
-        for (i, s) in enumerate(stage_data)
+    for t in 1:max_stages
+        stage_data = Float64[]
+        for (i, d) in enumerate(dataset)
+            if length(d) < t
+                # Skip this replication because it doesn't have enough stages
+                continue
+            end
+            s = stage_function(d[t])
             if !isfinite(s)
                 error(
-                    "Unable to plot `publication_plot` because stage $stage " *
+                    "Unable to plot `publication_plot` because stage $t " *
                     "of replication $i contains data that is not finite. " *
                     "The data function must return a finite real-valued " *
                     "scalar. Got: $s",
                 )
             end
+            push!(stage_data, s)
         end
-        output_array[:, stage] .= Statistics.quantile(stage_data, quantiles)
+        output_array[:, t] .= Statistics.quantile(stage_data, quantiles)
     end
     return output_array
 end

test/visualization/visualization.jl

@@ -78,6 +78,17 @@ function test_PublicationPlot()
     return
 end
 
+function test_PublicationPlot_different_lengths()
+    simulations = [
+        [Dict{Symbol,Any}(:x => 1), Dict{Symbol,Any}(:x => 5)],
+        [Dict{Symbol,Any}(:x => 2)],
+        [Dict{Symbol,Any}(:x => 3), Dict{Symbol,Any}(:x => 4)],
+    ]
+    data = SDDP.publication_data(simulations, [0.0, 0.25, 0.5, 1.0], d -> d[:x])
+    @test data == [1.0 4.0; 1.5 4.25; 2.0 4.5; 3.0 5.0]
+    return
+end
+
 end  # module
 
 TestVisualization.runtests()