repro-adam-zero-sum-games-ode-dynamics / official_code /experiments /source_identification /figure_upstream_downstream.jl
| """ | |
| Paper figure: Upstream vs Downstream observation comparison. | |
| Shows that posterior uncertainty correctly reflects information flow direction. | |
| With advection-dominated transport, downstream observations constrain the source | |
| well, while upstream observations provide little information. | |
| Usage: | |
| julia --project=../.. figure_upstream_downstream.jl | |
| """ | |
| using NPZ | |
| using CairoMakie | |
| function main() | |
| # Load results | |
| println("Loading results...") | |
| down_data = npzread("data/downstream_only.npz") | |
| down_gpfvm = npzread("results/downstream_only_gpfvm.npz") | |
| up_data = npzread("data/upstream_only.npz") | |
| up_gpfvm = npzread("results/upstream_only_gpfvm.npz") | |
| # Grid (same for both) | |
| xs = down_data["xs"] | |
| ys = down_data["ys"] | |
| # True source (same for both) | |
| s_true = down_data["s_int_true"] | |
| source_x = down_data["source_x"][1] | |
| source_y = down_data["source_y"][1] | |
| # Cell centers for true source | |
| cx = 0.5 .* (xs[1:end-1] .+ xs[2:end]) | |
| cy = 0.5 .* (ys[1:end-1] .+ ys[2:end]) | |
| Δx, Δy = xs[2] - xs[1], ys[2] - ys[1] | |
| s_true_density = s_true ./ (Δx * Δy) | |
| # Observations | |
| down_obs_x, down_obs_y = down_data["obs_x"], down_data["obs_y"] | |
| up_obs_x, up_obs_y = up_data["obs_x"], up_data["obs_y"] | |
| # GP-FVM results | |
| down_s_mean = down_gpfvm["s_mean"] | |
| down_s_std = down_gpfvm["s_std"] | |
| up_s_mean = up_gpfvm["s_mean"] | |
| up_s_std = up_gpfvm["s_std"] | |
| # Shared color ranges for fair comparison | |
| s_mean_range = ( | |
| min(minimum(down_s_mean), minimum(up_s_mean), 0), | |
| max(maximum(down_s_mean), maximum(up_s_mean)) | |
| ) | |
| s_std_range = ( | |
| min(minimum(down_s_std), minimum(up_s_std)), | |
| max(maximum(down_s_std), maximum(up_s_std)) | |
| ) | |
| # Create figure - 2x3 layout | |
| fig = Figure(size=(800, 500)) | |
| # Row 1: Downstream observations | |
| ax11 = Axis(fig[1, 1], xlabel="x", ylabel="y", title="True source", aspect=DataAspect()) | |
| hm11 = heatmap!(ax11, cx, cy, s_true_density, colormap=:hot) | |
| scatter!(ax11, [source_x], [source_y], color=:cyan, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax11, down_obs_x, down_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| # Flow arrow | |
| arrows!(ax11, [0.1], [0.1], [0.15], [0.0], color=:white, linewidth=2, arrowsize=10) | |
| text!(ax11, 0.18, 0.05, text="flow", color=:white, fontsize=10) | |
| ax12 = Axis(fig[1, 2], xlabel="x", ylabel="y", title="Inferred (downstream obs)", aspect=DataAspect()) | |
| hm12 = heatmap!(ax12, xs, ys, down_s_mean, colormap=:hot, colorrange=s_mean_range) | |
| scatter!(ax12, [source_x], [source_y], color=:cyan, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax12, down_obs_x, down_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| ax13 = Axis(fig[1, 3], xlabel="x", ylabel="y", title="Std (downstream obs)", aspect=DataAspect()) | |
| hm13 = heatmap!(ax13, xs, ys, down_s_std, colormap=:viridis, colorrange=s_std_range) | |
| scatter!(ax13, [source_x], [source_y], color=:white, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax13, down_obs_x, down_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| # Row 2: Upstream observations | |
| ax21 = Axis(fig[2, 1], xlabel="x", ylabel="y", aspect=DataAspect()) | |
| hm21 = heatmap!(ax21, cx, cy, s_true_density, colormap=:hot) | |
| scatter!(ax21, [source_x], [source_y], color=:cyan, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax21, up_obs_x, up_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| arrows!(ax21, [0.1], [0.1], [0.15], [0.0], color=:white, linewidth=2, arrowsize=10) | |
| text!(ax21, 0.18, 0.05, text="flow", color=:white, fontsize=10) | |
| ax22 = Axis(fig[2, 2], xlabel="x", ylabel="y", title="Inferred (upstream obs)", aspect=DataAspect()) | |
| hm22 = heatmap!(ax22, xs, ys, up_s_mean, colormap=:hot, colorrange=s_mean_range) | |
| scatter!(ax22, [source_x], [source_y], color=:cyan, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax22, up_obs_x, up_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| ax23 = Axis(fig[2, 3], xlabel="x", ylabel="y", title="Std (upstream obs)", aspect=DataAspect()) | |
| hm23 = heatmap!(ax23, xs, ys, up_s_std, colormap=:viridis, colorrange=s_std_range) | |
| scatter!(ax23, [source_x], [source_y], color=:white, markersize=10, | |
| marker=:star5, strokewidth=1, strokecolor=:black) | |
| scatter!(ax23, up_obs_x, up_obs_y, color=:white, markersize=8, | |
| strokewidth=1.5, strokecolor=:black, marker=:circle) | |
| # Colorbars | |
| Colorbar(fig[1:2, 4], hm12, label="s") | |
| Colorbar(fig[1:2, 5], hm13, label="σ") | |
| # Save | |
| output_path = "figures/paper_upstream_downstream.pdf" | |
| mkpath(dirname(output_path)) | |
| save(output_path, fig) | |
| println("Saved: $output_path") | |
| # Also save PNG for quick viewing | |
| save("figures/paper_upstream_downstream.png", fig, px_per_unit=3) | |
| println("Saved: figures/paper_upstream_downstream.png") | |
| # Print summary stats | |
| println("\n" * "="^50) | |
| println("Summary") | |
| println("="^50) | |
| println("Downstream observations: $(length(down_obs_x)) points at x ≈ $(round(mean(down_obs_x), digits=2))") | |
| println(" Source std range: [$(round(minimum(down_s_std), digits=3)), $(round(maximum(down_s_std), digits=3))]") | |
| println(" Mean std at source location: $(round(down_s_std[argmin(abs.(xs .- source_x)), argmin(abs.(ys .- source_y))], digits=3))") | |
| println("\nUpstream observations: $(length(up_obs_x)) points at x ≈ $(round(mean(up_obs_x), digits=2))") | |
| println(" Source std range: [$(round(minimum(up_s_std), digits=3)), $(round(maximum(up_s_std), digits=3))]") | |
| println(" Mean std at source location: $(round(up_s_std[argmin(abs.(xs .- source_x)), argmin(abs.(ys .- source_y))], digits=3))") | |
| return fig | |
| end | |
| # Need this for mean() | |
| using Statistics | |
| if abspath(PROGRAM_FILE) == | |
| main() | |
| end | |