diff --git a/Project.toml b/Project.toml index 4c84d8a..58dd56e 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/src/abstractitensor.jl b/src/abstractitensor.jl index ddcc882..d3e5f1a 100644 --- a/src/abstractitensor.jl +++ b/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)) diff --git a/test/test_vectorinterface.jl b/test/test_vectorinterface.jl new file mode 100644 index 0000000..9e5603c --- /dev/null +++ b/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