Some code refactoring (#24)

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.2.0"
+version = "0.2.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/LazyNamedDimsArrays/lazyinterface.jl

@@ -10,7 +10,9 @@ function walk(opmap, argmap, ex)
     if !iscall(ex)
         return argmap(ex)
     else
-        return mapfoldl((args...) -> walk(opmap, argmap, args...), opmap(operation(ex)), arguments(ex))
+        return mapfoldl(opmap(operation(ex)), arguments(ex)) do (args...)
+            return walk(opmap, argmap, args...)
+        end
     end
 end
 # Walk the expression `ex`, modifying the
@@ -21,13 +23,22 @@ opwalk(opmap, a) = walk(opmap, identity, a)
 argwalk(argmap, a) = walk(identity, argmap, a)
 
 # Generic lazy functionality.
+using DerivableInterfaces: AbstractArrayInterface, InterfaceFunction
+struct LazyInterface{N} <: AbstractArrayInterface{N} end
+LazyInterface() = LazyInterface{Any}()
+LazyInterface(::Val{N}) where {N} = LazyInterface{N}()
+LazyInterface{M}(::Val{N}) where {M, N} = LazyInterface{N}()
+const lazy_interface = LazyInterface()
+
+const maketerm_lazy = lazy_interface(maketerm)
 function maketerm_lazy(type::Type, head, args, metadata)
     if head ≡ *
         return type(maketerm(Mul, head, args, metadata))
     else
         return error("Only mul supported right now.")
     end
 end
+const getindex_lazy = lazy_interface(getindex)
 function getindex_lazy(a::AbstractArray, I...)
     u = unwrap(a)
     if !iscall(u)
@@ -36,6 +47,7 @@ function getindex_lazy(a::AbstractArray, I...)
         return error("Indexing into expression not supported.")
     end
 end
+const arguments_lazy = lazy_interface(arguments)
 function arguments_lazy(a)
     u = unwrap(a)
     if !iscall(u)
@@ -46,18 +58,18 @@ function arguments_lazy(a)
         return error("Variant not supported.")
     end
 end
-function children_lazy(a)
-    return arguments(a)
-end
-function head_lazy(a)
-    return operation(a)
-end
-function iscall_lazy(a)
-    return iscall(unwrap(a))
-end
-function isexpr_lazy(a)
-    return iscall(a)
-end
+using TermInterface: children
+const children_lazy = lazy_interface(children)
+children_lazy(a) = arguments(a)
+using TermInterface: head
+const head_lazy = lazy_interface(head)
+head_lazy(a) = operation(a)
+const iscall_lazy = lazy_interface(iscall)
+iscall_lazy(a) = iscall(unwrap(a))
+using TermInterface: isexpr
+const isexpr_lazy = lazy_interface(isexpr)
+isexpr_lazy(a) = iscall(a)
+const operation_lazy = lazy_interface(operation)
 function operation_lazy(a)
     u = unwrap(a)
     if !iscall(u)
@@ -68,6 +80,7 @@ function operation_lazy(a)
         return error("Variant not supported.")
     end
 end
+const sorted_arguments_lazy = lazy_interface(sorted_arguments)
 function sorted_arguments_lazy(a)
     u = unwrap(a)
     if !iscall(u)
@@ -78,27 +91,35 @@ function sorted_arguments_lazy(a)
         return error("Variant not supported.")
     end
 end
-function sorted_children_lazy(a)
-    return sorted_arguments(a)
-end
+using TermInterface: sorted_children
+const sorted_children_lazy = lazy_interface(sorted_children)
+sorted_children_lazy(a) = sorted_arguments(a)
+const ismul_lazy = lazy_interface(ismul)
 ismul_lazy(a) = ismul(unwrap(a))
+using AbstractTrees: AbstractTrees
+const abstracttrees_children_lazy = lazy_interface(AbstractTrees.children)
 function abstracttrees_children_lazy(a)
     if !iscall(a)
         return ()
     else
         return arguments(a)
     end
 end
+using AbstractTrees: nodevalue
+const nodevalue_lazy = lazy_interface(nodevalue)
 function nodevalue_lazy(a)
     if !iscall(a)
         return unwrap(a)
     else
         return operation(a)
     end
 end
-materialize_lazy(a) = argwalk(unwrap, a)
 using Base.Broadcast: materialize
+const materialize_lazy = lazy_interface(materialize)
+materialize_lazy(a) = argwalk(unwrap, a)
+const copy_lazy = lazy_interface(copy)
 copy_lazy(a) = materialize(a)
+const equals_lazy = lazy_interface(==)
 function equals_lazy(a1, a2)
     u1, u2 = unwrap.((a1, a2))
     if !iscall(u1) && !iscall(u2)
@@ -109,6 +130,7 @@ function equals_lazy(a1, a2)
         return false
     end
 end
+const isequal_lazy = lazy_interface(isequal)
 function isequal_lazy(a1, a2)
     u1, u2 = unwrap.((a1, a2))
     if !iscall(u1) && !iscall(u2)
@@ -119,11 +141,13 @@ function isequal_lazy(a1, a2)
         return false
     end
 end
+const hash_lazy = lazy_interface(hash)
 function hash_lazy(a, h::UInt64)
     h = hash(Symbol(unspecify_type_parameters(typeof(a))), h)
     # Use `_hash`, which defines a custom hash for NamedDimsArray.
     return _hash(unwrap(a), h)
 end
+const map_arguments_lazy = lazy_interface(map_arguments)
 function map_arguments_lazy(f, a)
     u = unwrap(a)
     if !iscall(u)
@@ -134,19 +158,22 @@ function map_arguments_lazy(f, a)
         return error("Variant not supported.")
     end
 end
+function substitute end
+const substitute_lazy = lazy_interface(substitute)
 function substitute_lazy(a, substitutions::AbstractDict)
     haskey(substitutions, a) && return substitutions[a]
     !iscall(a) && return a
     return map_arguments(arg -> substitute(arg, substitutions), a)
 end
-function substitute_lazy(a, substitutions)
-    return substitute(a, Dict(substitutions))
-end
+substitute_lazy(a, substitutions) = substitute(a, Dict(substitutions))
+using AbstractTrees: printnode
+const printnode_lazy = lazy_interface(printnode)
 function printnode_lazy(io, a)
     # Use `printnode_nameddims` to avoid type piracy,
     # since it overloads on `AbstractNamedDimsArray`.
     return printnode_nameddims(io, unwrap(a))
 end
+const show_lazy = lazy_interface(show)
 function show_lazy(io::IO, a)
     if !iscall(a)
         return show(io, unwrap(a))
@@ -160,9 +187,12 @@ function show_lazy(io::IO, mime::MIME"text/plain", a)
     !iscall(a) ? show(io, mime, unwrap(a)) : show(io, a)
     return nothing
 end
+const add_lazy = lazy_interface(+)
 add_lazy(a1, a2) = error("Not implemented.")
+const sub_lazy = lazy_interface(-)
 sub_lazy(a) = error("Not implemented.")
 sub_lazy(a1, a2) = error("Not implemented.")
+const mul_lazy = lazy_interface(*)
 function mul_lazy(a)
     u = unwrap(a)
     if !iscall(u)
@@ -186,6 +216,7 @@ mul_lazy(a1::Number, a2::Number) = a1 * a2
 div_lazy(a1, a2::Number) = error("Not implemented.")
 
 # NamedDimsArrays.jl interface.
+const inds_lazy = lazy_interface(inds)
 function inds_lazy(a)
     u = unwrap(a)
     if !iscall(u)
@@ -196,6 +227,7 @@ function inds_lazy(a)
         return error("Variant not supported.")
     end
 end
+const dename_lazy = lazy_interface(dename)
 function dename_lazy(a)
     u = unwrap(a)
     if !iscall(u)

src/TensorNetworkGenerators/ising_network.jl

@@ -1,5 +1,6 @@
 using DiagonalArrays: DiagonalArray
 using Graphs: degree, dst, edges, src
+using ..ITensorNetworksNext: @preserve_graph
 using LinearAlgebra: Diagonal, eigen
 using NamedDimsArrays: apply, dename, inds, operator, randname
 using NamedGraphs.GraphsExtensions: vertextype
@@ -42,8 +43,8 @@ function ising_network(
         deg2 = degree(tn, v2)
         m = sqrt_ising_bond(β; J, h, deg1, deg2)
         t = operator(m, (f̃(e),), (f(e),))
-        tn[v1] = apply(t, tn[v1])
-        tn[v2] = apply(t, tn[v2])
+        @preserve_graph tn[v1] = apply(t, tn[v1])
+        @preserve_graph tn[v2] = apply(t, tn[v2])
     end
     return tn
 end

src/contract_network.jl

@@ -56,8 +56,8 @@ function contraction_order(tn; alg = default_kwargs(contraction_order, tn).alg,
 end
 # Convert the tensor network to a flat symbolic multiplication expression.
 function contraction_order(alg::Algorithm"flat", tn)
-    syms = [symnameddims(i, Tuple(inds(tn[i]))) for i in keys(tn)]
     # Same as: `reduce((a, b) -> *(a, b; flatten = true), syms)`.
+    syms = vec([symnameddims(i, Tuple(inds(tn[i]))) for i in keys(tn)])
     return lazy(Mul(syms))
 end
 function contraction_order(alg::Algorithm"left_associative", tn)

test/test_tensornetworkgenerators.jl

@@ -8,46 +8,7 @@ using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges
 using NamedGraphs.NamedGraphGenerators: named_grid
 using Test: @test, @testset
 
-module TestUtils
-    using QuadGK: quadgk
-    # Exact critical inverse temperature for 2D square lattice Ising model.
-    βc_2d_ising(elt::Type{<:Number} = Float64) = elt(log(1 + √2) / 2)
-    # Exact infinite volume free energy density for 2D square lattice Ising model.
-    function f_2d_ising(β::Real; J::Real = one(β))
-        κ = 2sinh(2β * J) / cosh(2β * J)^2
-        integrand(θ) = log((1 + √(abs(1 - (κ * sin(θ))^2))) / 2)
-        integral, _ = quadgk(integrand, 0, π)
-        return (-log(2cosh(2β * J)) - (1 / (2π)) * integral) / β
-    end
-    function f_1d_ising(β::Real; J::Real = one(β), h::Real = zero(β))
-        λ⁺ = exp(β * J) * (cosh(β * h) + √(sinh(β * h)^2 + exp(-4β * J)))
-        return -(log(λ⁺) / β)
-    end
-    function f_1d_ising(β::Real, N::Integer; periodic::Bool = true, kwargs...)
-        return if periodic
-            f_1d_ising_periodic(β, N; kwargs...)
-        else
-            f_1d_ising_open(β, N; kwargs...)
-        end
-    end
-    function f_1d_ising_periodic(β::Real, N::Integer; J::Real = one(β), h::Real = zero(β))
-        r = √(sinh(β * h)^2 + exp(-4β * J))
-        λ⁺ = exp(β * J) * (cosh(β * h) + r)
-        λ⁻ = exp(β * J) * (cosh(β * h) - r)
-        Z = λ⁺^N + λ⁻^N
-        return -(log(Z) / (β * N))
-    end
-    function f_1d_ising_open(β::Real, N::Integer; J::Real = one(β), h::Real = zero(β))
-        isone(N) && return 2cosh(β * h)
-        T = [
-            exp(β * (J + h)) exp(-β * J);
-            exp(-β * J) exp(β * (J - h));
-        ]
-        b = [exp(β * h / 2), exp(-β * h / 2)]
-        Z = (b' * (T^(N - 1)) * b)[]
-        return -(log(Z) / (β * N))
-    end
-end
+!@isdefined(TestUtils) && include("utils.jl")
 
 @testset "TensorNetworkGenerators" begin
     @testset "Delta Network" begin

test/utils.jl

@@ -0,0 +1,40 @@
+module TestUtils
+using QuadGK: quadgk
+# Exact critical inverse temperature for 2D square lattice Ising model.
+βc_2d_ising(elt::Type{<:Number} = Float64) = elt(log(1 + √2) / 2)
+# Exact infinite volume free energy density for 2D square lattice Ising model.
+function f_2d_ising(β::Real; J::Real = one(β))
+    κ = 2sinh(2β * J) / cosh(2β * J)^2
+    integrand(θ) = log((1 + √(abs(1 - (κ * sin(θ))^2))) / 2)
+    integral, _ = quadgk(integrand, 0, π)
+    return (-log(2cosh(2β * J)) - (1 / (2π)) * integral) / β
+end
+function f_1d_ising(β::Real; J::Real = one(β), h::Real = zero(β))
+    λ⁺ = exp(β * J) * (cosh(β * h) + √(sinh(β * h)^2 + exp(-4β * J)))
+    return -(log(λ⁺) / β)
+end
+function f_1d_ising(β::Real, N::Integer; periodic::Bool = true, kwargs...)
+    return if periodic
+        f_1d_ising_periodic(β, N; kwargs...)
+    else
+        f_1d_ising_open(β, N; kwargs...)
+    end
+end
+function f_1d_ising_periodic(β::Real, N::Integer; J::Real = one(β), h::Real = zero(β))
+    r = √(sinh(β * h)^2 + exp(-4β * J))
+    λ⁺ = exp(β * J) * (cosh(β * h) + r)
+    λ⁻ = exp(β * J) * (cosh(β * h) - r)
+    Z = λ⁺^N + λ⁻^N
+    return -(log(Z) / (β * N))
+end
+function f_1d_ising_open(β::Real, N::Integer; J::Real = one(β), h::Real = zero(β))
+    isone(N) && return 2cosh(β * h)
+    T = [
+        exp(β * (J + h)) exp(-β * J);
+        exp(-β * J) exp(β * (J - h));
+    ]
+    b = [exp(β * h / 2), exp(-β * h / 2)]
+    Z = (b' * (T^(N - 1)) * b)[]
+    return -(log(Z) / (β * N))
+end
+end