Aleatoric Uncertainty with noise parameter sigma

[BernhardAhrens]

See Kendall, A., & Gal, Y. (2018). What uncertainty do we need in deep learning? doi:10.48550/arxiv.1703.04977

Equation 5 for heteroscedastic aleatoric uncertainty. sigma(x_i) is estimated with a neural network. For the homoscedastic case sigma is estimated per datastream as a global parameter or fixed a priori

[github-actions]

📚 Documentation preview 🚀

Preview URL: https://EarthyScience.github.io/EasyHybrid.jl/previews/PR284/

Note

The preview will be available once the documentation build completes successfully, and will reflect the last successful build for this PR.

[gemini-code-assist]

Code Review

This pull request introduces support for full-context training losses (such as Gaussian negative log-likelihood with a learned noise scale) by allowing custom loss functions to receive the full unmasked predictions, targets, masks, raw parameters, target names, and model parameters. It also updates the hybrid model to only forward declared keyword arguments to the mechanistic model, enabling loss-only parameters to be optimized without mechanistic model changes. The review feedback highlights three important issues: a guaranteed MethodError when converting a Symbol to a String using String(d) instead of string(d), a limitation in the auto-detection of 6-argument functions when arguments are type-annotated, and a potential MethodError when calling get on non-NamedTuple containers like ComponentArray.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In Julia, converting a Symbol to a String using String(d) throws a MethodError (e.g., MethodError: no method matching String(::Symbol)). The standard and correct way to convert a Symbol to a String is using string(d).

        for d in Base.kwarg_decl(mth)
            endswith(string(d), "...") && return nothing  # slurps kwargs → keep all
            push!(names, d)
        end

[gemini-code-assist]

Using hasmethod(f, NTuple{6, Any}) to detect the 6-argument signature will return false if the user has type-annotated any of the arguments in their custom loss function (e.g., gaussian_nll(ŷ::NamedTuple, ...)), because Tuple{Any, Any, Any, Any, Any, Any} is not a subtype of the annotated signature.

A more robust approach that supports type-annotated arguments is to inspect the methods of f directly and check if any method has exactly 6 positional arguments (which corresponds to 7 parameters in the method signature, including the function itself):

_accepts_params(f) = any(m -> length(Base.unwrap_unionall(m.sig).parameters) == 7, methods(f)) &&
                     !any(m -> length(Base.unwrap_unionall(m.sig).parameters) == 3, methods(f))

_accepts_params(f) = any(m -> length(Base.unwrap_unionall(m.sig).parameters) == 7, methods(f)) && !any(m -> length(Base.unwrap_unionall(m.sig).parameters) == 3, methods(f))

[gemini-code-assist]

Using get(ŷ, :parameters, (;)) can throw a MethodError if ŷ is not a standard NamedTuple (for example, if it is a ComponentArray or another custom container that does not support get with a Symbol key).

Using hasproperty and property access is much more robust and widely supported across different container types (including ComponentArray):

parameters = hasproperty(ŷ, :parameters) ? ŷ.parameters : (;)

        loss_value = logging.training_loss isa ParamLoss ?
            logging.training_loss.f(ŷ, y_t, y_nan, ps, targets, hasproperty(ŷ, :parameters) ? ŷ.parameters : (;)) :
            _compute_loss(ŷ, y_t, y_nan, targets, training_loss(logging), logging.agg)