more tests

docs/src/intro.qmd

@@ -37,7 +37,7 @@ using LaplaceRedux.Data: toy_data_regression
 using Flux.Optimise: update!, Adam
 
 # Data:
-n = 120           # number of observations
+n = 150           # number of observations
 σtrue = 0.3       # true observational noise
 x, y = toy_data_regression(100;noise=σtrue)
 xs = [[x] for x in x]

src/baselaplace.jl

@@ -30,12 +30,10 @@ function get_params(la::BaseLaplace)
     params = Flux.params(nn)
     n_elements = length(params)
     if la.subset_of_weights == :all
-        params = [θ for θ ∈ params] # get all parameters and constants in logitbinarycrossentropy
+        params = [θ for θ ∈ params]                         # get all parameters and constants in logitbinarycrossentropy
     elseif la.subset_of_weights == :last_layer
-        params = [params[n_elements-1],params[n_elements]] # only get last parameters and constants
-    else
-        @error "`subset_of_weights` of weights should be one of the following: `[:all, :last_layer]`"
-    end 
+        params = [params[n_elements-1],params[n_elements]]  # only get last parameters and constants
+    end
     return params
 end
 
@@ -123,11 +121,9 @@ function log_marginal_likelihood(la::BaseLaplace; P₀::Union{Nothing,AbstractFl
     if !isnothing(σ)
         @assert (la.likelihood==:regression || la.σ == σ) "Can only change observational noise σ for regression."
         la.σ = σ
-    end
-
-    mll = log_likelihood(la) - 0.5 * (log_det_ratio(la) + _weight_penalty(la))
+    end 
 
-    return mll
+    return log_likelihood(la) - 0.5 * (log_det_ratio(la) + _weight_penalty(la))
 end
 
 """

src/laplace.jl

@@ -45,7 +45,12 @@ function Laplace(model::Any; likelihood::Symbol, kwargs...)
 
     # Load hyperparameters:
     args = LaplaceParams(;kwargs...)
+
+    # Assertions:
     @assert !(args.σ != 1.0 && likelihood != :regression) "Observation noise σ ≠ 1 only available for regression."
+    @assert args.subset_of_weights ∈ [:all, :last_layer] "`subset_of_weights` of weights should be one of the following: `[:all, :last_layer]`"
+
+    # Setup:
     P₀ = isnothing(args.P₀) ? UniformScaling(args.λ) : args.P₀
     nn = model
     n_out = outdim(nn)

test/laplace.jl

@@ -10,6 +10,11 @@ using Statistics
 
     # One layer:
     nn = Chain(Dense(2,1))
+
+    # Expected error
+    @test_throws AssertionError Laplace(nn; likelihood=:classification, subset_of_weights=:last)
+
+    # Correct:
     la = Laplace(nn; likelihood=:classification)
     @test la.n_params == 3
 
@@ -179,8 +184,11 @@ end
 
             la = Laplace(nn; likelihood=likelihood, λ=λ, subset_of_weights=:last_layer)
             fit!(la, data)
-            optimize_prior!(la)
-            plt = plot(la, X, y)
+            optimize_prior!(la; verbose=true)
+            plot(la, X, y)                              # standard
+            plot(la, X, y; xlims=(-5,5), ylims=(-5,5))  # lims
+            plot(la, X, y; link_approx=:plugin)         # plugin approximation
+
         end
     end