Setting up multi-output Gaussian process with Surrogates.jl

Hi all,

I am trying to use Surrogates to train a Gaussian process surrogate to a model with 2 different outputs (i.e., a multi-output model).

According to the Surrogates documentation I can use custom kriging in Stheno, and I can see that AbstractGPs also provide ways to do multi-output models, but I can’t figure out how to set these up with Surrogates.

Below is a toy model to demonstrate roughly how this problem looks and where I have got so far. Training the surrogate model on one output works well enough - how can I adapt this to train on two outputs?

Would appreciate any advice on making this work and links to any pages that might put me on the right track.

FYI this is a stochastic population level epidemiological model of transmission of two competing strains of a pathogen, and I have data on overall prevalence and the frequency of each strain. I am planning to use the Surrogates package to do parameter tuning for this model.

Many thanks!

using Distributions, Plots, Surrogates

pois(x) = rand(Poisson(max(x, 0)))

function modelstep(comps, params)
    beta, mu, c, N, a, dt = params
    S,R = comps 
    X = N - S - R
    p = S + R
    if p > 0.
        ql = (1 .- ([S,R])./p .+ 1/2) .^a 
        comps[1] += pois(ql[1]*beta*(S/N)*X*dt) - pois(mu*S*dt)
        comps[2] += pois(ql[2]*beta*(R/N)*X*dt) - pois(c*mu*R*dt)

function runmodel(ps)
    params = [35., 9.6, ps[1], 10^5, ps[2], 4/52]
    n_it = 10000
    comps = [10,10]
    out = zeros(n_it, 2)
    for i in 1:n_it
        modelstep(comps, params)
        out[i,:] = comps
    return out

function getprev(ps)
    out = runmodel(ps)

lb = [1., 0.]
ub = [2., 2.]
x = Surrogates.sample(500, lb, ub, SobolSample())
y = getprev.(x)
surrogate = Kriging(x, y, lb, ub)

cs = [1.:0.1:2.0;]
as = [0.:0.2:2.0;]
prevorig = zeros(length(cs), length(as))
prevgp = zeros(length(cs), length(as))

for i in eachindex(cs)
    for j in eachindex(as)
        prevorig[i,j] = getprev([cs[i], as[j]])
        prevgp[i,j] = surrogate([cs[i], as[j]]) 

heatmap(as, cs, exp.(prevorig).-1)
heatmap(as, cs, exp.(prevgp).-1)

@SathvikBhagavan have you looked into this part recently?

I haven’t but I can look into this.