From 79a7c3e6180c976e128b1a52f057b423aa921275 Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Thu, 2 Jul 2026 21:45:14 +0200 Subject: [PATCH 01/12] rename and dispatches --- README.md | 2 +- .../research/synthetic_respiration.jl | 4 +- .../tutorials/example_synthetic_lstm.jl | 4 +- docs/literate/tutorials/folds.jl | 2 +- .../tutorials/synthetic_respiration_gpu.jl | 4 +- docs/src/get_started.md | 2 +- docs/src/tutorials/exponential_res.md | 2 +- docs/src/tutorials/hyperparameter_tuning.md | 2 +- projects/ExpoHybrid/ExpoHybridEstim.jl | 2 +- src/data/prepare_data.jl | 17 +- src/data/sequences.jl | 8 +- src/models/GenericHybridModel.jl | 289 +++++------------- src/models/show_generic.jl | 86 +++--- src/training/tune.jl | 6 +- test/test_autodiff_backend.jl | 2 +- test/test_compute_loss.jl | 4 +- test/test_generic_hybrid_model.jl | 80 ++--- test/test_show_generic_hybrid.jl | 10 +- test/test_split_data_train.jl | 8 +- 19 files changed, 196 insertions(+), 338 deletions(-) diff --git a/README.md b/README.md index 046001e8..f000e0e3 100644 --- a/README.md +++ b/README.md @@ -174,7 +174,7 @@ global_param_names = [:Q10] # Global parameters (same for all samples) neural_param_names = [:rb] # Neural network predicted parameters # Construct hybrid model -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/literate/research/synthetic_respiration.jl b/docs/literate/research/synthetic_respiration.jl index 1f4d6dbe..5a84fa97 100644 --- a/docs/literate/research/synthetic_respiration.jl +++ b/docs/literate/research/synthetic_respiration.jl @@ -73,7 +73,7 @@ neural_param_names = [:rb] # Neural network predicted parameters predictors_single_nn = [:sw_pot, :dsw_pot] # Predictor variables (solar radiation, and its derivative) -single_nn_hybrid_model = constructHybridModel( +single_nn_hybrid_model = HybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables @@ -150,7 +150,7 @@ mean(df.sw_pot) predictors_multi_nn = (rb = [:sw_pot, :dsw_pot],) -multi_nn_hybrid_model = constructHybridModel( +multi_nn_hybrid_model = HybridModel( predictors_multi_nn, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/literate/tutorials/example_synthetic_lstm.jl b/docs/literate/tutorials/example_synthetic_lstm.jl index 8bccb93c..222528c7 100644 --- a/docs/literate/tutorials/example_synthetic_lstm.jl +++ b/docs/literate/tutorials/example_synthetic_lstm.jl @@ -87,7 +87,7 @@ neural_param_names = [:rb] # ## 7. Construct LSTM Hybrid Model # Create LSTM hybrid model using the unified constructor -hlstm = constructHybridModel( +hlstm = HybridModel( predictors, forcing, target, @@ -216,7 +216,7 @@ first(out_lstm.val_obs_pred, 5) # ## 10. Train Single NN Hybrid Model (Optional) # For comparison, we can also train a hybrid model with a standard feed-forward neural network -hm = constructHybridModel( +hm = HybridModel( predictors, forcing, target, diff --git a/docs/literate/tutorials/folds.jl b/docs/literate/tutorials/folds.jl index aefc3b32..52b0c39d 100644 --- a/docs/literate/tutorials/folds.jl +++ b/docs/literate/tutorials/folds.jl @@ -61,7 +61,7 @@ neural_param_names = [:rb] # ## 6. Construct the Hybrid Model -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, forcing, target, diff --git a/docs/literate/tutorials/synthetic_respiration_gpu.jl b/docs/literate/tutorials/synthetic_respiration_gpu.jl index 88479d78..4bee0062 100644 --- a/docs/literate/tutorials/synthetic_respiration_gpu.jl +++ b/docs/literate/tutorials/synthetic_respiration_gpu.jl @@ -62,7 +62,7 @@ neural_param_names = [:rb] # Neural network predicted parameters # ## Single NN Hybrid Model Training predictors_single_nn = [:sw_pot, :dsw_pot] # Predictor variables (solar radiation, and its derivative) -small_nn_hybrid_model = constructHybridModel( +small_nn_hybrid_model = HybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables @@ -76,7 +76,7 @@ small_nn_hybrid_model = constructHybridModel( input_batchnorm = true # Apply batch normalization to inputs ) -large_nn_hybrid_model = constructHybridModel( +large_nn_hybrid_model = HybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/src/get_started.md b/docs/src/get_started.md index 0609da04..27ffebed 100644 --- a/docs/src/get_started.md +++ b/docs/src/get_started.md @@ -95,7 +95,7 @@ neural_param_names = [:rb] # Neural network predicted parameters Construct hybrid model ```@example quick_start_complete -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/src/tutorials/exponential_res.md b/docs/src/tutorials/exponential_res.md index 637515b2..08dbc2a9 100644 --- a/docs/src/tutorials/exponential_res.md +++ b/docs/src/tutorials/exponential_res.md @@ -117,7 +117,7 @@ global_param_names = [:k] ``` ```@example expo -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, forcings, targets, diff --git a/docs/src/tutorials/hyperparameter_tuning.md b/docs/src/tutorials/hyperparameter_tuning.md index 0bfc2505..b777104f 100644 --- a/docs/src/tutorials/hyperparameter_tuning.md +++ b/docs/src/tutorials/hyperparameter_tuning.md @@ -75,7 +75,7 @@ neural_param_names = [:rb] # Neural network predicted parameters Construct hybrid model ```@example hyperparameter_tuning -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables diff --git a/projects/ExpoHybrid/ExpoHybridEstim.jl b/projects/ExpoHybrid/ExpoHybridEstim.jl index f7ed6e49..b8bc7931 100644 --- a/projects/ExpoHybrid/ExpoHybridEstim.jl +++ b/projects/ExpoHybrid/ExpoHybridEstim.jl @@ -84,7 +84,7 @@ function Expo_resp_model(; T, Resp0, k) return (; Resp_obs, Resp0, k) end -hybrid_model = constructHybridModel( +hybrid_model = HybridModel( predictors, forcings, targets, diff --git a/src/data/prepare_data.jl b/src/data/prepare_data.jl index 00fe8e79..90067101 100644 --- a/src/data/prepare_data.jl +++ b/src/data/prepare_data.jl @@ -3,19 +3,14 @@ export prepare_data function prepare_data(hm, data::KeyedArray; cfg = DataConfig(), kwargs...) predictors, forcings, targets = get_prediction_target_names(hm) # KeyedArray: use () syntax for views that are differentiable - X_arr = Array(data(predictors)) - forcings_nt = NamedTuple([forcing => Array(data(forcing)) for forcing in forcings]) - targets_nt = NamedTuple([target => Array(data(target)) for target in targets]) - return ((X_arr, forcings_nt), targets_nt) -end - -function prepare_data(hm::MultiNNHybridModel, data::KeyedArray; cfg = DataConfig(), kwargs...) - predictors, forcings, targets = get_prediction_target_names(hm) - # KeyedArray: use () syntax for views that are differentiable - X_all = NamedTuple([name => Array(data(p)) for (name, p) in pairs(predictors)]) + if predictors isa NamedTuple + X = NamedTuple([name => Array(data(p)) for (name, p) in pairs(predictors)]) + else + X = Array(data(predictors)) + end forcings_nt = NamedTuple([forcing => Array(data(forcing)) for forcing in forcings]) targets_nt = NamedTuple([target => Array(data(target)) for target in targets]) - return ((X_all, forcings_nt), targets_nt) + return ((X, forcings_nt), targets_nt) end function prepare_data(hm, data::AbstractDimArray; kwargs...) diff --git a/src/data/sequences.jl b/src/data/sequences.jl index ba51c7aa..f44c3485 100644 --- a/src/data/sequences.jl +++ b/src/data/sequences.jl @@ -34,7 +34,7 @@ end _subset_y(y, valid) = y[:, :, valid] _subset_y(y::NamedTuple, valid) = map(v -> v[:, valid], y) -# NamedTuple predictors (MultiNNHybridModel): every branch is a 3D +# NamedTuple predictors (HybridModel): every branch is a 3D # `(feature, time, batch)` array sharing the batch axis with forcings/targets. # Drop a sequence if any branch has a NaN predictor or all targets are NaN. function filter_sequences(x_tuple::Tuple{<:NamedTuple, <:Any}, y) @@ -99,9 +99,9 @@ function split_into_sequences(x::Tuple, y; kwargs...) return (X_seq, forcings_seq), Y_seq end -# MultiNNHybridModel support. +# HybridModel support. # -# For `MultiNNHybridModel`, `prepare_data` returns the predictors as a +# For `HybridModel` with NamedTuple predictors, `prepare_data` returns the predictors as a # NamedTuple of per-branch `(feature, time)` arrays (one per neural network), # while forcings and targets stay shared across branches. Window every branch on # the common time axis and keep the NamedTuple structure so the model's forward @@ -123,7 +123,7 @@ function split_into_sequences(x::Tuple{<:NamedTuple, <:Any}, y::NamedTuple; kwar return (X_seq, forcings_seq), _array_to_nt(Y_seq, targetkeys) end -# Predictor branches built for `MultiNNHybridModel` are plain matrices (their +# Predictor branches built for `HybridModel` with NamedTuple predictors are plain matrices (their # axis keys are stripped in `prepare_data`). Wrap them with synthetic keys so # they flow through the keyed windowing path and produce keyed 3D outputs that # `collect_end_dim`/`view_end_dim` understand. Keyed/Dim inputs pass through. diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 5bf56e0c..f8cc94df 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -1,4 +1,4 @@ -export SingleNNHybridModel, MultiNNHybridModel, constructHybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid +export HybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid export HybridParams # Import necessary components for neural networks @@ -40,39 +40,10 @@ struct HybridParams{M <: Function} <: AbstractHybridModel end # ─────────────────────────────────────────────────────────────────────────── -# Single NN Hybrid Model Structure (optimized for performance) -struct SingleNNHybridModel <: LuxCore.AbstractLuxContainerLayer{ - ( - :NN, #:predictors, :forcing, :targets, - #:mechanistic_model, :parameters, :neural_param_names, :global_param_names, :fixed_param_names, - #:scale_nn_outputs, :start_from_default, - ), - } - NN::Chain - predictors::Vector{Symbol} - forcing::Vector{Symbol} - targets::Vector{Symbol} - mechanistic_model::Function - parameters::AbstractHybridModel - neural_param_names::Vector{Symbol} - global_param_names::Vector{Symbol} - fixed_param_names::Vector{Symbol} - scale_nn_outputs::Bool - start_from_default::Bool - config::NamedTuple -end - -# Multi-NN Hybrid Model Structure (optimized for performance) -struct MultiNNHybridModel <: LuxCore.AbstractLuxContainerLayer{ - ( - :NNs, #:predictors, :forcing, :targets, - # :mechanistic_model, :parameters, :neural_param_names, :global_param_names, :fixed_param_names, - # :scale_nn_outputs, :start_from_default, - ), - } - - NNs::NamedTuple - predictors::NamedTuple +# Unified Hybrid Model Structure (optimized for performance) +struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} + NNs::T + predictors::P forcing::Vector{Symbol} targets::Vector{Symbol} mechanistic_model::Function @@ -86,7 +57,7 @@ struct MultiNNHybridModel <: LuxCore.AbstractLuxContainerLayer{ end # Unified constructor that dispatches based on predictors type -function constructHybridModel( +function HybridModel( predictors::Vector{Symbol}, forcing, targets, @@ -136,10 +107,10 @@ function constructHybridModel( kwargs..., ) - return SingleNNHybridModel(NN, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) + return HybridModel(NN, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) end -function constructHybridModel( +function HybridModel( predictors::NamedTuple, forcing, targets, @@ -202,10 +173,10 @@ function constructHybridModel( kwargs..., ) - return MultiNNHybridModel(NNs, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) + return HybridModel(NNs, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) end -function constructHybridModel( +function HybridModel( ; predictors, forcing, targets, @@ -217,12 +188,12 @@ function constructHybridModel( ) if predictors isa Vector{Symbol} @assert neural_param_names !== nothing "Provide neural_param_names for Vector predictors" - return constructHybridModel( + return HybridModel( predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs... ) elseif predictors isa NamedTuple - return constructHybridModel( + return HybridModel( predictors, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs... ) @@ -232,40 +203,24 @@ function constructHybridModel( end # ─────────────────────────────────────────────────────────────────────────── -# Initial parameters for SingleNNHybridModel -function LuxCore.initialparameters(rng::AbstractRNG, m::SingleNNHybridModel) - ps_nn, _ = LuxCore.setup(rng, m.NN) - nt = (; ps = ps_nn) - - # Then append each global parameter as a 1-vector of Float32 - if !isempty(m.global_param_names) - if m.start_from_default - for g in m.global_param_names - default_val = scale_single_param_minmax(g, m.parameters) - nt = merge(nt, NamedTuple{(g,), Tuple{Vector{Float32}}}(([Float32(default_val)],))) - end - else - for g in m.global_param_names - random_val = rand(rng, Float32) - nt = merge(nt, NamedTuple{(g,), Tuple{Vector{Float32}}}(([random_val],))) - end - end - end - - return nt -end - -# Initial parameters for MultiNNHybridModel -function LuxCore.initialparameters(rng::AbstractRNG, m::MultiNNHybridModel) - # Setup parameters for each neural network +# Helper functions for NN initialization +function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) nn_params = NamedTuple() for (nn_name, nn) in pairs(m.NNs) ps_nn, _ = LuxCore.setup(rng, nn) nn_params = merge(nn_params, NamedTuple{(nn_name,), Tuple{typeof(ps_nn)}}((ps_nn,))) end + return (; nn_params...) +end + +function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) + ps_nn, _ = LuxCore.setup(rng, m.NNs) + return (; ps = ps_nn) +end - # Start with the NN weights - nt = (; nn_params...) +# Initial parameters for HybridModel +function LuxCore.initialparameters(rng::AbstractRNG, m::HybridModel) + nt = _init_nn_params(rng, m) # Then append each global parameter as a 1-vector of Float32 if !isempty(m.global_param_names) @@ -285,33 +240,23 @@ function LuxCore.initialparameters(rng::AbstractRNG, m::MultiNNHybridModel) return nt end -# Initial states for SingleNNHybridModel -function LuxCore.initialstates(rng::AbstractRNG, m::SingleNNHybridModel) - _, st_nn = LuxCore.setup(rng, m.NN) - nt = (;) - - # Then append each fixed parameter as a 1-vector of Float32 - if !isempty(m.fixed_param_names) - for f in m.fixed_param_names - default_val = default(m.parameters)[f] - nt = merge(nt, NamedTuple{(f,), Tuple{Vector{Float32}}}(([Float32(default_val)],))) - end - end - - nt = (; st_nn = st_nn, fixed = nt) - return nt -end - -# Initial states for MultiNNHybridModel -function LuxCore.initialstates(rng::AbstractRNG, m::MultiNNHybridModel) - # Setup states for each neural network +function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) nn_states = NamedTuple() for (nn_name, nn) in pairs(m.NNs) _, st_nn = LuxCore.setup(rng, nn) nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) end + return nn_states +end - # Start with the NN states +function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) + _, st_nn = LuxCore.setup(rng, m.NNs) + return (; st_nn = st_nn) +end + +# Initial states for HybridModel +function LuxCore.initialstates(rng::AbstractRNG, m::HybridModel) + nn_states_nt = _init_nn_states(rng, m) nt = (;) # Then append each fixed parameter as a 1-vector of Float32 @@ -322,8 +267,7 @@ function LuxCore.initialstates(rng::AbstractRNG, m::MultiNNHybridModel) end end - nt = (; nn_states..., fixed = nt) - return nt + return merge(nn_states_nt, (; fixed = nt)) end function default(p::AbstractHybridModel) @@ -366,34 +310,44 @@ end # ─────────────────────────────────────────────────────────────────────────── -# Forward pass for SingleNNHybridModel (optimized, no branching) -function (m::SingleNNHybridModel)(ds_k, ps, st) - # 1) get features - predictors = ds_k[1] #toArray(ds_k, m.predictors) +# ─────────────────────────────────────────────────────────────────────────── +function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) + nn_outputs = NamedTuple() + nn_states = NamedTuple() - parameters = m.parameters + for (nn_name, nn) in pairs(m.NNs) + nn_out, st_nn = LuxCore.apply(nn, ds_k[1][nn_name], ps[nn_name], st[nn_name]) + nn_outputs = merge(nn_outputs, NamedTuple{(nn_name,), Tuple{typeof(nn_out)}}((nn_out,))) + nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) + end - # 2) scale global parameters (handle empty case) - if !isempty(m.global_param_names) - global_vals = Tuple( - scale_single_param(g, ps[g], parameters) - for g in m.global_param_names - ) - global_params = NamedTuple{Tuple(m.global_param_names), Tuple{typeof.(global_vals)...}}(global_vals) - else - global_params = NamedTuple() + scaled_nn_params = NamedTuple() + for (nn_name, param_name) in zip(keys(m.NNs), m.neural_param_names) + nn_cols = eachslice(nn_outputs[nn_name]; dims = 1) + nn_param = NamedTuple{(param_name,), Tuple{typeof(nn_cols[1])}}((nn_cols[1],)) + + if m.scale_nn_outputs + scaled_nn_val = scale_single_param(param_name, nn_param[param_name], m.parameters) + else + scaled_nn_val = nn_param[param_name] + end + + nn_scaled_param = NamedTuple{(param_name,), Tuple{typeof(scaled_nn_val)}}((scaled_nn_val,)) + scaled_nn_params = merge(scaled_nn_params, nn_scaled_param) end - # 3) scale NN parameters (handle empty case) + return scaled_nn_params, (; nn_states...), (; nn_outputs = nn_outputs) +end + +function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) if !isempty(m.neural_param_names) - nn_out, st_nn = LuxCore.apply(m.NN, predictors, ps.ps, st.st_nn) + nn_out, st_nn = LuxCore.apply(m.NNs, ds_k[1], ps.ps, st.st_nn) nn_cols = eachslice(nn_out, dims = 1) nn_params = NamedTuple(zip(m.neural_param_names, nn_cols)) - # Use appropriate scaling based on setting if m.scale_nn_outputs scaled_nn_vals = Tuple( - scale_single_param(name, nn_params[name], parameters) + scale_single_param(name, nn_params[name], m.parameters) for name in m.neural_param_names ) else @@ -404,62 +358,14 @@ function (m::SingleNNHybridModel)(ds_k, ps, st) scaled_nn_params = NamedTuple() st_nn = st.st_nn end - - # 4) pick fixed parameters (handle empty case) - if !isempty(m.fixed_param_names) - fixed_vals = Tuple(st.fixed[f] for f in m.fixed_param_names) - fixed_params = NamedTuple{Tuple(m.fixed_param_names), Tuple{typeof.(fixed_vals)...}}(fixed_vals) - else - fixed_params = NamedTuple() - end - - # 5) unpack forcing data - forcing_data = ds_k[2] #toNamedTuple(ds_k, m.forcing) - - # 6) merge all parameters - all_params = merge(scaled_nn_params, global_params, fixed_params) - all_kwargs = merge(forcing_data, all_params) - # all_kwargs = merge(forcing_data, all_params) - - # 7) physics - y_pred = m.mechanistic_model(; all_kwargs...) - - out = (; y_pred..., parameters = all_params) - st_new = (; st_nn = st_nn, fixed = st.fixed) - - return out, st_new + return scaled_nn_params, (; st_nn = st_nn), (;) end -function (m::SingleNNHybridModel)(df::DataFrame, ps, st) - @warn "Only makes sense in test mode, not training!" - - - # Process numeric or missing-containing columns - for col in names(df) - what_type = eltype(df[!, col]) - if what_type <: Union{Missing, Real} || what_type <: Real - df[!, col] = Float32.(coalesce.(df[!, col], NaN)) - end - end - - all_data = to_keyedArray(df) - x, _ = prepare_data(m, all_data) - out, _ = m(x, ps, LuxCore.testmode(st)) - dfnew = copy(df) - for k in keys(out) - if length(out[k]) == size(x, 2) - dfnew[!, String(k) * "_pred"] = out[k] - end - end - return dfnew -end - -# Forward pass for MultiNNHybridModel (optimized, no branching) -function (m::MultiNNHybridModel)(ds_k::Tuple, ps, st) - +# Forward pass for HybridModel (optimized, using multiple dispatch for NN run) +function (m::HybridModel)(ds_k::Tuple, ps, st) parameters = m.parameters - # 2) Scale global parameters (handle empty case) + # 1) Scale global parameters (handle empty case) if !isempty(m.global_param_names) global_vals = Tuple( scale_single_param(g, ps[g], parameters) @@ -470,41 +376,10 @@ function (m::MultiNNHybridModel)(ds_k::Tuple, ps, st) global_params = NamedTuple() end - # 3) Run each neural network and collect outputs - nn_outputs = NamedTuple() - nn_states = NamedTuple() + # 2) Run neural network(s) + scaled_nn_params, st_new_nns, out_extra = _run_nn(m, ds_k, ps, st) - for (nn_name, nn) in pairs(m.NNs) - nn_out, st_nn = LuxCore.apply(nn, ds_k[1][nn_name], ps[nn_name], st[nn_name]) - nn_outputs = merge(nn_outputs, NamedTuple{(nn_name,), Tuple{typeof(nn_out)}}((nn_out,))) - nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) - end - - # 4) Scale neural network parameters using the mapping - scaled_nn_params = NamedTuple() - for (nn_name, param_name) in zip(keys(m.NNs), m.neural_param_names) - # `eachslice(...; dims = 1)` (instead of `eachrow`) so this works for both - # the feed-forward case (2D `(param, batch)` output) and the recurrent/LSTM - # case (3D `(param, time, batch)` sequence output). - nn_cols = eachslice(nn_outputs[nn_name]; dims = 1) - - # Create parameter for this NN - nn_param = NamedTuple{(param_name,), Tuple{typeof(nn_cols[1])}}((nn_cols[1],)) - - # Conditionally apply scaling based on scale_nn_outputs setting - if m.scale_nn_outputs - scaled_nn_val = scale_single_param(param_name, nn_param[param_name], parameters) - else - scaled_nn_val = nn_param[param_name] # Use raw NN output without scaling - end - - nn_scaled_param = NamedTuple{(param_name,), Tuple{typeof(scaled_nn_val)}}((scaled_nn_val,)) - - # Merge with existing scaled parameters - scaled_nn_params = merge(scaled_nn_params, nn_scaled_param) - end - - # 5) Pick fixed parameters (handle empty case) + # 3) Pick fixed parameters (handle empty case) if !isempty(m.fixed_param_names) fixed_vals = Tuple(st.fixed[f] for f in m.fixed_param_names) fixed_params = NamedTuple{Tuple(m.fixed_param_names), Tuple{typeof.(fixed_vals)...}}(fixed_vals) @@ -512,24 +387,28 @@ function (m::MultiNNHybridModel)(ds_k::Tuple, ps, st) fixed_params = NamedTuple() end + # 4) merge all parameters all_params = merge(scaled_nn_params, global_params, fixed_params) - # 6) unpack forcing data - + # 5) unpack forcing data forcing_data = ds_k[2] all_kwargs = merge(forcing_data, all_params) - # 7) Apply mechanistic model + # 6) Apply mechanistic model y_pred = m.mechanistic_model(; all_kwargs...) - out = (; y_pred..., parameters = all_params, nn_outputs = nn_outputs) - - st_new = (; nn_states..., fixed = st.fixed) + out = (; y_pred..., parameters = all_params, out_extra...) + st_new = (; st_new_nns..., fixed = st.fixed) return out, st_new end -function (m::MultiNNHybridModel)(df::DataFrame, ps, st) +function (m::HybridModel)(ds_k, ps, st) + # Forward pass fallback when ds_k is not explicitly typed as Tuple + return m(Tuple(ds_k), ps, st) +end + +function (m::HybridModel)(df::DataFrame, ps, st) @warn "Only makes sense in test mode, not training!" # Process numeric or missing-containing columns @@ -541,12 +420,12 @@ function (m::MultiNNHybridModel)(df::DataFrame, ps, st) end all_data = to_keyedArray(df) - x, _ = prepare_data(m, all_data) out, _ = m(x, ps, LuxCore.testmode(st)) dfnew = copy(df) + n_samples = x[1] isa NamedTuple ? size(first(values(x[1])), 2) : size(x[1], 2) for k in keys(out) - if length(out[k]) == size(x, 2) + if length(out[k]) == n_samples dfnew[!, String(k) * "_pred"] = out[k] end end diff --git a/src/models/show_generic.jl b/src/models/show_generic.jl index 49e5308b..fc42be88 100644 --- a/src/models/show_generic.jl +++ b/src/models/show_generic.jl @@ -100,62 +100,48 @@ function Base.show(io::IO, pc::ParameterContainer) return print(io, ")") end -function Base.show(io::IO, ::MIME"text/plain", hm::SingleNNHybridModel) - _print_header(io, "Hybrid Model (Single NN)") - - printstyled(io, "Neural Network: \n", color = :light_black) - show(IndentedIO(io), MIME"text/plain"(), hm.NN) - println(io) - - _print_header(io, "Configuration:", color = :light_blue, bold = false) - _print_field(io, "predictors", hm.predictors) - _print_field(io, "forcing", hm.forcing) - _print_field(io, "targets", hm.targets) - _print_field(io, "mechanistic_model", hm.mechanistic_model, value_color = :light_blue) - _print_field(io, "neural_param_names", hm.neural_param_names, value_color = :light_blue) - _print_field(io, "global_param_names", hm.global_param_names, value_color = :green) - _print_field(io, "fixed_param_names", hm.fixed_param_names, value_color = :yellow) - _print_field(io, "scale_nn_outputs", hm.scale_nn_outputs, value_color = hm.scale_nn_outputs ? :green : :red) - _print_field(io, "start_from_default", hm.start_from_default, value_color = hm.start_from_default ? :green : :red) - _print_field(io, "config", hm.config, value_color = :cyan) +function Base.show(io::IO, ::MIME"text/plain", hm::HybridModel) + if hm.predictors isa NamedTuple + _print_header(io, "Hybrid Model (Multi NN)") + + # Neural networks + _print_header(io, "Neural Networks:", color = :light_black, bold = false) + n_nns = length(hm.NNs) + idx = 0 + for (name, nn) in pairs(hm.NNs) + idx += 1 + printstyled(io, " ", color = :light_black) + printstyled(io, string(name), color = :cyan, bold = true) + println(io, ":") + show(IndentedIO(io; indent = " "), MIME"text/plain"(), nn) + if idx < n_nns + println(io) + end + end + println(io) - println(io) - _print_header(io, "Parameters:", color = :light_blue, bold = false) - io_full = IOContext(IndentedIO(io), :compact => false, :limit => false) - return show(io_full, MIME"text/plain"(), hm.parameters) -end + # Configuration + _print_header(io, "Configuration:", color = :light_blue, bold = false) -function Base.show(io::IO, ::MIME"text/plain", hm::MultiNNHybridModel) - _print_header(io, "Hybrid Model (Multi NN)") - - # Neural networks - _print_header(io, "Neural Networks:", color = :light_black, bold = false) - n_nns = length(hm.NNs) - idx = 0 - for (name, nn) in pairs(hm.NNs) - idx += 1 - printstyled(io, " ", color = :light_black) - printstyled(io, string(name), color = :cyan, bold = true) + # Predictors are per-network (NamedTuple) + printstyled(io, " predictors", color = :light_black) println(io, ":") - show(IndentedIO(io; indent = " "), MIME"text/plain"(), nn) - if idx < n_nns + for (name, preds) in pairs(hm.predictors) + printstyled(io, " ", color = :light_black) + printstyled(io, string(name), color = :yellow) + print(io, " = ") + printstyled(io, preds, color = :cyan) println(io) end - end - println(io) + else + _print_header(io, "Hybrid Model (Single NN)") - # Configuration - _print_header(io, "Configuration:", color = :light_blue, bold = false) - - # Predictors are per-network (NamedTuple) - printstyled(io, " predictors", color = :light_black) - println(io, ":") - for (name, preds) in pairs(hm.predictors) - printstyled(io, " ", color = :light_black) - printstyled(io, string(name), color = :yellow) - print(io, " = ") - printstyled(io, preds, color = :cyan) + printstyled(io, "Neural Network: \n", color = :light_black) + show(IndentedIO(io), MIME"text/plain"(), hm.NNs) println(io) + + _print_header(io, "Configuration:", color = :light_blue, bold = false) + _print_field(io, "predictors", hm.predictors) end _print_field(io, "forcing", hm.forcing) @@ -175,8 +161,6 @@ function Base.show(io::IO, ::MIME"text/plain", hm::MultiNNHybridModel) _print_field(io, "config", hm.config, value_color = :cyan) println(io) - - # Parameters _print_header(io, "Parameters:", color = :light_blue, bold = false) io_full = IOContext(IndentedIO(io), :compact => false, :limit => false) return show(io_full, MIME"text/plain"(), hm.parameters) diff --git a/src/training/tune.jl b/src/training/tune.jl index 45617066..3b51be51 100644 --- a/src/training/tune.jl +++ b/src/training/tune.jl @@ -26,7 +26,7 @@ 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) - hm = constructHybridModel(; kwargs_model...) + hm = HybridModel(; kwargs_model...) kwargs_train = merge((; kwargs...), mspec.hyper_train, (; target_names = hm.targets)) train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), kwargs_train) @@ -35,7 +35,7 @@ end function tune(hybrid_model, data; kwargs...) kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, (; kwargs...)) - hm = constructHybridModel(; kwargs_model...) + hm = HybridModel(; kwargs_model...) kwargs_train = merge((; kwargs...), (; target_names = hm.targets)) train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), kwargs_train) @@ -44,7 +44,7 @@ end function tune(hybrid_model, data, train_cfg::TrainConfig; data_cfg::DataConfig = DataConfig(), kwargs...) kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...)) - hm = constructHybridModel(; kwargs_model...) + hm = HybridModel(; kwargs_model...) 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) diff --git a/test/test_autodiff_backend.jl b/test/test_autodiff_backend.jl index 4eca0bfa..c09a22e1 100644 --- a/test/test_autodiff_backend.jl +++ b/test/test_autodiff_backend.jl @@ -11,7 +11,7 @@ using ForwardDiff global_param_names = [:Q10] neural_param_names = [:rb] - model = constructHybridModel( + model = HybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) diff --git a/test/test_compute_loss.jl b/test/test_compute_loss.jl index 81a7f64f..5972e480 100644 --- a/test/test_compute_loss.jl +++ b/test/test_compute_loss.jl @@ -1,6 +1,6 @@ using EasyHybrid: _compute_loss, PerTarget, _apply_loss, loss_fn using EasyHybrid: _get_target_nan, _get_target_y, _loss_name, compute_loss, LoggingLoss -using EasyHybrid: constructHybridModel, to_keyedArray +using EasyHybrid: HybridModel, to_keyedArray using Statistics using DimensionalData using Random @@ -223,7 +223,7 @@ end neural_param_names = [:a] global_param_names = [:b] - HM = constructHybridModel( + HM = HybridModel( predictors, forcing, targets, diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index 7e19e068..e8e5d6bf 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -126,15 +126,15 @@ end end end -@testset "GenericHybridModel - SingleNNHybridModel" begin - @testset "constructHybridModel with Vector predictors" begin +@testset "GenericHybridModel - HybridModel" begin + @testset "HybridModel with Vector predictors" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -144,7 +144,7 @@ end global_param_names ) - @test model isa SingleNNHybridModel + @test model isa HybridModel @test model.predictors == predictors @test model.forcing == forcing @test model.targets == targets @@ -153,17 +153,17 @@ end @test model.fixed_param_names == [:c, :d] @test model.scale_nn_outputs == false @test model.start_from_default == true - @test model.NN isa Chain + @test model.NNs isa Chain end - @testset "constructHybridModel with empty predictors" begin + @testset "HybridModel with empty predictors" begin predictors = Symbol[] forcing = [:x1] targets = [:obs] neural_param_names = Symbol[] global_param_names = [:a, :b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -173,20 +173,20 @@ end global_param_names ) - @test model isa SingleNNHybridModel + @test model isa HybridModel @test model.predictors == predictors - @test model.NN isa Chain + @test model.NNs isa Chain # @test typeof(model.NN.layers[1]) == Lux.NoOpLayer # Empty chain end - @testset "SingleNNHybridModel initialparameters" begin + @testset "HybridModel initialparameters" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -205,14 +205,14 @@ end @test ps.b[1] isa Float32 end - @testset "SingleNNHybridModel initialstates" begin + @testset "HybridModel initialstates" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -235,14 +235,14 @@ end @test st.fixed.d[1] isa Float32 end - @testset "SingleNNHybridModel forward pass" begin + @testset "HybridModel forward pass" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -269,14 +269,14 @@ end @test haskey(new_st, :fixed) end - @testset "SingleNNHybridModel with scale_nn_outputs=true" begin + @testset "HybridModel with scale_nn_outputs=true" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -301,14 +301,14 @@ end end end -@testset "GenericHybridModel - MultiNNHybridModel" begin - @testset "constructHybridModel with NamedTuple predictors" begin +@testset "GenericHybridModel - HybridModel" begin + @testset "HybridModel with NamedTuple predictors" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -317,7 +317,7 @@ end global_param_names ) - @test model isa MultiNNHybridModel + @test model isa HybridModel @test model.predictors == predictors @test model.forcing == forcing @test model.targets == targets @@ -330,13 +330,13 @@ end @test haskey(model.NNs, :d) end - @testset "MultiNNHybridModel with NamedTuple hidden_layers and activation" begin + @testset "HybridModel with NamedTuple hidden_layers and activation" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -347,18 +347,18 @@ end activation = (a = tanh, d = sigmoid) ) - @test model isa MultiNNHybridModel + @test model isa HybridModel @test haskey(model.NNs, :a) @test haskey(model.NNs, :d) end - @testset "MultiNNHybridModel initialparameters" begin + @testset "HybridModel initialparameters" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -377,13 +377,13 @@ end @test ps.b[1] isa Float32 end - @testset "MultiNNHybridModel initialstates" begin + @testset "HybridModel initialstates" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -403,13 +403,13 @@ end @test st.fixed.c[1] isa Float32 end - @testset "MultiNNHybridModel forward pass" begin + @testset "HybridModel forward pass" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -438,13 +438,13 @@ end @test haskey(new_st, :fixed) end - @testset "MultiNNHybridModel with scale_nn_outputs=true" begin + @testset "HybridModel with scale_nn_outputs=true" begin predictors = (a = [:x2, :x3], d = [:x1]) forcing = [:x1] targets = [:obs] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -481,7 +481,7 @@ end neural_param_names = Symbol[] global_param_names = Symbol[] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -491,7 +491,7 @@ end global_param_names ) - @test model isa SingleNNHybridModel + @test model isa HybridModel @test isempty(model.neural_param_names) @test isempty(model.global_param_names) @test model.fixed_param_names == [:a, :b, :c, :d] @@ -517,7 +517,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -548,7 +548,7 @@ end custom_chain = Chain(Dense(2, 16, tanh), Dense(16, 8, tanh)) - model = constructHybridModel( + model = HybridModel( predictors, forcing, targets, @@ -559,10 +559,10 @@ end hidden_layers = custom_chain ) - @test model isa SingleNNHybridModel - @test model.NN isa Chain + @test model isa HybridModel + @test model.NNs isa Chain # The chain should have the custom layers plus input and output layers - @test length(model.NN.layers) > length(custom_chain.layers) + @test length(model.NNs.layers) > length(custom_chain.layers) end end @@ -574,7 +574,7 @@ end neural_param_names = [:invalid_param] # Not in test_parameters global_param_names = [:b] - @test_throws AssertionError constructHybridModel( + @test_throws AssertionError HybridModel( predictors, forcing, targets, diff --git a/test/test_show_generic_hybrid.jl b/test/test_show_generic_hybrid.jl index 14fca091..0633cd7f 100644 --- a/test/test_show_generic_hybrid.jl +++ b/test/test_show_generic_hybrid.jl @@ -1,4 +1,4 @@ -using EasyHybrid: HybridParams, ParameterContainer, SingleNNHybridModel, MultiNNHybridModel +using EasyHybrid: HybridParams, ParameterContainer, HybridModel using EasyHybrid: _print_field, _print_header, IndentedIO @testset "show_generic.jl" begin @@ -44,12 +44,12 @@ using EasyHybrid: _print_field, _print_header, IndentedIO @test occursin("ParameterContainer(a, b)", result) end - @testset "SingleNNHybridModel show" begin + @testset "HybridModel show" begin function test_model(; x1, a, b) return (; y_pred = a .* x1 .+ b) end - model = constructHybridModel( + model = HybridModel( [:x1, :x2], [:x3], [:y], test_model, (a = (1.0, 0.0, 2.0), b = (2.0, 1.0, 3.0)), [:a], [:b]; @@ -71,12 +71,12 @@ using EasyHybrid: _print_field, _print_header, IndentedIO ) end - @testset "MultiNNHybridModel show" begin + @testset "HybridModel show" begin function test_model(; x1, x2, x3, a, b, c, d) return (; obs = a .* x2 .+ d .* x1 .+ b) end - model = constructHybridModel( + model = HybridModel( (a = [:x2, :x3], d = [:x1]), [:x1], [:obs], test_model, (a = (1.0, 0.0, 5.0), b = (2.0, 0.0, 10.0), c = (0.5, 0.0, 2.0), d = (0.5, 0.0, 2.0)), [:b]; # Only global_param_names, neural_param_names derived from predictors keys diff --git a/test/test_split_data_train.jl b/test/test_split_data_train.jl index a7176355..572fc745 100644 --- a/test/test_split_data_train.jl +++ b/test/test_split_data_train.jl @@ -57,11 +57,11 @@ const RbQ10_PARAMS = ( neural_param_names = [:rb] @testset "test DataFrame and thereby KeyedArray" begin - model = constructHybridModel( + model = HybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) - @test model isa SingleNNHybridModel + @test model isa HybridModel # prepare_data should produce something consumable by split_data ka = to_keyedArray(df) @test !isnothing(ka) @@ -135,11 +135,11 @@ const RbQ10_PARAMS = ( end @testset "test keep_history" begin - model = constructHybridModel( + model = HybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) - @test model isa SingleNNHybridModel + @test model isa HybridModel # prepare_data should produce something consumable by split_data ka = prepare_data(model, df) @test !isnothing(ka) From 5a17376126422ffee3973bf16fbaf7e8e04e1da8 Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Thu, 2 Jul 2026 22:42:32 +0200 Subject: [PATCH 02/12] performance review --- src/data/prepare_data.jl | 10 +++---- src/models/GenericHybridModel.jl | 51 ++++++++++---------------------- src/training/tune.jl | 6 ++-- 3 files changed, 23 insertions(+), 44 deletions(-) diff --git a/src/data/prepare_data.jl b/src/data/prepare_data.jl index 90067101..0c5bb080 100644 --- a/src/data/prepare_data.jl +++ b/src/data/prepare_data.jl @@ -4,12 +4,12 @@ function prepare_data(hm, data::KeyedArray; cfg = DataConfig(), kwargs...) predictors, forcings, targets = get_prediction_target_names(hm) # KeyedArray: use () syntax for views that are differentiable if predictors isa NamedTuple - X = NamedTuple([name => Array(data(p)) for (name, p) in pairs(predictors)]) + X = NamedTuple(name => Array(data(p)) for (name, p) in pairs(predictors)) else X = Array(data(predictors)) end - forcings_nt = NamedTuple([forcing => Array(data(forcing)) for forcing in forcings]) - targets_nt = NamedTuple([target => Array(data(target)) for target in targets]) + forcings_nt = NamedTuple(forcing => Array(data(forcing)) for forcing in forcings) + targets_nt = NamedTuple(target => Array(data(target)) for target in targets) return ((X, forcings_nt), targets_nt) end @@ -17,8 +17,8 @@ function prepare_data(hm, data::AbstractDimArray; kwargs...) predictors, forcings, targets = get_prediction_target_names(hm) # KeyedArray: use () syntax for views that are differentiable X_arr = data[variable = At(predictors)] - forcings_nt = NamedTuple([forcing => data[variable = At(forcing)] for forcing in forcings]) - targets_nt = NamedTuple([target => data[variable = At(target)] for target in targets]) + forcings_nt = NamedTuple(forcing => data[variable = At(forcing)] for forcing in forcings) + targets_nt = NamedTuple(target => data[variable = At(target)] for target in targets) # DimArray: use [] syntax (copies, but differentiable) return ((X_arr, forcings_nt), targets_nt) end diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index f8cc94df..aca9df48 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -205,12 +205,7 @@ end # ─────────────────────────────────────────────────────────────────────────── # Helper functions for NN initialization function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) - nn_params = NamedTuple() - for (nn_name, nn) in pairs(m.NNs) - ps_nn, _ = LuxCore.setup(rng, nn) - nn_params = merge(nn_params, NamedTuple{(nn_name,), Tuple{typeof(ps_nn)}}((ps_nn,))) - end - return (; nn_params...) + return map(nn -> LuxCore.setup(rng, nn)[1], m.NNs) end function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) @@ -241,12 +236,7 @@ function LuxCore.initialparameters(rng::AbstractRNG, m::HybridModel) end function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) - nn_states = NamedTuple() - for (nn_name, nn) in pairs(m.NNs) - _, st_nn = LuxCore.setup(rng, nn) - nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) - end - return nn_states + return map(nn -> LuxCore.setup(rng, nn)[2], m.NNs) end function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) @@ -312,31 +302,20 @@ end # ─────────────────────────────────────────────────────────────────────────── # ─────────────────────────────────────────────────────────────────────────── function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) - nn_outputs = NamedTuple() - nn_states = NamedTuple() - - for (nn_name, nn) in pairs(m.NNs) - nn_out, st_nn = LuxCore.apply(nn, ds_k[1][nn_name], ps[nn_name], st[nn_name]) - nn_outputs = merge(nn_outputs, NamedTuple{(nn_name,), Tuple{typeof(nn_out)}}((nn_out,))) - nn_states = merge(nn_states, NamedTuple{(nn_name,), Tuple{typeof(st_nn)}}((st_nn,))) - end - - scaled_nn_params = NamedTuple() - for (nn_name, param_name) in zip(keys(m.NNs), m.neural_param_names) - nn_cols = eachslice(nn_outputs[nn_name]; dims = 1) - nn_param = NamedTuple{(param_name,), Tuple{typeof(nn_cols[1])}}((nn_cols[1],)) - - if m.scale_nn_outputs - scaled_nn_val = scale_single_param(param_name, nn_param[param_name], m.parameters) - else - scaled_nn_val = nn_param[param_name] - end - - nn_scaled_param = NamedTuple{(param_name,), Tuple{typeof(scaled_nn_val)}}((scaled_nn_val,)) - scaled_nn_params = merge(scaled_nn_params, nn_scaled_param) - end + applied = map(LuxCore.apply, m.NNs, ds_k[1], ps, st) + nn_outputs = map(first, applied) + nn_states = map(last, applied) + + scaled_vals = Tuple( + begin + val = eachslice(nn_outputs[nn_name]; dims = 1)[1] + m.scale_nn_outputs ? scale_single_param(param_name, val, m.parameters) : val + end + for (nn_name, param_name) in zip(keys(m.NNs), m.neural_param_names) + ) + scaled_nn_params = NamedTuple{Tuple(m.neural_param_names)}(scaled_vals) - return scaled_nn_params, (; nn_states...), (; nn_outputs = nn_outputs) + return scaled_nn_params, nn_states, (; nn_outputs = nn_outputs) end function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) diff --git a/src/training/tune.jl b/src/training/tune.jl index 3b51be51..e0cfca1d 100644 --- a/src/training/tune.jl +++ b/src/training/tune.jl @@ -25,7 +25,7 @@ Construct a new hybrid model from `hybrid_model` plus hyperparameters, then call Returns a [`TrainResults`](@ref) (or `nothing` if data preparation fails, as in `train`). """ function tune(hybrid_model, data, mspec::ModelSpec; kwargs...) - kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, (; kwargs...), mspec.hyper_model) + kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, (; kwargs...), mspec.hyper_model) hm = HybridModel(; kwargs_model...) kwargs_train = merge((; kwargs...), mspec.hyper_train, (; target_names = hm.targets)) @@ -34,7 +34,7 @@ function tune(hybrid_model, data, mspec::ModelSpec; kwargs...) end function tune(hybrid_model, data; kwargs...) - kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, (; kwargs...)) + kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, (; kwargs...)) hm = HybridModel(; kwargs_model...) kwargs_train = merge((; kwargs...), (; target_names = hm.targets)) @@ -43,7 +43,7 @@ function tune(hybrid_model, data; kwargs...) end function tune(hybrid_model, data, train_cfg::TrainConfig; data_cfg::DataConfig = DataConfig(), kwargs...) - kwargs_model = merge(to_namedtuple(hybrid_model), hybrid_model.config, to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...)) + kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...)) hm = HybridModel(; kwargs_model...) kwargs_train = merge(to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...), (; target_names = hm.targets)) From 934177aab7f1788ca7c0155b11efa0e06caea5c9 Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Thu, 2 Jul 2026 22:59:14 +0200 Subject: [PATCH 03/12] keys --- src/models/GenericHybridModel.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index aca9df48..f9cbf05d 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -302,7 +302,8 @@ end # ─────────────────────────────────────────────────────────────────────────── # ─────────────────────────────────────────────────────────────────────────── function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) - applied = map(LuxCore.apply, m.NNs, ds_k[1], ps, st) + nn_names = keys(m.NNs) + applied = map(LuxCore.apply, m.NNs, ds_k[1][nn_names], ps[nn_names], st[nn_names]) nn_outputs = map(first, applied) nn_states = map(last, applied) From 8b1bb616f81067b7a1c63a4b8624ad2433259e3a Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Thu, 2 Jul 2026 23:52:43 +0200 Subject: [PATCH 04/12] map over keys --- src/models/GenericHybridModel.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index f9cbf05d..b829f709 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -303,9 +303,11 @@ end # ─────────────────────────────────────────────────────────────────────────── function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) nn_names = keys(m.NNs) - applied = map(LuxCore.apply, m.NNs, ds_k[1][nn_names], ps[nn_names], st[nn_names]) - nn_outputs = map(first, applied) - nn_states = map(last, applied) + applied = map(nn_names) do nn_name + LuxCore.apply(m.NNs[nn_name], ds_k[1][nn_name], ps[nn_name], st[nn_name]) + end + nn_outputs = NamedTuple{nn_names}(map(first, applied)) + nn_states = NamedTuple{nn_names}(map(last, applied)) scaled_vals = Tuple( begin From 09e9fde19e8e4c81c5903b55c000631818d4ac37 Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Fri, 3 Jul 2026 12:06:56 +0200 Subject: [PATCH 05/12] mv legacy and docstrings --- src/config/config_yaml.jl | 11 +- src/models/GenericHybridModel.jl | 143 +++++++++--------- src/models/helpers_for_HybridModel.jl | 31 +--- .../{ => legacy}/FluxPartModel_Q10_Lux.jl | 0 src/models/{ => legacy}/LinearHM.jl | 0 src/models/{ => legacy}/Respiration_Rb_Q10.jl | 0 src/models/{ => legacy}/Rs_components.jl | 0 src/models/{ => legacy}/simple_Rb_Q10_PBM.jl | 0 src/models/models.jl | 9 +- src/models/show_generic.jl | 13 -- src/models/utils.jl | 51 +++++++ test/test_generic_hybrid_model.jl | 31 ++-- test/test_show_generic_hybrid.jl | 15 +- 13 files changed, 144 insertions(+), 160 deletions(-) rename src/models/{ => legacy}/FluxPartModel_Q10_Lux.jl (100%) rename src/models/{ => legacy}/LinearHM.jl (100%) rename src/models/{ => legacy}/Respiration_Rb_Q10.jl (100%) rename src/models/{ => legacy}/Rs_components.jl (100%) rename src/models/{ => legacy}/simple_Rb_Q10_PBM.jl (100%) create mode 100644 src/models/utils.jl diff --git a/src/config/config_yaml.jl b/src/config/config_yaml.jl index 9bb78643..7916eaa9 100644 --- a/src/config/config_yaml.jl +++ b/src/config/config_yaml.jl @@ -18,20 +18,19 @@ function get_hybrid_config(hm::LuxCore.AbstractLuxContainerLayer) end _yaml_field_value(x) = x -_yaml_field_value(p::AbstractHybridModel) = get_parameters_config(p) +_yaml_field_value(p::ParameterContainer) = get_parameters_config(p) _yaml_field_value(f::Function) = get_mechanistic_model_config(f) """ - get_parameters_config(p::AbstractHybridModel) + get_parameters_config(pc::ParameterContainer) Serialize the parameter table (`default`, `lower`, `upper` per parameter) of a -`HybridParams`/`ParameterContainer` into a nested `OrderedDict` suitable for +`ParameterContainer` into a nested `OrderedDict` suitable for YAML output. Without this, only the compact `show` of the struct -(e.g. `HybridParams(ParameterContainer(RUE, Rb, Q10))`) would be written, which +(e.g. `ParameterContainer(RUE, Rb, Q10)`) would be written, which drops all of the actual default and bound values. """ -function get_parameters_config(p::AbstractHybridModel) - pc = hasfield(typeof(p), :hybrid) ? p.hybrid : p +function get_parameters_config(pc::ParameterContainer) out = OrderedDict{String, Any}() for name in keys(pc.values) d, l, u = pc.values[name] diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index b829f709..733d504f 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -1,25 +1,10 @@ -export HybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid -export HybridParams +export HybridModel, ParameterContainer # Import necessary components for neural networks using Lux: BatchNorm using Lux: sigmoid -# Define the hard sigmoid activation function -function hard_sigmoid(x) - return clamp.(0.2 .* x .+ 0.5, 0.0, 1.0) -end - -# Inverse of `hard_sigmoid` on the linear region (0, 1). -# Saturated inputs (y ≤ 0 or y ≥ 1) are extrapolated linearly since the -# clamp makes the forward map non-invertible there. -function inv_hard_sigmoid(y) - return (y .- 0.5) ./ 0.2 -end - -abstract type AbstractHybridModel end - -mutable struct ParameterContainer{NT <: NamedTuple, T} <: AbstractHybridModel +mutable struct ParameterContainer{NT <: NamedTuple, T} values::NT table::T @@ -30,33 +15,64 @@ mutable struct ParameterContainer{NT <: NamedTuple, T} <: AbstractHybridModel end """ - HybridParams{M<:Function} + HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} -A little parametric stub for “the params of function `M`.” -All of your function‐based models become `HybridParams{typeof(f)}`. +A unified hybrid model struct that handles both single and multi neural network architectures. +It combines predictive neural networks (`NNs`) with a `mechanistic_model` to form a differentiable hybrid model. """ -struct HybridParams{M <: Function} <: AbstractHybridModel - hybrid::ParameterContainer -end - -# ─────────────────────────────────────────────────────────────────────────── -# Unified Hybrid Model Structure (optimized for performance) struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} + "Neural network(s) used to predict parameters. Can be a single `Chain` or a `NamedTuple` of `Chain`s." NNs::T + + "Predictor variables for the neural networks. Can be a `Vector{Symbol}` or a `NamedTuple`." predictors::P + + "Forcing variables passed directly to the mechanistic model." forcing::Vector{Symbol} + + "Target variables the model will output/be trained against." targets::Vector{Symbol} + + "The core process-based or mechanistic model function." mechanistic_model::Function - parameters::AbstractHybridModel + + "Base parameters of the model (encapsulated in a `ParameterContainer`)." + parameters::ParameterContainer + + "Names of the parameters predicted by the neural network(s)." neural_param_names::Vector{Symbol} + + "Names of the globally optimized (constant) parameters." global_param_names::Vector{Symbol} + + "Names of the fixed (non-optimized) parameters." fixed_param_names::Vector{Symbol} + + "Whether to scale neural network outputs to the parameter bounds." scale_nn_outputs::Bool + + "Whether to initialize global parameters from their default values." start_from_default::Bool + + "Configuration named tuple capturing the hyperparameters used for initialization." config::NamedTuple end -# Unified constructor that dispatches based on predictors type +""" + HybridModel(predictors::Vector{Symbol}, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs...) + +Construct a `HybridModel` with a single neural network architecture predicting all `neural_param_names` from the `predictors`. + +# Arguments: +- `predictors::Vector{Symbol}`: Variables used as inputs to the neural network. +- `forcing`: Variables passed directly to the mechanistic model. +- `targets`: The target variables to predict. +- `mechanistic_model`: A function implementing the process-based model. +- `parameters`: A parameter container defining defaults, lowers, and uppers. +- `neural_param_names`: Names of the parameters to be predicted by the neural network. +- `global_param_names`: Names of the parameters to be globally optimized. +- `kwargs`: Additional configuration like `hidden_layers`, `activation`, `scale_nn_outputs`, etc. +""" function HybridModel( predictors::Vector{Symbol}, forcing, @@ -73,8 +89,8 @@ function HybridModel( kwargs... ) - if !isa(parameters, AbstractHybridModel) - parameters = build_parameters(parameters, mechanistic_model) + if !isa(parameters, ParameterContainer) + parameters = ParameterContainer(parameters) end all_names = pnames(parameters) @@ -110,6 +126,20 @@ function HybridModel( return HybridModel(NN, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) end +""" + HybridModel(predictors::NamedTuple, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs...) + +Construct a `HybridModel` with multiple neural network architectures. A separate neural network is built for each key in the `predictors` NamedTuple. + +# Arguments: +- `predictors::NamedTuple`: A NamedTuple where keys are network names, and values are vectors of predictor variables for that network. +- `forcing`: Variables passed directly to the mechanistic model. +- `targets`: The target variables to predict. +- `mechanistic_model`: A function implementing the process-based model. +- `parameters`: A parameter container defining defaults, lowers, and uppers. +- `global_param_names`: Names of the parameters to be globally optimized. +- `kwargs`: Additional configuration. `hidden_layers` and `activation` can also be NamedTuples to configure each network independently. +""" function HybridModel( predictors::NamedTuple, forcing, @@ -125,8 +155,8 @@ function HybridModel( kwargs... ) - if !isa(parameters, AbstractHybridModel) - parameters = build_parameters(parameters, mechanistic_model) + if !isa(parameters, ParameterContainer) + parameters = ParameterContainer(parameters) end all_names = pnames(parameters) @@ -260,47 +290,6 @@ function LuxCore.initialstates(rng::AbstractRNG, m::HybridModel) return merge(nn_states_nt, (; fixed = nt)) end -function default(p::AbstractHybridModel) - return p.hybrid.table[:, :default] -end - -function lower(p::AbstractHybridModel) - return p.hybrid.table[:, :lower] -end - -function upper(p::AbstractHybridModel) - return p.hybrid.table[:, :upper] -end - -pnames(p::AbstractHybridModel) = keys(p.hybrid.table.axes[1]) - -""" - scale_single_param(name, raw_val, parameters) - -Scale a single parameter using the sigmoid scaling function. -""" -function scale_single_param(name, raw_val, hm::AbstractHybridModel) - ℓ = lower(hm)[name] - u = upper(hm)[name] - return ℓ .+ (u .- ℓ) .* sigmoid.(raw_val) -end - -inv_sigmoid(y) = log.(y ./ (1 .- y)) - -""" - scale_single_param_minmax(name, hm::AbstractHybridModel) - -Scale a single parameter using the minmax scaling function. -""" -function scale_single_param_minmax(name, hm::AbstractHybridModel) - ℓ = lower(hm)[name] - u = upper(hm)[name] - return inv_sigmoid.((default(hm)[name] .- ℓ) ./ (u .- ℓ)) -end - - -# ─────────────────────────────────────────────────────────────────────────── -# ─────────────────────────────────────────────────────────────────────────── function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) nn_names = keys(m.NNs) applied = map(nn_names) do nn_name @@ -343,7 +332,13 @@ function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) return scaled_nn_params, (; st_nn = st_nn), (;) end -# Forward pass for HybridModel (optimized, using multiple dispatch for NN run) +""" + (m::HybridModel)(ds_k::Tuple, ps, st) + +Forward pass of the hybrid model. +Evaluates the neural networks to predict parameters, merges them with scaled global parameters and fixed parameters, and executes the mechanistic model. +Returns a tuple `(out, st_new)`. +""" function (m::HybridModel)(ds_k::Tuple, ps, st) parameters = m.parameters diff --git a/src/models/helpers_for_HybridModel.jl b/src/models/helpers_for_HybridModel.jl index ef440a23..70b25368 100644 --- a/src/models/helpers_for_HybridModel.jl +++ b/src/models/helpers_for_HybridModel.jl @@ -1,4 +1,4 @@ -export display_parameter_bounds, build_parameters, construct_dispatch_functions, build_parameter_matrix +export display_parameter_bounds, construct_dispatch_functions, build_parameter_matrix function construct_dispatch_functions(f) function new_f end # Create a new generic function @@ -6,13 +6,13 @@ function construct_dispatch_functions(f) println("constructing on KeyedArray function for $f") function new_f(forcing_data::KeyedArray, parameters::NamedTuple, forcing_names::Vector{Symbol}) forcing = toNamedTuple(forcing_data, forcing_names) - parameter_container = build_parameters(parameters, f) + parameter_container = ParameterContainer(parameters) return f(; forcing..., values(default(parameter_container))...) end function new_f(forcing_data::DataFrame, parameters::NamedTuple, forcing_names::Vector{Symbol}) forcing = (; (name => forcing_data[!, name] for name in forcing_names)...) - parameter_container = build_parameters(parameters, f) + parameter_container = ParameterContainer(parameters) return f(; forcing..., values(default(parameter_container))...) end @@ -24,32 +24,7 @@ function construct_dispatch_functions(f) return new_f end -""" - build_parameters(parameters::NamedTuple, f::DataType) -> AbstractHybridModel - -Constructs a parameter container from a named tuple of parameter bounds and wraps it in a user-defined subtype of `AbstractHybridModel`. - -# Arguments -- `parameters::NamedTuple`: Named tuple where each entry is a tuple of (default, lower, upper) bounds for a parameter. -- `f::DataType`: A constructor for a subtype of `AbstractHybridModel` that takes a `ParameterContainer` as its argument. -# Returns -- An instance of the user-defined `AbstractHybridModel` subtype containing the parameter container. -""" -function build_parameters(parameters::NamedTuple, f::DataType) - ca = ParameterContainer(parameters) - return f(ca) -end - -# the “core” build_parameters that knows how to turn a NamedTuple + a -# subtype of AbstractHybridModel into an actual model instance: -function build_parameters(params::NamedTuple, ::Type{P}) where {P <: AbstractHybridModel} - return P(ParameterContainer(params)) -end - -function build_parameters(params::NamedTuple, f::M) where {M <: Function} - return build_parameters(params, HybridParams{M}) -end """ build_parameter_matrix(parameter_defaults_and_bounds::NamedTuple) diff --git a/src/models/FluxPartModel_Q10_Lux.jl b/src/models/legacy/FluxPartModel_Q10_Lux.jl similarity index 100% rename from src/models/FluxPartModel_Q10_Lux.jl rename to src/models/legacy/FluxPartModel_Q10_Lux.jl diff --git a/src/models/LinearHM.jl b/src/models/legacy/LinearHM.jl similarity index 100% rename from src/models/LinearHM.jl rename to src/models/legacy/LinearHM.jl diff --git a/src/models/Respiration_Rb_Q10.jl b/src/models/legacy/Respiration_Rb_Q10.jl similarity index 100% rename from src/models/Respiration_Rb_Q10.jl rename to src/models/legacy/Respiration_Rb_Q10.jl diff --git a/src/models/Rs_components.jl b/src/models/legacy/Rs_components.jl similarity index 100% rename from src/models/Rs_components.jl rename to src/models/legacy/Rs_components.jl diff --git a/src/models/simple_Rb_Q10_PBM.jl b/src/models/legacy/simple_Rb_Q10_PBM.jl similarity index 100% rename from src/models/simple_Rb_Q10_PBM.jl rename to src/models/legacy/simple_Rb_Q10_PBM.jl diff --git a/src/models/models.jl b/src/models/models.jl index 328ed5f4..a1c89446 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,8 +1,9 @@ -include("LinearHM.jl") -include("Respiration_Rb_Q10.jl") -include("Rs_components.jl") -include("simple_Rb_Q10_PBM.jl") +# include("LinearHM.jl") +# include("Respiration_Rb_Q10.jl") +# include("Rs_components.jl") +# include("simple_Rb_Q10_PBM.jl") include("GenericHybridModel.jl") +include("utils.jl") include("helpers_for_HybridModel.jl") include("NNModels.jl") include("show_generic.jl") diff --git a/src/models/show_generic.jl b/src/models/show_generic.jl index fc42be88..6d5ec2ec 100644 --- a/src/models/show_generic.jl +++ b/src/models/show_generic.jl @@ -52,19 +52,6 @@ function _print_header(io::IO, text::String; color = :blue, bold = true) printstyled(io, text, color = color, bold = bold) return println(io) end -function Base.show(io::IO, ::MIME"text/plain", hp::HybridParams) - _print_header(io, "Hybrid Parameters", color = :blue) - - # Delegate to the contained ParameterContainer - io_full = IOContext(IndentedIO(io), :compact => false, :limit => false) - return show(io_full, MIME"text/plain"(), hp.hybrid) -end - -function Base.show(io::IO, hp::HybridParams) - print(io, "HybridParams(") - show(io, hp.hybrid) - return print(io, ")") -end function Base.show(io::IO, ::MIME"text/plain", pc::ParameterContainer) table = pc.table diff --git a/src/models/utils.jl b/src/models/utils.jl new file mode 100644 index 00000000..db463cfa --- /dev/null +++ b/src/models/utils.jl @@ -0,0 +1,51 @@ +export scale_single_param, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid, scale_single_param_minmax + +# Define the hard sigmoid activation function +function hard_sigmoid(x) + return clamp.(0.2 .* x .+ 0.5, 0.0, 1.0) +end + +# Inverse of `hard_sigmoid` on the linear region (0, 1). +# Saturated inputs (y ≤ 0 or y ≥ 1) are extrapolated linearly since the +# clamp makes the forward map non-invertible there. +function inv_hard_sigmoid(y) + return (y .- 0.5) ./ 0.2 +end + +function default(p::ParameterContainer) + return p.table[:, :default] +end + +function lower(p::ParameterContainer) + return p.table[:, :lower] +end + +function upper(p::ParameterContainer) + return p.table[:, :upper] +end + +pnames(p::ParameterContainer) = keys(p.table.axes[1]) + +""" + scale_single_param(name, raw_val, parameters) + +Scale a single parameter using the sigmoid scaling function. +""" +function scale_single_param(name, raw_val, hm::ParameterContainer) + ℓ = lower(hm)[name] + u = upper(hm)[name] + return ℓ .+ (u .- ℓ) .* sigmoid.(raw_val) +end + +inv_sigmoid(y) = log.(y ./ (1 .- y)) + +""" + scale_single_param_minmax(name, hm::AbstractHybridModel) + +Scale a single parameter using the minmax scaling function. +""" +function scale_single_param_minmax(name, hm::ParameterContainer) + ℓ = lower(hm)[name] + u = upper(hm)[name] + return inv_sigmoid.((default(hm)[name] .- ℓ) ./ (u .- ℓ)) +end diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index e8e5d6bf..a6c261d2 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -34,11 +34,7 @@ test_parameters = ( @test result[5] ≈ 1.0f0 # hard_sigmoid(5) = 1.0 (clamped) end - @testset "AbstractHybridModel types" begin - @test AbstractHybridModel <: Any - @test ParameterContainer <: AbstractHybridModel - @test HybridParams <: AbstractHybridModel - end + @testset "ParameterContainer construction" begin params = (a = (1.0f0, 0.0f0, 2.0f0), b = (2.0f0, 1.0f0, 3.0f0)) @@ -69,59 +65,52 @@ test_parameters = ( @test result == expected end - @testset "HybridParams construction" begin - params = (a = (1.0f0, 0.0f0, 2.0f0),) - pc = ParameterContainer(params) - hp = HybridParams{typeof(test_mechanistic_model)}(pc) - @test hp isa HybridParams - @test hp.hybrid == pc - end end @testset "GenericHybridModel - Parameter Functions" begin params = (a = (1.0f0, 0.0f0, 2.0f0), b = (2.0f0, 1.0f0, 3.0f0)) - hp = HybridParams{typeof(test_mechanistic_model)}(ParameterContainer(params)) + pc = ParameterContainer(params) @testset "default function" begin - defaults = default(hp) + defaults = default(pc) @test defaults.a == 1.0f0 @test defaults.b == 2.0f0 end @testset "lower function" begin - lowers = lower(hp) + lowers = lower(pc) @test lowers.a == 0.0f0 @test lowers.b == 1.0f0 end @testset "upper function" begin - uppers = upper(hp) + uppers = upper(pc) @test uppers.a == 2.0f0 @test uppers.b == 3.0f0 end @testset "pnames function" begin - names = EasyHybrid.pnames(hp) + names = EasyHybrid.pnames(pc) @test collect(names) == [:a, :b] end @testset "scale_single_param function" begin # Test sigmoid scaling - scaled_a = scale_single_param(:a, [0.0f0], hp) + scaled_a = scale_single_param(:a, [0.0f0], pc) @test length(scaled_a) == 1 @test scaled_a[1] ≈ 1.0f0 # sigmoid(0) = 0.5, so 0 + 0.5*(2-0) = 1.0 - scaled_b = scale_single_param(:b, [0.0f0], hp) + scaled_b = scale_single_param(:b, [0.0f0], pc) @test scaled_b[1] ≈ 2.0f0 # sigmoid(0) = 0.5, so 1 + 0.5*(3-1) = 2.0 end @testset "scale_single_param_minmax function" begin # Test inverse sigmoid scaling - inv_scaled_a = EasyHybrid.scale_single_param_minmax(:a, hp) + inv_scaled_a = EasyHybrid.scale_single_param_minmax(:a, pc) @test inv_scaled_a ≈ 0.0f0 # inverse sigmoid of 0.5 is 0.0 - inv_scaled_b = EasyHybrid.scale_single_param_minmax(:b, hp) + inv_scaled_b = EasyHybrid.scale_single_param_minmax(:b, pc) @test inv_scaled_b ≈ 0.0f0 # inverse sigmoid of 0.5 is 0.0 end end diff --git a/test/test_show_generic_hybrid.jl b/test/test_show_generic_hybrid.jl index 0633cd7f..f7233dcb 100644 --- a/test/test_show_generic_hybrid.jl +++ b/test/test_show_generic_hybrid.jl @@ -1,4 +1,4 @@ -using EasyHybrid: HybridParams, ParameterContainer, HybridModel +using EasyHybrid: ParameterContainer, HybridModel using EasyHybrid: _print_field, _print_header, IndentedIO @testset "show_generic.jl" begin @@ -22,19 +22,6 @@ using EasyHybrid: _print_field, _print_header, IndentedIO @test occursin("function", result) && occursin("sum", result) end - @testset "HybridParams show" begin - params = (a = (1.0, 0.0, 2.0), b = (2.0, 1.0, 3.0)) - pc = ParameterContainer(params) - hp = HybridParams{typeof(sum)}(pc) - - # Compact show - result_compact = sprint(show, hp, context = :color => false) - @test occursin("HybridParams", result_compact) && occursin("ParameterContainer", result_compact) - - # Text/plain show - result_full = sprint(show, MIME"text/plain"(), hp, context = :color => false) - @test occursin("Hybrid Parameters", result_full) # only check for the header - end @testset "ParameterContainer compact show" begin params = (a = (1.0, 0.0, 2.0), b = (2.0, 1.0, 3.0)) From 798c90a25c9b21907730d199d8766dec258b8fb4 Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Fri, 3 Jul 2026 14:22:59 +0200 Subject: [PATCH 06/12] to use constructHybridModel --- .../tutorials/example_synthetic_lstm.jl | 4 +- src/EasyHybrid.jl | 2 +- src/models/GenericHybridModel.jl | 42 ++++++++++++------- test/test_generic_hybrid_model.jl | 1 - 4 files changed, 30 insertions(+), 19 deletions(-) diff --git a/docs/literate/tutorials/example_synthetic_lstm.jl b/docs/literate/tutorials/example_synthetic_lstm.jl index 222528c7..8bccb93c 100644 --- a/docs/literate/tutorials/example_synthetic_lstm.jl +++ b/docs/literate/tutorials/example_synthetic_lstm.jl @@ -87,7 +87,7 @@ neural_param_names = [:rb] # ## 7. Construct LSTM Hybrid Model # Create LSTM hybrid model using the unified constructor -hlstm = HybridModel( +hlstm = constructHybridModel( predictors, forcing, target, @@ -216,7 +216,7 @@ first(out_lstm.val_obs_pred, 5) # ## 10. Train Single NN Hybrid Model (Optional) # For comparison, we can also train a hybrid model with a standard feed-forward neural network -hm = HybridModel( +hm = constructHybridModel( predictors, forcing, target, diff --git a/src/EasyHybrid.jl b/src/EasyHybrid.jl index e29d53a4..d10e1a69 100644 --- a/src/EasyHybrid.jl +++ b/src/EasyHybrid.jl @@ -28,7 +28,7 @@ using Downloads: Downloads using Hyperopt: Hyperopt, Hyperoptimizer using JLD2: JLD2, jldopen using LuxCore: LuxCore -using Lux: Lux +using Lux: Lux, BatchNorm, sigmoid using MLJ: partition using MLUtils: MLUtils, DataLoader, kfolds, numobs, rpad, splitobs using NCDatasets: NCDatasets, NCDataset, close, name diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 733d504f..9da59ea5 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -1,11 +1,15 @@ -export HybridModel, ParameterContainer +export HybridModel, ParameterContainer, constructHybridModel -# Import necessary components for neural networks -using Lux: BatchNorm -using Lux: sigmoid +""" + ParameterContainer{NT <: NamedTuple, T} +A container for holding the parameter definitions of a model, including their default values, lower bounds, and upper bounds. +""" mutable struct ParameterContainer{NT <: NamedTuple, T} + "The raw parameter definitions. A `NamedTuple` where each entry is a tuple of `(default, lower, upper)` bounds for a parameter." values::NT + + "A `ComponentArray` matrix representation of the parameter bounds, organized for efficient access by name and bound type." table::T function ParameterContainer(values::NT) where {NT <: NamedTuple} @@ -23,37 +27,37 @@ It combines predictive neural networks (`NNs`) with a `mechanistic_model` to for struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} "Neural network(s) used to predict parameters. Can be a single `Chain` or a `NamedTuple` of `Chain`s." NNs::T - + "Predictor variables for the neural networks. Can be a `Vector{Symbol}` or a `NamedTuple`." predictors::P - + "Forcing variables passed directly to the mechanistic model." forcing::Vector{Symbol} - + "Target variables the model will output/be trained against." targets::Vector{Symbol} - + "The core process-based or mechanistic model function." mechanistic_model::Function - + "Base parameters of the model (encapsulated in a `ParameterContainer`)." parameters::ParameterContainer - + "Names of the parameters predicted by the neural network(s)." neural_param_names::Vector{Symbol} - + "Names of the globally optimized (constant) parameters." global_param_names::Vector{Symbol} - + "Names of the fixed (non-optimized) parameters." fixed_param_names::Vector{Symbol} - + "Whether to scale neural network outputs to the parameter bounds." scale_nn_outputs::Bool - + "Whether to initialize global parameters from their default values." start_from_default::Bool - + "Configuration named tuple capturing the hyperparameters used for initialization." config::NamedTuple end @@ -408,3 +412,11 @@ function (m::HybridModel)(df::DataFrame, ps, st) end return dfnew end + +""" + constructHybridModel(args...; kwargs...) + +Alias for `HybridModel`. Provided for backward compatibility with earlier versions of EasyHybrid.jl. +Users are encouraged to use `HybridModel` directly. +""" +const constructHybridModel = HybridModel diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index a6c261d2..851cc842 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -35,7 +35,6 @@ test_parameters = ( end - @testset "ParameterContainer construction" begin params = (a = (1.0f0, 0.0f0, 2.0f0), b = (2.0f0, 1.0f0, 3.0f0)) pc = ParameterContainer(params) From d6c76b7a0d0d89fbc4349a77be5ed138aa8f18bd Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Fri, 3 Jul 2026 14:57:37 +0200 Subject: [PATCH 07/12] remove old code, but keep their core ideas in the docs --- .gitignore | 1 + docs/literate/tutorials/building_models.jl | 167 +++++++++++++++++++++ docs/make.jl | 1 + src/models/legacy/FluxPartModel_Q10_Lux.jl | 79 ---------- src/models/legacy/LinearHM.jl | 68 --------- src/models/legacy/Respiration_Rb_Q10.jl | 72 --------- src/models/legacy/Rs_components.jl | 57 ------- src/models/legacy/simple_Rb_Q10_PBM.jl | 44 ------ src/models/models.jl | 4 - test/runtests.jl | 9 -- 10 files changed, 169 insertions(+), 333 deletions(-) create mode 100644 docs/literate/tutorials/building_models.jl delete mode 100644 src/models/legacy/FluxPartModel_Q10_Lux.jl delete mode 100644 src/models/legacy/LinearHM.jl delete mode 100644 src/models/legacy/Respiration_Rb_Q10.jl delete mode 100644 src/models/legacy/Rs_components.jl delete mode 100644 src/models/legacy/simple_Rb_Q10_PBM.jl diff --git a/.gitignore b/.gitignore index 0aa43f4e..9e822727 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ Manifest.toml *.err *.out /docs/src/tutorials/folds.md +/docs/src/tutorials/building_models.md *.jl.cov *.jl.*.cov coverage/lcov.info diff --git a/docs/literate/tutorials/building_models.jl b/docs/literate/tutorials/building_models.jl new file mode 100644 index 00000000..b49f5954 --- /dev/null +++ b/docs/literate/tutorials/building_models.jl @@ -0,0 +1,167 @@ +# # Building Models Examples +# +# `EasyHybrid.jl` allows constructing diverse modeling architectures using the unified `HybridModel` struct. +# Previously, users defined bespoke structs (like `LinearHM`, `RespirationRbQ10`) for different configurations. +# Here we demonstrate how those legacy model architectures can be trivially constructed via `HybridModel`. +# +# ## Setup +# First, let's load our required packages: +using EasyHybrid + +# ## 1. Linear Hybrid Model +# This is a basic model with one neural network predicting a coefficient `α`, and an explicit global parameter `β`. +# The equation is: `ŷ = α * x + β` +# +# ### Process-Based Definition +linear_mechanistic(; x, α, β) = (; obs = α .* x .+ β) + +# ### Parameter Setup +params_linear = ( + α = (1.0f0, 0.0f0, 2.0f0), + β = (1.5f0, -1.0f0, 3.0f0), +) + +# ### HybridModel Construction +# We use `x` as forcing data, predict `α` with a neural network based on some predictors `a` and `b`, +# and leave `β` as a globally optimized constant parameter. +lhm = HybridModel( + [:a, :b], # predictors for the NN (predicts α) + [:x], # forcing variable + [:obs], # targets + linear_mechanistic, # our mechanistic model + params_linear, # parameter container + [:α], # parameters predicted by the NN + [:β]; # globally optimized constant parameters + hidden_layers = [4, 4], + activation = tanh +) + + +# ## 2. Respiration Rb Q10 +# A single NN predicting `Rb` for a Q10 temperature-sensitive respiration formulation. +# The equation is: `R_soil = Rb * Q10^(0.1 * (Temp - 15))` +# +# ### Process-Based Definition +function mRbQ10(; Temp, Rb, Q10) + R_soil = @. Rb * Q10^(0.1f0 * (Temp - 15.0f0)) + return (; R_soil) +end + +# ### Parameter Setup +params_rbq10 = ( + Rb = (1.0f0, 0.0f0, 5.0f0), + Q10 = (1.5f0, 1.0f0, 3.0f0), +) + +# ### HybridModel Construction +m_rbq10 = HybridModel( + [:SWC, :TA], # predictors for Rb + [:Temp], # forcing variable + [:R_soil], # targets + mRbQ10, # mechanistic model + params_rbq10, + [:Rb], # predicted by NN + [:Q10]; # globally optimized + hidden_layers = [8, 8] +) + + +# ## 3. Respiration Components +# A single NN outputting 3 distinct parameters (`Rb_het`, `Rb_root`, `Rb_myc`). +# +# ### Process-Based Definition +function rs_comp(; Temp, Rb_het, Rb_root, Rb_myc, Q10_het, Q10_root, Q10_myc) + R_het = @. Rb_het * Q10_het^(0.1f0 * (Temp - 15.0f0)) + R_root = @. Rb_root * Q10_root^(0.1f0 * (Temp - 15.0f0)) + R_myc = @. Rb_myc * Q10_myc^(0.1f0 * (Temp - 15.0f0)) + R_soil = R_het .+ R_root .+ R_myc + return (; R_soil, R_het, R_root, R_myc) +end + +# ### Parameter Setup +params_rs_comp = ( + Rb_het = (1.0f0, 0.0f0, 5.0f0), + Rb_root = (1.0f0, 0.0f0, 5.0f0), + Rb_myc = (1.0f0, 0.0f0, 5.0f0), + Q10_het = (1.5f0, 1.0f0, 3.0f0), + Q10_root = (1.5f0, 1.0f0, 3.0f0), + Q10_myc = (1.5f0, 1.0f0, 3.0f0), +) + +# ### HybridModel Construction +m_rs_comp = HybridModel( + [:SWC, :TA], # predictors for all 3 Rb parameters + [:Temp], + [:R_soil], + rs_comp, + params_rs_comp, + [:Rb_het, :Rb_root, :Rb_myc], + [:Q10_het, :Q10_root, :Q10_myc]; + hidden_layers = [16, 16] +) + + +# ## 4. Flux Partitioning with Multiple NNs +# A multi-NN architecture predicting `RUE` (Radiation Use Efficiency) and `Rb` from different sets of predictors. +# +# ### Process-Based Definition +function flux_part(; SW_IN, TA, RUE, Rb, Q10) + GPP = @. SW_IN * RUE / 12.011f0 + RECO = @. Rb * Q10^(0.1f0 * (TA - 15.0f0)) + NEE = RECO .- GPP + return (; NEE, GPP, RECO) +end + +# ### Parameter Setup +params_flux = ( + RUE = (1.0f0, 0.0f0, 5.0f0), + Rb = (1.0f0, 0.0f0, 5.0f0), + Q10 = (1.5f0, 1.0f0, 3.0f0), +) + +# ### HybridModel Construction +# By passing a `NamedTuple` to `predictors`, `HybridModel` automatically provisions +# an independent Neural Network for each key. +predictors_multi = ( + RUE = [:SWC, :TA, :SW_IN], + Rb = [:SWC, :TA], +) + +m_flux = HybridModel( + predictors_multi, # Triggers Multi-NN construction + [:SW_IN, :TA], # Forcing variables + [:NEE], # Targets + flux_part, # Mechanistic model + params_flux, + [:Q10]; # Global parameter + hidden_layers = (RUE = [8, 8], Rb = [4, 4]), # Custom architectures per NN + activation = (RUE = Lux.sigmoid, Rb = tanh) +) + + +# ## 5. Process-Based Model (Zero NNs) +# A purely process-based configuration where all parameters are optimized globally, and no Neural Networks are built. +# +# ### Process-Based Definition +function mRbQ10_0(; Temp, Rb, Q10) + R_soil = @. Rb * Q10^(0.1f0 * (Temp - 0.0f0)) + return (; R_soil) +end + +# ### HybridModel Construction +# Passing an empty `Symbol[]` array to `predictors` prevents any Neural Networks from being created. +m_pbm = HybridModel( + Symbol[], # No predictors -> No Neural Network + [:Temp], # Forcing + [:R_soil], # Target + mRbQ10_0, + params_rbq10, + Symbol[], # No neural params + [:Rb, :Q10] # Both are optimized as global parameters +) + +# ## Summary +# +# As demonstrated above, `HybridModel` provides a highly flexible, unified interface. +# By simply modifying the `predictors` argument and your mechanistic function, you can rapidly scale from a purely +# process-based model, to a single Neural Network hybrid model, all the way up to complex multi-Neural Network architectures! diff --git a/docs/make.jl b/docs/make.jl index 3ebf3c4b..82274fdc 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -59,6 +59,7 @@ makedocs(; "Get Started" => "get_started.md", "Tutorial" => [ "Overview" => "tutorials/overview.md", + "Building Models Examples" => "tutorials/building_models.md", "Dashboard" => "tutorials/dashboard.md", "Exponential Response" => "tutorials/exponential_res.md", "Hyperparameter Tuning" => "tutorials/hyperparameter_tuning.md", diff --git a/src/models/legacy/FluxPartModel_Q10_Lux.jl b/src/models/legacy/FluxPartModel_Q10_Lux.jl deleted file mode 100644 index 89b12d7e..00000000 --- a/src/models/legacy/FluxPartModel_Q10_Lux.jl +++ /dev/null @@ -1,79 +0,0 @@ -export FluxPartModelQ10Lux - -""" - FluxPartModelQ10Lux(RUE_NN, Rb_NN, RUE_predictors, Rb_predictors, forcing, targets, Q10) - -A flux partitioning model with separate neural networks for RUE (Radiation Use Efficiency) and Rb (basal respiration), -using Q10 temperature sensitivity for respiration calculations. -""" -struct FluxPartModelQ10Lux{D1, D2, T1, T2, T3, T4, T5} <: LuxCore.AbstractLuxContainerLayer{(:RUE_NN, :Rb_NN, :RUE_predictors, :Rb_predictors, :forcing, :targets, :Q10)} - RUE_NN - Rb_NN - RUE_predictors - Rb_predictors - forcing - targets - Q10 - function FluxPartModelQ10Lux(RUE_NN::D1, Rb_NN::D2, RUE_predictors::T1, Rb_predictors::T2, forcing::T3, targets::T4, Q10::T5) where {D1, D2, T1, T2, T3, T4, T5} - return new{D1, D2, T1, T2, T3, T4, T5}(RUE_NN, Rb_NN, collect(RUE_predictors), collect(Rb_predictors), collect(forcing), collect(targets), [Q10]) - end -end - -function LuxCore.initialparameters(::AbstractRNG, layer::FluxPartModelQ10Lux) - ps_RUE, _ = LuxCore.setup(Random.default_rng(), layer.RUE_NN) - ps_Rb, _ = LuxCore.setup(Random.default_rng(), layer.Rb_NN) - return (; RUE = ps_RUE, Rb = ps_Rb, Q10 = layer.Q10) -end - -function LuxCore.initialstates(::AbstractRNG, layer::FluxPartModelQ10Lux) - _, st_RUE = LuxCore.setup(Random.default_rng(), layer.RUE_NN) - _, st_Rb = LuxCore.setup(Random.default_rng(), layer.Rb_NN) - return (; RUE = st_RUE, Rb = st_Rb) -end - -""" - FluxPartModelQ10Lux(RUE_NN, Rb_NN, RUE_predictors, Rb_predictors, forcing, targets, Q10)(ds_k, ps, st) - -# Model definition -- GPP = SW_IN * RUE(αᵢ(t)) / 12.011 # µmol/m²/s = J/s/m² * g/MJ / g/mol -- Reco = Rb(βᵢ(t)) * Q10^((T(t) - T_ref)/10) -- NEE = Reco - GPP - -where: -- RUE(αᵢ(t)) is the radiation use efficiency predicted by neural network -- Rb(βᵢ(t)) is the basal respiration predicted by neural network -- SW_IN is incoming shortwave radiation -- T(t) is air temperature -- T_ref is reference temperature (15°C) -- Q10 is temperature sensitivity factor -""" -function (hm::FluxPartModelQ10Lux)(ds_k, ps, st::NamedTuple) - # Get inputs for RUE neural network - RUE_input = ds_k(hm.RUE_predictors) - - # Get inputs for Rb neural network - Rb_input = ds_k(hm.Rb_predictors) - - # Get forcing variables - forcing_data = ds_k(hm.forcing) # don't propagate names after this - sw_in = Array(forcing_data([:SW_IN])) # SW_IN - ta = Array(forcing_data([:TA])) # TA - - - RUE, st_RUE = LuxCore.apply(hm.RUE_NN, RUE_input, ps.RUE, st.RUE) # TODO could be simplified if we move diagnostics out of the tuple with st - Rb, st_Rb = LuxCore.apply(hm.Rb_NN, Rb_input, ps.Rb, st.Rb) - - # Scale outputs - RUE_scaled = 1.0f0 .* RUE - Rb_scaled = 1.0f0 .* Rb - - # Calculate fluxes - GPP = sw_in .* RUE_scaled ./ 12.011f0 # µmol/m²/s - RECO = Rb_scaled .* ps.Q10 .^ (0.1f0 .* (ta .- 15.0f0)) - NEE = RECO .- GPP - - # Update states - new_st = (; RUE = st_RUE, Rb = st_Rb) - - return (; NEE, RUE = RUE_scaled, Rb = Rb_scaled, GPP, RECO), (; st = new_st) -end diff --git a/src/models/legacy/LinearHM.jl b/src/models/legacy/LinearHM.jl deleted file mode 100644 index c07c9617..00000000 --- a/src/models/legacy/LinearHM.jl +++ /dev/null @@ -1,68 +0,0 @@ -export LinearHM - -""" - LinearHM(NN, predictors, forcing, targets, β) - -A linear hybrid model with a neural network `NN`, `predictors`, `forcing` and `targets` terms. -""" -struct LinearHM{D, T1, T2, T3, T4} <: LuxCore.AbstractLuxContainerLayer{(:NN, :predictors, :forcing, :targets, :β)} - NN - predictors - forcing - targets - β - function LinearHM(NN::D, predictors::T1, forcing::T2, targets::T3, β::T4) where {D, T1, T2, T3, T4} - return new{D, T1, T2, T3, T4}(NN, collect(predictors), collect(forcing), collect(targets), [β]) - end -end - -# ? β is a parameter, so expand the initialparameters! -function LuxCore.initialparameters(::AbstractRNG, layer::LinearHM) - ps, _ = LuxCore.setup(Random.default_rng(), layer.NN) - return (; ps, β = layer.β) -end - -function LuxCore.initialstates(::AbstractRNG, layer::LinearHM) - _, st = LuxCore.setup(Random.default_rng(), layer.NN) - return (; st) -end - -""" - LinearHM(NN, predictors, forcing, β)(ds_k) - -# Model definition `ŷ = α x + β` - -Apply the linear hybrid model to the input data `ds_k` (a `KeyedArray` with proper dimensions). -The model uses the neural network `NN` to compute new `α`'s based on the `predictors` and then computes `ŷ` using the `forcing` term `x`. - -Returns: - -A tuple containing the predicted values `ŷ` and a named tuple with the computed values of `α` and the state `st`. - -## Example: -````julia -using Lux -using EasyHybrid -using Random -using AxisKeys - -ds_k = KeyedArray(rand(Float32, 3,4); data=[:a, :b, :c], sample=1:4) -m = Lux.Chain(Dense(2, 5), Dense(5, 1)) -# Instantiate the model -# Note: The model is initialized with a neural network and the predictors and forcing terms -lh_model = LinearHM(m, (:a, :b), (:c,), 1.5f0) -rng = Random.default_rng() -Random.seed!(rng, 0) -ps, st = LuxCore.setup(rng, lh_model) -# Apply the model to the data -ŷ, αst = LuxCore.apply(lh_model, ds_k, ps, st) -```` -""" -function (lhm::LinearHM)(ds_k, ps, st::NamedTuple) - p = ds_k(lhm.predictors) - x = Array(ds_k(lhm.forcing)) # don't propagate name dims - α, st = LuxCore.apply(lhm.NN, p, ps.ps, st.st) - ŷ = α .* x .+ ps.β - - return (; obs = ŷ), (; st = (; st)) -end diff --git a/src/models/legacy/Respiration_Rb_Q10.jl b/src/models/legacy/Respiration_Rb_Q10.jl deleted file mode 100644 index d9be3825..00000000 --- a/src/models/legacy/Respiration_Rb_Q10.jl +++ /dev/null @@ -1,72 +0,0 @@ -export RespirationRbQ10 -export mRbQ10 - -""" - RespirationRbQ10(NN, predictors, forcing, targets, Q10) - -A linear hybrid model with a neural network `NN`, `predictors`, `targets` and `forcing` terms. -""" -struct RespirationRbQ10{D, T1, T2, T3, T4} <: LuxCore.AbstractLuxContainerLayer{(:NN, :predictors, :forcing, :targets, :Q10)} - NN - predictors - forcing #TODO order is messed up compared to new - targets - Q10 - function RespirationRbQ10(NN::D, predictors::T1, forcing::T2, targets::T3, Q10::T4) where {D, T1, T2, T3, T4} - return new{D, T1, T2, T3, T4}(NN, collect(predictors), collect(forcing), collect(targets), [Q10]) - end -end - -# ? Q10 is a parameter, so expand the initialparameters! -function LuxCore.initialparameters(::AbstractRNG, layer::RespirationRbQ10) - ps, _ = LuxCore.setup(Random.default_rng(), layer.NN) - return (; ps, Q10 = layer.Q10) -end -# TODO: trainable vs non-trainable! set example! -# see: https://lux.csail.mit.edu/stable/manual/migrate_from_flux#Implementing-Custom-Layers -function LuxCore.initialstates(::AbstractRNG, layer::RespirationRbQ10) - _, st = LuxCore.setup(Random.default_rng(), layer.NN) - return (; st) -end - -""" - mRbQ10(Rb, Q10, Temp, Tref) - - Rb base respiration, Q10 temperature sensitivity, Temp current temperature, Tref reference temperature - -""" - -function mRbQ10(Rb, Q10, Temp, Tref) - return @. Rb * Q10^(0.1f0 * (Temp - Tref)) -end - -""" - RespirationRbQ10(NN, predictors, forcing, targets, Q10)(ds_k) - -# Model definition `ŷ = Rb(αᵢ(t)) * Q10^((T(t) - T_ref)/10)` - -ŷ (respiration rate) is computed as a function of the neural network output `Rb(αᵢ(t))` and the temperature `T(t)` adjusted by the reference temperature `T_ref` (default 15°C) using the Q10 temperature sensitivity factor. -```` -""" -function (hm::RespirationRbQ10)(ds_k, ps, st::NamedTuple) - p = ds_k(hm.predictors) - x = Array(ds_k(hm.forcing)) # don't propagate names after this - Rb, stQ10 = LuxCore.apply(hm.NN, p, ps.ps, st.st) #! NN(αᵢ(t)) ≡ Rb(T(t), M(t)) - - #TODO output name flexible - could be R_soil, heterotrophic, autotrophic, etc. - R_soil = mRbQ10(Rb, ps.Q10, x, 15.0f0) # ? should 15°C be the reference temperature also an input variable? - - return (; R_soil, Rb), (; st = stQ10) -end - -function (hm::RespirationRbQ10)(ds_k::AbstractDimArray, ps, st::NamedTuple) - p = ds_k[variable = At(hm.predictors)] # No @view - needs to be differentiable - x = Array(ds_k[variable = At(hm.forcing)]) # No @view - needs to be differentiable - - Rb, stQ10 = LuxCore.apply(hm.NN, p, ps.ps, st.st) #! NN(αᵢ(t)) ≡ Rb(T(t), M(t)) - - #TODO output name flexible - could be R_soil, heterotrophic, autotrophic, etc. - R_soil = mRbQ10(Rb, ps.Q10, x, 15.0f0) # ? should 15°C be the reference temperature also an input variable? - - return (; R_soil, Rb), (; st = stQ10) -end diff --git a/src/models/legacy/Rs_components.jl b/src/models/legacy/Rs_components.jl deleted file mode 100644 index aa7df244..00000000 --- a/src/models/legacy/Rs_components.jl +++ /dev/null @@ -1,57 +0,0 @@ -export Rs_components - -""" - Rs_components(NN, predictors, forcing, targets, Q10_het, Q10_root, Q10_myc) - -A linear hybrid model with a neural network `NN`, `predictors`, `targets` and `forcing` terms. -""" -struct Rs_components{D, T1, T2, T3, T4} <: LuxCore.AbstractLuxContainerLayer{(:NN, :predictors, :forcing, :targets, :Q10_het, :Q10_root, :Q10_myc)} - NN - predictors - forcing - targets - Q10_het - Q10_root - Q10_myc - function Rs_components(NN::D, predictors::T1, forcing::T2, targets::T3, Q10_het::T4, Q10_root::T4, Q10_myc::T4) where {D, T1, T2, T3, T4} - return new{D, T1, T2, T3, T4}(NN, collect(predictors), collect(forcing), collect(targets), [Q10_het], [Q10_root], [Q10_myc]) - end -end - -# ? Q10 is a parameter, so expand the initialparameters! -function LuxCore.initialparameters(::AbstractRNG, layer::Rs_components) - ps, _ = LuxCore.setup(Random.default_rng(), layer.NN) - return (; ps, Q10_het = layer.Q10_het, Q10_root = layer.Q10_root, Q10_myc = layer.Q10_myc) -end - -function LuxCore.initialstates(::AbstractRNG, layer::Rs_components) - _, st = LuxCore.setup(Random.default_rng(), layer.NN) - return (; st) -end - -""" - Rs_components(NN, predictors, forcing, targets, Q10)(ds_k) - -# Model definition `ŷ = Rb(αᵢ(t)) * Q10^((T(t) - T_ref)/10)` - -ŷ (respiration rate) is computed as a function of the neural network output `Rb(αᵢ(t))` and the temperature `T(t)` adjusted by the reference temperature `T_ref` (default 15°C) using the Q10 temperature sensitivity factor. -```` -""" -function (hm::Rs_components)(ds_k, ps, st::NamedTuple) - p = ds_k(hm.predictors) - x = Array(ds_k(hm.forcing)) # don't propagate names after this - - out, stRs = LuxCore.apply(hm.NN, p, ps.ps, st.st) - - Rb_het = out[1, :] - Rb_root = out[2, :] - Rb_myc = out[3, :] - - R_het = mRbQ10(Rb_het, ps.Q10_het, x[1, :], 15.0f0) - R_root = mRbQ10(Rb_root, ps.Q10_root, x[1, :], 15.0f0) - R_myc = mRbQ10(Rb_myc, ps.Q10_myc, x[1, :], 15.0f0) - - R_soil = R_het .+ R_root .+ R_myc - - return (; R_soil, R_het, R_root, R_myc), (; st = (; st = stRs)) -end diff --git a/src/models/legacy/simple_Rb_Q10_PBM.jl b/src/models/legacy/simple_Rb_Q10_PBM.jl deleted file mode 100644 index 10e8a0b2..00000000 --- a/src/models/legacy/simple_Rb_Q10_PBM.jl +++ /dev/null @@ -1,44 +0,0 @@ -export RbQ10_2p - -""" - RbQ10_2p(forcing, targets, Q10) -""" -struct RbQ10_2p{T2, T3, T4} <: LuxCore.AbstractLuxContainerLayer{(:forcing, :targets, :Q10, :Rb)} - forcing - targets - Q10 - Rb - function RbQ10_2p(forcing::T2, targets::T3, Q10::T4, Rb::T4) where {T2, T3, T4} - return new{T2, T3, T4}(collect(forcing), collect(targets), [Q10], [Rb]) - end -end - -# ? Q10 is a parameter, so expand the initialparameters! -function LuxCore.initialparameters(::AbstractRNG, layer::RbQ10_2p) - #ps, _ = LuxCore.setup(Random.default_rng(), layer.NN) - return (; Q10 = layer.Q10, Rb = layer.Rb) -end -# TODO: trainable vs non-trainable! set example! -# see: https://lux.csail.mit.edu/stable/manual/migrate_from_flux#Implementing-Custom-Layers -function LuxCore.initialstates(::AbstractRNG, layer::RbQ10_2p) - #_, st = LuxCore.setup(Random.default_rng(), layer.NN) - return NamedTuple() -end - - -""" - RbQ10_2p(NN, predictors, forcing, targets, Q10)(ds_k) - -# Model definition `ŷ = Rb(αᵢ(t)) * Q10^((T(t) - T_ref)/10)` - -ŷ (respiration rate) is computed as a function of the neural network output `Rb(αᵢ(t))` and the temperature `T(t)` adjusted by the reference temperature `T_ref` (default 15°C) using the Q10 temperature sensitivity factor. -```` -""" -function (hm::RbQ10_2p)(ds_k, ps, st) - - x = Array(ds_k(hm.forcing)) # don't propagate names after this - - R_soil = mRbQ10(ps.Rb, ps.Q10, x, 0.0f0) # ? should 15°C be the reference temperature also an input variable? - - return (; R_soil), (; st) -end diff --git a/src/models/models.jl b/src/models/models.jl index a1c89446..b2046a4b 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,7 +1,3 @@ -# include("LinearHM.jl") -# include("Respiration_Rb_Q10.jl") -# include("Rs_components.jl") -# include("simple_Rb_Q10_PBM.jl") include("GenericHybridModel.jl") include("utils.jl") include("helpers_for_HybridModel.jl") diff --git a/test/runtests.jl b/test/runtests.jl index 210088f2..3b10e0fd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,12 +14,3 @@ include("test_show_train.jl") include("test_show_generic_hybrid.jl") include("test_wrap_tuples.jl") include("test_extract_weights.jl") - -@testset "LinearHM" begin - # test model instantiation - NN = Lux.Chain(Lux.Dense(2, 5), Lux.Dense(5, 1)) - lhm = LinearHM(NN, (:x2, :x3), (:x1,), (:obs,), 1.5f0) - @test lhm.forcing == [:x1] - @test lhm.β == [1.5f0] - @test lhm.predictors == [:x2, :x3] -end From e46532f51f4d99cfe32ba3369d544f4146a0806c Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Sat, 4 Jul 2026 10:44:09 +0200 Subject: [PATCH 08/12] keep name --- docs/literate/tutorials/folds.jl | 2 +- docs/literate/tutorials/synthetic_respiration_gpu.jl | 4 ++-- docs/src/get_started.md | 2 +- docs/src/tutorials/exponential_res.md | 2 +- docs/src/tutorials/hyperparameter_tuning.md | 2 +- projects/ExpoHybrid/ExpoHybridEstim.jl | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/literate/tutorials/folds.jl b/docs/literate/tutorials/folds.jl index 52b0c39d..aefc3b32 100644 --- a/docs/literate/tutorials/folds.jl +++ b/docs/literate/tutorials/folds.jl @@ -61,7 +61,7 @@ neural_param_names = [:rb] # ## 6. Construct the Hybrid Model -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, forcing, target, diff --git a/docs/literate/tutorials/synthetic_respiration_gpu.jl b/docs/literate/tutorials/synthetic_respiration_gpu.jl index 4bee0062..88479d78 100644 --- a/docs/literate/tutorials/synthetic_respiration_gpu.jl +++ b/docs/literate/tutorials/synthetic_respiration_gpu.jl @@ -62,7 +62,7 @@ neural_param_names = [:rb] # Neural network predicted parameters # ## Single NN Hybrid Model Training predictors_single_nn = [:sw_pot, :dsw_pot] # Predictor variables (solar radiation, and its derivative) -small_nn_hybrid_model = HybridModel( +small_nn_hybrid_model = constructHybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables @@ -76,7 +76,7 @@ small_nn_hybrid_model = HybridModel( input_batchnorm = true # Apply batch normalization to inputs ) -large_nn_hybrid_model = HybridModel( +large_nn_hybrid_model = constructHybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/src/get_started.md b/docs/src/get_started.md index 27ffebed..0609da04 100644 --- a/docs/src/get_started.md +++ b/docs/src/get_started.md @@ -95,7 +95,7 @@ neural_param_names = [:rb] # Neural network predicted parameters Construct hybrid model ```@example quick_start_complete -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/src/tutorials/exponential_res.md b/docs/src/tutorials/exponential_res.md index 08dbc2a9..637515b2 100644 --- a/docs/src/tutorials/exponential_res.md +++ b/docs/src/tutorials/exponential_res.md @@ -117,7 +117,7 @@ global_param_names = [:k] ``` ```@example expo -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, forcings, targets, diff --git a/docs/src/tutorials/hyperparameter_tuning.md b/docs/src/tutorials/hyperparameter_tuning.md index b777104f..0bfc2505 100644 --- a/docs/src/tutorials/hyperparameter_tuning.md +++ b/docs/src/tutorials/hyperparameter_tuning.md @@ -75,7 +75,7 @@ neural_param_names = [:rb] # Neural network predicted parameters Construct hybrid model ```@example hyperparameter_tuning -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables diff --git a/projects/ExpoHybrid/ExpoHybridEstim.jl b/projects/ExpoHybrid/ExpoHybridEstim.jl index b8bc7931..f7ed6e49 100644 --- a/projects/ExpoHybrid/ExpoHybridEstim.jl +++ b/projects/ExpoHybrid/ExpoHybridEstim.jl @@ -84,7 +84,7 @@ function Expo_resp_model(; T, Resp0, k) return (; Resp_obs, Resp0, k) end -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, forcings, targets, From 413b2303a5031f8effe79b68c3c7df0d791a303a Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Sat, 4 Jul 2026 11:03:09 +0200 Subject: [PATCH 09/12] break lines --- src/models/GenericHybridModel.jl | 70 ++++++++++++++++++++++++++++++-- 1 file changed, 66 insertions(+), 4 deletions(-) diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 9da59ea5..77e0a913 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -127,7 +127,20 @@ function HybridModel( kwargs..., ) - return HybridModel(NN, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) + return HybridModel( + NN, + predictors, + forcing, + targets, + mechanistic_model, + parameters, + neural_param_names, + global_param_names, + fixed_param_names, + scale_nn_outputs, + start_from_default, + config + ) end """ @@ -207,7 +220,20 @@ function HybridModel( kwargs..., ) - return HybridModel(NNs, predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names, fixed_param_names, scale_nn_outputs, start_from_default, config) + return HybridModel( + NNs, + predictors, + forcing, + targets, + mechanistic_model, + parameters, + neural_param_names, + global_param_names, + fixed_param_names, + scale_nn_outputs, + start_from_default, + config + ) end function HybridModel( @@ -236,12 +262,22 @@ function HybridModel( end end -# ─────────────────────────────────────────────────────────────────────────── -# Helper functions for NN initialization +""" + _init_nn_params(rng, m::HybridModel{<:Any, <:NamedTuple}) + +Initialize parameters for a multi-neural network architecture. +Returns a `NamedTuple` containing the initialized parameters for each sub-network. +""" function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) return map(nn -> LuxCore.setup(rng, nn)[1], m.NNs) end +""" + _init_nn_params(rng, m::HybridModel{<:Any, <:Vector}) + +Initialize parameters for a single-neural network architecture. +Returns a `NamedTuple` containing a single `ps` field with the network's parameters. +""" function _init_nn_params(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) ps_nn, _ = LuxCore.setup(rng, m.NNs) return (; ps = ps_nn) @@ -269,10 +305,22 @@ function LuxCore.initialparameters(rng::AbstractRNG, m::HybridModel) return nt end +""" + _init_nn_states(rng, m::HybridModel{<:Any, <:NamedTuple}) + +Initialize states for a multi-neural network architecture. +Returns a `NamedTuple` containing the initialized states for each sub-network. +""" function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:NamedTuple}) return map(nn -> LuxCore.setup(rng, nn)[2], m.NNs) end +""" + _init_nn_states(rng, m::HybridModel{<:Any, <:Vector}) + +Initialize states for a single-neural network architecture. +Returns a `NamedTuple` containing a single `st_nn` field with the network's states. +""" function _init_nn_states(rng::AbstractRNG, m::HybridModel{<:Any, <:Vector}) _, st_nn = LuxCore.setup(rng, m.NNs) return (; st_nn = st_nn) @@ -294,6 +342,13 @@ function LuxCore.initialstates(rng::AbstractRNG, m::HybridModel) return merge(nn_states_nt, (; fixed = nt)) end +""" + _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) + +Execute the forward pass for a multi-neural network architecture. +Applies each sub-network to its specific predictors, and applies scaling to the outputs if required. +Returns scaled parameter values, updated states, and raw network outputs. +""" function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) nn_names = keys(m.NNs) applied = map(nn_names) do nn_name @@ -314,6 +369,13 @@ function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) return scaled_nn_params, nn_states, (; nn_outputs = nn_outputs) end +""" + _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) + +Execute the forward pass for a single-neural network architecture. +Applies the neural network to the given predictors, slices the output for multiple predicted parameters, and scales them if required. +Returns scaled parameter values, updated states, and raw network outputs. +""" function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) if !isempty(m.neural_param_names) nn_out, st_nn = LuxCore.apply(m.NNs, ds_k[1], ps.ps, st.st_nn) From ab0b5ddf5b80f8f9c64ba421f092ff046827862f Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Mon, 6 Jul 2026 13:26:14 +0200 Subject: [PATCH 10/12] function definition --- .../research/synthetic_respiration.jl | 4 +-- docs/literate/tutorials/building_models.jl | 10 +++---- src/models/GenericHybridModel.jl | 22 +++++--------- src/training/tune.jl | 6 ++-- test/test_autodiff_backend.jl | 2 +- test/test_compute_loss.jl | 4 +-- test/test_generic_hybrid_model.jl | 30 +++++++++---------- test/test_show_generic_hybrid.jl | 6 ++-- test/test_split_data_train.jl | 4 +-- 9 files changed, 40 insertions(+), 48 deletions(-) diff --git a/docs/literate/research/synthetic_respiration.jl b/docs/literate/research/synthetic_respiration.jl index 5a84fa97..1f4d6dbe 100644 --- a/docs/literate/research/synthetic_respiration.jl +++ b/docs/literate/research/synthetic_respiration.jl @@ -73,7 +73,7 @@ neural_param_names = [:rb] # Neural network predicted parameters predictors_single_nn = [:sw_pot, :dsw_pot] # Predictor variables (solar radiation, and its derivative) -single_nn_hybrid_model = HybridModel( +single_nn_hybrid_model = constructHybridModel( predictors_single_nn, # Input features forcing, # Forcing variables target, # Target variables @@ -150,7 +150,7 @@ mean(df.sw_pot) predictors_multi_nn = (rb = [:sw_pot, :dsw_pot],) -multi_nn_hybrid_model = HybridModel( +multi_nn_hybrid_model = constructHybridModel( predictors_multi_nn, # Input features forcing, # Forcing variables target, # Target variables diff --git a/docs/literate/tutorials/building_models.jl b/docs/literate/tutorials/building_models.jl index b49f5954..f6af3d3e 100644 --- a/docs/literate/tutorials/building_models.jl +++ b/docs/literate/tutorials/building_models.jl @@ -24,7 +24,7 @@ params_linear = ( # ### HybridModel Construction # We use `x` as forcing data, predict `α` with a neural network based on some predictors `a` and `b`, # and leave `β` as a globally optimized constant parameter. -lhm = HybridModel( +lhm = constructHybridModel( [:a, :b], # predictors for the NN (predicts α) [:x], # forcing variable [:obs], # targets @@ -54,7 +54,7 @@ params_rbq10 = ( ) # ### HybridModel Construction -m_rbq10 = HybridModel( +m_rbq10 = constructHybridModel( [:SWC, :TA], # predictors for Rb [:Temp], # forcing variable [:R_soil], # targets @@ -89,7 +89,7 @@ params_rs_comp = ( ) # ### HybridModel Construction -m_rs_comp = HybridModel( +m_rs_comp = constructHybridModel( [:SWC, :TA], # predictors for all 3 Rb parameters [:Temp], [:R_soil], @@ -127,7 +127,7 @@ predictors_multi = ( Rb = [:SWC, :TA], ) -m_flux = HybridModel( +m_flux = constructHybridModel( predictors_multi, # Triggers Multi-NN construction [:SW_IN, :TA], # Forcing variables [:NEE], # Targets @@ -150,7 +150,7 @@ end # ### HybridModel Construction # Passing an empty `Symbol[]` array to `predictors` prevents any Neural Networks from being created. -m_pbm = HybridModel( +m_pbm = constructHybridModel( Symbol[], # No predictors -> No Neural Network [:Temp], # Forcing [:R_soil], # Target diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 77e0a913..b9da61f6 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -63,7 +63,7 @@ struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} end """ - HybridModel(predictors::Vector{Symbol}, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs...) + constructHybridModel(predictors::Vector{Symbol}, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs...) Construct a `HybridModel` with a single neural network architecture predicting all `neural_param_names` from the `predictors`. @@ -77,7 +77,7 @@ Construct a `HybridModel` with a single neural network architecture predicting a - `global_param_names`: Names of the parameters to be globally optimized. - `kwargs`: Additional configuration like `hidden_layers`, `activation`, `scale_nn_outputs`, etc. """ -function HybridModel( +function constructHybridModel( predictors::Vector{Symbol}, forcing, targets, @@ -144,7 +144,7 @@ function HybridModel( end """ - HybridModel(predictors::NamedTuple, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs...) + constructHybridModel(predictors::NamedTuple, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs...) Construct a `HybridModel` with multiple neural network architectures. A separate neural network is built for each key in the `predictors` NamedTuple. @@ -157,7 +157,7 @@ Construct a `HybridModel` with multiple neural network architectures. A separate - `global_param_names`: Names of the parameters to be globally optimized. - `kwargs`: Additional configuration. `hidden_layers` and `activation` can also be NamedTuples to configure each network independently. """ -function HybridModel( +function constructHybridModel( predictors::NamedTuple, forcing, targets, @@ -236,7 +236,7 @@ function HybridModel( ) end -function HybridModel( +function constructHybridModel( ; predictors, forcing, targets, @@ -248,12 +248,12 @@ function HybridModel( ) if predictors isa Vector{Symbol} @assert neural_param_names !== nothing "Provide neural_param_names for Vector predictors" - return HybridModel( + return constructHybridModel( predictors, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs... ) elseif predictors isa NamedTuple - return HybridModel( + return constructHybridModel( predictors, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs... ) @@ -474,11 +474,3 @@ function (m::HybridModel)(df::DataFrame, ps, st) end return dfnew end - -""" - constructHybridModel(args...; kwargs...) - -Alias for `HybridModel`. Provided for backward compatibility with earlier versions of EasyHybrid.jl. -Users are encouraged to use `HybridModel` directly. -""" -const constructHybridModel = HybridModel diff --git a/src/training/tune.jl b/src/training/tune.jl index e0cfca1d..36de25fc 100644 --- a/src/training/tune.jl +++ b/src/training/tune.jl @@ -26,7 +26,7 @@ Returns a [`TrainResults`](@ref) (or `nothing` if data preparation fails, as in """ function tune(hybrid_model, data, mspec::ModelSpec; kwargs...) kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, (; kwargs...), mspec.hyper_model) - hm = HybridModel(; kwargs_model...) + 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) @@ -35,7 +35,7 @@ end function tune(hybrid_model, data; kwargs...) kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, (; kwargs...)) - hm = HybridModel(; kwargs_model...) + hm = constructHybridModel(; kwargs_model...) kwargs_train = merge((; kwargs...), (; target_names = hm.targets)) train_cfg, data_cfg = EasyHybrid.kwargs_to_configs((), kwargs_train) @@ -44,7 +44,7 @@ end function tune(hybrid_model, data, train_cfg::TrainConfig; data_cfg::DataConfig = DataConfig(), kwargs...) kwargs_model = merge(Base.structdiff(to_namedtuple(hybrid_model), NamedTuple{(:config,)}), hybrid_model.config, to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...)) - hm = HybridModel(; kwargs_model...) + hm = constructHybridModel(; kwargs_model...) 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) diff --git a/test/test_autodiff_backend.jl b/test/test_autodiff_backend.jl index c09a22e1..4eca0bfa 100644 --- a/test/test_autodiff_backend.jl +++ b/test/test_autodiff_backend.jl @@ -11,7 +11,7 @@ using ForwardDiff global_param_names = [:Q10] neural_param_names = [:rb] - model = HybridModel( + model = constructHybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) diff --git a/test/test_compute_loss.jl b/test/test_compute_loss.jl index 5972e480..42bb57ff 100644 --- a/test/test_compute_loss.jl +++ b/test/test_compute_loss.jl @@ -1,6 +1,6 @@ using EasyHybrid: _compute_loss, PerTarget, _apply_loss, loss_fn using EasyHybrid: _get_target_nan, _get_target_y, _loss_name, compute_loss, LoggingLoss -using EasyHybrid: HybridModel, to_keyedArray +using EasyHybrid: HybridModel, constructHybridModel, to_keyedArray using Statistics using DimensionalData using Random @@ -223,7 +223,7 @@ end neural_param_names = [:a] global_param_names = [:b] - HM = HybridModel( + HM = constructHybridModel( predictors, forcing, targets, diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index 851cc842..e7fc696f 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -122,7 +122,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -151,7 +151,7 @@ end neural_param_names = Symbol[] global_param_names = [:a, :b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -174,7 +174,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -200,7 +200,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -230,7 +230,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -264,7 +264,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -296,7 +296,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -324,7 +324,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -346,7 +346,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -371,7 +371,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -397,7 +397,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -432,7 +432,7 @@ end targets = [:obs] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -469,7 +469,7 @@ end neural_param_names = Symbol[] global_param_names = Symbol[] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -505,7 +505,7 @@ end neural_param_names = [:a] global_param_names = [:b] - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, @@ -536,7 +536,7 @@ end custom_chain = Chain(Dense(2, 16, tanh), Dense(16, 8, tanh)) - model = HybridModel( + model = constructHybridModel( predictors, forcing, targets, diff --git a/test/test_show_generic_hybrid.jl b/test/test_show_generic_hybrid.jl index f7233dcb..672ec6ee 100644 --- a/test/test_show_generic_hybrid.jl +++ b/test/test_show_generic_hybrid.jl @@ -1,4 +1,4 @@ -using EasyHybrid: ParameterContainer, HybridModel +using EasyHybrid: ParameterContainer, HybridModel, constructHybridModel using EasyHybrid: _print_field, _print_header, IndentedIO @testset "show_generic.jl" begin @@ -36,7 +36,7 @@ using EasyHybrid: _print_field, _print_header, IndentedIO return (; y_pred = a .* x1 .+ b) end - model = HybridModel( + model = constructHybridModel( [:x1, :x2], [:x3], [:y], test_model, (a = (1.0, 0.0, 2.0), b = (2.0, 1.0, 3.0)), [:a], [:b]; @@ -63,7 +63,7 @@ using EasyHybrid: _print_field, _print_header, IndentedIO return (; obs = a .* x2 .+ d .* x1 .+ b) end - model = HybridModel( + model = constructHybridModel( (a = [:x2, :x3], d = [:x1]), [:x1], [:obs], test_model, (a = (1.0, 0.0, 5.0), b = (2.0, 0.0, 10.0), c = (0.5, 0.0, 2.0), d = (0.5, 0.0, 2.0)), [:b]; # Only global_param_names, neural_param_names derived from predictors keys diff --git a/test/test_split_data_train.jl b/test/test_split_data_train.jl index 572fc745..e5d6d05c 100644 --- a/test/test_split_data_train.jl +++ b/test/test_split_data_train.jl @@ -57,7 +57,7 @@ const RbQ10_PARAMS = ( neural_param_names = [:rb] @testset "test DataFrame and thereby KeyedArray" begin - model = HybridModel( + model = constructHybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) @@ -135,7 +135,7 @@ const RbQ10_PARAMS = ( end @testset "test keep_history" begin - model = HybridModel( + model = constructHybridModel( predictors, forcing, target, RbQ10, RbQ10_PARAMS, neural_param_names, global_param_names ) From 2ca3cd3f2d010f1e0f0d19f4ea15d62d3227c96e Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Mon, 6 Jul 2026 13:45:31 +0200 Subject: [PATCH 11/12] also --- test/test_generic_hybrid_model.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index e7fc696f..228a2ca8 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -562,7 +562,7 @@ end neural_param_names = [:invalid_param] # Not in test_parameters global_param_names = [:b] - @test_throws AssertionError HybridModel( + @test_throws AssertionError constructHybridModel( predictors, forcing, targets, From c179347478a9b102b8ce9c6e70ea17a754173f7d Mon Sep 17 00:00:00 2001 From: Lazaro Alonso Date: Mon, 6 Jul 2026 15:06:15 +0200 Subject: [PATCH 12/12] back --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index f000e0e3..046001e8 100644 --- a/README.md +++ b/README.md @@ -174,7 +174,7 @@ global_param_names = [:Q10] # Global parameters (same for all samples) neural_param_names = [:rb] # Neural network predicted parameters # Construct hybrid model -hybrid_model = HybridModel( +hybrid_model = constructHybridModel( predictors, # Input features forcing, # Forcing variables target, # Target variables