From 497db8ea5ba60daf0c5b49e04995f10cebcbfdcb Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 16:41:39 +0200
Subject: [PATCH 01/10] allow arbitrary out indices

---
 src/mokernels/moinput.jl | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index 2c2fd1f62..d16aab2e9 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -24,10 +24,17 @@ The first `out_dim` elements represent all outputs for the first input, the seco
 
 See [Inputs for Multiple Outputs](@ref) in the docs for more info.
 """
-struct MOInputIsotopicByFeatures{S,T<:AbstractVector{S},Tout_dim<:Integer} <:
-       AbstractVector{Tuple{S,Int}}
+struct MOInputIsotopicByFeatures{
+    S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
+} <: AbstractVector{Tuple{S,IdxType}}
     x::T
-    out_dim::Tout_dim
+    out_dim::ToutIndices
+end
+
+function MOInputIsotopicByFeatures(
+    x::T, out_dim::Tout_dim
+) where {S,T<:AbstractVector{S},Tout_dim<:Integer}
+    return MOInputIsotopicByFeatures{S,T,Int,Base.OneTo{Tout_dim}}(x, Base.OneTo(out_dim))
 end
 
 """
@@ -54,10 +61,17 @@ As shown above, an `MOInputIsotopicByOutputs` represents a vector of tuples.
 The first `length(x)` elements represent the inputs for the first output, the second
 `length(x)` elements represent the inputs for the second output, etc.
 """
-struct MOInputIsotopicByOutputs{S,T<:AbstractVector{S},Tout_dim<:Integer} <:
-       AbstractVector{Tuple{S,Int}}
+struct MOInputIsotopicByOutputs{
+    S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
+} <: AbstractVector{Tuple{S,IdxType}}
     x::T
-    out_dim::Tout_dim
+    out_dim::ToutIndices
+end
+
+function MOInputIsotopicByOutputs(
+    x::T, out_dim::Tout_dim
+) where {S,T<:AbstractVector{S},Tout_dim<:Integer}
+    return MOInputIsotopicByOutputs{S,T,Int,Base.OneTo{Tout_dim}}(x, Base.OneTo(out_dim))
 end
 
 const IsotopicMOInputsUnion = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}
@@ -66,17 +80,17 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    return feature, output_index
+    return feature, @inbounds inp.out_dim[output_index]
 end
 
 function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
-    feature_index, output_index = fldmod1(ind, inp.out_dim)
+    feature_index, output_index = fldmod1(ind, length(inp.out_dim))
     feature = @inbounds inp.x[feature_index]
-    return feature, output_index
+    return feature, @inbounds inp.out_dim[output_index]
 end
 
-Base.size(inp::IsotopicMOInputsUnion) = (inp.out_dim * length(inp.x),)
+Base.size(inp::IsotopicMOInputsUnion) = (length(inp.out_dim) * length(inp.x),)
 
 function Base.vcat(x::MOInputIsotopicByFeatures, y::MOInputIsotopicByFeatures)
     x.out_dim == y.out_dim || throw(DimensionMismatch("out_dim mismatch"))

From fc1cbfc8ae99493e515d6e811e0da6d18282e39c Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 17:40:21 +0200
Subject: [PATCH 02/10] green tests

---
 src/matrix/kernelkroneckermat.jl   |  8 ++++----
 src/mokernels/independent.jl       | 14 +++++++-------
 src/mokernels/intrinsiccoregion.jl | 14 +++++++-------
 src/mokernels/moinput.jl           | 20 ++++++++++----------
 src/mokernels/slfm.jl              |  4 ++--
 5 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/src/matrix/kernelkroneckermat.jl b/src/matrix/kernelkroneckermat.jl
index 113d0f53d..49d2def1c 100644
--- a/src/matrix/kernelkroneckermat.jl
+++ b/src/matrix/kernelkroneckermat.jl
@@ -82,10 +82,10 @@ efficiently inverted or decomposed. See also [`kernelmatrix`](@ref).
 function kronecker_kernelmatrix(
     k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel}, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.out_dim == y.out_dim ||
-        throw(DimensionMismatch("`x` and `y` must have the same `out_dim`"))
+    x.outIndices == y.outIndices ||
+        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, x.out_dim)
+    Koutputs = _mo_output_covariance(k, length(x.outIndices))
     return _kernelmatrix_kroneckerjl_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -93,7 +93,7 @@ function kronecker_kernelmatrix(
     k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel}, x::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
     Kfeatures = kernelmatrix(k.kernel, x.x)
-    Koutputs = _mo_output_covariance(k, x.out_dim)
+    Koutputs = _mo_output_covariance(k, length(x.outIndices))
     return _kernelmatrix_kroneckerjl_helper(MOI, Kfeatures, Koutputs)
 end
 
diff --git a/src/mokernels/independent.jl b/src/mokernels/independent.jl
index 1f7811b14..1d7f4d6eb 100644
--- a/src/mokernels/independent.jl
+++ b/src/mokernels/independent.jl
@@ -27,15 +27,15 @@ function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,I
     return κ.kernel(x, y) * (px == py)
 end
 
-_mo_output_covariance(k::IndependentMOKernel, out_dim) = Eye{Bool}(out_dim)
+_mo_output_covariance(k::IndependentMOKernel, out_dim::Integer) = Eye{Bool}(out_dim)
 
 function kernelmatrix(
     k::IndependentMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.out_dim == y.out_dim ||
-        throw(DimensionMismatch("`x` and `y` must have the same `out_dim`"))
+    x.outIndices == y.outIndices ||
+        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, x.out_dim)
+    Koutputs = _mo_output_covariance(k, length(x.outIndices))
     return _kernelmatrix_kron_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -43,10 +43,10 @@ if VERSION >= v"1.6"
     function kernelmatrix!(
         K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
     ) where {MOI<:IsotopicMOInputsUnion}
-        x.out_dim == y.out_dim ||
-            throw(DimensionMismatch("`x` and `y` must have the same `out_dim`"))
+        x.outIndices == y.outIndices ||
+            throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
         Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-        Koutputs = _mo_output_covariance(k, x.out_dim)
+        Koutputs = _mo_output_covariance(k, length(x.outIndices))
         return _kernelmatrix_kron_helper!(K, MOI, Kfeatures, Koutputs)
     end
 end
diff --git a/src/mokernels/intrinsiccoregion.jl b/src/mokernels/intrinsiccoregion.jl
index 0a940796b..1f79bca11 100644
--- a/src/mokernels/intrinsiccoregion.jl
+++ b/src/mokernels/intrinsiccoregion.jl
@@ -42,7 +42,7 @@ function (k::IntrinsicCoregionMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{
     return k.B[px, py] * k.kernel(x, y)
 end
 
-function _mo_output_covariance(k::IntrinsicCoregionMOKernel, out_dim)
+function _mo_output_covariance(k::IntrinsicCoregionMOKernel, out_dim::Integer)
     @assert size(k.B) == (out_dim, out_dim)
     return k.B
 end
@@ -50,10 +50,10 @@ end
 function kernelmatrix(
     k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.out_dim == y.out_dim ||
-        throw(DimensionMismatch("`x` and `y` must have the same `out_dim`"))
+    x.outIndices == y.outIndices ||
+        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, x.out_dim)
+    Koutputs = _mo_output_covariance(k, length(x.outIndices))
     return _kernelmatrix_kron_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -61,10 +61,10 @@ if VERSION >= v"1.6"
     function kernelmatrix!(
         K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
     ) where {MOI<:IsotopicMOInputsUnion}
-        x.out_dim == y.out_dim ||
-            throw(DimensionMismatch("`x` and `y` must have the same `out_dim`"))
+        x.outIndices == y.outIndices ||
+            throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
         Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-        Koutputs = _mo_output_covariance(k, x.out_dim)
+        Koutputs = _mo_output_covariance(k, length(x.outIndices))
         return _kernelmatrix_kron_helper!(K, MOI, Kfeatures, Koutputs)
     end
 end
diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index d16aab2e9..07d014f52 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -28,7 +28,7 @@ struct MOInputIsotopicByFeatures{
     S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
 } <: AbstractVector{Tuple{S,IdxType}}
     x::T
-    out_dim::ToutIndices
+    outIndices::ToutIndices
 end
 
 function MOInputIsotopicByFeatures(
@@ -65,7 +65,7 @@ struct MOInputIsotopicByOutputs{
     S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
 } <: AbstractVector{Tuple{S,IdxType}}
     x::T
-    out_dim::ToutIndices
+    outIndices::ToutIndices
 end
 
 function MOInputIsotopicByOutputs(
@@ -80,26 +80,26 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.out_dim[output_index]
+    return feature, @inbounds inp.outIndices[output_index]
 end
 
 function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
-    feature_index, output_index = fldmod1(ind, length(inp.out_dim))
+    feature_index, output_index = fldmod1(ind, length(inp.outIndices))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.out_dim[output_index]
+    return feature, @inbounds inp.outIndices[output_index]
 end
 
-Base.size(inp::IsotopicMOInputsUnion) = (length(inp.out_dim) * length(inp.x),)
+Base.size(inp::IsotopicMOInputsUnion) = (length(inp.outIndices) * length(inp.x),)
 
 function Base.vcat(x::MOInputIsotopicByFeatures, y::MOInputIsotopicByFeatures)
-    x.out_dim == y.out_dim || throw(DimensionMismatch("out_dim mismatch"))
-    return MOInputIsotopicByFeatures(vcat(x.x, y.x), x.out_dim)
+    x.outIndices == y.outIndices || throw(DimensionMismatch("outIndices mismatch"))
+    return MOInputIsotopicByFeatures(vcat(x.x, y.x), x.outIndices)
 end
 
 function Base.vcat(x::MOInputIsotopicByOutputs, y::MOInputIsotopicByOutputs)
-    x.out_dim == y.out_dim || throw(DimensionMismatch("out_dim mismatch"))
-    return MOInputIsotopicByOutputs(vcat(x.x, y.x), x.out_dim)
+    x.outIndices == y.outIndices || throw(DimensionMismatch("outIndices mismatch"))
+    return MOInputIsotopicByOutputs(vcat(x.x, y.x), x.outIndices)
 end
 
 """
diff --git a/src/mokernels/slfm.jl b/src/mokernels/slfm.jl
index 521d07de7..692764cf0 100644
--- a/src/mokernels/slfm.jl
+++ b/src/mokernels/slfm.jl
@@ -34,8 +34,8 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
 end
 
 function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
-    x.out_dim == y.out_dim || error("`x` and `y` should have the same output dimension")
-    x.out_dim == size(k.A, 1) ||
+    x.outIndices == y.outIndices || error("`x` and `y` should have the same output dimension")
+    length(x.outIndices) == size(k.A, 1) ||
         error("Kernel not compatible with the given multi-output inputs")
 
     # Weights matrix ((out_dim x out_dim) x length(k.g))

From bc7e55302a6b106f27b19ec4fb00e5b24ac5ca04 Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 17:58:29 +0200
Subject: [PATCH 03/10] formatter

---
 src/mokernels/slfm.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/mokernels/slfm.jl b/src/mokernels/slfm.jl
index 692764cf0..8eb051092 100644
--- a/src/mokernels/slfm.jl
+++ b/src/mokernels/slfm.jl
@@ -34,7 +34,8 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
 end
 
 function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
-    x.outIndices == y.outIndices || error("`x` and `y` should have the same output dimension")
+    x.outIndices == y.outIndices ||
+        error("`x` and `y` should have the same output dimension")
     length(x.outIndices) == size(k.A, 1) ||
         error("Kernel not compatible with the given multi-output inputs")
 

From 133599ee3988bbde4988cc0f71b59105f955ed9a Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 19:28:39 +0200
Subject: [PATCH 04/10] implement most review comments

---
 src/matrix/kernelkroneckermat.jl   |  8 +++---
 src/mokernels/independent.jl       | 12 ++++-----
 src/mokernels/intrinsiccoregion.jl | 12 ++++-----
 src/mokernels/moinput.jl           | 39 +++++++++++++-----------------
 src/mokernels/slfm.jl              |  4 +--
 5 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/src/matrix/kernelkroneckermat.jl b/src/matrix/kernelkroneckermat.jl
index 49d2def1c..90ca1d2da 100644
--- a/src/matrix/kernelkroneckermat.jl
+++ b/src/matrix/kernelkroneckermat.jl
@@ -82,10 +82,10 @@ efficiently inverted or decomposed. See also [`kernelmatrix`](@ref).
 function kronecker_kernelmatrix(
     k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel}, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.outIndices == y.outIndices ||
-        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
+    x.out_axis == y.out_axis ||
+        throw(DimensionMismatch("`x` and `y` must have the same `out_axis`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, length(x.outIndices))
+    Koutputs = _mo_output_covariance(k, length(x.out_axis))
     return _kernelmatrix_kroneckerjl_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -93,7 +93,7 @@ function kronecker_kernelmatrix(
     k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel}, x::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
     Kfeatures = kernelmatrix(k.kernel, x.x)
-    Koutputs = _mo_output_covariance(k, length(x.outIndices))
+    Koutputs = _mo_output_covariance(k, length(x.out_axis))
     return _kernelmatrix_kroneckerjl_helper(MOI, Kfeatures, Koutputs)
 end
 
diff --git a/src/mokernels/independent.jl b/src/mokernels/independent.jl
index 1d7f4d6eb..a6e5e2811 100644
--- a/src/mokernels/independent.jl
+++ b/src/mokernels/independent.jl
@@ -32,10 +32,10 @@ _mo_output_covariance(k::IndependentMOKernel, out_dim::Integer) = Eye{Bool}(out_
 function kernelmatrix(
     k::IndependentMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.outIndices == y.outIndices ||
-        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
+    x.out_axis == y.out_axis ||
+        throw(DimensionMismatch("`x` and `y` must have the same `out_axis`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, length(x.outIndices))
+    Koutputs = _mo_output_covariance(k, length(x.out_axis))
     return _kernelmatrix_kron_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -43,10 +43,10 @@ if VERSION >= v"1.6"
     function kernelmatrix!(
         K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
     ) where {MOI<:IsotopicMOInputsUnion}
-        x.outIndices == y.outIndices ||
-            throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
+        x.out_axis == y.out_axis ||
+            throw(DimensionMismatch("`x` and `y` must have the same `out_axis`"))
         Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-        Koutputs = _mo_output_covariance(k, length(x.outIndices))
+        Koutputs = _mo_output_covariance(k, length(x.out_axis))
         return _kernelmatrix_kron_helper!(K, MOI, Kfeatures, Koutputs)
     end
 end
diff --git a/src/mokernels/intrinsiccoregion.jl b/src/mokernels/intrinsiccoregion.jl
index 1f79bca11..89b81434e 100644
--- a/src/mokernels/intrinsiccoregion.jl
+++ b/src/mokernels/intrinsiccoregion.jl
@@ -50,10 +50,10 @@ end
 function kernelmatrix(
     k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
-    x.outIndices == y.outIndices ||
-        throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
+    x.out_axis == y.out_axis ||
+        throw(DimensionMismatch("`x` and `y` must have the same `out_axis`"))
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    Koutputs = _mo_output_covariance(k, length(x.outIndices))
+    Koutputs = _mo_output_covariance(k, length(x.out_axis))
     return _kernelmatrix_kron_helper(MOI, Kfeatures, Koutputs)
 end
 
@@ -61,10 +61,10 @@ if VERSION >= v"1.6"
     function kernelmatrix!(
         K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
     ) where {MOI<:IsotopicMOInputsUnion}
-        x.outIndices == y.outIndices ||
-            throw(DimensionMismatch("`x` and `y` must have the same `outIndices`"))
+        x.out_axis == y.out_axis ||
+            throw(DimensionMismatch("`x` and `y` must have the same `out_axis`"))
         Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-        Koutputs = _mo_output_covariance(k, length(x.outIndices))
+        Koutputs = _mo_output_covariance(k, length(x.out_axis))
         return _kernelmatrix_kron_helper!(K, MOI, Kfeatures, Koutputs)
     end
 end
diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index 07d014f52..f7cb026b3 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -25,16 +25,14 @@ The first `out_dim` elements represent all outputs for the first input, the seco
 See [Inputs for Multiple Outputs](@ref) in the docs for more info.
 """
 struct MOInputIsotopicByFeatures{
-    S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
+    S,T<:AbstractVector{S},IdxType,Tout_axis<:AbstractVector{IdxType}
 } <: AbstractVector{Tuple{S,IdxType}}
     x::T
-    outIndices::ToutIndices
+    out_axis::Tout_axis
 end
 
-function MOInputIsotopicByFeatures(
-    x::T, out_dim::Tout_dim
-) where {S,T<:AbstractVector{S},Tout_dim<:Integer}
-    return MOInputIsotopicByFeatures{S,T,Int,Base.OneTo{Tout_dim}}(x, Base.OneTo(out_dim))
+function MOInputIsotopicByFeatures(x::AbstractVector, out_dim::Integer)
+    return MOInputIsotopicByFeatures(x, Base.OneTo(out_dim))
 end
 
 """
@@ -61,17 +59,14 @@ As shown above, an `MOInputIsotopicByOutputs` represents a vector of tuples.
 The first `length(x)` elements represent the inputs for the first output, the second
 `length(x)` elements represent the inputs for the second output, etc.
 """
-struct MOInputIsotopicByOutputs{
-    S,T<:AbstractVector{S},IdxType,ToutIndices<:AbstractVector{IdxType}
-} <: AbstractVector{Tuple{S,IdxType}}
+struct MOInputIsotopicByOutputs{S,I,T<:AbstractVector{S},Tout_axis<:AbstractVector{I}} <:
+       AbstractVector{Tuple{S,I}}
     x::T
-    outIndices::ToutIndices
+    out_axis::Tout_axis
 end
 
-function MOInputIsotopicByOutputs(
-    x::T, out_dim::Tout_dim
-) where {S,T<:AbstractVector{S},Tout_dim<:Integer}
-    return MOInputIsotopicByOutputs{S,T,Int,Base.OneTo{Tout_dim}}(x, Base.OneTo(out_dim))
+function MOInputIsotopicByOutputs(x::AbstractVector, out_dim::Integer)
+    return MOInputIsotopicByOutputs(x, Base.OneTo(out_dim))
 end
 
 const IsotopicMOInputsUnion = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}
@@ -80,26 +75,26 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.outIndices[output_index]
+    return feature, @inbounds inp.out_axis[output_index]
 end
 
 function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
-    feature_index, output_index = fldmod1(ind, length(inp.outIndices))
+    feature_index, output_index = fldmod1(ind, length(inp.out_axis))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.outIndices[output_index]
+    return feature, @inbounds inp.out_axis[output_index]
 end
 
-Base.size(inp::IsotopicMOInputsUnion) = (length(inp.outIndices) * length(inp.x),)
+Base.size(inp::IsotopicMOInputsUnion) = (length(inp.out_axis) * length(inp.x),)
 
 function Base.vcat(x::MOInputIsotopicByFeatures, y::MOInputIsotopicByFeatures)
-    x.outIndices == y.outIndices || throw(DimensionMismatch("outIndices mismatch"))
-    return MOInputIsotopicByFeatures(vcat(x.x, y.x), x.outIndices)
+    x.out_axis == y.out_axis || throw(DimensionMismatch("out_axis mismatch"))
+    return MOInputIsotopicByFeatures(vcat(x.x, y.x), x.out_axis)
 end
 
 function Base.vcat(x::MOInputIsotopicByOutputs, y::MOInputIsotopicByOutputs)
-    x.outIndices == y.outIndices || throw(DimensionMismatch("outIndices mismatch"))
-    return MOInputIsotopicByOutputs(vcat(x.x, y.x), x.outIndices)
+    x.out_axis == y.out_axis || throw(DimensionMismatch("out_axis mismatch"))
+    return MOInputIsotopicByOutputs(vcat(x.x, y.x), x.out_axis)
 end
 
 """
diff --git a/src/mokernels/slfm.jl b/src/mokernels/slfm.jl
index 8eb051092..f1e035b26 100644
--- a/src/mokernels/slfm.jl
+++ b/src/mokernels/slfm.jl
@@ -34,9 +34,9 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
 end
 
 function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
-    x.outIndices == y.outIndices ||
+    x.out_axis == y.out_axis ||
         error("`x` and `y` should have the same output dimension")
-    length(x.outIndices) == size(k.A, 1) ||
+    length(x.out_axis) == size(k.A, 1) ||
         error("Kernel not compatible with the given multi-output inputs")
 
     # Weights matrix ((out_dim x out_dim) x length(k.g))

From 913ed384c68ae815b7139f44aeca0911b29b2e35 Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 19:34:07 +0200
Subject: [PATCH 05/10] actually ensure @inbounds

---
 src/mokernels/moinput.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index f7cb026b3..19c00edd1 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -75,14 +75,16 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.out_axis[output_index]
+    out_idx = axes(inp.out_axis,1)[output_index]
+    return feature, @inbounds inp.out_axis[out_idx]
 end
 
 function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     feature_index, output_index = fldmod1(ind, length(inp.out_axis))
     feature = @inbounds inp.x[feature_index]
-    return feature, @inbounds inp.out_axis[output_index]
+    out_idx = axes(inp.out_axis,1)[output_index]
+    return feature, @inbounds inp.out_axis[out_idx]
 end
 
 Base.size(inp::IsotopicMOInputsUnion) = (length(inp.out_axis) * length(inp.x),)

From d6fd996b10988599548ce258e2c861b2a02fbd9d Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 19:47:27 +0200
Subject: [PATCH 06/10] missed suggestion

---
 src/mokernels/moinput.jl | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index 19c00edd1..bfe9abd29 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -24,9 +24,8 @@ The first `out_dim` elements represent all outputs for the first input, the seco
 
 See [Inputs for Multiple Outputs](@ref) in the docs for more info.
 """
-struct MOInputIsotopicByFeatures{
-    S,T<:AbstractVector{S},IdxType,Tout_axis<:AbstractVector{IdxType}
-} <: AbstractVector{Tuple{S,IdxType}}
+struct MOInputIsotopicByFeatures{S,I,T<:AbstractVector{S},Tout_axis<:AbstractVector{I}} <:
+       AbstractVector{Tuple{S,IdxType}}
     x::T
     out_axis::Tout_axis
 end
@@ -75,7 +74,7 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    out_idx = axes(inp.out_axis,1)[output_index]
+    out_idx = axes(inp.out_axis, 1)[output_index]
     return feature, @inbounds inp.out_axis[out_idx]
 end
 
@@ -83,7 +82,7 @@ function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     feature_index, output_index = fldmod1(ind, length(inp.out_axis))
     feature = @inbounds inp.x[feature_index]
-    out_idx = axes(inp.out_axis,1)[output_index]
+    out_idx = axes(inp.out_axis, 1)[output_index]
     return feature, @inbounds inp.out_axis[out_idx]
 end
 

From be23f44bf61649411836eb876ad1e1b6a9e4203f Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 20:22:16 +0200
Subject: [PATCH 07/10] fix typo

---
 src/mokernels/moinput.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index bfe9abd29..7a5ee40ea 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -25,7 +25,7 @@ The first `out_dim` elements represent all outputs for the first input, the seco
 See [Inputs for Multiple Outputs](@ref) in the docs for more info.
 """
 struct MOInputIsotopicByFeatures{S,I,T<:AbstractVector{S},Tout_axis<:AbstractVector{I}} <:
-       AbstractVector{Tuple{S,IdxType}}
+       AbstractVector{Tuple{S,I}}
     x::T
     out_axis::Tout_axis
 end

From bdee7823f0dd330d094ffb58b8d96b9a3f85a3ff Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Thu, 25 May 2023 20:36:33 +0200
Subject: [PATCH 08/10] formatter

---
 src/mokernels/slfm.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/mokernels/slfm.jl b/src/mokernels/slfm.jl
index f1e035b26..70ab2007f 100644
--- a/src/mokernels/slfm.jl
+++ b/src/mokernels/slfm.jl
@@ -34,8 +34,7 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
 end
 
 function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
-    x.out_axis == y.out_axis ||
-        error("`x` and `y` should have the same output dimension")
+    x.out_axis == y.out_axis || error("`x` and `y` should have the same output dimension")
     length(x.out_axis) == size(k.A, 1) ||
         error("Kernel not compatible with the given multi-output inputs")
 

From 1e790456ab3aed1c12143c0517908762e112d481 Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Fri, 26 May 2023 17:51:01 +0200
Subject: [PATCH 09/10] ensure offset arrays work with tests

---
 src/mokernels/moinput.jl  |  6 ++++--
 test/Project.toml         |  1 +
 test/mokernels/moinput.jl | 22 ++++++++++++++++++++++
 test/runtests.jl          |  1 +
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index 7a5ee40ea..201a6379e 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -74,7 +74,8 @@ function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     output_index, feature_index = fldmod1(ind, length(inp.x))
     feature = @inbounds inp.x[feature_index]
-    out_idx = axes(inp.out_axis, 1)[output_index]
+    out_indices = axes(inp.out_axis, 1)
+    out_idx = out_indices[begin + output_index - 1]
     return feature, @inbounds inp.out_axis[out_idx]
 end
 
@@ -82,7 +83,8 @@ function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     @boundscheck checkbounds(inp, ind)
     feature_index, output_index = fldmod1(ind, length(inp.out_axis))
     feature = @inbounds inp.x[feature_index]
-    out_idx = axes(inp.out_axis, 1)[output_index]
+    out_indices = axes(inp.out_axis, 1)
+    out_idx = out_indices[begin + output_index - 1]
     return feature, @inbounds inp.out_axis[out_idx]
 end
 
diff --git a/test/Project.toml b/test/Project.toml
index e16f39a6c..2a54c4331 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -9,6 +9,7 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Kronecker = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
diff --git a/test/mokernels/moinput.jl b/test/mokernels/moinput.jl
index ce1a4757e..e56ef3393 100644
--- a/test/mokernels/moinput.jl
+++ b/test/mokernels/moinput.jl
@@ -4,6 +4,8 @@
     type_2 = AbstractVector{Tuple{AbstractVector{Vector{Float64}},Int}}
 
     @testset "isotopicbyoutputs" begin
+        @test KernelFunctions.MOInputIsotopicByOutputs(x, []) == []
+
         ibo = KernelFunctions.MOInputIsotopicByOutputs(x, 3)
         ibo2 = KernelFunctions.MOInputIsotopicByOutputs(x, 2)
 
@@ -27,9 +29,20 @@
         @test ibo[7] == (x[3], 2)
         @test all([(x_, i) for i in 1:3 for x_ in x] .== ibo)
         @inferred getindex(ibo, 1)
+
+        # test non-standard array
+        ibo3 = KernelFunctions.MOInputIsotopicByOutputs(x, OA.Origin(0)([:a, :b, :c]))
+        ibo4 = KernelFunctions.MOInputIsotopicByOutputs(x, [:a, :b, :c])
+        @test ibo3 == ibo4
+        @test length(ibo3) == 12
+        @test lastindex(ibo3) == 12
+        @test firstindex(ibo3) == 1
+        @test_throws BoundsError ibo3[0]
+        @test_throws BoundsError ibo3[13]
     end
 
     @testset "isotopicbyfeatures" begin
+        @test KernelFunctions.MOInputIsotopicByFeatures([], [1.0, 2.0]) == []
         ibf = KernelFunctions.MOInputIsotopicByFeatures(x, 3)
 
         @test isa(ibf, type_1) == true
@@ -48,6 +61,15 @@
         @test ibf[7] == (x[3], 1)
         @test all([(x_, i) for x_ in x for i in 1:3] .== ibf)
         @inferred getindex(ibf, 1)
+
+        # test non-standard array
+        ibf2 = KernelFunctions.MOInputIsotopicByFeatures(x, OA.Origin(0)(["a", "b", "c"]))
+        ibf3 = KernelFunctions.MOInputIsotopicByFeatures(x, ["a", "b", "c"])
+        @test ibf2 == ibf3
+        @test lastindex(ibf2) == 12
+        @test firstindex(ibf2) == 1
+        @test_throws BoundsError ibf2[0]
+        @test_throws BoundsError ibf2[13]
     end
 
     @testset "prepare_isotopic_multi_output_data" begin
diff --git a/test/runtests.jl b/test/runtests.jl
index caf43cb91..d6ea88b06 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,6 +18,7 @@ using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
 using FiniteDifferences: FiniteDifferences
 using Compat: only
+import OffsetArrays as OA
 
 using KernelFunctions: SimpleKernel, metric, kappa, ColVecs, RowVecs, TestUtils
 

From 12a1e058214665810a10a74859b03e67ef7fb32a Mon Sep 17 00:00:00 2001
From: Felix Benning <felix.benning@gmail.com>
Date: Fri, 26 May 2023 18:01:16 +0200
Subject: [PATCH 10/10] fix jldoctest

---
 src/mokernels/moinput.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/mokernels/moinput.jl b/src/mokernels/moinput.jl
index 201a6379e..cf81b9f2f 100644
--- a/src/mokernels/moinput.jl
+++ b/src/mokernels/moinput.jl
@@ -7,7 +7,7 @@
 julia> x = [1, 2, 3];
 
 julia> KernelFunctions.MOInputIsotopicByFeatures(x, 2)
-6-element KernelFunctions.MOInputIsotopicByFeatures{Int64, Vector{Int64}, Int64}:
+6-element KernelFunctions.MOInputIsotopicByFeatures{Int64, Int64, Vector{Int64}, Base.OneTo{Int64}}:
  (1, 1)
  (1, 2)
  (2, 1)
@@ -43,7 +43,7 @@ end
 julia> x = [1, 2, 3];
 
 julia> KernelFunctions.MOInputIsotopicByOutputs(x, 2)
-6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}, Int64}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Int64, Vector{Int64}, Base.OneTo{Int64}}:
  (1, 1)
  (2, 1)
  (3, 1)
@@ -110,7 +110,7 @@ A data type to accommodate modelling multi-dimensional output data.
 julia> x = [1, 2, 3];
 
 julia> MOInput(x, 2)
-6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}, Int64}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Int64, Vector{Int64}, Base.OneTo{Int64}}:
  (1, 1)
  (2, 1)
  (3, 1)
@@ -150,7 +150,7 @@ julia> Y = [1.1 2.1 3.1; 1.2 2.2 3.2]
 julia> inputs, outputs = prepare_isotopic_multi_output_data(x, ColVecs(Y));
 
 julia> inputs
-6-element KernelFunctions.MOInputIsotopicByFeatures{Float64, Vector{Float64}, Int64}:
+6-element KernelFunctions.MOInputIsotopicByFeatures{Float64, Int64, Vector{Float64}, Base.OneTo{Int64}}:
  (1.0, 1)
  (1.0, 2)
  (2.0, 1)
@@ -198,7 +198,7 @@ julia> Y = [1.1 1.2; 2.1 2.2; 3.1 3.2]
 julia> inputs, outputs = prepare_isotopic_multi_output_data(x, RowVecs(Y));
 
 julia> inputs
-6-element KernelFunctions.MOInputIsotopicByOutputs{Float64, Vector{Float64}, Int64}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Float64, Int64, Vector{Float64}, Base.OneTo{Int64}}:
  (1.0, 1)
  (2.0, 1)
  (3.0, 1)