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..f6af3d3e --- /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 = constructHybridModel( + [: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 = constructHybridModel( + [: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 = constructHybridModel( + [: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 = constructHybridModel( + 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 = constructHybridModel( + 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/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/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/data/prepare_data.jl b/src/data/prepare_data.jl index 00fe8e79..0c5bb080 100644 --- a/src/data/prepare_data.jl +++ b/src/data/prepare_data.jl @@ -3,27 +3,22 @@ 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)]) - 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) + 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, forcings_nt), targets_nt) end 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/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/FluxPartModel_Q10_Lux.jl b/src/models/FluxPartModel_Q10_Lux.jl deleted file mode 100644 index 89b12d7e..00000000 --- a/src/models/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/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 5bf56e0c..b9da61f6 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -1,26 +1,15 @@ -export SingleNNHybridModel, MultiNNHybridModel, constructHybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid -export HybridParams +export HybridModel, ParameterContainer, constructHybridModel -# 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 +""" + ParameterContainer{NT <: NamedTuple, T} -mutable struct ParameterContainer{NT <: NamedTuple, T} <: AbstractHybridModel +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} @@ -30,62 +19,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 +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 -# ─────────────────────────────────────────────────────────────────────────── -# 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 + "Predictor variables for the neural networks. Can be a `Vector{Symbol}` or a `NamedTuple`." + predictors::P -# 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 + "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 +""" + 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`. + +# 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 constructHybridModel( predictors::Vector{Symbol}, forcing, @@ -102,8 +93,8 @@ function constructHybridModel( kwargs... ) - if !isa(parameters, AbstractHybridModel) - parameters = build_parameters(parameters, mechanistic_model) + if !isa(parameters, ParameterContainer) + parameters = ParameterContainer(parameters) end all_names = pnames(parameters) @@ -136,9 +127,36 @@ 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 +""" + 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. + +# 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 constructHybridModel( predictors::NamedTuple, forcing, @@ -154,8 +172,8 @@ function constructHybridModel( kwargs... ) - if !isa(parameters, AbstractHybridModel) - parameters = build_parameters(parameters, mechanistic_model) + if !isa(parameters, ParameterContainer) + parameters = ParameterContainer(parameters) end all_names = pnames(parameters) @@ -202,7 +220,20 @@ 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( @@ -231,41 +262,30 @@ function constructHybridModel( end 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 +""" + _init_nn_params(rng, m::HybridModel{<:Any, <:NamedTuple}) - return nt +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 -# Initial parameters for MultiNNHybridModel -function LuxCore.initialparameters(rng::AbstractRNG, m::MultiNNHybridModel) - # Setup parameters for each neural network - 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 +""" + _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) +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 +305,30 @@ 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 +""" + _init_nn_states(rng, m::HybridModel{<:Any, <:NamedTuple}) - nt = (; st_nn = st_nn, fixed = nt) - return nt +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 -# Initial states for MultiNNHybridModel -function LuxCore.initialstates(rng::AbstractRNG, m::MultiNNHybridModel) - # Setup states for each neural network - 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 +""" + _init_nn_states(rng, m::HybridModel{<:Any, <:Vector}) - # Start with the NN states +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) +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,78 +339,52 @@ function LuxCore.initialstates(rng::AbstractRNG, m::MultiNNHybridModel) end end - nt = (; nn_states..., fixed = nt) - return nt -end - -function default(p::AbstractHybridModel) - return p.hybrid.table[:, :default] -end - -function lower(p::AbstractHybridModel) - return p.hybrid.table[:, :lower] + return merge(nn_states_nt, (; fixed = nt)) 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) + _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) -Scale a single parameter using the sigmoid scaling function. +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 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)) +function _run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st) + nn_names = keys(m.NNs) + 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)) -""" - scale_single_param_minmax(name, hm::AbstractHybridModel) + 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) -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 .- ℓ)) + return scaled_nn_params, nn_states, (; nn_outputs = nn_outputs) end +""" + _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) -# ─────────────────────────────────────────────────────────────────────────── -# 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) - - parameters = m.parameters - - # 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() - end - - # 3) scale NN parameters (handle empty case) +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.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 +395,20 @@ 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 -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 + return scaled_nn_params, (; st_nn = st_nn), (;) end -# Forward pass for MultiNNHybridModel (optimized, no branching) -function (m::MultiNNHybridModel)(ds_k::Tuple, ps, st) +""" + (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 - # 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 +419,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() - - 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,)) + # 2) Run neural network(s) + scaled_nn_params, st_new_nns, out_extra = _run_nn(m, ds_k, ps, st) - # 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 +430,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 +463,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/LinearHM.jl b/src/models/LinearHM.jl deleted file mode 100644 index c07c9617..00000000 --- a/src/models/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/Respiration_Rb_Q10.jl b/src/models/Respiration_Rb_Q10.jl deleted file mode 100644 index d9be3825..00000000 --- a/src/models/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/Rs_components.jl b/src/models/Rs_components.jl deleted file mode 100644 index aa7df244..00000000 --- a/src/models/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/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/models.jl b/src/models/models.jl index 328ed5f4..b2046a4b 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,8 +1,5 @@ -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 49e5308b..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 @@ -100,62 +87,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 +148,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/models/simple_Rb_Q10_PBM.jl b/src/models/simple_Rb_Q10_PBM.jl deleted file mode 100644 index 10e8a0b2..00000000 --- a/src/models/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/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/src/training/tune.jl b/src/training/tune.jl index 45617066..36de25fc 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 = constructHybridModel(; 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 = constructHybridModel(; 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 = constructHybridModel(; kwargs_model...) kwargs_train = merge(to_namedtuple(train_cfg), to_namedtuple(data_cfg), (; kwargs...), (; target_names = hm.targets)) 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 diff --git a/test/test_compute_loss.jl b/test/test_compute_loss.jl index 81a7f64f..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: constructHybridModel, to_keyedArray +using EasyHybrid: HybridModel, constructHybridModel, to_keyedArray using Statistics using DimensionalData using Random diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index 7e19e068..228a2ca8 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -34,11 +34,6 @@ 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,65 +64,58 @@ 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 -@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] @@ -144,7 +132,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,10 +141,10 @@ 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] @@ -173,13 +161,13 @@ 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] @@ -205,7 +193,7 @@ end @test ps.b[1] isa Float32 end - @testset "SingleNNHybridModel initialstates" begin + @testset "HybridModel initialstates" begin predictors = [:x2, :x3] forcing = [:x1] targets = [:obs] @@ -235,7 +223,7 @@ 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] @@ -269,7 +257,7 @@ 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] @@ -301,8 +289,8 @@ 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] @@ -317,7 +305,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,7 +318,7 @@ 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] @@ -347,12 +335,12 @@ 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] @@ -377,7 +365,7 @@ 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] @@ -403,7 +391,7 @@ 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] @@ -438,7 +426,7 @@ 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] @@ -491,7 +479,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] @@ -559,10 +547,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 diff --git a/test/test_show_generic_hybrid.jl b/test/test_show_generic_hybrid.jl index 14fca091..672ec6ee 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: ParameterContainer, HybridModel, constructHybridModel 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)) @@ -44,7 +31,7 @@ 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 @@ -71,7 +58,7 @@ 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 diff --git a/test/test_split_data_train.jl b/test/test_split_data_train.jl index a7176355..e5d6d05c 100644 --- a/test/test_split_data_train.jl +++ b/test/test_split_data_train.jl @@ -61,7 +61,7 @@ const RbQ10_PARAMS = ( 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) @@ -139,7 +139,7 @@ const RbQ10_PARAMS = ( 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)