Heterotopic MO kernel helper (#353)

* Heterotopic helper

* Bump patch

* Tweak docstring

* Update docs

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Bump patch

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.12"
+version = "0.10.13"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/api.md

@@ -58,13 +58,19 @@ For an explanation of this design choice, see [the design notes on multi-output
 An input to a multi-output `Kernel` should be a `Tuple{T, Int}`, whose first element specifies a location in the domain of the multi-output GP, and whose second element specifies which output the inputs corresponds to.
 The type of collections of inputs for multi-output GPs is therefore `AbstractVector{<:Tuple{T, Int}}`.
 
-KernelFunctions.jl provides the following helper function for situations in which all outputs are observed all of the time:
+KernelFunctions.jl provides the following helper functions to reduce the cognitive load
+associated with working with multi-output kernels by dealing with transforming data from the
+formats in which it is commonly found into the format required by KernelFunctions.
+The intention is that users can pass their data to these functions, and use the returned
+values throughout their code, without having to worry further about correctly formatting
+their data for KernelFunctions' sake:
 ```@docs
 prepare_isotopic_multi_output_data(x::AbstractVector, y::ColVecs)
 prepare_isotopic_multi_output_data(x::AbstractVector, y::RowVecs)
+prepare_heterotopic_multi_output_data
 ```
 
-The input types that it constructs can also be constructed manually:
+The input types returned by `prepare_isotopic_multi_output_data` can also be constructed manually:
 ```@docs
 MOInput
 ```

src/KernelFunctions.jl

@@ -36,7 +36,7 @@ export spectral_mixture_kernel, spectral_mixture_product_kernel
 
 export ColVecs, RowVecs
 
-export MOInput, prepare_isotopic_multi_output_data
+export MOInput, prepare_isotopic_multi_output_data, prepare_heterotopic_multi_output_data
 export IndependentMOKernel,
     LatentFactorMOKernel, IntrinsicCoregionMOKernel, LinearMixingModelKernel
 

src/mokernels/moinput.jl

@@ -143,6 +143,8 @@ julia> outputs
  3.1
  3.2
 ```
+
+See also [`prepare_heterotopic_multi_output_data`](@ref).
 """
 function prepare_isotopic_multi_output_data(x::AbstractVector, y::ColVecs)
     length(x) == length(y) || throw(ArgumentError("length(x) not equal to length(y)."))
@@ -189,8 +191,65 @@ julia> outputs
  2.2
  3.2
 ```
+
+See also [`prepare_heterotopic_multi_output_data`](@ref).
 """
 function prepare_isotopic_multi_output_data(x::AbstractVector, y::RowVecs)
     length(x) == length(y) || throw(ArgumentError("length(x) not equal to length(y)."))
     return MOInputIsotopicByOutputs(x, size(y.X, 2)), vec(y.X)
 end
+
+"""
+    prepare_heterotopic_multi_output_data(
+        x::AbstractVector, y::AbstractVector{<:Real}, output_indices::AbstractVector{Int},
+    )
+
+Utility functionality to convert a collection of inputs `x`, observations `y`, and
+`output_indices` into a format suitable for use with multi-output kernels.
+Handles the situation in which only one (or a subset) of outputs are observed at each
+feature.
+Ensures that all arguments are compatible with one another, and returns a vector of inputs
+and a vector of outputs.
+
+`y[n]` should be the observed value associated with output `output_indices[n]` at feature
+`x[n]`.
+
+```jldoctest
+julia> x = [1.0, 2.0, 3.0];
+
+julia> y = [-1.0, 0.0, 1.0];
+
+julia> output_indices = [3, 2, 1];
+
+julia> inputs, outputs = prepare_heterotopic_multi_output_data(x, y, output_indices);
+
+julia> inputs
+3-element Vector{Tuple{Float64, Int64}}:
+ (1.0, 3)
+ (2.0, 2)
+ (3.0, 1)
+
+julia> outputs
+3-element Vector{Float64}:
+ -1.0
+  0.0
+  1.0
+```
+
+See also [`prepare_isotopic_multi_output_data`](@ref).
+"""
+function prepare_heterotopic_multi_output_data(
+    x::AbstractVector, y::AbstractVector{<:Real}, output_indices::AbstractVector{Int}
+)
+    # Ensure validity of arguments.
+    if length(x) != length(y)
+        throw(ArgumentError("length(x) != length(y)"))
+    end
+    if length(x) != length(output_indices)
+        throw(ArgumentError("length(x) != length(output_indices)"))
+    end
+
+    # Construct inputs and outputs for multi-output kernel.
+    x_mogp = map(tuple, x, output_indices)
+    return x_mogp, y
+end

test/mokernels/moinput.jl

@@ -75,4 +75,17 @@
             @test length(y_canon) == length(x_canon)
         end
     end
+    @testset "prepare_heterotopic_multi_output_data" begin
+        x = randn(3)
+        y = randn(3)
+        output_indices = [3, 1, 2]
+        x_canon, y_canon = prepare_heterotopic_multi_output_data(x, y, output_indices)
+        @test x_canon isa AbstractVector{<:Tuple{<:Real,Int}}
+        @test y isa AbstractVector{<:Real}
+
+        @test_throws ArgumentError prepare_heterotopic_multi_output_data(x, y, [3, 1])
+        @test_throws(
+            ArgumentError, prepare_heterotopic_multi_output_data(x, [1.0], output_indices),
+        )
+    end
 end