|
| 1 | +function _cov_row!(cK, k, X::AbstractMatrix, data, j, dim) |
| 2 | + cK[j,j] = cov_ij(k, X, X, data, j, j, dim) |
| 3 | + @inbounds for i in 1:j-1 |
| 4 | + cK[i,j] = cov_ij(k, X, X, data, i, j, dim) |
| 5 | + cK[j,i] = cK[i,j] |
| 6 | + end |
| 7 | +end |
| 8 | +function cov!(cK::Matrix, k::Kernel, X::Matrix, data::KernelData=EmptyData()) |
| 9 | + dim, nobs = size(X) |
| 10 | + (nobs,nobs) == size(cK) || throw(ArgumentError("cK has size $(size(cK)) and X has size $(size(X))")) |
| 11 | + kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] # in case k is not threadsafe (e.g. ADkernel) |
| 12 | + @inbounds Threads.@threads for j in 1:nobs |
| 13 | + kthread = kcopies[Threads.threadid()] |
| 14 | + _cov_row!(cK, k, X, data, j, dim) |
| 15 | + end |
| 16 | + return cK |
| 17 | +end |
| 18 | +function _cov_row!(cK, k, X1::AbstractMatrix, X2::AbstractMatrix, data, i, dim, nobs2) |
| 19 | + @inbounds for j in 1:nobs2 |
| 20 | + cK[i,j] = cov_ij(k, X1, X2, data, i, j, dim) |
| 21 | + end |
| 22 | +end |
| 23 | +""" |
| 24 | + cov!(cK::AbstractMatrix, k::Kernel, X1::AbstractMatrix, X2::AbstractMatrix, data::KernelData=EmptyData()) |
| 25 | +
|
| 26 | +Like [`cov(k, X1, X2)`](@ref), but stores the result in `cK` rather than a new matrix. |
| 27 | +""" |
| 28 | +function cov!(cK::Matrix, k::Kernel, X1::Matrix, X2::Matrix, data::KernelData=EmptyData()) |
| 29 | + if X1 === X2 |
| 30 | + return cov!(cK, k, X1, data) |
| 31 | + end |
| 32 | + dim1, nobs1 = size(X1) |
| 33 | + dim2, nobs2 = size(X2) |
| 34 | + dim1==dim2 || throw(ArgumentError("X1 and X2 must have same dimension")) |
| 35 | + dim = size(X1, 1) |
| 36 | + (nobs1,nobs2) == size(cK) || throw(ArgumentError("cK has size $(size(cK)) X1 $(size(X1)) and X2 $(size(X2))")) |
| 37 | + kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] |
| 38 | + @inbounds Threads.@threads for i in 1:nobs1 |
| 39 | + kthread = kcopies[Threads.threadid()] |
| 40 | + _cov_row!(cK, kthread, X1, X2, data, i, dim, nobs2) |
| 41 | + end |
| 42 | + return cK |
| 43 | +end |
| 44 | + |
| 45 | +function _grad_slice_row!(dK, k, X::AbstractMatrix, data, j, p, dim) |
| 46 | + dK[j,j] = dKij_dθp(k,X,X,data,j,j,p,dim) |
| 47 | + @inbounds @simd for i in 1:(j-1) |
| 48 | + dK[i,j] = dKij_dθp(k,X,X,data,i,j,p,dim) |
| 49 | + dK[j,i] = dK[i,j] |
| 50 | + end |
| 51 | +end |
| 52 | +function grad_slice!(dK::AbstractMatrix, k::Kernel, X::Matrix, data::KernelData, p::Int) |
| 53 | + dim, nobs = size(X) |
| 54 | + (nobs,nobs) == size(dK) || throw(ArgumentError("dK has size $(size(dK)) and X has size $(size(X))")) |
| 55 | + kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] |
| 56 | + @inbounds Threads.@threads for j in 1:nobs |
| 57 | + kthread = kcopies[Threads.threadid()] |
| 58 | + _grad_slice_row!(dK, kthread, X, data, j, p, dim) |
| 59 | + end |
| 60 | + return dK |
| 61 | +end |
| 62 | +function _grad_slice_row!(dK, k, X1::AbstractMatrix, X2::AbstractMatrix, data, i, p, dim, nobs2) |
| 63 | + @inbounds @simd for j in 1:nobs2 |
| 64 | + dK[i,j] = dKij_dθp(k,X1,X2,data,i,j,p,dim) |
| 65 | + end |
| 66 | +end |
| 67 | +function grad_slice!(dK::AbstractMatrix, k::Kernel, X1::Matrix, X2::Matrix, data::KernelData, p::Int) |
| 68 | + if X1 === X2 |
| 69 | + return grad_slice!(dK, k, X1, data, p) |
| 70 | + end |
| 71 | + dim1, nobs1 = size(X1) |
| 72 | + dim2, nobs2 = size(X2) |
| 73 | + dim1==dim2 || throw(ArgumentError("X1 and X2 must have same dimension")) |
| 74 | + (nobs1,nobs2) == size(dK) || throw(ArgumentError("dK has size $(size(dK)) X1 $(size(X1)) and X2 $(size(X2))")) |
| 75 | + dim=dim1 |
| 76 | + kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] |
| 77 | + @inbounds Threads.@threads for i in 1:nobs1 |
| 78 | + kthread = kcopies[Threads.threadid()] |
| 79 | + _grad_slice_row!(dK, kthread, X1, X2, data, i, p, dim, nobs2) |
| 80 | + end |
| 81 | + return dK |
| 82 | +end |
| 83 | + |
| 84 | +function _dmll_kern_row!(dmll, buf, k, ααinvcKI, X, data, j, dim, nparams) |
| 85 | + # diagonal |
| 86 | + dKij_dθ!(buf, k, X, X, data, j, j, dim, nparams) |
| 87 | + @inbounds for iparam in 1:nparams |
| 88 | + dmll[iparam] += buf[iparam] * ααinvcKI[j, j] / 2.0 |
| 89 | + end |
| 90 | + # off-diagonal |
| 91 | + @inbounds for i in 1:j-1 |
| 92 | + dKij_dθ!(buf, k, X, X, data, i, j, dim, nparams) |
| 93 | + @simd for iparam in 1:nparams |
| 94 | + dmll[iparam] += buf[iparam] * ααinvcKI[i, j] |
| 95 | + end |
| 96 | + end |
| 97 | +end |
| 98 | +""" |
| 99 | + dmll_kern!((dmll::AbstractVector, k::Kernel, X::AbstractMatrix, data::KernelData, ααinvcKI::AbstractMatrix)) |
| 100 | +
|
| 101 | +Derivative of the marginal log likelihood log p(Y|θ) with respect to the kernel hyperparameters. |
| 102 | +""" |
| 103 | +function dmll_kern!(dmll::AbstractVector, k::Kernel, X::Matrix, data::KernelData, |
| 104 | + ααinvcKI::Matrix{Float64}, covstrat::CovarianceStrategy) |
| 105 | + dim, nobs = size(X) |
| 106 | + nparams = num_params(k) |
| 107 | + @assert nparams == length(dmll) |
| 108 | + dK_buffer = Vector{Float64}(undef, nparams) |
| 109 | + dmll[:] .= 0.0 |
| 110 | + # make a copy per thread for objects that are potentially not thread-safe: |
| 111 | + kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] |
| 112 | + buffercopies = [similar(dK_buffer) for _ in 1:Threads.nthreads()] |
| 113 | + dmllcopies = [deepcopy(dmll) for _ in 1:Threads.nthreads()] |
| 114 | + |
| 115 | + @inbounds Threads.@threads for j in 1:nobs |
| 116 | + kthread = kcopies[Threads.threadid()] |
| 117 | + bufthread = buffercopies[Threads.threadid()] |
| 118 | + dmllthread = dmllcopies[Threads.threadid()] |
| 119 | + _dmll_kern_row!(dmllthread, bufthread, kthread, |
| 120 | + ααinvcKI, X, data, j, dim, nparams) |
| 121 | + end |
| 122 | + |
| 123 | + dmll[:] = sum(dmllcopies) # sum up the results from all threads |
| 124 | + return dmll |
| 125 | +end |
0 commit comments