refactor generic hybrid

.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

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!

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",

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

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]

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

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.