diff --git a/Project.toml b/Project.toml index 595f7120..0aba9c24 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrameMacros = "75880514-38bc-4a95-a458-c2aea5a3a702" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -51,6 +52,7 @@ ChainRulesCore = "1.25.1" ComponentArrays = "0.15.28" DataFrameMacros = "0.4" DataFrames = "1" +DataStructures = "0.19.4" DimensionalData = "0.29.24, 0.30" Downloads = "1.6.0" ForwardDiff = "1" diff --git a/docs/Project.toml b/docs/Project.toml index 9b98d244..4eb8bd2b 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -3,6 +3,7 @@ AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" Chain = "8be319e6-bccf-4806-a6f7-6fae938471bc" +DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365" diff --git a/docs/literate/research/synthetic_respiration.jl b/docs/literate/research/synthetic_respiration.jl index a9c57651..1f4d6dbe 100644 --- a/docs/literate/research/synthetic_respiration.jl +++ b/docs/literate/research/synthetic_respiration.jl @@ -98,7 +98,7 @@ single_nn_out = train( single_nn_hybrid_model, df, (); - nepochs = 10, # Number of training epochs + nepochs = 100, # Number of training epochs batchsize = 512, # Batch size for training opt = AdamW(0.1), # Optimizer and learning rate monitor_names = [:rb, :Q10], # Parameters to monitor during training diff --git a/docs/literate/tutorials/example_synthetic_lstm.jl b/docs/literate/tutorials/example_synthetic_lstm.jl index db8dc153..8bccb93c 100644 --- a/docs/literate/tutorials/example_synthetic_lstm.jl +++ b/docs/literate/tutorials/example_synthetic_lstm.jl @@ -11,6 +11,7 @@ using EasyHybrid using AxisKeys using DimensionalData +using CairoMakie # ## 2. Data Loading and Preprocessing @@ -195,18 +196,22 @@ out_lstm = train( opt = RMSProp(0.01), # Optimizer and learning rate monitor_names = [:rb, :Q10], # Parameters to monitor during training yscale = identity, # Scaling for outputs - shuffleobs = true, + shuffleobs = false, training_loss = :nseLoss, - loss_types = [:nse], + loss_types = [:nse, :nseLoss], sequence_kwargs = (; input_window = input_window, output_window = output_window, output_shift = output_shift, lead_time = 0), - plotting = false, + plotting = true, show_progress = false, input_batchnorm = false, array_type = pref_array_type, model_name = "RbQ10_synthetic_lstm" ); -out_lstm.val_obs_pred +# ```@raw html +# +# ``` + +first(out_lstm.val_obs_pred, 5) # ## 10. Train Single NN Hybrid Model (Optional) @@ -234,15 +239,18 @@ single_nn_out = train( opt = RMSProp(0.01), # Optimizer and learning rate monitor_names = [:rb, :Q10], # Parameters to monitor during training yscale = identity, # Scaling for outputs - shuffleobs = true, + shuffleobs = false, training_loss = :nseLoss, - loss_types = [:nse], + loss_types = [:nse, :nseLoss], array_type = :DimArray, - plotting = false, + plotting = true, show_progress = false, model_name = "RbQ10_synthetic_single_nn" ); +# ```@raw html +# +# ``` + # Close enough -out_lstm.best_loss -single_nn_out.best_loss +out_lstm.best_loss, single_nn_out.best_loss diff --git a/docs/make.jl b/docs/make.jl index 845afe57..cd1fa436 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -63,9 +63,9 @@ makedocs(; "Hyperparameter Tuning" => "tutorials/hyperparameter_tuning.md", "GPU Acceleration" => "tutorials/gpu.md", "Synthetic Respiration on GPU" => "tutorials/synthetic_respiration_gpu.md", - # "Slurm" => "tutorials/slurm.md", + "Slurm" => "tutorials/slurm.md", "Cross-validation" => "tutorials/folds.md", - # "LSTM Hybrid Model" => "tutorials/example_synthetic_lstm.md", + "LSTM Hybrid Model" => "tutorials/example_synthetic_lstm.md", "Loss Functions" => "tutorials/losses.md", ], "Research" => [ diff --git a/docs/src/tutorials/exponential_res.md b/docs/src/tutorials/exponential_res.md index d7acbffa..637515b2 100644 --- a/docs/src/tutorials/exponential_res.md +++ b/docs/src/tutorials/exponential_res.md @@ -135,7 +135,7 @@ hybrid_model = constructHybridModel( ```@example expo out = train(hybrid_model, df, (:k,); nepochs=300, batchsize=64, - opt=AdamW(0.01, (0.9, 0.999), 0.01), loss_types=[:mse, :nse], + opt=AdamW(0.01, (0.9, 0.999), 0.01), loss_types=[:mse, :nse, :nseLoss], training_loss=:nseLoss, random_seed=123, yscale = identity, monitor_names=[:Resp0, :k], show_progress=false, diff --git a/docs/test_expo.jl b/docs/test_expo.jl new file mode 100644 index 00000000..7fd5468d --- /dev/null +++ b/docs/test_expo.jl @@ -0,0 +1,55 @@ +using EasyHybrid +using GLMakie + +using Random +Random.seed!(2314) + +T = rand(500) .* 40 .- 10 # Random temperature +SM = rand(500) .* 0.8 .+ 0.1 # Random soil moisture +SM_fac = exp.(-8.0 * (SM .- 0.6) .^ 2) +Resp0 = 1.1 .* SM_fac # Base respiration dependent on soil moisture +Resp = Resp0 .* exp.(0.07 .* T) +Resp_obs = Resp .+ randn(length(Resp)) .* 0.05 .* mean(Resp); # Add some noise + +df = DataFrame(; T, SM, SM_fac, Resp0, Resp, Resp_obs); + +function Expo_resp_model(; T, Resp0, k) + Resp_obs = Resp0 .* exp.(k .* T) + return (; Resp_obs, Resp0, k) +end; + +parameters = ( + # name: (default, lower_bound, upper_bound) # Description + k = (0.01f0, 0.0f0, 0.2f0), # Exponent + Resp0 = (2.0f0, 0.0f0, 8.0f0), # Basal respiration [μmol/m²/s] +); + +targets = [:Resp_obs] +forcings = [:T] +predictors = (Resp0 = [:SM],); + +global_param_names = [:k] + +hybrid_model = constructHybridModel( + predictors, + forcings, + targets, + Expo_resp_model, + parameters, + global_param_names, + scale_nn_outputs = false, # TODO `true` also works with good lower and upper bounds + hidden_layers = [16, 16], + activation = sigmoid, + input_batchnorm = true +) + +out = train( + hybrid_model, df, (:k,); nepochs = 500, batchsize = 64, + opt = AdamW(0.01, (0.9, 0.999), 0.01), loss_types = [:mse, :nse, :nseLoss], + training_loss = :nseLoss, random_seed = 123, yscale = identity, + monitor_names = [:Resp0, :k], + show_progress = false, + hybrid_name = "expo_response" +); + +# predictionplot(out.train_obs_pred[!, :Resp_obs], out.train_obs_pred[!, :Resp_obs_pred]; color = :tomato) diff --git a/docs/test_extMakie.jl b/docs/test_extMakie.jl new file mode 100644 index 00000000..79a1e8a0 --- /dev/null +++ b/docs/test_extMakie.jl @@ -0,0 +1,127 @@ +using EasyHybrid +using GLMakie +using GLMakie.Makie.GeometryBasics: AbstractPoint +using DataStructures: CircularBuffer + +n_epochs = [0.9] +t_arr = sin.(rand(1)) +v_arr = cos.(rand(1)) +fig, ax, plt = lossplot(n_epochs, t_arr, v_arr; axis = (; xlabel = "Epochs", ylabel = "Loss")) +# axislegend(ax, plt) + +Legend(fig[1, 1, Top()], ax, plt) +hidespines!(ax, :r, :t) +fig + +# ! do buffer! +zoom_epochs = 50 +n_epochs_buffer = CircularBuffer{Int64}(zoom_epochs) +fill!(n_epochs_buffer, 0) +t_arr_buffer = CircularBuffer{Float64}(zoom_epochs) +fill!(t_arr_buffer, t_arr[1]) +v_arr_buffer = CircularBuffer{Float64}(zoom_epochs) +fill!(v_arr_buffer, v_arr[1]) + +ax_z = Axis( + fig[1, 1], + width = Relative(0.35), + height = Relative(0.35), + halign = 0.95, + valign = 1, + xlabel = "", + ylabel = "", + rightspinecolor = :dodgerblue, + leftspinecolor = :dodgerblue, + topspinecolor = :dodgerblue, + bottomspinecolor = :dodgerblue, + title = "Zoomed View" +) + +plt_z = lossplot!(ax_z, n_epochs_buffer, t_arr_buffer, v_arr_buffer) +# hidespines!(ax_z, :l, :t) +translate!(ax_z.blockscene, 0, 0, 150) +fig + +o_ax_z = ax_z.scene.viewport[].origin + +function current_rect2(n_epochs_buffer, t_arr_buffer, v_arr_buffer, zoom_epochs, epoch) + xzoom_rect = epoch < zoom_epochs ? epoch : zoom_epochs + mn_tv = minimum(map(minimum, [t_arr_buffer, v_arr_buffer])) + mx_tv = maximum(map(maximum, [t_arr_buffer, v_arr_buffer])) + z_rect = Rect2(minimum(n_epochs_buffer), 0.95 * mn_tv, xzoom_rect, 1.05 * (mx_tv - mn_tv)) + + return z_rect +end + +z_rect = current_rect2(n_epochs_buffer, t_arr_buffer, v_arr_buffer, zoom_epochs, 0) + +plt_b = lines!(ax, z_rect, color = :dodgerblue, linewidth = 1) +fig + +# scatter!(fig.scene, Point2f(o_ax_z)) +# scatter!(fig.scene, Point2f(o_ax_z) + Point2f(first(ax_z.scene.viewport[].widths), 0)) + +function _project_points_to_figure(ax, p::AbstractPoint) + return ax.scene.viewport[].origin + Makie.project(ax.scene, p) +end + +function _axis_bottom_points(ax_z) + left_point = Point2f(ax_z.scene.viewport[].origin) + x_right = first(ax_z.scene.viewport[].widths) + right_point = left_point + Point2f(x_right, 0) + return [left_point, right_point] +end + +_axis_bottom_points(ax_z) + +Legend(fig[1, 1, Top()], ax, plt; nbanks = 3, framewidth = 0.25, halign = 0) +fig + +for epoch in 1:1000 + # push a new data point + n_tv = sin(rand()) / epoch + n_vv = cos(rand()) / epoch + push!(n_epochs, epoch) + push!(t_arr, n_tv) + push!(v_arr, n_vv) + #! now the buffers + push!(n_epochs_buffer, epoch) + push!(t_arr_buffer, n_tv) + push!(v_arr_buffer, n_vv) + + new_z_rect = current_rect2(n_epochs_buffer, t_arr_buffer, v_arr_buffer, zoom_epochs, epoch) + + #? now that all are updated and synchronized we can update the plot + + update!(plt, n_epochs, t_arr, v_arr) + update!(plt_z, n_epochs_buffer, t_arr_buffer, v_arr_buffer) + update!(plt_b, arg1 = new_z_rect) + autolimits!(ax) + autolimits!(ax_z) + sleep(0.002) +end +fig + + +# oo = _project_points_to_figure(ax, Point2f(1000, 0.01)) +# scatter!(fig.scene, Point2f(oo); color = :olive, markersize=15) +# fig + +ax.yscale = log10 + +# oo2 = _project_points_to_figure(ax, Point2f(1000, 0.02)) +# scatter!(fig.scene, Point2f(oo2); color = :orange, markersize=15) + +ax.xscale = log10 +fig + + +fig, ax, plt = lossplot(rand(10), rand(10)) +scatter!(rand(10), label = "some dots") +Legend(fig[0, 1], ax, plt; position = :ct, nbanks = 3, tellheight = true, tellwidth = false) +fig + +fig, ax, plt = lossplot(rand(10), rand(10); validation_label = "validate me") +scatter!(rand(10), label = "some dots") +Legend(fig[0, 1], ax, plt; position = :ct, nbanks = 3, tellheight = true, tellwidth = false) +fig diff --git a/docs/test_monitor.jl b/docs/test_monitor.jl new file mode 100644 index 00000000..e4a53fe6 --- /dev/null +++ b/docs/test_monitor.jl @@ -0,0 +1,73 @@ +using EasyHybrid +using GLMakie +using GLMakie.Makie.GeometryBasics: AbstractPoint +using DataStructures: CircularBuffer + +epochs = 1:20 +# Build fake monitor data matching the expected structure +# Scalar monitors are expected to be tuples of the form (scalar = ,) +training_monitor = ( + loss = (scalar = rand(20) .* 0.5 .+ 0.1,), + accuracy = (scalar = rand(20) .* 0.3 .+ 0.6,), +) +validation_monitor = ( + loss = (scalar = rand(20) .* 0.5 .+ 0.2,), + accuracy = (scalar = rand(20) .* 0.3 .+ 0.5,), +) + +# Quantile monitors are expected to be tuples of the form (quantile = (q25 = , q50 = , q75 = ),) +training_monitor_q = ( + loss = ( + quantile = ( + q25 = rand(20) .* 0.3 .+ 0.05, + q50 = rand(20) .* 0.3 .+ 0.15, + q75 = rand(20) .* 0.3 .+ 0.25, + ), + ), +) +validation_monitor_q = ( + loss = ( + quantile = ( + q25 = rand(20) .* 0.3 .+ 0.1, + q50 = rand(20) .* 0.3 .+ 0.2, + q75 = rand(20) .* 0.3 .+ 0.3, + ), + ), +) + +# 1. Scalar, standalone figure +fig1, ax1, plt1 = monitorplot(epochs, training_monitor, validation_monitor, :loss; axis = (xlabel = "Epoch", ylabel = "Loss")) +# axislegend(ax1, plt1) +hidespines!(ax1, :r, :t) +Legend(fig1[1, 1, Top()], ax1, plt1; orientation = :horizontal, titleposition = :left) +fig1 + +# 2. Quantile, standalone figure +fig2, ax2, plt2 = monitorplot(epochs, training_monitor_q, validation_monitor_q, :loss) +# axislegend(ax2, plt2) +hidespines!(ax2, :r, :t) +Legend(fig2[1, 1, Top()], ax2, plt2; orientation = :horizontal) +fig2 + +# 3. Mutating form + attribute overrides +fig3 = Figure() +ax3 = Axis(fig3[1, 1], title = "Loss (custom style)", xlabel = "Epoch", ylabel = "Loss") +plt3 = monitorplot!( + ax3, epochs, training_monitor, validation_monitor, :loss; + training_color = :steelblue, + validation_color = :crimson, + linewidth = 3, + training_label = "Train", + validation_label = "Val", +) +axislegend(ax3, plt3) +fig3 + +# 4. Multi-panel figure +fig4 = Figure(size = (900, 400)) +for (col, name) in enumerate([:loss, :accuracy]) + ax = Axis(fig4[1, col], title = string(name), xlabel = "Epoch") + plt = monitorplot!(ax, epochs, training_monitor, validation_monitor, name) + axislegend(ax, plt) +end +fig4 diff --git a/ext/EasyHybridMakie.jl b/ext/EasyHybridMakie.jl index 1356fda5..5fe1e237 100644 --- a/ext/EasyHybridMakie.jl +++ b/ext/EasyHybridMakie.jl @@ -6,8 +6,9 @@ using Makie.Colors using DataFrames import Makie import EasyHybrid -import EasyHybrid: get_loss_value +import EasyHybrid: get_loss_value, get_monitor_values, collect_monitor_history using Statistics +using DataStructures: CircularBuffer include("HybridTheme.jl") @@ -32,6 +33,11 @@ function _series(wt::WrappedTuples, attributes) return data_matrix, merged_attributes end +include("recipes/LossPlot.jl") +include("recipes/MonitorPlot.jl") +include("recipes/PredictionPlot.jl") +include("recipes/TimeSeriesPlot.jl") + # ============================================================================= # Prediction vs Observed Plotting Functions # ============================================================================= @@ -82,6 +88,22 @@ function EasyHybrid.poplot!(fig, pred, obs, title_prefix, row::Int, col::Int; xl return EasyHybrid.plot_pred_vs_obs!(ax, pred, obs, title_prefix; xlabel, ylabel) end +""" + plot_pred_vs_obs!(ax, pred, obs, title_prefix; xlabel="Predicted", ylabel="Observed") + +Add a scatter plot comparing predicted vs observed values with performance metrics on an existing axis. + +# Arguments +- `ax`: Makie axis to plot on +- `pred`: Vector of predicted values +- `obs`: Vector of observed values +- `title_prefix`: Title prefix for the plot +- `xlabel`: Label for the x-axis (default: "Predicted") +- `ylabel`: Label for the y-axis (default: "Observed") + +# Returns +- A `Legend` object containing the 1:1 line legend entry. +""" function EasyHybrid.plot_pred_vs_obs!(ax, pred, obs, title_prefix; xlabel = "Predicted", ylabel = "Observed") ss_res = sum((obs .- pred) .^ 2) ss_tot = sum((obs .- mean(obs)) .^ 2) @@ -200,6 +222,18 @@ end # Original Observable-based Loss Plotting (for live training updates) # ============================================================================= +""" + plot_loss(loss, yscale) + +Create an observable-based loss plot for live training updates. + +# Arguments +- `loss`: Observable containing the training loss history +- `yscale`: Y-axis scale function (e.g. `log10`) + +# Returns +- A Makie `Figure` object containing the loss plot +""" function EasyHybrid.plot_loss(loss, yscale) fig = Makie.Figure() ax = Makie.Axis(fig[1, 1]; yscale = yscale, xlabel = "epoch", ylabel = "loss") @@ -210,6 +244,17 @@ function EasyHybrid.plot_loss(loss, yscale) return display(fig; title = "EasyHybrid.jl", focus_on_show = true) end +""" + plot_loss!(loss) + +Add a validation loss line to the current observable-based loss plot. + +# Arguments +- `loss`: Observable containing the validation loss history + +# Returns +- The axis legend added to the current plot +""" function EasyHybrid.plot_loss!(loss) if nameof(Makie.current_backend()) == :WGLMakie # TODO for our CPU cluster - alternatives? sleep(2.0) @@ -219,269 +264,419 @@ function EasyHybrid.plot_loss!(loss) return Makie.axislegend(ax; position = :rt) end +""" + log_tick_formatter(values) + +Format logarithmic axis ticks as superscript powers of 10. + +# Arguments +- `values`: Array of numeric values to format + +# Returns +- Array of formatted string labels (e.g., "10²") +""" function log_tick_formatter(values) return map(v -> "10" * Makie.UnicodeFun.to_superscript(round(Int64, v)), values) end -# ============================================================================= -# Multi‑Target Live Training Dashboard with Monitors -# ============================================================================= - """ - train_board(train_loss, val_loss, - train_preds, train_obs, - val_preds, val_obs, - train_monitor, val_monitor, - yscale, zoom_epochs, - target_names; - monitor_names) + _extract_monitor(monitor, name) -Create a live‑updating dashboard showing per‑target scatter plots for training and validation, -loss curves, and time‑series for additional monitored outputs. +Extract monitor values for a specific parameter name. # Arguments -- `train_loss`, `val_loss`: Observables of loss history vectors -- `train_preds`, `val_preds`: NamedTuples of Observables for per‑target predictions -- `train_obs`, `val_obs`: NamedTuples of Observables for per‑target observations -- `train_monitor`, `val_monitor`: NamedTuples of Observables for extra model outputs -- `yscale`: Y‑axis scale function (e.g. `log10`) -- `target_names`: Symbols of targets to plot -- `monitor_names`: Symbols of extra outputs to monitor -- `zoom_epochs`: Number of epochs to zoom in on loss curve -""" -function EasyHybrid.train_board( - train_loss, - val_loss, - train_preds, - val_preds, - train_monitor, - val_monitor, - train_obs, - val_obs, - yscale, - target_names, - loss_type; - monitor_names, - zoom_epochs - ) - n_targets = length(target_names) - n_monitors = length(monitor_names) - # total_rows = max(n_targets, n_monitors) - total_gds = (n_targets + n_monitors) - j_max = Int(floor(total_gds / 2)) + 1 - - fig = Makie.Figure(; size = (1200, 400 * j_max)) - # let's do a GridLayout per topic, Per‑target scatter subplots (side by side) - gd_losses = GridLayout(fig[1, 1]) - gd_t1 = GridLayout(fig[1, 2]) - gd_tm = [gd_t1] - # create more grid layouts to accommodate additional targets and monitor_names. - max_counter = 1 - for j in 2:j_max - for i in 1:2 - push!(gd_tm, GridLayout(fig[j, i])) - max_counter += 1 - if max_counter >= total_gds - break - end - end +- `monitor`: Dictionary or NamedTuple containing monitor values +- `name`: Symbol of the monitor parameter to extract + +# Returns +- Tuple containing: + - The extracted values (either a scalar or a quantile dictionary) + - Boolean indicating if the values are quantiles + - Array of quantile keys (empty if scalar) +""" +function _extract_monitor(monitor, name) + entry = monitor[name] + if haskey(entry, :quantile) + q = entry[:quantile] + return q, true, collect(keys(q)) + else + return entry[:scalar], false, Symbol[] end +end - # gd_losses - ax_loss = Makie.Axis(gd_losses[1, 1]; yscale = yscale, xlabel = "Epoch", ylabel = "$(string(loss_type))", aspect = 1) - Makie.lines!(ax_loss, train_loss; color = :grey25, label = "Training", linewidth = 2) - Makie.lines!(ax_loss, val_loss; color = :tomato, label = "Validation", linewidth = 2) - # Zoomed loss in last zoom_epochs - ax_zoom = Makie.Axis( - gd_losses[1, 2], - xlabel = "Epoch", ylabel = "", aspect = 1, - title = "Zoomed View", titlefont = :regular - ) +""" + build_dashboards(history, cfg, y_train, y_val) - zoom_idx = @lift(max(1, length($train_loss) - zoom_epochs)) - tlz = @lift($train_loss[$zoom_idx:end]) - vlz = @lift($val_loss[$zoom_idx:end]) - Makie.lines!(ax_zoom, tlz; color = :grey25, label = "Training (Zoom)", linewidth = 2) - Makie.lines!(ax_zoom, vlz; color = :tomato, label = "Validation (Zoom)", linewidth = 2) - # Makie.axislegend(ax_zoom; position = :rt, nbanks = 2) - on(train_loss) do _ - autolimits!(ax_loss); autolimits!(ax_zoom) - end +Initialize the training dashboards with static and live-updating components. - Makie.Legend( - gd_losses[0, 1:2], ax_loss; nbanks = 2, - framewidth = 0, backgroundcolor = (:grey25, 0.1), tellheight = true, - ) - hidespines!(ax_loss, :r, :t) +# Arguments +- `history`: `TrainingHistory` object containing metrics +- `cfg`: `TrainConfig` object containing plotting configuration +- `y_train`: Training targets +- `y_val`: Validation targets +# Returns +- Tuple containing: + - `figures`: Dict mapping component name to Figure + - `axes_dict`: Dict mapping component name to named tuple of axes + - `plots_dict`: Dict mapping component name to named tuple of plots +""" +function EasyHybrid.build_dashboards(history, cfg, y_train, y_val) + components = cfg.dashboard_components + split = cfg.split_dashboard + + figures = Dict{Symbol, Any}() + axes_dict = Dict{Symbol, Any}() + plots_dict = Dict{Symbol, Any}() + + n_epochs = get_epochs(history) + + if split + for comp in components + fig = Makie.Figure(size = (800, 600)) + figures[comp] = fig + _build_component!(comp, fig, fig[1, 1], history, cfg, y_train, y_val, n_epochs, axes_dict, plots_dict) + display(fig) + end + else + fig = Makie.Figure(size = (1200, 800)) + figures[:dashboard] = fig - for (i, t) in enumerate(target_names) - # Data - p_tr = getfield(train_preds, t) - o_tr = getfield(train_obs, t) + n_comp = length(components) + rows = n_comp > 2 ? 2 : 1 + cols = ceil(Int, n_comp / rows) - maxpoints = 10_000 - idx = @lift begin #TODO better with density plot? - n = length($p_tr) - if n > maxpoints - randperm(n)[1:maxpoints] - else - 1:n - end + for (i, comp) in enumerate(components) + r = ceil(Int, i / cols) + c = ((i - 1) % cols) + 1 + _build_component!(comp, fig, fig[r, c], history, cfg, y_train, y_val, n_epochs, axes_dict, plots_dict) end + display(fig) + end - p_tr_sub = @lift($p_tr[$idx]) # Observable - o_tr_sub = @lift(o_tr[$idx]) # o_val captured as constant + return figures, axes_dict, plots_dict +end - mn, mx = extrema(filter(!isnan, o_tr)) - δd = 0.1 - # Training scatter plot - ax_tr = Makie.Axis( - gd_tm[i][1, 1]; aspect = 1, xlabel = "Predicted", ylabel = "", - limits = (mn - δd, mx + δd, mn - δd, mx + δd) +function _build_component!(comp, fig, layout, history, cfg, y_train, y_val, n_epochs, axes_dict, plots_dict) + return if comp == :loss + vals_train = get_loss_value_t(history, cfg.training_loss, Symbol("$(cfg.agg)")) + vals_val = get_loss_value_v(history, cfg.training_loss, Symbol("$(cfg.agg)")) + + ax, plt = lossplot( + layout, + n_epochs, vals_train, vals_val; + axis = (; + xlabel = "Epochs", ylabel = "Loss", yscale = log10, + xtrimspine = (true, false), ytrimspine = true, + ) ) - hidespines!(ax_tr, :r, :t) + Legend(layout[1, 1, Top()], ax, plt; orientation = :horizontal, halign = :left, framevisible = false) + hidespines!(ax, :r, :t) + z_rect = z_Rect2(n_epochs, vals_train, vals_val) + plt_rect = lines!(ax, z_rect, color = :dodgerblue, linewidth = 1) + + ax_z = Axis( + layout, + width = Relative(0.35), height = Relative(0.35), + halign = 0.95, valign = 1, + xlabel = "", ylabel = "", + rightspinecolor = :dodgerblue, leftspinecolor = :dodgerblue, + topspinecolor = :dodgerblue, bottomspinecolor = :dodgerblue, + title = "Zoomed View" + ) + plt_z = lossplot!(ax_z, n_epochs, vals_train, vals_val) + translate!(ax_z.blockscene, 0, 0, 150) + + axes_dict[:loss] = (; ax, ax_z) + plots_dict[:loss] = (; plt, plt_rect, plt_z) + + elseif comp == :prediction + y_pred_train = get_prediction_values(history, cfg.target_names[1], :train) + y_pred_val = get_prediction_values(history, cfg.target_names[1], :validation) + y_obs_train = getfield(y_train, cfg.target_names[1]) + y_obs_val = getfield(y_val, cfg.target_names[1]) + + gd_pred = GridLayout(layout) + ax_pred_train = Axis( + gd_pred[1, 1]; xlabel = "", ylabel = "Observed", title = "Training", + xtrimspine = true, ytrimspine = true, aspect = 1 + ) + hidespines!(ax_pred_train, :r, :t) + plt_pred_train = predictionplot!(ax_pred_train, y_pred_train, y_obs_train) - Box(gd_tm[i][1, 1:2, Top()]; color = (:grey25, 0.1), strokevisible = false) - Label(gd_tm[i][1, 1:2, Top()], "$(t)") + ax_pred_val = Axis( + gd_pred[1, 2]; xlabel = "", ylabel = "", title = "Validation", + xtrimspine = true, ytrimspine = true, aspect = 1 + ) + hidespines!(ax_pred_val, :l, :r, :t) + plt_pred_val = predictionplot!(ax_pred_val, y_pred_val, y_obs_val; color = :tomato) + hideydecorations!(ax_pred_val, grid = false, ticks = false) + linkyaxes!(ax_pred_train, ax_pred_val) + + Box(gd_pred[1, 1:2, Top()]; color = (:grey25, 0.1), strokevisible = false) + Label(gd_pred[1, 1:2, Top()], "$(cfg.target_names[1])") + + Box(gd_pred[1, 1:2, Bottom()]; color = (:grey45, 0.1), strokevisible = false) + Label(gd_pred[1, 1:2, Bottom()], "Predicted") + + axes_dict[:prediction] = (; ax_pred_train, ax_pred_val) + plots_dict[:prediction] = (; plt_pred_train, plt_pred_val) + + elseif comp == :timeseries + y_pred_train = get_prediction_values(history, cfg.target_names[1], :train) + y_pred_val = get_prediction_values(history, cfg.target_names[1], :validation) + y_obs_train = getfield(y_train, cfg.target_names[1]) + y_obs_val = getfield(y_val, cfg.target_names[1]) + + gd_ts = GridLayout(layout) + ax_ts_train = Axis( + gd_ts[1, 1]; xlabel = "Time Index", ylabel = "Value", title = "Training (Time Series)", + xtrimspine = true, ytrimspine = true + ) + hidespines!(ax_ts_train, :r, :t) + plt_ts_train = timeseriesplot!(ax_ts_train, y_pred_train, y_obs_train) - Makie.scatter!(ax_tr, p_tr_sub, o_tr_sub; color = :grey25, alpha = 0.6, markersize = 6) - Makie.lines!(ax_tr, sort(o_tr), sort(o_tr); color = :black, linestyle = :dash) - # Validation scatter plot - ax_val = Makie.Axis( - gd_tm[i][1, 2]; aspect = 1, xlabel = "Predicted", ylabel = "", - limits = (mn - δd, mx + δd, mn - δd, mx + δd) + ax_ts_val = Axis( + gd_ts[1, 2]; xlabel = "Time Index", ylabel = "", title = "Validation (Time Series)", + xtrimspine = true, ytrimspine = true ) - hideydecorations!(ax_val, grid = false) - hidespines!(ax_val, :l, :t) + hidespines!(ax_ts_val, :l, :r, :t) + plt_ts_val = timeseriesplot!(ax_ts_val, y_pred_val, y_obs_val) + hideydecorations!(ax_ts_val, grid = false, ticks = false) + linkyaxes!(ax_ts_train, ax_ts_val) + + axes_dict[:timeseries] = (; ax_ts_train, ax_ts_val) + plots_dict[:timeseries] = (; plt_ts_train, plt_ts_val) + + elseif comp == :monitor + if !isempty(cfg.monitor_names) + gl_m, axes_m, plts_m = setup_monitor_panel!(fig, layout, history, cfg) + axes_dict[:monitor] = (; axes_m) + plots_dict[:monitor] = (; plts_m) + end + end +end - p_val = getfield(val_preds, t) - o_val = getfield(val_obs, t) +""" + z_Rect2(z_n_epochs, train_zoom, val_zoom) - val_idx = @lift begin - n = length($p_val) - if n > maxpoints - randperm(n)[1:maxpoints] - else - 1:n - end - end +Create a bounding rectangle for the zoomed-in view of the loss curve. - p_val_sub = @lift($p_val[$val_idx]) # Observable - o_val_sub = @lift(o_val[$val_idx]) +# Arguments +- `z_n_epochs`: Array of epoch indices for the zoomed window +- `train_zoom`: Array of training loss values in the zoomed window +- `val_zoom`: Array of validation loss values in the zoomed window + +# Returns +- A `Rect2` representing the bounding box for the zoomed region +""" +function z_Rect2(z_n_epochs, train_zoom, val_zoom) + mn_epoch = minimum(z_n_epochs) + mx_epoch = maximum(z_n_epochs) + xzoom_rect = mx_epoch - mn_epoch + 1 + mn_tv = minimum(map(minimum, [train_zoom, val_zoom])) + mx_tv = maximum(map(maximum, [train_zoom, val_zoom])) + z_rect = Rect2(mn_epoch - 0.5, 0.95 * mn_tv, xzoom_rect, 1.05 * (mx_tv - mn_tv)) + + return z_rect +end + +""" + setup_monitor_panel!(fig, grid_position, history, cfg) - Makie.scatter!(ax_val, p_val_sub, o_val_sub; color = :tomato, alpha = 0.6, markersize = 6) - Makie.lines!(ax_val, sort(o_val), sort(o_val); color = :black, linestyle = :dash) +Initialize the monitor plot panel within a specific grid layout. + +# Arguments +- `fig`: The parent Makie figure +- `grid_position`: Tuple representing the position in the GridLayout +- `history`: `TrainingHistory` object containing metrics +- `cfg`: `TrainConfig` object containing plotting configuration + +# Returns +- Tuple containing: + - `gl`: The created GridLayout for the panel + - `axes`: Array of initialized axes + - `plts`: Array of initialized plots +""" +function setup_monitor_panel!(fig, grid_position, history, cfg) + monitor_names = cfg.monitor_names + n = length(monitor_names) + + raw_train = get_monitor_values(history, monitor_names, :train) + training_mon = collect_monitor_history(raw_train, monitor_names) + raw_val = get_monitor_values(history, monitor_names, :validation) + validation_mon = collect_monitor_history(raw_val, monitor_names) + + n_epochs = get_epochs(history) + + # Use a nested GridLayout at the given position + gl = GridLayout(grid_position) + + axes = Vector{Axis}(undef, n) + plts = Vector{MonitorPlot}(undef, n) + + for (i, name) in enumerate(monitor_names) + y_train, is_q, qkeys = _extract_monitor(training_mon, name) + y_val, _, _ = _extract_monitor(validation_mon, name) + + ax = Axis( + gl[1, i]; + xlabel = "Epochs", + ylabel = string(name), + xtrimspine = (true, false), + ytrimspine = true, + ) + hidespines!(ax, :r, :t) + + plt = monitorplot!( + ax, n_epochs, y_train, y_val; + is_quantile = is_q, + quantile_keys = qkeys, + ) + Legend( + gl[1, i, Top()], ax, plt; + orientation = :horizontal, + titleposition = :left, + framevisible = false, + ) + + axes[i] = ax + plts[i] = plt end - Label(gd_tm[1][1:end, 0], "Observed", tellheight = false, rotation = pi / 2) - # Label(gd_tm[end+1,1:end], "Predicted") - Label(gd_tm[1][0, 1], "Training"; color = :grey25, tellwidth = false) - Label(gd_tm[1][0, 2], "Validation"; color = :tomato, tellwidth = false) - - # Columns 5-6: Additional monitored outputs - for (j, m) in enumerate(monitor_names) - ax_mt = Makie.Axis(gd_tm[j + n_targets][1, 1]; xlabel = "Epoch", ylabel = string(m), title = "Monitor: $(m)") - m_tr = getfield(train_monitor, m) - m_val = getfield(val_monitor, m) - - if length(m_tr) > 1 - for (qi, q) in enumerate([0.75, 0.5, 0.25]) - qntl = Symbol("q", string(Int(q * 100))) - m_tr_ex = getfield(m_tr, qntl) - m_val_ex = getfield(m_val, qntl) - lw = q == 0.5 ? 3 : 1 # thickest for q50, thin for q25 and q75 - Makie.lines!(ax_mt, m_tr_ex; color = :grey25, linewidth = lw, label = String(qntl)) - Makie.lines!(ax_mt, m_val_ex; color = :tomato, linewidth = lw, linestyle = :dash) - on(m_val_ex) do _ - autolimits!(ax_mt) - end - Makie.linkxaxes!(ax_loss, ax_mt) - end - Makie.axislegend(ax_mt; position = :lt) - else - m_tr_ex = getfield(m_tr, :scalar) - m_val_ex = getfield(m_val, :scalar) - Makie.lines!(ax_mt, m_tr_ex; color = :grey25, linewidth = 2, label = "Training") - Makie.lines!(ax_mt, m_val_ex; color = :tomato, linewidth = 2, linestyle = :dash, label = "Validation") - #Makie.axislegend(ax_mt; position = :rt) - on(m_val_ex) do _ - autolimits!(ax_mt) - end - Makie.linkxaxes!(ax_loss, ax_mt) - end + + return gl, axes, plts +end + +""" + update_monitor_panel!(axes, plts, history, cfg) + +Update the monitor panel plots with new values from the training history. + +# Arguments +- `axes`: Array of axes in the monitor panel +- `plts`: Array of plot objects in the monitor panel +- `history`: `TrainingHistory` object containing updated metrics +- `cfg`: `TrainConfig` object +""" +function update_monitor_panel!(axes, plts, history, cfg) + monitor_names = cfg.monitor_names + n_epochs = get_epochs(history) + + raw_train = get_monitor_values(history, monitor_names, :train) + training_mon = collect_monitor_history(raw_train, monitor_names) + raw_val = get_monitor_values(history, monitor_names, :validation) + validation_mon = collect_monitor_history(raw_val, monitor_names) + + for (i, name) in enumerate(monitor_names) + y_train, _, _ = _extract_monitor(training_mon, name) + y_val, _, _ = _extract_monitor(validation_mon, name) + update!(plts[i], n_epochs, y_train, y_val) + autolimits!(axes[i]) end - return display(fig) + return end """ - update_plotting_observables(ext, train_h_obs, val_h_obs, train_preds, val_preds, train_monitor, val_monitor, hybridModel, x_train, x_val, ps, st, l_train, l_val, training_loss, agg, epoch, monitor_names) - -Update plotting observables during training if the Makie extension is loaded. -""" -function EasyHybrid.update_plotting_observables( - train_h_obs, - val_h_obs, - train_preds, - val_preds, - train_monitor, - val_monitor, - l_train, - l_val, - training_loss, - agg, - current_ŷ_train, - current_ŷ_val, - target_names, - epoch; - monitor_names - ) + update_step_dashboards!(dashboard, history, cfg) - l_value = get_loss_value(l_train, training_loss, Symbol("$agg")) - new_p = Point2f(epoch, l_value) - push!(train_h_obs[], new_p) - notify(train_h_obs) - - l_value_val = get_loss_value(l_val, training_loss, Symbol("$agg")) - new_p_val = Point2f(epoch, l_value_val) - push!(val_h_obs[], new_p_val) - - for t in target_names - # replace the array stored in the Observable: - train_preds[t][] = vec(getfield(current_ŷ_train, t)) - val_preds[t][] = vec(getfield(current_ŷ_val, t)) - # and notify Makie that it changed: - notify(train_preds[t]) - notify(val_preds[t]) +Update all plots in the training dashboard with the latest epoch data. + +# Arguments +- `dashboard`: The `TrainDashboard` object +- `history`: `TrainingHistory` object containing updated metrics +- `cfg`: `TrainConfig` object +""" +function EasyHybrid.update_step_dashboards!(dashboard, history, cfg) + n_epochs = get_epochs(history) + + if haskey(dashboard.plots, :loss) + zoom_epochs = 50 + vals_train = get_loss_value_t(history, cfg.training_loss, Symbol("$(cfg.agg)")) + vals_val = get_loss_value_v(history, cfg.training_loss, Symbol("$(cfg.agg)")) + + update!(dashboard.plots[:loss].plt, n_epochs, vals_train, vals_val) + autolimits!(dashboard.axes[:loss].ax) + + zoom_idx = max(1, length(vals_train) - zoom_epochs) + train_zoom = vals_train[zoom_idx:end] + val_zoom = vals_val[zoom_idx:end] + z_n_epochs = n_epochs[zoom_idx:end] + + updatedRect2 = z_Rect2(z_n_epochs, train_zoom, val_zoom) + update!(dashboard.plots[:loss].plt_rect, arg1 = updatedRect2) + + update!(dashboard.plots[:loss].plt_z, z_n_epochs, val_zoom, train_zoom) + autolimits!(dashboard.axes[:loss].ax_z) end - if !isempty(monitor_names) - for m in monitor_names - v_tr = vec(getfield(current_ŷ_val, m)) # ? it was set to train before? bug? - m_tr = vec(getfield(current_ŷ_train, m)) - - if length(v_tr) > 1 - for q in [0.25, 0.5, 0.75] - push!(val_monitor[m][Symbol("q", string(Int(q * 100)))][], Point2f(epoch, quantile(v_tr, q))) - push!(train_monitor[m][Symbol("q", string(Int(q * 100)))][], Point2f(epoch, quantile(m_tr, q))) - notify(val_monitor[m][Symbol("q", string(Int(q * 100)))]) - notify(train_monitor[m][Symbol("q", string(Int(q * 100)))]) - end - else - push!(val_monitor[m][:scalar][], Point2f(epoch, v_tr[1])) - push!(train_monitor[m][:scalar][], Point2f(epoch, m_tr[1])) - notify(val_monitor[m][:scalar]) - notify(train_monitor[m][:scalar]) - end - end + if haskey(dashboard.plots, :prediction) + y_pred_train = get_prediction_values(history, cfg.target_names[1], :train) + y_pred_val = get_prediction_values(history, cfg.target_names[1], :validation) + update!(dashboard.plots[:prediction].plt_pred_train, y_pred_train) + update!(dashboard.plots[:prediction].plt_pred_val, y_pred_val) + autolimits!(dashboard.axes[:prediction].ax_pred_train) + autolimits!(dashboard.axes[:prediction].ax_pred_val) + end + + if haskey(dashboard.plots, :timeseries) + y_pred_train = get_prediction_values(history, cfg.target_names[1], :train) + y_pred_val = get_prediction_values(history, cfg.target_names[1], :validation) + update!(dashboard.plots[:timeseries].plt_ts_train, y_pred_train) + update!(dashboard.plots[:timeseries].plt_ts_val, y_pred_val) + autolimits!(dashboard.axes[:timeseries].ax_ts_train) + autolimits!(dashboard.axes[:timeseries].ax_ts_val) + end + + if haskey(dashboard.plots, :monitor) && !isempty(cfg.monitor_names) + update_monitor_panel!(dashboard.axes[:monitor].axes_m, dashboard.plots[:monitor].plts_m, history, cfg) end - return notify(val_h_obs) + + return nothing end + +""" + dashboard_figure() + +Get the current Makie figure for the dashboard. +""" EasyHybrid.dashboard_figure() = Makie.current_figure() + +""" + record_history(args...; kargs...) + +Record a video of the dashboard using `Makie.record`. +""" EasyHybrid.record_history(args...; kargs...) = Makie.record(args...; backend = Makie.current_backend(), kargs...) + +""" + VideoStream(fig; kargs...) + +Create a video stream using `Makie.VideoStream`. +""" +EasyHybrid.VideoStream(fig; kargs...) = Makie.VideoStream(fig; kargs...) + +""" + recordframe!(io) + +Record a single frame to the video stream. +""" EasyHybrid.recordframe!(io) = Makie.recordframe!(io) + +""" + save_fig(args...) + +Save the current figure to a file. +""" EasyHybrid.save_fig(args...) = Makie.save(args...) +""" + save_video(path, io) + +Save the recorded video stream to a file. +""" +EasyHybrid.save_video(path, io) = Makie.save(path, io) + # ============================================================================= # Generic Dispatch Methods for Loss and Parameter Plotting # ============================================================================= @@ -734,10 +929,20 @@ function EasyHybrid.plot_training_summary(results::EasyHybrid.TrainResults; loss return fig end +""" + to_obs(o) + +Convert a value to a Makie Observable. +""" function EasyHybrid.to_obs(o) return Makie.Observable(o) end +""" + to_point2f(i, p) + +Create a Point2f from an index and a value. +""" function EasyHybrid.to_point2f(i, p) return Makie.Point2f(i, p) end diff --git a/ext/HybridTheme.jl b/ext/HybridTheme.jl index ec75daea..08afa907 100644 --- a/ext/HybridTheme.jl +++ b/ext/HybridTheme.jl @@ -53,8 +53,8 @@ function theme_easy_hybrid() ylabel = "y", xtickalign = 1, ytickalign = 1, - yticksize = 10, - xticksize = 10, + # yticksize = 5, + # xticksize = 5, xgridstyle = :dash, ygridstyle = :dash, xminorgridstyle = :dash, yminorgridstyle = :dash, xminorgridvisible = true, @@ -63,7 +63,11 @@ function theme_easy_hybrid() # ytrimspine = true, # rightspinevisible = false, # topspinevisible = false - spinewidth = 0.5, + spinewidth = 1, + leftspinecolor = :black, + bottomspinecolor = :black, + topspinecolor = :black, + rightspinecolor = :black, titlefont = :regular, ), Legend = ( diff --git a/ext/recipes/LossPlot.jl b/ext/recipes/LossPlot.jl new file mode 100644 index 00000000..90fcf3c4 --- /dev/null +++ b/ext/recipes/LossPlot.jl @@ -0,0 +1,68 @@ +import EasyHybrid: lossplot, lossplot! + +@recipe LossPlot (epochs_range, training_loss, validation_loss) begin + "Y-axis scale function, e.g. `log10`" + yscale = identity + "Number of recent epochs shown in the zoom panel" + zoom_epochs = 50 + "Training label" + training_label = "Training" + "Validation label" + validation_label = "Validation" + "Colour for training curves" + training_color = :grey25 + "Colour for validation curves" + validation_color = :tomato + "Line width for both curves" + linewidth = 2 +end + +function Makie.plot!(p::LossPlot) + Makie.lines!( + p, p[:epochs_range], p[:training_loss]; + color = p.training_color, + linewidth = p.linewidth, + label = p.training_label + ) + Makie.lines!( + p, p[:epochs_range], p[:validation_loss]; + color = p.validation_color, + linewidth = p.linewidth, + label = p.validation_label + ) + return p +end + +function _legend_entries(ax::Makie.Axis, plt::LossPlot) + loss_plots = collect(plt.plots) + loss_labels = [plt.training_label[], plt.validation_label[]] + + other_plots = filter(ax.scene.plots) do p + haskey(p.attributes, :label) && + p.label[] != "" && + p ∉ loss_plots + end + other_labels = String[p.label[] for p in other_plots] + + return [loss_plots; other_plots], [loss_labels; other_labels] +end + +function Makie.axislegend(ax::Makie.Axis, plt::LossPlot; title = nothing, kwargs...) + plots, labels = _legend_entries(ax, plt) + return Makie.axislegend(ax, plots, labels, title; kwargs...) +end + +function Makie.Legend(gp, ax::Makie.Axis, plt::LossPlot; title = nothing, kwargs...) + plots, labels = _legend_entries(ax, plt) + return Makie.Legend(gp, plots, labels, title; kwargs...) +end + +function Makie.update!(plt::LossPlot, epochs_range, training_loss, validation_loss) + Makie.update!( + plt, + arg1 = epochs_range, + arg2 = training_loss, + arg3 = validation_loss + ) + return nothing +end diff --git a/ext/recipes/MonitorPlot.jl b/ext/recipes/MonitorPlot.jl new file mode 100644 index 00000000..66770e4b --- /dev/null +++ b/ext/recipes/MonitorPlot.jl @@ -0,0 +1,125 @@ +import EasyHybrid: monitorplot, monitorplot! + +@recipe MonitorPlot (epochs_range, y_train, y_val) begin + "Colour for training curves" + training_color = :grey25 + "Colour for validation curves" + validation_color = :tomato + "Training label" + training_label = "Training" + "Validation label" + validation_label = "Validation" + "Line width for both curves" + linewidth = 2 + "Whether the monitor data is quantile-based (i.e. contains multiple quantiles per monitor) or scalar-based (one value per monitor)" + is_quantile = false + "If `is_quantile` is true, the keys of the quantiles to plot, e.g. `[:q25, :q50, :q75]`. Ignored if `is_quantile` is false." + quantile_keys = Symbol[] +end + +function Makie.plot!(p::MonitorPlot) + if p.is_quantile[] + _plot_monitor_quantiles!(p) + else + _plot_monitor_lines!(p) + end + return p +end + +function _plot_monitor_lines!(p) + Makie.lines!( + p, p[:epochs_range], p[:y_train]; + color = p.training_color, linewidth = p.linewidth, label = p.training_label + ) + Makie.lines!( + p, p[:epochs_range], p[:y_val]; + color = p.validation_color, linewidth = p.linewidth, + linestyle = :dash, label = p.validation_label + ) + return p +end + +function _plot_monitor_quantiles!(p) + qkeys = p.quantile_keys[] + n = length(qkeys) + mid_idx = something(findfirst(==(:q50), qkeys), n ÷ 2 + 1) + + for (i, qntl) in enumerate(qkeys) + distance = abs(i - mid_idx) + alpha = 1.0 - 0.25 * distance + lw = p.linewidth[] / (1 + 0.5 * distance) + + # Create a plain Observable for this slot's data + tr_obs = Observable(p[:y_train][][qntl]) + val_obs = Observable(p[:y_val][][qntl]) + + # Wire them to update whenever y_train / y_val change + on(p[:y_train]) do yt + tr_obs[] = yt[qntl] + end + on(p[:y_val]) do yv + val_obs[] = yv[qntl] + end + + Makie.lines!( + p, p[:epochs_range], tr_obs; + color = (p.training_color[], alpha), linewidth = lw, label = string(qntl) + ) + Makie.lines!( + p, p[:epochs_range], val_obs; + color = (p.validation_color[], alpha), linewidth = lw, + linestyle = :dash, label = "" + ) + end + return p +end + +function _legend_entries(ax::Makie.Axis, plt::MonitorPlot) + child_plots_with_labels = [ + p for p in plt.plots + if haskey(p.attributes, :label) && p.label[] != "" + ] + child_labels = [p.label[] for p in child_plots_with_labels] + + other_plots = filter(ax.scene.plots) do p + haskey(p.attributes, :label) && p.label[] != "" && p ∉ plt.plots + end + other_labels = String[p.label[] for p in other_plots] + + # Only add solid/dashed explanation dummies in the quantile case, + # where training_label/validation_label are not already in the child labels + has_quantiles = !any(==(plt.training_label[]), child_labels) + extras = has_quantiles ? [ + LineElement(color = plt.training_color[], linestyle = :solid), + LineElement(color = plt.validation_color[], linestyle = :dash), + ] : [] + extra_labels = has_quantiles ? [plt.training_label[], plt.validation_label[]] : String[] + + return ( + plots = [child_plots_with_labels; other_plots], + labels = [child_labels; other_labels], + extras = extras, + extra_labels = extra_labels, + ) +end + +function Makie.axislegend(ax::Makie.Axis, plt::MonitorPlot; title = nothing, kwargs...) + leg_entry = _legend_entries(ax, plt) + return Makie.axislegend(ax, [leg_entry.extras; leg_entry.plots], [leg_entry.extra_labels; leg_entry.labels], title; kwargs...) +end + +function Makie.Legend(gp, ax::Makie.Axis, plt::MonitorPlot; title = nothing, kwargs...) + leg_entry = _legend_entries(ax, plt) + return Makie.Legend(gp, [leg_entry.extras; leg_entry.plots], [leg_entry.extra_labels; leg_entry.labels], title; kwargs...) +end + +function Makie.update!(plt::MonitorPlot, epochs_range, training_monitor, validation_monitor, monitor_name) + Makie.update!( + plt, + arg1 = epochs_range, + arg2 = training_monitor, + arg3 = validation_monitor, + arg4 = monitor_name, + ) + return nothing +end diff --git a/ext/recipes/PredictionPlot.jl b/ext/recipes/PredictionPlot.jl new file mode 100644 index 00000000..3efec38d --- /dev/null +++ b/ext/recipes/PredictionPlot.jl @@ -0,0 +1,64 @@ +import EasyHybrid: predictionplot, predictionplot! + +@recipe PredictionPlot (predictions, observations) begin + "Label shown in the panel header" + target_name = "target" + "Maximum scatter points drawn (random subsample when exceeded)" + maxpoints = 10_000 + "Marker colour for training scatter" + color = :grey25 + "Marker size" + markersize = 8 + "Marker alpha (transparency)" + alpha = 0.6 + "Line style for the 1:1 reference line" + linestyle = :solid + "Line width for the 1:1 reference line" + linewidth = 0.85 +end + +function _align_predictions(preds::AbstractArray, obs::AbstractArray) + if length(preds) != length(obs) + if ndims(preds) == 2 + batch_size = size(preds, 2) + if ndims(obs) == 2 && size(obs, 2) == batch_size + nout = size(obs, 1) + preds = preds[(end - nout + 1):end, :] + elseif ndims(obs) == 1 && length(obs) == batch_size + preds = preds[end:end, :] + end + end + end + return vec(preds), vec(obs) +end + +Makie.convert_arguments(::Type{<:PredictionPlot}, preds::AbstractArray, obs::AbstractArray) = _align_predictions(preds, obs) + +function Makie.plot!(p::PredictionPlot) + Makie.scatter!( + p, p[:predictions], p[:observations]; + color = p.color, + markersize = p.markersize, + alpha = p.alpha, + ) + + Makie.ablines!( + p, 0, 1; + color = :black, linestyle = p.linestyle, linewidth = p.linewidth + ) + + return p +end + +function Makie.update!(plt::PredictionPlot, predictions) + obs = plt[:observations][] + preds_aligned, _ = _align_predictions(predictions, obs) + Makie.update!(plt, arg1 = preds_aligned) + return nothing +end + +function Makie.update!(plt::PredictionPlot, predictions, observations) + preds_aligned, obs_aligned = _align_predictions(predictions, observations) + Makie.update!(plt, arg1 = preds_aligned, arg2 = obs_aligned) + return nothing +end diff --git a/ext/recipes/TimeSeriesPlot.jl b/ext/recipes/TimeSeriesPlot.jl new file mode 100644 index 00000000..89602ac5 --- /dev/null +++ b/ext/recipes/TimeSeriesPlot.jl @@ -0,0 +1,65 @@ +import EasyHybrid: timeseriesplot, timeseriesplot! + +@recipe TimeSeriesPlot (predictions, observations) begin + "Marker colour for predictions" + pred_color = :tomato + "Marker colour for observations" + obs_color = :grey25 + "Marker size" + markersize = 6 + "Marker alpha (transparency)" + alpha = 0.6 +end + +function _align_timeseries(preds::AbstractArray, obs::AbstractArray) + if length(preds) != length(obs) + if ndims(preds) == 2 + batch_size = size(preds, 2) + if ndims(obs) == 2 && size(obs, 2) == batch_size + nout = size(obs, 1) + preds = preds[(end - nout + 1):end, :] + elseif ndims(obs) == 1 && length(obs) == batch_size + preds = preds[end:end, :] + end + end + end + return vec(preds), vec(obs) +end + +Makie.convert_arguments(::Type{<:TimeSeriesPlot}, preds::AbstractArray, obs::AbstractArray) = _align_timeseries(preds, obs) + +function Makie.plot!(p::TimeSeriesPlot) + preds = p[:predictions] + obs = p[:observations] + + Makie.scatter!( + p, obs; + color = p.obs_color, + markersize = p.markersize, + alpha = p.alpha, + label = "Observed" + ) + + Makie.scatter!( + p, preds; + color = p.pred_color, + markersize = p.markersize, + alpha = p.alpha, + label = "Predicted" + ) + + return p +end + +function Makie.update!(plt::TimeSeriesPlot, predictions) + obs = plt[:observations][] + preds_aligned, _ = _align_timeseries(predictions, obs) + Makie.update!(plt, arg1 = preds_aligned) + return nothing +end + +function Makie.update!(plt::TimeSeriesPlot, predictions, observations) + preds_aligned, obs_aligned = _align_timeseries(predictions, observations) + Makie.update!(plt, arg1 = preds_aligned, arg2 = obs_aligned) + return nothing +end diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl index cb3a19cc..6caa8922 100644 --- a/src/config/TrainingConfig.jl +++ b/src/config/TrainingConfig.jl @@ -109,6 +109,9 @@ loss computation, data handling, output, and visualization. "Vector of monitor names to track during training. Default: `[]`." monitor_names::Vector = [] + "Vector of target names for plotting. Default: `[]`." + target_names::Vector = [] + "Additional folder name string appended to the output path. Default: empty string." output_folder::String = "" @@ -118,6 +121,15 @@ loss computation, data handling, output, and visualization. "Whether to show progress bars during training. Default: `true`." show_progress::Bool = true + "Which components to show in the dashboard. Options: `:loss`, `:prediction`, `:timeseries`, `:monitor`. Default: `[:loss, :prediction, :timeseries, :monitor]`." + dashboard_components::Vector{Symbol} = [:loss, :prediction, :timeseries, :monitor] + + "If `true`, each component requested will be rendered in its own separate figure instead of one unified dashboard. Default: `false`." + split_dashboard::Bool = false + + "Which components to save as animations. By default `[:all]` saves the unified dashboard if `split_dashboard=false` or all individual components if `true`." + save_animations::Vector{Symbol} = [:all] + "Scale applied to the y-axis for plotting. Default: `identity`." yscale = identity diff --git a/src/config/TrainingPaths.jl b/src/config/TrainingPaths.jl index 37b5052f..63474d70 100644 --- a/src/config/TrainingPaths.jl +++ b/src/config/TrainingPaths.jl @@ -16,4 +16,10 @@ struct TrainingPaths "Training animation saved as an `.mp4` file." history_video::String + + "Base directory for training outputs." + base_dir::String + + "Suffix for output files." + suffix::String end diff --git a/src/data/split_data.jl b/src/data/split_data.jl index b561f74b..4c58defb 100644 --- a/src/data/split_data.jl +++ b/src/data/split_data.jl @@ -168,9 +168,7 @@ function collect_end_dim(x_all::Union{KeyedArray{Float32, 3}, AbstractDimArray{F return collect(getindex(x_all, :, :, idx)) end -# 2D (feature, time): split along time; 3D (feature, time, batch): split along batch -_num_samples(x::AbstractArray{<:Any, 3}) = size(x, 3) -_num_samples(x::AbstractArray) = size(x, 2) +_num_samples(x::AbstractArray) = size(x, ndims(x)) _num_samples(x::NamedTuple) = _num_samples(first(values(x))) function _split_and_pack(x_all, forcings_all, y_all, train_idx, val_idx) diff --git a/src/io/paths.jl b/src/io/paths.jl index f4358a3b..0eacf016 100644 --- a/src/io/paths.jl +++ b/src/io/paths.jl @@ -10,5 +10,7 @@ function resolve_paths(cfg::TrainConfig) joinpath(folder, "config_settings$(suffix).yaml"), joinpath(folder, "train_history$(suffix).png"), joinpath(folder, "training_history$(suffix).mp4"), + folder, + suffix, ) end diff --git a/src/training/dashboard.jl b/src/training/dashboard.jl index 1d50271b..1d716377 100644 --- a/src/training/dashboard.jl +++ b/src/training/dashboard.jl @@ -1,85 +1,57 @@ struct TrainDashboard - observables - fixed_observations - eval_metric - agg - target_names - monitor_names + figures::Dict{Symbol, Any} + axes::Dict{Symbol, Any} + plots::Dict{Symbol, Any} end -function init_dashboard(ext, init::EpochSnapshot, cfg::TrainConfig, y_train, y_val, target_names) +function init_dashboard(ext, history::TrainingHistory, cfg::TrainConfig, y_train, y_val, target_names) isnothing(ext) && return nothing - observables, fixed_observations = initialize_plotting_observables( - init.ŷ_train, - init.ŷ_val, - y_train, - y_val, - init.l_train, - init.l_val, - cfg.loss_types[1], - cfg.agg, - target_names; - monitor_names = cfg.monitor_names # ← was missing - ) - - zoom_epochs = min(cfg.patience, 50) - EasyHybrid.train_board( - observables..., - fixed_observations..., - cfg.yscale, - target_names, - string(cfg.loss_types[1]); - monitor_names = cfg.monitor_names, - zoom_epochs - ) - - return TrainDashboard( - observables, - fixed_observations, - cfg.loss_types[1], - cfg.agg, - target_names, - cfg.monitor_names - ) + figures, axes, plots = build_dashboards(history, cfg, y_train, y_val) + return TrainDashboard(figures, axes, plots) end -function update_dashboard!(dashboard, ext, snapshot::EpochSnapshot, epoch::Int, io, cfg::TrainConfig) +function update_dashboard!(dashboard, ext, history::TrainingHistory, streams, cfg::TrainConfig) isnothing(ext) && !cfg.save_training && return isnothing(dashboard) && return - update_plotting_observables( - dashboard.observables..., - snapshot.l_train, - snapshot.l_val, - dashboard.eval_metric, - dashboard.agg, - snapshot.ŷ_train, - snapshot.ŷ_val, - dashboard.target_names, - epoch; - monitor_names = dashboard.monitor_names - ) + update_step_dashboards!(dashboard, history, cfg) - if io !== nothing - recordframe!(io) + if streams !== nothing + for stream in values(streams) + recordframe!(stream) + end end return nothing end function save_dashboard_img!(dashboard, ext, paths::TrainingPaths, cfg::TrainConfig, best_epoch::Int) return if !isnothing(ext) && cfg.save_training - save_fig(paths.history_img, dashboard_figure()) - @info "Dashboard saved to $(paths.history_img)" + for (name, fig) in pairs(dashboard.figures) + path = name == :dashboard ? paths.history_img : joinpath(paths.base_dir, "$(name)_history$(paths.suffix).png") + save_fig(path, fig) + @info "Dashboard ($name) saved to $(path)" + end else nothing end end -function record_or_run(f, ext, paths::TrainingPaths, cfg::TrainConfig) - return if !isnothing(ext) && cfg.save_training - record_history(dashboard_figure(), paths.history_video; framerate = 24) do io - f(io) +function record_or_run(f, ext, dashboard, paths::TrainingPaths, cfg::TrainConfig) + return if !isnothing(ext) && !isnothing(dashboard) && cfg.save_training + streams = Dict{Symbol, Any}() + for (name, fig) in pairs(dashboard.figures) + if :all in cfg.save_animations || name in cfg.save_animations + streams[name] = VideoStream(fig; framerate = 24) + end + end + + f(streams) + + for (name, stream) in pairs(streams) + path = name == :dashboard ? paths.history_video : joinpath(paths.base_dir, "$(name)_history$(paths.suffix).mp4") + save_video(path, stream) + @info "Animation ($name) saved to $(path)" end else f(nothing) diff --git a/src/training/early_stopping.jl b/src/training/early_stopping.jl index aae960c3..d28250bf 100644 --- a/src/training/early_stopping.jl +++ b/src/training/early_stopping.jl @@ -13,9 +13,9 @@ function EarlyStopping(init_loss, ps, st, cfg) return EarlyStopping(best_loss, deepcopy(cfg.cdev(ps)), deepcopy(cfg.cdev(st)), 0, 0, cfg.patience, false) end -function update!(es::EarlyStopping, history::TrainingHistory, snapshot::EpochSnapshot, ps, st, epoch, cfg::TrainConfig) +function update!(es::EarlyStopping, history::TrainingHistory, snapshot::EpochSnapshot, ps, st, cfg::TrainConfig) current_loss = extract_agg_loss(snapshot.l_val) - new_snapshot = EpochSnapshot(snapshot.l_train, snapshot.l_val, deepcopy(snapshot.ŷ_train), deepcopy(snapshot.ŷ_val)) + new_snapshot = EpochSnapshot(snapshot.l_train, snapshot.l_val, deepcopy(snapshot.ŷ_train), deepcopy(snapshot.ŷ_val), snapshot.epoch) if cfg.keep_history push!(history.snapshots, new_snapshot) @@ -25,7 +25,7 @@ function update!(es::EarlyStopping, history::TrainingHistory, snapshot::EpochSna es.best_loss = current_loss es.best_ps = deepcopy(cfg.cdev(ps)) es.best_st = deepcopy(cfg.cdev(st)) - es.best_epoch = epoch + es.best_epoch = snapshot.epoch es.counter = 0 if !cfg.keep_history history.snapshots[1] = new_snapshot @@ -36,7 +36,7 @@ function update!(es::EarlyStopping, history::TrainingHistory, snapshot::EpochSna return if es.counter >= es.patience metric_name = first(keys(snapshot.l_val)) - @warn "Early stopping at epoch $epoch, best validation loss wrt $metric_name: $(es.best_loss) at epoch $(es.best_epoch)" + @warn "Early stopping at epoch $(snapshot.epoch), best validation loss wrt $metric_name: $(es.best_loss) at epoch $(es.best_epoch)" es.done = true end end diff --git a/src/training/epoch.jl b/src/training/epoch.jl index 7a2686f1..aaf4b43d 100644 --- a/src/training/epoch.jl +++ b/src/training/epoch.jl @@ -50,7 +50,7 @@ function build_loss_fn(model, cfg::TrainConfig) ) end -function evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, init::EpochSnapshot, cfg::TrainConfig) +function evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, epoch::Int, init::EpochSnapshot, cfg::TrainConfig) ps_cpu = ps |> cfg.cdev st_cpu = st |> cfg.cdev l_train, _, ŷ_train = evaluate_acc( @@ -62,5 +62,5 @@ function evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_v ps_cpu, st_cpu, cfg.loss_types, cfg.training_loss, cfg.extra_loss, cfg.agg ) - return EpochSnapshot(l_train, l_val, ŷ_train, ŷ_val) + return EpochSnapshot(l_train, l_val, ŷ_train, ŷ_val, epoch) end diff --git a/src/training/history.jl b/src/training/history.jl index 3e09a1af..57a40278 100644 --- a/src/training/history.jl +++ b/src/training/history.jl @@ -7,3 +7,65 @@ TrainingHistory(init::EpochSnapshot) = TrainingHistory([init]) train_losses(history::TrainingHistory) = [s.l_train for s in history.snapshots] val_losses(history::TrainingHistory) = [s.l_val for s in history.snapshots] predictions(history::TrainingHistory) = [s.ŷ_train for s in history.snapshots] +get_epochs(history::TrainingHistory) = [s.epoch for s in history.snapshots] + +function get_loss_value_t(history::TrainingHistory, loss_type::Symbol, agg::Symbol) + return [get_loss_value(s.l_train, loss_type, agg) for s in history.snapshots] +end + +function get_loss_value_v(history::TrainingHistory, loss_type::Symbol, agg::Symbol) + return [get_loss_value(s.l_val, loss_type, agg) for s in history.snapshots] +end + +function get_prediction_values(history::TrainingHistory, target_name::Symbol, dataset_split::Symbol = :train) + ŷ_now = last(history.snapshots) + if dataset_split == :train + return getfield(ŷ_now.ŷ_train, target_name) + elseif dataset_split == :validation + return getfield(ŷ_now.ŷ_val, target_name) + else + throw(ArgumentError("Invalid dataset_split specified. Use :train or :validation.")) + end +end +function get_monitor_values(history::TrainingHistory, monitor_names::Vector{Symbol}, dataset_split::Symbol = :train) + if dataset_split == :train + return [get_monitor_values(s.ŷ_train, monitor_names) for s in history.snapshots] + elseif dataset_split == :validation + return [get_monitor_values(s.ŷ_val, monitor_names) for s in history.snapshots] + else + throw(ArgumentError("Invalid dataset_split specified. Use :train or :validation.")) + end +end + +function collect_monitor_history(history_vec::Vector, monitor_names::Vector{Symbol}) + return (; + (m => _collect_monitor_field(history_vec, m) for m in monitor_names)..., + ) +end + +function _collect_monitor_field(history_vec::Vector, name::Symbol) + entries = [getfield(snap, name) for snap in history_vec] + first_entry = first(entries) + + if haskey(first_entry, :scalar) + # scalar case: collect into a single vector + return (; :scalar => [e.scalar for e in entries]) + + elseif haskey(first_entry, :quantile) + # quantile case: collect each quantile level into its own vector + qlabels = keys(first_entry.quantile) + return (; + :quantile => (; + (q => [e.quantile[q] for e in entries] for q in qlabels)..., + ), + ) + else + error("Unknown monitor entry format for field $name") + end +end + +export get_loss_value_t, get_loss_value_v +export collect_monitor_history +export get_epochs +export get_monitor_values +export get_prediction_values diff --git a/src/training/initialization.jl b/src/training/initialization.jl index d2feb6b4..1b5dfc78 100644 --- a/src/training/initialization.jl +++ b/src/training/initialization.jl @@ -10,6 +10,10 @@ function load_makie_extension(cfg::TrainConfig) @info "Plotting disabled." return nothing end + if !cfg.keep_history + @warn "Plotting enabled but keep_history is false. Plots will not be generated." + return nothing + end return ext end @@ -55,6 +59,7 @@ struct EpochSnapshot l_val ŷ_train ŷ_val + epoch end function compute_initial_state(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, cfg::TrainConfig) @@ -69,5 +74,5 @@ function compute_initial_state(model, x_train, forcings_train, y_train, mask_tra @debug "Initial train loss: $(l_train) | val loss: $(l_val)" - return EpochSnapshot(l_train, l_val, ŷ_train, ŷ_val) + return EpochSnapshot(l_train, l_val, ŷ_train, ŷ_val, 0.9) end diff --git a/src/training/train.jl b/src/training/train.jl index a7f53989..3ede5dff 100644 --- a/src/training/train.jl +++ b/src/training/train.jl @@ -108,29 +108,35 @@ function _train(model, data, train_cfg::TrainConfig, data_cfg::DataConfig) stopper = EarlyStopping(init.l_val, ps, st, train_cfg) paths = resolve_paths(train_cfg) prog = build_progress(train_cfg) - dashboard = init_dashboard(ext, init, train_cfg, y_train, y_val, model.targets) + + # @show train_cfg.agg + # @show train_cfg.training_loss + # @show get_loss_value_t(history, train_cfg.training_loss, Symbol("$(train_cfg.agg)")) + # @show get_loss_value_v(history, train_cfg.val_loss, Symbol("$(train_cfg.agg)")) + + dashboard = init_dashboard(ext, history, train_cfg, y_train, y_val, model.targets) save_initial_state!(paths, model, ps, st, train_cfg) ps = ps |> train_cfg.gdev st = st |> train_cfg.gdev train_state = train_state |> train_cfg.gdev - record_or_run(ext, paths, train_cfg) do io + record_or_run(ext, dashboard, paths, train_cfg) do streams for epoch in 1:train_cfg.nepochs ps, st, train_state = run_epoch!(loader, model, ps, st, train_state, train_cfg) - snapshot = evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, init, train_cfg) + snapshot = evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, epoch, init, train_cfg) - update!(stopper, history, snapshot, ps, st, epoch, train_cfg) - save_epoch!(paths, model, ps, st, snapshot, epoch, train_cfg) - update_dashboard!(dashboard, ext, snapshot, epoch, io, train_cfg) - log_progress!(prog, init, snapshot, epoch, train_cfg) + update!(stopper, history, snapshot, ps, st, train_cfg) + # save_epoch!(paths, model, ps, st, snapshot, train_cfg) + update_dashboard!(dashboard, ext, history, streams, train_cfg) + # log_progress!(prog, init, snapshot, train_cfg) is_done(stopper) && break end end - save_dashboard_img!(dashboard, ext, paths, train_cfg, stopper.best_epoch) + # save_dashboard_img!(dashboard, ext, paths, train_cfg, stopper.best_epoch) ps, st = best_or_final(stopper, ps, st, train_cfg) - save_final!(paths, model, ps, st, x_train, forcings_train, y_train, x_val, forcings_val, y_val, stopper, train_cfg) + # save_final!(paths, model, ps, st, x_train, forcings_train, y_train, x_val, forcings_val, y_val, stopper, train_cfg) return build_results(model, history, stopper, ps, st, x_train, forcings_train, y_train, x_val, forcings_val, y_val, train_cfg) end @@ -232,7 +238,9 @@ function valid_mask(y) end function train(model, data, save_ps; kwargs...) - train_cfg, data_cfg, solve_kwargs = kwargs_to_configs(save_ps, kwargs) + target_names = model.targets + merge_kwargs = (; kwargs..., target_names) + train_cfg, data_cfg, solve_kwargs = kwargs_to_configs(save_ps, merge_kwargs) return _train(model, data, train_cfg, data_cfg, solve_kwargs) end diff --git a/src/training/train_optimization.jl b/src/training/train_optimization.jl index 5dc9ffae..a2358478 100644 --- a/src/training/train_optimization.jl +++ b/src/training/train_optimization.jl @@ -54,7 +54,7 @@ function _train_optimization(model, data, train_cfg::TrainConfig, data_cfg::Data stopper = EarlyStopping(init.l_val, ps_ca, st, train_cfg) paths = resolve_paths(train_cfg) prog = build_progress(train_cfg) - dashboard = init_dashboard(ext, init, train_cfg, y_train, y_val, model.targets) + dashboard = init_dashboard(ext, history, train_cfg, y_train, y_val, model.targets) save_initial_state!(paths, model, ps_ca, st, train_cfg) @@ -62,7 +62,7 @@ function _train_optimization(model, data, train_cfg::TrainConfig, data_cfg::Data opt_func = OptimizationFunction(loss_fn, train_cfg.autodiff_backend) final_ps = Ref{Any}(ps_ca) - record_or_run(ext, paths, train_cfg) do io + record_or_run(ext, dashboard, paths, train_cfg) do streams if train_cfg.full_batch # Convert once (not per objective/gradient evaluation): the full # training set is the fixed `p` for the entire solve. @@ -74,7 +74,7 @@ function _train_optimization(model, data, train_cfg::TrainConfig, data_cfg::Data model, st, init, history, stopper, dashboard, ext, prog, paths, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, - io, train_cfg, final_ps, + streams, train_cfg, final_ps, ) res = solve(opt_prob, train_cfg.opt; callback = cb, solve_kwargs...) final_ps[] = res.u @@ -150,7 +150,7 @@ function _run_minibatch!( dashboard, ext, prog, paths, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, - io, cfg::TrainConfig, solve_kwargs::NamedTuple, + streams, cfg::TrainConfig, solve_kwargs::NamedTuple, ) loader = build_loader(x_train, forcings_train, y_train, mask_train, cfg) inner_kwargs = Base.structdiff(solve_kwargs, (; maxiters = nothing, epochs = nothing)) @@ -183,7 +183,7 @@ function _run_minibatch!( ) update!(stopper, history, snapshot, ps, st, epoch, cfg) save_epoch!(paths, model, ps, st, snapshot, epoch, cfg) - update_dashboard!(dashboard, ext, snapshot, epoch, io, cfg) + update_dashboard!(dashboard, ext, snapshot, epoch, streams, cfg) log_progress!(prog, init, snapshot, epoch, cfg) is_done(stopper) && break @@ -196,7 +196,7 @@ function _optim_callback( model, st, init, history, stopper, dashboard, ext, prog, paths, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, - io, cfg::TrainConfig, final_ps, + streams, cfg::TrainConfig, final_ps, ) return function (state, _loss) iter = state.iter @@ -211,7 +211,7 @@ function _optim_callback( ) update!(stopper, history, snapshot, ps_cur, st, iter, cfg) save_epoch!(paths, model, ps_cur, st, snapshot, iter, cfg) - update_dashboard!(dashboard, ext, snapshot, iter, io, cfg) + update_dashboard!(dashboard, ext, snapshot, state.iter, streams, cfg) log_progress!(prog, init, snapshot, iter, cfg) end diff --git a/src/training/tune.jl b/src/training/tune.jl index 4d819700..45617066 100644 --- a/src/training/tune.jl +++ b/src/training/tune.jl @@ -26,20 +26,19 @@ Returns a [`TrainResults`](@ref) (or `nothing` if data preparation fails, as in """ function tune(hybrid_model, data, mspec::ModelSpec; kwargs...) kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, (; kwargs...), mspec.hyper_model) - kwargs_train = merge((; kwargs...), mspec.hyper_train) - hm = constructHybridModel(; kwargs_model...) + kwargs_train = merge((; kwargs...), mspec.hyper_train, (; target_names = hm.targets)) train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), kwargs_train) return train(hm, data; train_cfg, data_cfg) end function tune(hybrid_model, data; kwargs...) kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, (; kwargs...)) - hm = constructHybridModel(; kwargs_model...) - train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), (; kwargs...)) + kwargs_train = merge((; kwargs...), (; target_names = hm.targets)) + train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), kwargs_train) return train(hm, data; train_cfg, data_cfg) end @@ -47,9 +46,10 @@ function tune(hybrid_model, data, train_cfg::TrainConfig; data_cfg::DataConfig = kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...)) hm = constructHybridModel(; kwargs_model...) - train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), merge(to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...))) + kwargs_train = merge(to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...), (; target_names = hm.targets)) + train_cfg_new, data_cfg_new = EasyHybrid.kwargs_to_configs((), kwargs_train) - return train(hm, data; train_cfg, data_cfg) + return train(hm, data; train_cfg = train_cfg_new, data_cfg = data_cfg_new) end function best_hyperparams(ho::Hyperoptimizer) diff --git a/src/utils/plotrecipes.jl b/src/utils/plotrecipes.jl index af3b8cfd..4a234a82 100644 --- a/src/utils/plotrecipes.jl +++ b/src/utils/plotrecipes.jl @@ -1,3 +1,10 @@ +export lossplot, lossplot! +export monitorplot, monitorplot! +export predictionplot, predictionplot! +export timeseriesplot, timeseriesplot! +export train_dashboard, update_step_dashboard! +export build_dashboards, update_step_dashboards! + function poplot() return @error("Please load `Makie.jl` and then call this function. If Makie is loaded, then you can't call `poplot` with no arguments!") end @@ -42,6 +49,12 @@ function record_history end function dashboard_figure end function recordframe! end function save_fig end +function VideoStream end +function save_video end +function train_dashboard end +function update_step_dashboard! end +function build_dashboards end +function update_step_dashboards! end """ initialize_plotting_observables(init_ŷ_train, init_ŷ_val, y_train, y_val, l_init_train, l_init_val, training_loss, agg, monitor_names, target_names) @@ -132,3 +145,28 @@ function monitor_to_obs(ŷ, monitor_names; cuts = (0.25, 0.5, 0.75)) )..., ) end + +function get_monitor_values(ŷ, monitor_names; cuts = (0.25, 0.5, 0.75)) + labels = map(q -> Symbol("q$(Int(q * 100))"), cuts) + return (; (m => _monitor_entry(getfield(ŷ, m), cuts, labels) for m in monitor_names)...) +end + +function _monitor_entry(v, cuts, labels) + v = vec(v) + if length(v) > 1 + return (; :quantile => (; zip(labels, quantile(v, collect(cuts)))...)) + else + return (; :scalar => only(v)) + end +end + +# for recipes +function lossplot end +function lossplot! end +function monitorplot end +function monitorplot! end +function predictionplot end +function predictionplot! end +function timeseriesplot end +function timeseriesplot! end +#