Vector-output GEKPLS via PLS2 + multi-output BLUP

[ChrisRackauckas-Claude]

Note

Draft — please ignore until reviewed by @ChrisRackauckas.

Summary

GEKPLS previously assumed a scalar response: y_matrix was hardcoded to (nt, 1), gradients to (nt, dim), and the kriging BLUP collapsed to a single output via transpose(Q) ⋅ Yt. This PR generalizes the model to vector-valued targets f : R^dim -> R^ny while keeping the legacy scalar API bit-for-bit unchanged.

Motivated by a Discourse thread asking for derivative-informed surrogates that handle vector responses: https://discourse.julialang.org/t/surrogate-with-derivatives/137277/5

Math

Taylor augmentation (_ge_compute_pls) — the local FOTA response is now a (n_bb, ny) matrix:

_y[bb, k] = y[i, k] + Σⱼ bb_dx[bb, j] · ∂y_k/∂x_j

PLS step (_get_first_singular_vectors_power_method) — the scalar response case (PLS1) is exact after one power iteration; that path is preserved bit-for-bit. For vector response (PLS2) the same body is iterated to convergence following Wold's NIPALS algorithm. Initial y_score is the column of Y with the largest variance, matching scikit-learn's PLSRegression default.

Multi-output BLUP (_reduced_likelihood_function) — under shared θ across outputs the kriging correlation matrix R and its Cholesky factor are shared; β (1, ny) and γ (nt, ny) are then obtained from a single multi-RHS solve. The previous transpose(Q) ⋅ Yt (dot product, scalar) becomes transpose(Q) * Yt (matrix product).

References:

• Bouhlel & Martins, Gradient-enhanced kriging for high-dimensional problems, Engineering with Computers 35 (2019), arXiv:1708.02663.
• Wold, Sjöström & Eriksson, PLS-regression: a basic tool of chemometrics, Chemometrics and Intelligent Laboratory Systems 58 (2001).

API

The constructor auto-detects shape from eltype(y_vec):

• Vector{<:Number} — legacy scalar API, unchanged.
• Vector{<:AbstractVector{<:Number}} — vector output; grads_vec entries must be (ny, dim) Jacobian matrices (or 1-tuples wrapping them, matching Zygote.jacobian output).

g(x) returns a scalar for scalar y and a length-ny Vector for vector y. update! mirrors the same auto-detection.

Internal representation

• y_matrix::Matrix{T} (nt, ny)
• grads::Array{T,3} (nt, dim, ny) — stores ∂y_k/∂x_j
• beta::Matrix{T}, gamma::Matrix{T}, y_mean::Matrix{T}, y_std::Matrix{T} — one column per output
• new scalar_output::Bool flag so predict returns the same type the user supplied

Tests

• Existing 11 GEKPLS regression tests pass unchanged (bit-for-bit on the PLS1 path; per-output BLUP results match modulo float reordering).
• Test 12 — vector-output fit on a 6D 3-output polynomial problem: each output's relative RMSE < 5%.
• Test 13 — packaging a scalar response as a length-1 vector via the new path reproduces the legacy scalar surrogate to rtol=1e-10 over 30 test points.
• Test 14 — update! on a vector-output model strictly improves per-output RMSE after additional samples.

Also verified locally that the Radials and Kriging algorithm test suites pass (unrelated, sanity check).

Test plan

• Pkg.test GEKPLS — 49 tests pass
• Radials sanity — 30 tests pass
• Kriging sanity — 26 tests pass
• Runic formatting applied
• CI green

🤖 Generated with Claude Code