Add name-aware VectorInterface methods for ITensors

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.8.1"
+version = "0.8.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/abstractitensor.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra: LinearAlgebra
 using Random: Random
-using TensorAlgebra: permuteddims
+using TensorAlgebra: TensorAlgebra, permuteddims, zero!
 
 # Some of the interface is inspired by:
 # https://github.com/ITensor/ITensors.jl
@@ -226,9 +226,60 @@ Base.ndims(a::AbstractITensor) = ndims(unnamed(a))
 # Circumvent issue when eltype isn't known at compile time.
 Base.eltype(a::AbstractITensor) = eltype(unnamed(a))
 
-using VectorInterface: VectorInterface, scalartype
-# Circumvent issue when eltype isn't known at compile time.
+# In-place zero of an ITensor, delegating to its unnamed parent array.
+TensorAlgebra.zero!(a::AbstractITensor) = (zero!(unnamed(a)); a)
+
+# Name-aware `VectorInterface` methods so that ITensors can be used directly as the vectors
+# in iterative solvers such as `KrylovKit.eigsolve`, which drive their Krylov vectors through
+# `VectorInterface`; the generic `AbstractArray` fallbacks are not name-aware. The `!` methods
+# operate in place via broadcasting; each `!!` method does so too when the result fits the
+# destination's element type, and otherwise allocates. `scalartype` is computed in the value
+# domain because an ITensor's element type is not encoded in its type.
+using VectorInterface: VectorInterface, add, add!, scalartype, scale, scale!, zerovector!
 VectorInterface.scalartype(a::AbstractITensor) = scalartype(unnamed(a))
+function VectorInterface.scalartype(a::AbstractArray{<:AbstractITensor})
+    return mapreduce(scalartype, promote_type, a; init = Bool)
+end
+
+function VectorInterface.zerovector(a::AbstractITensor, ::Type{S}) where {S <: Number}
+    return zerovector!(similar(a, S))
+end
+VectorInterface.zerovector!(a::AbstractITensor) = zero!(a)
+VectorInterface.zerovector!!(a::AbstractITensor) = zerovector!(a)
+
+VectorInterface.scale(a::AbstractITensor, α::Number) = a * α
+function VectorInterface.scale!(a::AbstractITensor, α::Number)
+    @. a = a * α
+    return a
+end
+function VectorInterface.scale!(b::AbstractITensor, a::AbstractITensor, α::Number)
+    @. b = a * α
+    return b
+end
+function VectorInterface.scale!!(a::AbstractITensor, α::Number)
+    promote_type(scalartype(a), typeof(α)) <: scalartype(a) || return scale(a, α)
+    return scale!(a, α)
+end
+function VectorInterface.scale!!(b::AbstractITensor, a::AbstractITensor, α::Number)
+    promote_type(scalartype(b), scalartype(a), typeof(α)) <: scalartype(b) ||
+        return scale(a, α)
+    return scale!(b, a, α)
+end
+
+function VectorInterface.add(y::AbstractITensor, x::AbstractITensor, α::Number, β::Number)
+    return @. y * β + x * α
+end
+function VectorInterface.add!(y::AbstractITensor, x::AbstractITensor, α::Number, β::Number)
+    @. y = y * β + x * α
+    return y
+end
+function VectorInterface.add!!(y::AbstractITensor, x::AbstractITensor, α::Number, β::Number)
+    promote_type(scalartype(y), scalartype(x), typeof(α), typeof(β)) <: scalartype(y) ||
+        return add(y, x, α, β)
+    return add!(y, x, α, β)
+end
+
+VectorInterface.inner(x::AbstractITensor, y::AbstractITensor) = (conj(x) * y)[]
 
 Base.axes(a::AbstractITensor, dimname::Name) = axes(a, dim(a, dimname))
 Base.size(a::AbstractITensor, dimname::Name) = size(a, dim(a, dimname))

test/test_vectorinterface.jl

@@ -0,0 +1,51 @@
+import VectorInterface as VI
+using ITensorBase: dimnames, named, unnamed
+using Test: @test, @testset
+
+# These name-aware methods are what let an ITensor be used directly as a vector in
+# iterative solvers such as `KrylovKit.eigsolve`, which drive their Krylov vectors
+# through `VectorInterface`.
+@testset "VectorInterface (eltype=$(elt))" for elt in
+    (Float32, Float64, Complex{Float32})
+    i, j = named.(2, (:i, :j))
+    a = randn(elt, i, j)
+    b = randn(elt, j, i)
+    ua = unnamed(a)
+    ub = unnamed(b, dimnames(a))
+
+    @test VI.scalartype(a) === elt
+    @test VI.scalartype([a, b]) === elt
+    @test VI.scalartype([a, randn(ComplexF64, i, j)]) === ComplexF64
+
+    # zerovector / zerovector! / zerovector!!
+    z = VI.zerovector(a, ComplexF64)
+    @test VI.scalartype(z) === ComplexF64
+    @test iszero(unnamed(z))
+    @test dimnames(z) == dimnames(a)
+    z = VI.zerovector!(copy(a))
+    @test VI.scalartype(z) === elt
+    @test iszero(unnamed(z))
+    @test VI.zerovector!!(a) ≈ VI.zerovector(a)
+
+    # scale / scale! / scale!!
+    @test unnamed(VI.scale(a, 2)) ≈ 2 * ua
+    @test unnamed(VI.scale!(copy(a), 2)) ≈ 2 * ua
+    @test unnamed(VI.scale!(similar(a), a, 2)) ≈ 2 * ua
+    @test unnamed(VI.scale!!(copy(a), 2)) ≈ 2 * ua
+    @test unnamed(VI.scale!!(similar(a), a, 2)) ≈ 2 * ua
+    # `!!` allocates a new array when the scalar promotes beyond the element type.
+    s = VI.scale!!(copy(a), 2im)
+    @test VI.scalartype(s) === complex(elt)
+    @test unnamed(s) ≈ 2im * ua
+
+    # add / add! / add!!
+    @test unnamed(VI.add(b, a), dimnames(a)) ≈ ub + ua
+    @test unnamed(VI.add(b, a, 2, 3), dimnames(a)) ≈ 3 * ub + 2 * ua
+    @test unnamed(VI.add!(copy(b), a, 2, 3), dimnames(a)) ≈ 3 * ub + 2 * ua
+    @test unnamed(VI.add!!(copy(b), a, 2, 3), dimnames(a)) ≈ 3 * ub + 2 * ua
+    r = VI.add!!(copy(b), a, 2im, 3)
+    @test VI.scalartype(r) === complex(elt)
+    @test unnamed(r, dimnames(a)) ≈ 3 * ub + 2im * ua
+
+    @test VI.inner(a, b) ≈ VI.inner(ua, ub)
+end