diff --git a/docs/literate/research/synthetic_respiration.jl b/docs/literate/research/synthetic_respiration.jl index 1f4d6dbe..a8cbe52e 100644 --- a/docs/literate/research/synthetic_respiration.jl +++ b/docs/literate/research/synthetic_respiration.jl @@ -110,6 +110,111 @@ single_nn_out = train( model_name = "RbQ10_synthetic1" ) +# ## Learning the observation noise σ (Gaussian NLL) +# +# To fit a Gaussian negative log-likelihood +# ``\tfrac12 \sum_i (r_i/σ_i)^2 + \log σ_i`` we need a *learned* noise scale ``σ``. +# The mechanistic model needs **no changes at all**: `σ` is just declared as a +# parameter. The framework only forwards to the mechanistic model the kwargs it +# declares, so `σ` (which `RbQ10` does not accept) is skipped there, yet it is +# still optimized and exposed to the loss. +# +# Full-context training losses use the signature +# `loss(ŷ, y, y_nan, ps, targets, parameters)` (auto-detected — just pass such a +# function to `training_loss`). The last argument, `parameters`, is `ŷ.parameters`: +# all model parameters (NN-predicted, global and fixed). We read `σ` from there. +# Parameters the mechanistic model doesn't consume (like `σ`) are also surfaced as +# top-level fields of `ŷ`, so they can be listed in `monitor_names` for plots. + +parameters_σ = ( + rb = (3.0f0, 0.0f0, 13.0f0), # Basal respiration [μmol/m²/s] + Q10 = (2.0f0, 1.0f0, 4.0f0), # Temperature sensitivity factor [-] + sigma = (1.0f0, 0.01f0, 5.0f0), # Learned noise scale (bounds keep σ > 0) +) + +# One loss for both cases: `σ` comes from `parameters.sigma` — a single value +# (global parameter) or one value per observation (NN output). Min–max scaling of +# global/NN outputs keeps it positive. +function gaussian_nll(ŷ, y, y_nan, ps, targets, parameters) + total = zero(eltype(ŷ.reco)) + for t in targets + m = y_nan[t] + r = ŷ[t][m] .- y[t][m] + σ = length(parameters.sigma) == 1 ? parameters.sigma[1] : parameters.sigma[m] + total += sum(@. 0.5f0 * (r / σ)^2 + log(σ)) + end + return total +end + +# ### Per-target σ: σ is a global parameter (one value, read from `parameters`) +# We reuse the unchanged `RbQ10` model defined above. + +model_σ_global = constructHybridModel( + predictors_single_nn, + forcing, + target, + RbQ10, + parameters_σ, + [:rb], # NN-predicted parameters + [:Q10, :sigma], # σ is a global (per-target) learned parameter + hidden_layers = [16, 16], + activation = sigmoid, + scale_nn_outputs = true, + input_batchnorm = true +) + +nll_global_out = train( + model_σ_global, + df; + nepochs = 100, + batchsize = 512, + opt = AdamW(0.1), + monitor_names = [:rb, :Q10, :sigma], # σ is surfaced top-level in `ŷ` (and in `ŷ.parameters`) + yscale = identity, + shuffleobs = true, + loss_types = [:mse, :nse], # metrics for logging (still masked f(ŷ, y)) + training_loss = gaussian_nll, # full-context loss, auto-detected + show_progress = false, + model_name = "RbQ10_nll_global" +) + +# ### Per-observation σ: σ is predicted by the neural network (one value per obs) +# +# Only the construction changes — `sigma` moves from the global parameters to the +# NN-predicted ones, so `parameters.sigma` becomes a heteroscedastic +# per-observation vector aligned with the predictions. The model (`RbQ10`) and the +# loss are exactly the same as above. + +model_σ_nn = constructHybridModel( + predictors_single_nn, + forcing, + target, + RbQ10, + parameters_σ, + [:rb, :sigma], # NN predicts both rb and a per-observation σ + [:Q10], # global parameters + hidden_layers = [16, 16], + activation = sigmoid, + scale_nn_outputs = true, + input_batchnorm = true +) + +nll_nn_out = train( + model_σ_nn, + df, + (); + nepochs = 100, + batchsize = 512, + opt = AdamW(0.1), + monitor_names = [:rb, :Q10, :sigma], # σ is surfaced top-level in `ŷ` (and in `ŷ.parameters`) + yscale = identity, + shuffleobs = true, + loss_types = [:mse, :nse], + training_loss = gaussian_nll, + show_progress = false, + model_name = "RbQ10_nll_nn" +) + # ### train on KeyedArray single_nn_out = train( diff --git a/docs/src/tutorials/losses.md b/docs/src/tutorials/losses.md index f445119e..cd667665 100644 --- a/docs/src/tutorials/losses.md +++ b/docs/src/tutorials/losses.md @@ -157,4 +157,51 @@ train(...; training_loss = (fn_args, (args...,), (kwargs...,)), loss_types = [(fn_args, (args...,), (kwargs...,)), :mae, :nse] ) -``` \ No newline at end of file +``` + +### full-context losses (access to the whole `ŷ`, masks, `ps` and `parameters`) + +A 'standard' custom loss `f(ŷ, y)` only receives *one target's masked* predictions +and observations. When the loss needs more — the full prediction NamedTuple, the +NaN masks, the raw parameters, or the model parameters — pass a function with the +6-argument signature `f(ŷ, y, y_nan, ps, targets, parameters)`. It is +**auto-detected** (no wrapper needed) and called with the *full, unmasked* `ŷ` and +`y`, the masks `y_nan`, the raw parameters `ps`, the target names, and +`parameters` (i.e. `ŷ.parameters`: every NN-predicted, global and fixed value). +You do the masking and aggregation yourself and return a scalar. + +The typical use case is a Gaussian negative log-likelihood with a *learned* noise +scale `σ`. You do **not** need the mechanistic model to know about `σ`: just +declare it as a parameter and read `parameters.sigma` in the loss (the mechanistic +model only receives the kwargs it declares, so a loss-only `σ` is skipped there): + +```julia +function gaussian_nll(ŷ, y, y_nan, ps, targets, parameters) + total = zero(eltype(ŷ.reco)) + for t in targets + m = y_nan[t] + r = ŷ[t][m] .- y[t][m] + σ = length(parameters.sigma) == 1 ? parameters.sigma[1] : parameters.sigma[m] + total += sum(@. 0.5f0 * (r / σ)^2 + log(σ)) + end + return total +end + +train(...; training_loss = gaussian_nll) +``` + +The same loss works whether `σ` is declared as a global parameter (one value per +target) or an NN-predicted parameter (one value per observation) — only the +model construction changes. See the synthetic respiration tutorial for both. + +::: warning + +- The 6-argument function is detected only when it has no 2-argument method; + classic `f(ŷ, y)` losses are unaffected. +- Full-context losses are only used for `training_loss`. Entries in `loss_types` + (logging/metrics) still use the masked `f(ŷ_masked, y_masked)` form. +- Because the loss needs the *full* `ŷ`/`parameters`, it cannot use the bare + 2-argument `f(ŷ, y)` signature (that one is per-target and masked). Use the + 6-argument form and simply ignore the arguments you don't need. + +::: \ No newline at end of file diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl index 6caa8922..7d4bd20d 100644 --- a/src/config/TrainingConfig.jl +++ b/src/config/TrainingConfig.jl @@ -60,8 +60,13 @@ loss computation, data handling, output, and visualization. "Set the `cpu_device`, useful for sending back to the cpu model parameters" cdev = cpu_device() - "Loss type to use during training. Default: `:mse`." - training_loss::Symbol = :mse + """ + Loss to use during training. Default: `:mse`. Accepts a predefined metric + (`Symbol`), a custom loss (`Function`) `f(ŷ_masked, y_masked)`, a + `(f, args[, kwargs])` tuple, or a full-context loss `f(ŷ, y, y_nan, ps, targets)` + (auto-detected) that additionally receives the model parameters `ps`. + """ + training_loss = :mse """ Vector of loss types to compute during training. Default: `[:mse, :r2]`. @@ -190,10 +195,14 @@ function validate_config(cfg::TrainConfig) 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 || - @warn "training_loss :$(cfg.training_loss) is not in loss_types $(cfg.loss_types), it won't appear in plots" + # Direction/plotting checks only make sense for predefined (Symbol) metrics. + # Function / tuple / full-context losses are validated at call time instead. + if cfg.training_loss isa Symbol + check_training_loss(cfg.training_loss) # TODO: revisit implementation + cfg.training_loss in cfg.loss_types || + @warn "training_loss :$(cfg.training_loss) is not in loss_types $(cfg.loss_types), it won't appear in plots" + end + return nothing end """ diff --git a/src/losses/compute_loss.jl b/src/losses/compute_loss.jl index 74aa59da..18314a9e 100644 --- a/src/losses/compute_loss.jl +++ b/src/losses/compute_loss.jl @@ -26,7 +26,13 @@ function compute_loss( ext_loss = extra_loss(logging) if logging.train_mode ŷ, st = HM((x, forcings), ps, st) - loss_value = _compute_loss(ŷ, y_t, y_nan, targets, training_loss(logging), logging.agg) + # Full-context losses (auto-detected `f(ŷ, y, y_nan, ps, targets, parameters)`) + # get the predictions, targets, masks, raw params `ps`, target names and the + # model parameters (`ŷ.parameters`), and do their own masking/aggregation; + # everything else uses the per-target machinery. + loss_value = logging.training_loss isa ParamLoss ? + logging.training_loss.f(ŷ, y_t, y_nan, ps, targets, get(ŷ, :parameters, (;))) : + _compute_loss(ŷ, y_t, y_nan, targets, training_loss(logging), logging.agg) # Add extra_loss if provided if ext_loss !== nothing extra_loss_value = ext_loss(ŷ, ps) diff --git a/src/losses/compute_loss_types.jl b/src/losses/compute_loss_types.jl index 8245d56e..83208663 100644 --- a/src/losses/compute_loss_types.jl +++ b/src/losses/compute_loss_types.jl @@ -25,6 +25,25 @@ struct ExtraLoss <: LossSpec f::Union{Function, Nothing} end +""" + ParamLoss(f) + +Wrapper for a *full-context* training loss with signature +`f(ŷ, y, y_nan, ps, targets, parameters)`. Unlike a classic custom loss +`f(ŷ_masked, y_masked)`, it receives the full (unmasked) predictions `ŷ`, targets +`y`, NaN masks `y_nan`, the raw parameters `ps`, the `targets` names, and the model +`parameters` (i.e. `ŷ.parameters`: NN-predicted, global and fixed values), and must +return a scalar loss. + +Users normally do not construct this directly: passing a function with this +6-argument signature to `training_loss` auto-detects it (see [`_accepts_params`](@ref)). +It is the natural choice when the loss needs a learned quantity such as a noise +scale `parameters.sigma`, e.g. a Gaussian negative log-likelihood. +""" +struct ParamLoss <: LossSpec + f::Function +end + """ PerTarget(losses) @@ -111,9 +130,21 @@ end _to_loss_spec(s::Symbol) = SymbolicLoss(s) -_to_loss_spec(f::Function) = FunctionLoss(f) +_to_loss_spec(f::Function) = _accepts_params(f) ? ParamLoss(f) : FunctionLoss(f) _to_loss_spec(ls::LossSpec) = ls +""" + _accepts_params(f) -> Bool + +Auto-detect whether a custom loss `f` uses the full-context signature +`f(ŷ, y, y_nan, ps, targets, parameters)` (6 positional args) instead of the +classic masked signature `f(ŷ_masked, y_masked)` (2 positional args). Returns +`true` only when `f` has a 6-argument method and no 2-argument method, so existing +2-arg losses keep their behavior. Functions that are ambiguous (e.g. varargs +matching both arities) fall back to the classic 2-arg path. +""" +_accepts_params(f) = hasmethod(f, NTuple{6, Any}) && !hasmethod(f, NTuple{2, Any}) + _to_loss_spec(t::Tuple{<:Function, <:Tuple}) = ParameterizedLoss(t[1], t[2]) _to_loss_spec(t::Tuple{<:Function, <:NamedTuple}) = ParameterizedLoss(t[1], (), t[2]) _to_loss_spec(t::Tuple{<:Function, <:Tuple, <:NamedTuple}) = ParameterizedLoss(t[1], t[2], t[3]) @@ -137,6 +168,7 @@ loss_spec(ls::ParameterizedLoss) = (ls.f, ls.args, ls.kwargs) loss_spec(el::ExtraLoss) = el.f +loss_spec(pl::ParamLoss) = pl.f loss_spec(pt::PerTarget) = PerTarget(map(loss_spec, pt.losses)) loss_types(logging::LoggingLoss) = map(loss_spec, logging.loss_types) diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index b9da61f6..8e96a8cb 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -437,10 +437,17 @@ function (m::HybridModel)(ds_k::Tuple, ps, st) forcing_data = ds_k[2] all_kwargs = merge(forcing_data, all_params) - # 6) Apply mechanistic model - y_pred = m.mechanistic_model(; all_kwargs...) - - out = (; y_pred..., parameters = all_params, out_extra...) + # 6) Apply mechanistic model. Only forward the kwargs it actually declares, so + # "loss-only" parameters (e.g. a learned noise scale used only in the loss) + # can be defined without the mechanistic model having to accept them. They + # still live in `all_params` and are exposed below under `parameters`. + y_pred = m.mechanistic_model(; _mechanistic_kwargs(m.mechanistic_model, all_kwargs)...) + + # Parameters the mechanistic model does not consume (e.g. loss-only ones such as + # a learned noise scale) are surfaced at the top level so they can be monitored + # and plotted, in addition to always being available under `parameters`. + extra_params = _extra_params(m.mechanistic_model, all_params) + out = (; y_pred..., extra_params..., parameters = all_params, out_extra...) st_new = (; st_new_nns..., fixed = st.fixed) return out, st_new @@ -451,6 +458,52 @@ function (m::HybridModel)(ds_k, ps, st) return m(Tuple(ds_k), ps, st) end +""" + _mechanistic_kwargs(f, all_kwargs::NamedTuple) + +Select from `all_kwargs` only the keyword arguments the mechanistic model `f` +declares, so parameters used solely by the loss (e.g. a learned noise scale) do +not need to be accepted by `f`. Falls back to passing everything when `f` slurps +`kwargs...` or its keyword signature cannot be introspected. +""" +function _mechanistic_kwargs(f, all_kwargs::NamedTuple) + keep = ChainRulesCore.ignore_derivatives() do + _accepted_kwarg_names(f, keys(all_kwargs)) + end + keep === nothing && return all_kwargs + return NamedTuple{keep}(map(k -> all_kwargs[k], keep)) +end + +""" + _extra_params(f, all_params::NamedTuple) + +The parameters not consumed by the mechanistic model `f` (e.g. loss-only ones such +as a learned noise scale). They are surfaced at the top level of the model output +so they can be monitored/plotted, in addition to always being available under +`parameters`. Returns an empty `NamedTuple` when `f` consumes everything. +""" +function _extra_params(f, all_params::NamedTuple) + keep = ChainRulesCore.ignore_derivatives() do + acc = _accepted_kwarg_names(f, keys(all_params)) + acc === nothing ? () : Tuple(k for k in keys(all_params) if !(k in acc)) + end + return NamedTuple{keep}(map(k -> all_params[k], keep)) +end + +# Returns the tuple of `all_kwargs` names accepted by `f`, or `nothing` to signal +# "pass everything" (the model slurps `kwargs...`, or has no introspectable kwargs). +function _accepted_kwarg_names(f, available::Tuple) + names = Symbol[] + for mth in methods(f) + for d in Base.kwarg_decl(mth) + endswith(string(d), "...") && return nothing # slurps kwargs → keep all + push!(names, d) + end + end + isempty(names) && return nothing + return Tuple(k for k in available if k in names) +end + function (m::HybridModel)(df::DataFrame, ps, st) @warn "Only makes sense in test mode, not training!"