Optim algorithms like BFGS, multi-LSTM hybrids, yaml with bounds and PBM, regularization - L2, per parameter optimisers

Project.toml

@@ -23,6 +23,7 @@ MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
@@ -63,6 +64,7 @@ MLUtils = "0.4.8"
 Makie = "0.22, 0.23, 0.24"
 NCDatasets = "0.14.8"
 NamedDims = "1.2.3"
+Optimization = "4, 5"
 OptimizationOptimisers = "0.3.7"
 OrderedCollections = "1.8.1"
 PrettyTables = "2.4.0, 3.1.2"

docs/literate/research/synthetic_respiration.jl

@@ -89,7 +89,7 @@ single_nn_hybrid_model = constructHybridModel(
 
 # ### train on DataFrame
 
-extra_loss = function (ŷ)
+function rb_prior_extra_loss(ŷ, ps)
     return (; a = sum((5.0 .- ŷ.rb) .^ 2))
 end
 
@@ -106,7 +106,7 @@ single_nn_out = train(
     shuffleobs = true,
     loss_types = [:mse, :nse],
     show_progress = false,
-    extra_loss = extra_loss,
+    extra_loss = rb_prior_extra_loss,
     model_name = "RbQ10_synthetic1"
 )
 

src/EasyHybrid.jl

@@ -32,6 +32,7 @@ using Lux: Lux
 using MLJ: partition
 using MLUtils: MLUtils, DataLoader, kfolds, numobs, rpad, splitobs
 using NCDatasets: NCDatasets, NCDataset, close, name
+using Optimization: Optimization, OptimizationFunction, OptimizationProblem, solve, remake
 using OptimizationOptimisers: OptimizationOptimisers, AdamW, Adam, Optimisers
 using OrderedCollections: OrderedDict
 using PrettyTables: PrettyTables
@@ -54,15 +55,16 @@ using Static: False, True
     using DataFrames
     using CSV
     using AxisKeys: axiskeys
+    using Optimization: Optimization, OptimizationFunction, OptimizationProblem, solve, remake
     using OptimizationOptimisers: OptimizationOptimisers, Optimisers, Adam, AdamW, RMSProp
     using ComponentArrays: ComponentArrays, ComponentArray
 end
 
 abstract type EasyHybridModels end
 
-include("config/config.jl")
 include("utils/utils.jl")
 include("models/models.jl")
+include("config/config.jl")
 include("data/data.jl")
 include("losses/losses.jl")
 include("training/training.jl")

src/config/TrainingConfig.jl

@@ -13,7 +13,33 @@ loss computation, data handling, output, and visualization.
     "Size of the training batches. Default: 64."
     batchsize::Int = 64
 
-    "Optimizer to use for training. Default: Adam(0.01)."
+    """
+    Optimizer to use for training. Default: `Adam(0.01)`.
+
+    On the `Optimisers.jl` path (selected when `opt` originates from
+    `Optimisers.jl`) three forms are accepted:
+
+    1. A single `Optimisers.AbstractRule` (default), applied to the whole
+       parameter tree, e.g. `Adam(0.01)`. `ps` is wrapped in a
+       `ComponentArray`.
+    2. A `NamedTuple` of rules — one per top-level branch of the parameter
+       tree — e.g. for an `RbQ10` hybrid model:
+
+       ```julia
+       opt = (; Rb = Adam(1e-3), Q10 = Descent(1e-2))
+       ```
+
+       The framework calls `Optimisers.setup` per branch; branches not
+       listed fall back to `Adam()`. `ps` is kept as a nested `NamedTuple`
+       (no `ComponentArray` wrap) because `Optimisers.jl` treats a
+       `ComponentArray` as a single leaf and would collapse the per-branch
+       state tree.
+    3. A `NamedTuple` of pre-built state trees (already returned by
+       `Optimisers.setup`). Useful when one branch needs an
+       `Optimisers.OptimiserChain` or a frozen leaf built up by hand.
+
+    Forms 2 and 3 can be freely mixed in the same `NamedTuple`.
+    """
     opt = Adam(0.01)
 
     """
@@ -44,7 +70,7 @@ loss computation, data handling, output, and visualization.
     """
     loss_types::Vector{Symbol} = [:mse, :r2]
 
-    "Additional loss term to add to the training loss. Default: `nothing`."
+    "Additional loss `(ŷ, ps; kwargs...) -> NamedTuple` added to the training loss. Default: `nothing`."
     extra_loss = nothing
 
     "Aggregation function applied to computed losses. Default: `sum`."
@@ -97,6 +123,40 @@ loss computation, data handling, output, and visualization.
 
     "Tuple of parameter names to track across epochs. Default: `()`."
     tracked_params::Tuple = ()
+
+    """
+    `Optimization.jl` path only: pass the full training set as a single tuple
+    to `OptimizationProblem` instead of an `MLUtils.DataLoader`. Has no effect
+    on the `Optimisers.jl` path. Default: `false`.
+    """
+    full_batch::Bool = false
+
+    """
+    `Optimization.jl` path only: promote `ps` to `Float64` before optimization.
+    Mirrors the workaround documented in `projects/RbQ10/Q10_lbfgs.jl` for
+    [Lux.jl#1260](https://github.com/LuxDL/Lux.jl/issues/1260). Has no effect
+    on the `Optimisers.jl` path. Default: `false`.
+    """
+    promote_f64::Bool = false
+
+    """
+    `Optimization.jl` full-batch path only (`full_batch = true`): build a
+    validation `EpochSnapshot` (driving history / early-stopping / dashboard)
+    every `eval_every` solver iterations. Default: `1`.
+    """
+    eval_every::Int = 1
+
+    """
+    `Optimization.jl` minibatch path only (`full_batch = false`): number of
+    optimizer iterations to run on each *fixed* minibatch before resampling the
+    next one. Implements the "several L-BFGS steps per minibatch, then resample"
+    scheme of Le et al., 2011 (*On Optimization Methods for Deep Learning*, ICML,
+    §4.2): holding the minibatch fixed for a few iterations keeps the objective
+    (and L-BFGS curvature pairs / line search) consistent. `nepochs` controls the
+    number of outer resampling passes; `batchsize` the minibatch size. Setting
+    `inner_maxiters = 1` reduces to one optimizer step per minibatch. Default: `4`.
+    """
+    inner_maxiters::Int = 4
 end
 
 function validate_config(cfg::TrainConfig)
@@ -112,6 +172,12 @@ function validate_config(cfg::TrainConfig)
     cfg.patience > 0 ||
         throw(ArgumentError("patience must be positive, got $(cfg.patience)"))
 
+    cfg.eval_every > 0 ||
+        throw(ArgumentError("eval_every must be positive, got $(cfg.eval_every)"))
+
+    cfg.inner_maxiters > 0 ||
+        throw(ArgumentError("inner_maxiters must be positive, got $(cfg.inner_maxiters)"))
+
     check_training_loss(cfg.training_loss) # TODO: revisit implementation
 
     return cfg.training_loss in cfg.loss_types ||

src/config/config_yaml.jl

@@ -1,5 +1,5 @@
 export load_hybrid_config, save_hybrid_config
-export get_hybrid_config, get_train_config, get_full_config
+export get_hybrid_config, get_train_config, get_full_config, get_parameters_config, get_mechanistic_model_config
 
 function load_hybrid_config(path::String; dicttype = OrderedDict{String, Any})
     return load_file(path; dicttype)
@@ -12,11 +12,108 @@ end
 function get_hybrid_config(hm::LuxCore.AbstractLuxContainerLayer)
     hm_config = OrderedDict{String, Any}()
     for field in fieldnames(typeof(hm))
-        hm_config[string(field)] = getfield(hm, field)
+        hm_config[string(field)] = _yaml_field_value(getfield(hm, field))
     end
     return hm_config
 end
 
+_yaml_field_value(x) = x
+_yaml_field_value(p::AbstractHybridModel) = get_parameters_config(p)
+_yaml_field_value(f::Function) = get_mechanistic_model_config(f)
+
+"""
+    get_parameters_config(p::AbstractHybridModel)
+
+Serialize the parameter table (`default`, `lower`, `upper` per parameter) of a
+`HybridParams`/`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
+drops all of the actual default and bound values.
+"""
+function get_parameters_config(p::AbstractHybridModel)
+    pc = hasfield(typeof(p), :hybrid) ? p.hybrid : p
+    out = OrderedDict{String, Any}()
+    for name in keys(pc.values)
+        d, l, u = pc.values[name]
+        out[string(name)] = OrderedDict{String, Any}(
+            "default" => d,
+            "lower" => l,
+            "upper" => u,
+        )
+    end
+    return out
+end
+
+"""
+    get_mechanistic_model_config(f::Function)
+
+Build an `OrderedDict` describing a function for YAML output: its `name`, and
+for each `Method`, its source `file`, starting `line`, and the full `source`
+text extracted from disk by parsing one complete expression starting at that
+line. The single-method case is flattened so the YAML stays compact.
+
+Used to record `mechanistic_model` in the saved config so the exact function
+definition (not just its name) is preserved alongside the run.
+"""
+function get_mechanistic_model_config(f::Function)
+    out = OrderedDict{String, Any}()
+    out["name"] = string(nameof(f))
+    method_entries = OrderedDict{String, Any}[]
+    for m in methods(f)
+        entry = OrderedDict{String, Any}()
+        try
+            file, line = Base.functionloc(m)
+            entry["file"] = string(file)
+            entry["line"] = line
+            src = _try_extract_function_source(string(file), line)
+            if src !== nothing
+                entry["source"] = src
+            end
+        catch err
+            entry["error"] = sprint(showerror, err)
+        end
+        push!(method_entries, entry)
+    end
+    if length(method_entries) == 1
+        merge!(out, method_entries[1])
+    elseif !isempty(method_entries)
+        out["methods"] = method_entries
+    end
+    return out
+end
+
+# Read `file` and return the source text of the first complete top-level
+# expression starting at `line`. Returns `nothing` if the file can't be read or
+# no expression can be parsed (e.g. function defined at the REPL or inside an
+# eval'd string).
+function _try_extract_function_source(file::AbstractString, line::Integer)
+    isfile(file) || return nothing
+    text = try
+        read(file, String)
+    catch
+        return nothing
+    end
+    idx = firstindex(text)
+    cur = 1
+    while cur < line
+        nl = findnext(==('\n'), text, idx)
+        nl === nothing && return nothing
+        idx = nextind(text, nl)
+        cur += 1
+    end
+    expr_and_next = try
+        Meta.parse(text, idx; greedy = true, raise = false)
+    catch
+        return nothing
+    end
+    expr_and_next === nothing && return nothing
+    _, next_idx = expr_and_next
+    last_idx = lastindex(text)
+    endidx = next_idx > last_idx ? last_idx : prevind(text, next_idx)
+    endidx < idx && return nothing
+    return rstrip(text[idx:endidx])
+end
+
 function get_train_config(train_args::NamedTuple)
     train_config = OrderedDict{String, Any}()
     for field in fieldnames(typeof(train_args))

src/data/sequences.jl

@@ -34,6 +34,37 @@ 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
+# `(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)
+    x, forcings = x_tuple
+    valid = _valid_seq_indices(x, y)
+    nseq = _nseq(x)
+    length(valid) < nseq && @info "Dropped $(nseq - length(valid)) / $nseq sequences with NaN predictors or all-NaN targets"
+    x_filtered = map(b -> b[:, :, valid], x)
+    forcings_filtered = map(v -> v[:, valid], forcings)
+    return (x_filtered, forcings_filtered), _subset_y(y, valid)
+end
+
+_nseq(x::NamedTuple) = size(first(values(x)), 3)
+
+_branches_finite(x::NamedTuple, ii) = all(b -> !any(isnan, @view(b[:, :, ii])), values(x))
+_targets_present(y, ii) = any(!isnan, @view(y[:, :, ii]))
+_targets_present(y::NamedTuple, ii) = any(v -> any(!isnan, @view(v[:, ii])), values(y))
+
+# `(x::NamedTuple, y::NamedTuple)` is defined explicitly to disambiguate against
+# `_valid_seq_indices(x, y::NamedTuple)` above.
+function _valid_seq_indices(x::NamedTuple, y::NamedTuple)
+    n = _nseq(x)
+    return findall(ii -> _branches_finite(x, ii) && _targets_present(y, ii), 1:n)
+end
+
+function _valid_seq_indices(x::NamedTuple, y)
+    n = _nseq(x)
+    return findall(ii -> _branches_finite(x, ii) && _targets_present(y, ii), 1:n)
+end
+
 """
     split_into_sequences(x, y; input_window=5, output_window=1, output_shift=1, lead_time=1)
 
@@ -68,6 +99,44 @@ function split_into_sequences(x::Tuple, y; kwargs...)
     return (X_seq, forcings_seq), Y_seq
 end
 
+# MultiNNHybridModel support.
+#
+# For `MultiNNHybridModel`, `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
+# pass can index `ds_k[1][nn_name]`.
+function split_into_sequences(x::Tuple{<:NamedTuple, <:Any}, y::NamedTuple; kwargs...)
+    x_branches, forcings = x
+    isempty(x_branches) && throw(ArgumentError("predictor NamedTuple is empty"))
+    targetkeys = keys(y)
+
+    # All branches share the same time axis; use the first as the reference for
+    # windowing the (shared) targets and forcings.
+    ref = _as_keyed_time(first(values(x_branches)))
+    y_arr = _nt_to_array(y, ref)
+
+    X_seq = map(b -> first(split_into_sequences(_as_keyed_time(b), y_arr; kwargs...)), x_branches)
+    _, Y_seq = split_into_sequences(ref, y_arr; kwargs...)
+    forcings_seq = _window_forcings(forcings, ref; kwargs...)
+
+    return (X_seq, forcings_seq), _array_to_nt(Y_seq, targetkeys)
+end
+
+# Predictor branches built for `MultiNNHybridModel` 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.
+function _as_keyed_time(b)
+    (b isa KeyedArray || b isa AbstractDimArray) && return b
+    ndims(b) == 2 || throw(ArgumentError("expected a 2D (feature, time) predictor branch; got ndims = $(ndims(b))"))
+    return KeyedArray(
+        Float32.(b);
+        variable = Symbol.("f", 1:size(b, 1)),
+        time = 1:size(b, 2),
+    )
+end
+
 function split_into_sequences(x, y::NamedTuple; kwargs...)
     targetkeys = keys(y)
     y_arr = _nt_to_array(y, x)

src/data/split_data.jl

@@ -168,6 +168,8 @@ function collect_end_dim(x_all::Union{KeyedArray{Float32, 3}, AbstractDimArray{F
     return collect(getindex(x_all, :, :, idx))
 end
 
+# 2D (feature, time): split along time; 3D (feature, time, batch): split along batch
+_num_samples(x::AbstractArray{<:Any, 3}) = size(x, 3)
 _num_samples(x::AbstractArray) = size(x, 2)
 _num_samples(x::NamedTuple) = _num_samples(first(values(x)))
 

src/io/checkpoints.jl

@@ -20,7 +20,7 @@ function save_final!(paths::TrainingPaths, model, ps, st, x_train, forcings_trai
     if cfg.save_training
         target_names = model.targets
         save_epoch = stopper.best_epoch == 0 ? 0 : stopper.best_epoch
-        save_ps_st!(paths.best_model, model, cfg.cdev(ps), cfg.cdev(st), cfg.tracked_params, save_epoch)
+        save_ps_st(paths.best_model, model, cfg.cdev(ps), cfg.cdev(st), cfg.tracked_params, save_epoch)
 
         ŷ_train, αst_train = model((cfg.cdev(x_train), cfg.cdev(forcings_train)), cfg.cdev(ps), LuxCore.testmode(cfg.cdev(st)))
         ŷ_val, αst_val = model((cfg.cdev(x_val), cfg.cdev(forcings_val)), cfg.cdev(ps), LuxCore.testmode(cfg.cdev(st)))

src/io/save.jl

@@ -1,6 +1,6 @@
 export get_all_groups
 export load_group
-function save_ps_st(file_name, hm, ps, st, save_ps)
+function save_ps_st(file_name, hm, ps, st, save_ps, epoch = 0)
     hm_name = string(nameof(typeof(hm)))
     # split physical parameters
     tmp_e = if !isempty(save_ps)
@@ -10,9 +10,9 @@ function save_ps_st(file_name, hm, ps, st, save_ps)
 
     isfile(file_name) && rm(file_name)
     return jldopen(file_name, "w") do file
-        file["HybridModel_$hm_name/epoch_0"] = (ps, st)
+        file["HybridModel_$hm_name/epoch_$epoch"] = (ps, st)
         if !isempty(save_ps)
-            file["physical_params/epoch_0"] = tmp_e
+            file["physical_params/epoch_$epoch"] = tmp_e
         end
     end
 end