Aleatoric Uncertainty with noise parameter sigma

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(

docs/src/tutorials/losses.md

@@ -157,4 +157,51 @@ train(...;
     training_loss = (fn_args, (args...,), (kwargs...,)),
     loss_types = [(fn_args, (args...,), (kwargs...,)), :mae, :nse]
     )
-```
+```
+
+### 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.
+
+:::

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

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)

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)

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