inferring non-symmetric Toeplitz structure automatically

Project.toml

@@ -1,7 +1,7 @@
 name = "CovarianceFunctions"
 uuid = "b3329135-7958-41d4-bb02-e084c5a526bf"
 authors = ["sebastianament"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 BesselK = "432ab697-7a72-484f-bc4a-bc531f5c819b"

src/gramian.jl

@@ -161,13 +161,22 @@ gramian(k, x, ::InputTrait) = gramian(k, x)
 
 # 1D stationary kernel on equi-spaced grid yields Toeplitz structure
 # works if input_trait(k) <: Union{IsotropicInput, StationaryInput}
-gramian(k, x::StepRangeLen{<:Real}) = gramian(k, x, input_trait(k))
+gramian(k, x::StepRangeLen{<:Real}, y::StepRangeLen{<:Real}) = gramian(k, x, y, input_trait(k))
+gramian(k, x::StepRangeLen{<:Real}, y::StepRangeLen{<:Real}, ::GenericInput) = Gramian(k, x, y)
 
 # IDEA: should this be in factorization? since dft still costs linear amount of information
 # while gramian is usually lazy and O(1) in structure construction
-function gramian(k, x::StepRangeLen{<:Real}, ::Union{IsotropicInput, StationaryInput})
-    k1 = k.(x[1], x)
-    SymmetricToeplitz(k1)
+function gramian(k, x::StepRangeLen{<:Real}, y::StepRangeLen{<:Real}, ::Union{IsotropicInput, StationaryInput})
+    if x === y
+        k1 = k.(x[1], x)
+        SymmetricToeplitz(k1)
+    elseif x.step == y.step
+        k1 = k.(x, y[1])
+        k2 = k.(x[1], y)
+        Toeplitz(k1, k2)
+    else
+        Gramian(k, x, y)
+    end
 end
 
 # 1D stationary kernel on equi-spaced grid with periodic boundary conditions
@@ -228,7 +237,7 @@ end
 function BlockFactorizations.blockmul!(y::AbstractVecOfVecOrMat, G::Gramian, x::AbstractVecOfVecOrMat, α::Real = 1, β::Real = 0)
     Gijs = [G[1, 1] for _ in 1:Base.Threads.nthreads()] # pre-allocate storage for elements
     IT = input_trait(G.k)
-    @threads for i in eachindex(y)
+    @threads for i in eachindex(y) # NOTE: without threading, this does not allocate anything
         @. y[i] = β * y[i]
         Gij = Gijs[Base.Threads.threadid()]
         for j in eachindex(x)

test/gramian.jl

@@ -128,7 +128,29 @@ end
     @test G isa Circulant
     @test size(G) == (n, n)
     @test G*x ≈ Matrix(G)*x
-    # @test Matrix(G) ≈ Matrix(Gramian(k, x))
+
+    k = (x, y)->CovarianceFunctions.EQ()(x, y)
+    G = gramian(k, x)
+    @test G isa Gramian # if we can't infer the input trait of k, falls back to lazy representation
+    @test Matrix(G) ≈ Matrix(Gramian(k, x))
+
+    # testing that we can hook anonymous kernel into the system
+    CovarianceFunctions.input_trait(::typeof(k)) = CovarianceFunctions.IsotropicInput()
+    G = gramian(k, x)
+    @test G isa SymmetricToeplitz
+    @test Matrix(G) ≈ Matrix(Gramian(k, x))
+
+    # testing non-symmetric Toeplitz systems
+    y = x .+ randn() # shifted version of x with same step length
+    G = gramian(k, x, y)
+    @test G isa Toeplitz
+    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))
+
+    y = range(-1, 1, n÷2) # if y has different step length, defaults to gramian
+    G = gramian(k, x, y)
+    @test G isa Gramian
+    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))
+
 end # test solves and multiplications
 
 end # TestGramian