Define substitute

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.8"
+version = "0.1.9"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/lazynameddimsarrays.jl

@@ -39,6 +39,16 @@ function Base.show(io::IO, m::Mul)
     print(io, "(", join(args, " $(operation(m)) "), ")")
     return nothing
 end
+function Base.hash(m::Mul, h::UInt64)
+    h = hash(:Mul, h)
+    for arg in arguments(m)
+        h = hash(arg, h)
+    end
+    return h
+end
+function map_arguments(f, m::Mul)
+    return Mul(map(f, arguments(m)))
+end
 
 @wrapped struct LazyNamedDimsArray{
         T, A <: AbstractNamedDimsArray{T},
@@ -65,6 +75,18 @@ function NamedDimsArrays.dename(a::LazyNamedDimsArray)
     end
 end
 
+function getindex_lazy(a::AbstractArray, I...)
+    u = unwrap(a)
+    if !iscall(u)
+        return u[I...]
+    else
+        return error("Indexing into expression not supported.")
+    end
+end
+function Base.getindex(a::LazyNamedDimsArray, I::Int...)
+    return getindex_lazy(a, I...)
+end
+
 function TermInterface.arguments(a::LazyNamedDimsArray)
     u = unwrap(a)
     if !iscall(u)
@@ -158,6 +180,49 @@ function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray)
     end
 end
 
+# Defined to avoid type piracy.
+# TODO: Define a proper hash function
+# in NamedDimsArrays.jl, maybe one that is
+# independent of the order of dimensions.
+function _hash(a::NamedDimsArray, h::UInt64)
+    h = hash(:NamedDimsArray, h)
+    h = hash(dename(a), h)
+    for i in inds(a)
+        h = hash(i, h)
+    end
+    return h
+end
+function _hash(x, h::UInt64)
+    return hash(x, h)
+end
+function Base.hash(a::LazyNamedDimsArray, h::UInt64)
+    h = hash(:LazyNamedDimsArray, h)
+    # Use `_hash`, which defines a custom hash for NamedDimsArray.
+    return _hash(unwrap(a), h)
+end
+
+generic_map(f, v) = map(f, v)
+generic_map(f, v::AbstractDict) = Dict(eachindex(v) .=> map(f, values(v)))
+generic_map(f, v::AbstractSet) = Set([f(x) for x in v])
+function map_arguments(f, a::LazyNamedDimsArray)
+    u = unwrap(a)
+    if !iscall(u)
+        return error("No arguments to map.")
+    elseif ismul(u)
+        return LazyNamedDimsArray(map_arguments(f, u))
+    else
+        return error("Variant not supported.")
+    end
+end
+function substitute(a::LazyNamedDimsArray, substitutions::AbstractDict)
+    haskey(substitutions, a) && return substitutions[a]
+    !iscall(a) && return a
+    return map_arguments(arg -> substitute(arg, substitutions), a)
+end
+function substitute(a::LazyNamedDimsArray, substitutions)
+    return substitute(a, Dict(substitutions))
+end
+
 function printnode(io::IO, a::LazyNamedDimsArray)
     return printnode(io, unwrap(a))
 end
@@ -280,6 +345,17 @@ Base.size(a::SymbolicArray) = length.(axes(a))
 function Base.:(==)(a::SymbolicArray, b::SymbolicArray)
     return symname(a) == symname(b) && axes(a) == axes(b)
 end
+function Base.hash(a::SymbolicArray, h::UInt64)
+    h = hash(:SymbolicArray, h)
+    h = hash(symname(a), h)
+    return hash(size(a), h)
+end
+function Base.getindex(a::SymbolicArray, I...)
+    return error("Indexing into SymbolicArray not supported.")
+end
+function Base.setindex!(a::SymbolicArray, value, I...)
+    return error("Indexing into SymbolicArray not supported.")
+end
 function Base.show(io::IO, mime::MIME"text/plain", a::SymbolicArray)
     Base.summary(io, a)
     println(io, ":")

test/test_lazynameddimsarrays.jl

@@ -1,8 +1,8 @@
 using AbstractTrees: AbstractTrees, print_tree, printnode
 using Base.Broadcast: materialize
 using ITensorNetworksNext.LazyNamedDimsArrays:
-    LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, symnameddims
-using NamedDimsArrays: NamedDimsArray, dename, dimnames, inds, nameddims
+    LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, substitute, symnameddims
+using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims
 using TermInterface:
     arguments,
     arity,
@@ -88,22 +88,32 @@ using WrappedUnions: unwrap
     end
 
     @testset "symnameddims" begin
-        a = symnameddims(:a)
-        b = symnameddims(:b)
-        c = symnameddims(:c)
-        @test a isa LazyNamedDimsArray
-        @test unwrap(a) isa NamedDimsArray
-        @test dename(a) isa SymbolicArray
-        @test dename(unwrap(a)) isa SymbolicArray
-        @test dename(unwrap(a)) == SymbolicArray(:a)
+        a1, a2, a3 = symnameddims.((:a1, :a2, :a3))
+        @test a1 isa LazyNamedDimsArray
+        @test unwrap(a1) isa NamedDimsArray
+        @test dename(a1) isa SymbolicArray
+        @test dename(unwrap(a1)) isa SymbolicArray
+        @test dename(unwrap(a1)) == SymbolicArray(:a1)
         @test inds(a) == ()
-        @test dimnames(a) == ()
+        @test dimnames(a1) == ()
 
-        ex = a * b * c
+        ex = a1 * a2 * a3
         @test copy(ex) == ex
-        @test arguments(ex) == [a * b, c]
+        @test arguments(ex) == [a1 * a2, a3]
         @test operation(ex) ≡ *
-        @test sprint(show, ex) == "((a * b) * c)"
-        @test sprint(show, MIME"text/plain"(), ex) == "((a * b) * c)"
+        @test sprint(show, ex) == "((a1 * a2) * a3)"
+        @test sprint(show, MIME"text/plain"(), ex) == "((a1 * a2) * a3)"
+    end
+
+    @testset "substitute" begin
+        s = symnameddims.((:a1, :a2, :a3))
+        i = @names i[1:4]
+        a = (randn(2, 2)[i[1], i[2]], randn(2, 2)[i[2], i[3]], randn(2, 2)[i[3], i[4]])
+        l = lazy.(a)
+
+        seq = s[1] * (s[2] * s[3])
+        net = substitute(seq, s .=> l)
+        @test net == l[1] * (l[2] * l[3])
+        @test arguments(net) == [l[1], l[2] * l[3]]
     end
 end