add ITensor interface

AFeuerpfeil · AFeuerpfeil · commit fe7cac9ba7c8 · 2025-12-17T14:09:36.000-05:00
diff --git a/ext/ITensorsExt/ITensorsExt.jl b/ext/ITensorsExt/ITensorsExt.jl
@@ -0,0 +1,20 @@
+module ITensorsExt 
+
+using ITensors 
+using TensorKit
+
+function TensorKit.TensorMap(src::ITensors.ITensor)
+    if hasqns(src)
+        return _blocksparse_tensormap(src)
+    end 
+    return _dense_tensormap(src)
+end
+
+
+include("itensor_utility.jl")
+include("ToTensorKit/dense.jl")
+include("ToTensorKit/blocksparse.jl")
+
+include("ToITensor/itensor.jl")
+
+end
diff --git a/ext/ITensorsExt/ToITensor/itensor.jl b/ext/ITensorsExt/ToITensor/itensor.jl
@@ -0,0 +1,144 @@
+const tk_tensor = Union{TensorMap, BraidingTensor}
+
+function ITensors.ITensor(t::tk_tensor;
+    tags::AbstractVector{<:AbstractString}=itensor_standard_tags(numind(t)),
+    ids=itensor_standard_ids(numind(t)),
+    plevs=itensor_standard_plevs(numind(t)),
+    qn_names::Union{<:AbstractVector{<:AbstractString},Missing}=missing,
+)::ITensor
+    T = eltype(t)
+    ## First we take t to be a tensor from a 0-dimensional codomain to a rank N-domain
+    t = repartition(t, numind(t), 0)
+    sp = TensorKit.space(t)
+    indices = [[(isdual(V) ? dual(c) : c) => TensorKit.dim(V, c) for c in sectors(V)] for V in sp.codomain]
+    direction = [isdual(V) ? ITensors.In : ITensors.Out for V in sp.codomain]
+    direction_sign = [isdual(V) ? -1 : 1 for V in sp.codomain]
+    if sectortype(sp) <: Trivial
+        sizes = last.(last.(indices))
+        index = map(eachindex(indices)) do i
+            id = ids[i]
+            dir = direction[i]
+            tag = tags[i]
+            plev = plevs[i]
+            return Index(id, sizes[i], dir, tag, plev)
+        end
+        iten = ITensors.DenseTensor(T, undef, index)
+        copyto!(iten.storage.data, t.data)
+        return ITensors.itensor(iten)
+    end
+    isa_product = sectortype(sp) <: ProductSector
+    if isa_product
+        moduli = convert_to_modulus.(fieldtypes(sectortype(sp).parameters[1]))
+    else
+        moduli = convert_to_modulus.([sectortype(sp)])
+    end
+
+    ismissing(qn_names) && (qn_names=itensor_standard_qn_names(length(moduli)))
+
+    all_u1s = findall(x -> x == 1, moduli)
+    multiply_by_two=false
+    for i in eachindex(indices)
+        space = indices[i]
+        for (c, _) in space
+            for i in all_u1s
+                charge = get_charge(isa_product ? c[i] : c; rescale=true)
+                if mod(charge,2)!=0
+                    multiply_by_two=true
+                    break
+                end
+            end
+            if multiply_by_two
+                break
+            end
+        end
+        if multiply_by_two
+            break
+        end
+    end
+
+    index_dict=[Dict{Sector, Int}() for _ in eachindex(indices)]
+
+    indices = map(eachindex(indices)) do i
+        space_tensorkit = indices[i]
+        space = map(x->construct_QN(x[1],direction_sign[i],qn_names,moduli, multiply_by_two, isa_product)=>x[2], space_tensorkit)
+        index_dict[i] = Dict(zip(first.(space_tensorkit),eachindex(space)))
+
+        id = ids[i]
+        dir = direction[i]
+        tag = tags[i]
+        plev = plevs[i]
+        return Index(id, space, dir, tag, plev)
+    end
+
+    N=numind(t)
+
+    all_blocks = ITensors.Block{N}[]
+    for (f1,f2) in fusiontrees(t)
+        push!(all_blocks, get_block(f1,f2,index_dict, Val(N)))
+    end
+    ITens = ITensors.BlockSparseTensor(T, undef, all_blocks, indices)
+    for (f1, f2) in fusiontrees(t)
+        b = get_block(f1,f2,index_dict, Val(N))
+        ITensors.blockview(ITens, b) .= t[f1, f2]
+    end
+    ITens = ITensors.itensor(ITens)
+    check_flux(ITens)
+    return combineblocks_itensor(ITens)
+end
+
+function construct_QN(x::Sector, sign::Int, qn_names, moduli, multiply_by_two::Bool=false, isa_product::Bool=false)
+    charges = map(eachindex(moduli)) do i
+        #charg = sign*
+        
+        charg = get_charge(isa_product ? x[i] : x; rescale=multiply_by_two)
+        if !isone(moduli[i])
+            charg = mod(charg, moduli[i])
+        end
+        return charg
+    end
+    return QN(map(i->(qn_names[i],charges[i],moduli[i]),eachindex(moduli))...)
+end
+
+function get_block(f1::FusionTree, f2::FusionTree, index_dict::Vector{Dict{Sector,Int}},::Val{N}) where {N}
+    charges= (
+        (flag ? dual(c) : c for (flag, c) in zip(f1.isdual, f1.uncoupled))...,
+        (flag ? dual(c) : c for (flag, c) in zip(f2.isdual, f2.uncoupled))...,
+    )
+    blockinds=ntuple(i-> index_dict[i][charges[i]], N)
+    return ITensors.Block{N}(blockinds)
+end
+
+function convert_to_modulus(::Type{U1Irrep})
+    return 1
+end
+function convert_to_modulus(::Type{ZNIrrep{N}}) where {N}
+    return N
+end
+function get_charge(x::ZNIrrep{N}; rescale::Bool=false) where {N}
+    return Int(x.n)
+end
+function get_charge(x::U1Irrep; rescale::Bool=false)
+    if !rescale
+        return Int(x.charge*1)
+    else
+        return Int(x.charge*2)
+    end
+end
+
+function itensor_standard_ids(N::Int)
+    return [rand(index_id_rng(), ITensors.IDType) for _ in 1:N]
+end
+function itensor_standard_plevs(N::Int)
+    return fill(0, N)
+end
+function itensor_standard_tags(N::Int)
+    return ["" for _ in 1:N]
+end
+function itensor_standard_qn_names(N::Int)
+    return [string(Char('A' + i - 1)) for i in 1:N]
+end
+function check_flux(T::ITensor)
+    if !iszero(T)
+        @assert iszero(flux(T))
+    end
+end
diff --git a/ext/ITensorsExt/ToTensorKit/blocksparse.jl b/ext/ITensorsExt/ToTensorKit/blocksparse.jl
@@ -0,0 +1,69 @@
+function _blocksparse_tensormap(tsrc::ITensors.ITensor)
+    N = ndims(tsrc)
+    T = eltype(tsrc)
+    indices = map(index, inds(tsrc))
+    allequal(Base.Fix2(getindex, 2), indices) || error("All QNs in the indices must have consistent symmetry types")
+    allequal(Base.Fix2(getindex, 3), indices) || error("All QNs in the indices must have consistent QN names")
+
+    symmetry = indices[1][2] 
+
+    charge_and_dims = getindex.(indices, 1)
+    directions = getindex.(indices, 4)
+    all_spaces_axes = map((x, y) -> canonicalize_space(x, y; symmetry), charge_and_dims, directions)
+    all_spaces = first.(all_spaces_axes)
+    all_axes = last.(all_spaces_axes)
+
+    tdst = zeros(T, prod(all_spaces))
+    t_tensor = ITensors.tensor(tsrc)
+    for b in eachnzblock(t_tensor)
+        key = b.data .% Int
+        block = Array(t_tensor[b])
+        charges = ntuple(N) do i
+            c = first(all_axes[i][key[i]])
+            return isdual(all_spaces[i]) ? dual(c) : c
+        end
+        axs = ntuple(N) do i
+            return last(all_axes[i][key[i]])
+        end
+        @assert isapprox(norm(tdst[charges][axs...]), 0; atol = 1.0e-10)
+        tdst[charges][axs...] .= block
+    end
+
+    return tdst
+end
+
+# little bit of annoyance: ITensors allows repeated charges, so have to merge them here
+# this helper function merges repeated charges and obtains the total space, while also
+# returning the axes for each original charge
+function canonicalize_space(inds, arrow; symmetry)
+    # determine the locations of each entry
+    axs = map(inds) do (label, dim)
+        length(label) == 1 && return (convert(symmetry, label[1]) => 1:dim)
+        return (convert(symmetry, label) => 1:dim)
+    end
+    for c in unique(first.(axs))
+        firstid = true
+        dlast = 0
+        for j in findall(==(c), first.(axs))
+            if firstid
+                dlast = last(last(axs[j]))
+                firstid = false
+            else
+                axs[j] = (c => (axs[j][2] .+ dlast))
+                dlast = last(axs[j][2])
+            end
+        end
+    end
+
+    # combine everything
+    combiner = Dict()
+    for (label, dim) in inds
+        label′ = length(label) == 1 ? convert(symmetry, label[1]) : convert(symmetry, label)
+        combiner[label′] = get(combiner, label′, 0) + dim
+    end
+    V = Vect[symmetry](label => dim for (label, dim) in combiner; dual = (arrow == -1))
+
+    return V, axs
+end
+
+
diff --git a/ext/ITensorsExt/ToTensorKit/dense.jl b/ext/ITensorsExt/ToTensorKit/dense.jl
@@ -0,0 +1,13 @@
+function _dense_tensormap(tsrc::ITensors.ITensor)
+    indices = map(index, inds(tsrc))
+
+    dims = getindex.(indices, 1)
+    directions = getindex.(indices, 4)
+    all_spaces = map(dims, directions) do dim, arrow
+        ComplexSpace(dim; dual = (arrow == -1))
+    end
+
+    tdest = zeros(eltype(tsrc), prod(all_spaces))
+    copyto!(tdest.data, ITensors.tensor(tsrc).storage.data)
+    return tdest
+end
diff --git a/ext/ITensorsExt/itensor_utility.jl b/ext/ITensorsExt/itensor_utility.jl
@@ -0,0 +1,62 @@
+function index(ind::Index{T}) where {T<:Integer}
+    direction = dir(ind) == ITensors.Out || dir(ind) == ITensors.Neither ? 1 : -1
+    tag = itensor_tag(ind)
+    return ind.space, direction, plev(ind), tag, ITensors.id(ind)
+end
+
+function index(ind::Index{<:AbstractVector{Pair{QN, T}}}) where {T<:Integer}
+    space = ind.space
+    blockdims = ITensors.blockdim.(space)
+    qns = first.(space)
+    qn_names = [[string(x[i].name) for i in eachindex(x)] for x in qns]
+    qn_names = filter(x->x!="",unique(reduce(vcat,qn_names)))
+    len = length(qn_names)
+    qn_moduli = [[x[i].modulus for i in 1:len] for x in qns]
+    qn_moduli = [filter(!iszero,unique(getindex.(qn_moduli,i))) for i in 1:len]
+    qn_moduli = map(x->isempty(x) ? 1 : only(x), qn_moduli)
+    index_output = [Tuple(x[j].val for j in 1:len)=>blockdims[i] for (i,x) in enumerate(qns)]
+
+    symmetry = symmetry_type(Tuple(qn_moduli))
+    direction = dir(ind)==ITensors.Out ? 1 : -1
+    
+    tag = itensor_tag(ind)
+    ids = ITensors.id(ind)
+    return index_output, symmetry, qn_names, direction, plev(ind), tag, ids
+end
+
+function itensor_tag(ind::Index)
+    ind_tag = tags(ind)
+    if string(ind_tag[1]) != ""
+        ind_tag = join(string.(ind_tag), ",")
+    else 
+        ind_tag = ""
+    end
+    return ind_tag 
+end
+
+function symmetry_type(t::NTuple{N, Int}) where {N}
+    factors = map(t) do sym
+        sym == 1 && return U1Irrep
+        return Irrep[TensorKit.ℤ{sym}]
+    end
+
+    if N == 1
+        return factors[1]
+    else
+        return ⊠(factors...)
+    end
+end
+
+function combineblocks_itensor(T::ITensor,inds_old=inds(T))
+    for i in eachindex(inds_old)
+        ind_old = inds_old[i]
+        C = combiner(ind_old; tags=tags(ind_old))
+        new_ind = uniqueind(C,T)
+        if inds_old[i].space != new_ind.space
+            T *= C
+            new_ind_old_id=Index(ITensors.id(ind_old), new_ind.space, dir(ind_old), tags(ind_old), plev(inds_old[i]))
+            T *= delta(dag(new_ind),new_ind_old_id)
+        end
+    end
+    return T
+end
diff --git a/test/test_itensor.jl b/test/test_itensor.jl
@@ -0,0 +1,73 @@
+using TensorKitAdapters
+using Test
+using TestExtras
+using ITensors
+using TensorKit 
+
+@testset "ITensor" begin
+    for T in [ComplexF64,ComplexF64, Float64]
+        for tag in [missing, "Link","Site,c=2,n=1"]
+            for plev in Iterators.product(0:2,0:4)
+                plev[1] == plev[2] && continue
+
+                tags = ismissing(tag) ? missing : [tag, tag]
+                if ismissing(tag) 
+                    i = Index(10;dir=ITensors.Out)
+                    iqn = ITensors.combineblocks(Index([QN(("N",2))=>5, QN("N",0)=>5]))[1]
+                    iqn2 = ITensors.combineblocks(Index([QN(("N",2))=>3, QN(("N",0),("P",3))=>7]))[1]
+                    ## This guarantees consistent sorting of the blocks for the comparison
+                    ## Here, we test combining blocks with the same QN:
+                    iqn3 = Index([QN("Sz",-1)=>2, QN("Sz",1)=>14,QN("Sz",-1)=>4])
+                else
+                    i = Index(10;dir=ITensors.Out,tags=tag)
+                    iqn = ITensors.combineblocks(Index([QN(("N",2))=>5, QN("N",0)=>5];tags=tag))[1]
+                    iqn2 = ITensors.combineblocks(Index([QN(("N",2))=>3, QN(("N",0),("P",3))=>7];tags=tag))[1]
+                    iqn3 = Index([QN("Sz",-1)=>2, QN("Sz",1)=>14,QN("Sz",-1)=>4]; tags=tag)
+                end
+
+                i_p1 = prime(i, plev[1])
+                iqn_p1 = prime(iqn, plev[1])
+                iqn2_p1 = prime(iqn2, plev[1]) 
+                iqn3_p1 = prime(iqn3, plev[1])
+
+                i_p2 = dag(prime(i, plev[2]))
+                iqn_p2 = dag(prime(iqn, plev[2]))
+                iqn2_p2 = dag(prime(iqn2, plev[2]))
+                iqn3_p2 = dag(prime(iqn3, plev[2]))
+
+                A = random_itensor(T,i_p1, i_p2)
+                B = random_itensor(T,iqn_p1, iqn_p2)
+                C = random_itensor(T,iqn2_p1, iqn2_p2)
+                D = random_itensor(T,iqn3_p1, iqn3_p2)
+
+                x = TensorMap(A)
+                y = TensorMap(B)
+                z = TensorMap(C)
+                w = TensorMap(D)
+
+                if ismissing(tag)
+                    A2 = ITensor(x;ids=ITensors.id.(inds(A)),plevs=plev,qn_names=["N","P"])
+                    B2 = ITensor(y;ids=ITensors.id.(inds(B)),plevs=plev,qn_names=["N"])
+                    C2 = ITensor(z;ids=ITensors.id.(inds(C)),plevs=plev,qn_names=["N","P"])
+                else
+                    A2 = ITensor(x;ids=ITensors.id.(inds(A)),tags=tags,plevs=plev,qn_names=["N","P"])
+                    B2 = ITensor(y;ids=ITensors.id.(inds(B)),tags=tags,plevs=plev,qn_names=["N"])
+                    C2 = ITensor(z;ids=ITensors.id.(inds(C)),tags=tags,plevs=plev,qn_names=["N","P"])
+                end
+                @test isapprox(norm(A2-A),0;atol=1.0e-10)
+                @test isapprox(norm(B2-B),0;atol=1.0e-10)
+                @test isapprox(norm(C2-C),0;atol=1.0e-10)
+
+
+                x2 = TensorMap(A2)
+                y2 = TensorMap(B2)
+                z2 = TensorMap(C2)
+                w2 = TensorMap(ITensor(w;ids=ITensors.id.(inds(D)),plevs=plev,qn_names=["N","P"]))
+                @test isapprox(norm(x2 - x), 0; atol=1.0e-10)
+                @test isapprox(norm(y2 - y), 0; atol=1.0e-10)
+                @test isapprox(norm(z2 - z), 0; atol=1.0e-10)
+                @test isapprox(norm(w2 - w), 0; atol=1.0e-10)
+            end
+        end
+    end
+end
