willtebbutt
commented
2 months ago Hmmm I'm really not sure what Zygote is doing here -- it has an annoying habit of breaking when you don't expect it to...
I would suggest trying with Tapir.jl. Could you please see if the following works for you?
using Stheno, Tapir, Random, Optim, ParameterHandling
l1 = 0.4
s1 = 0.2
l2 = 5.0
s2 = 1.0
g = @gppp let
f1 = s1 * stretch(GP(Matern52Kernel()), 1 / l1)
f2 = s2 * stretch(GP(SEKernel()), 1 / l2)
f3 = f1 + f2
end;
x = GPPPInput(:f3, collect(range(-5.0, 5.0; length=100)));
σ²_n = 0.02;
fx = g(x, σ²_n);
y = rand(fx);
θ = (
# Short length-scale and small variance.
l1 = positive(0.4),
s1 = positive(0.2),
# Long length-scale and larger variance.
l2 = positive(5.0),
s2 = positive(1.0),
# Observation noise variance -- we'll be learning this as well. Constrained to be
# at least 1e-3.
s_noise = positive(0.1, exp, 1e-3),
)
θ_flat_init, unflatten = flatten(θ);
function build_model(θ::NamedTuple)
return @gppp let
f1 = θ.s1 * stretch(GP(SEKernel()), 1 / θ.l1)
f2 = θ.s2 * stretch(GP(SEKernel()), 1 / θ.l2)
f3 = f1 + f2
end
end
function nlml(θ::NamedTuple)
f = build_model(θ)
return -logpdf(f(x, θ.s_noise + 1e-6), y)
end
# Define objective function, check it runs, and compute a gradient to check that works.
obj(x) = nlml(ParameterHandling.value(unflatten(x)))
obj(θ_flat_init)
rule = Tapir.build_rrule(obj, θ_flat_init)
Tapir.value_and_gradient!!(rule, obj, θ_flat_init)
# Run optimisation.
results = Optim.optimize(
obj,
θ->Tapir.value_and_gradient!!(rule, obj, θ)[2][2],
θ_flat_init + randn(length(θ_flat_init)),
BFGS(),
Optim.Options(show_trace=true,);
inplace=false,
)
Jarrod-Angove
I'm quite new to this package, so please forgive me if this is a mistake on my end, but I've been having some difficulties getting Optim AD methods working with Stheno. Even after directly copying the BFGS example from the docs, I end up with the following error:
This occurs if I copy and past the tutorial in the docs directly into the repl, or if I paste everything as a single function:
This issue does not occur if I do not provide a gradient to optim, as I believe it defaults to a finite different method to calculate the gradient. Similarly, using
NelderMead()allows it to run without issue.Here is the stack trace: