From 92c36d5d318e5b064f6f31e0bf34fc7a404f4482 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Thu, 28 Mar 2024 18:16:45 -0400
Subject: [PATCH 01/10] Operators code

---
 src/TensorNetworkFunctionals.jl             | 15 ++++-
 src/itensornetworkfunction.jl               |  2 -
 src/itensornetworks_elementary_functions.jl | 16 ++---
 src/itensornetworks_elementary_operators.jl | 73 +++++++++++++++++++++
 test/test_itensorfunction.jl                | 26 ++++----
 test/test_operators.jl                      | 46 +++++++++++++
 6 files changed, 151 insertions(+), 27 deletions(-)
 create mode 100644 src/itensornetworks_elementary_operators.jl
 create mode 100644 test/test_operators.jl

diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index d1c4cfe..3793f60 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -4,13 +4,19 @@ using ITensors
 using ITensorNetworks
 using NamedGraphs
 using EllipsisNotation
+using Dictionaries
 using Graphs
 
+using SplitApplyCombine: group
 using ITensorNetworks: delta_network
 using NamedGraphs: add_edges, random_bfs_tree, rem_edges
 
-#include("QTT_utils.jl")
+include("bitmaps.jl")
+include("itensornetworkfunction.jl")
+include("itensornetworks_elementary_functions.jl")
+include("itensornetworks_elementary_operators.jl")
 include("itensornetworksutils.jl")
+
 export ITensorNetworkFunction
 export BitMap,
   default_dimension_map, vertex, calculate_xyz, calculate_x, calculate_bit_values, dimension
@@ -22,8 +28,11 @@ export const_itensornetwork,
   cos_itensornetwork,
   sin_itensornetwork,
   get_edge_toward_root,
-  polynomial_itensornetwork
+  polynomial_itensornetwork,
+  Laplacian_operator,
+  derivative_operator
 export const_itn, poly_itn, cosh_itn, sinh_itn, tanh_itn, exp_itn, sin_itn, cos_itn
-export calculate_fx
+export calculate_fx, calculate_fxyz
+export operate
 
 end
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index c60690a..ac6e0b4 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -1,5 +1,3 @@
-include("bitmaps.jl")
-
 struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},BM<:BitMap} <:
        AbstractITensorNetwork{V}
   itensornetwork::TN
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index 2dc5ac7..cd7dac3 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -1,5 +1,3 @@
-include("itensornetworkfunction.jl")
-
 default_c_value() = 1.0
 default_a_value() = 0.0
 default_k_value() = 1.0
@@ -28,13 +26,13 @@ function exp_itensornetwork(
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ = copy(itensornetwork(const_itensornetwork(s, bit_map)))
+  ψ = const_itensornetwork(s, bit_map)
   Lx = length(vertices(bit_map, dimension))
   for v in vertices(bit_map, dimension)
     ψ[v] = ITensor([exp(a / Lx), exp(a / Lx) * exp(k / (2^bit(bit_map, v)))], inds(ψ[v]))
   end
 
-  return ITensorNetworkFunction(ψ, bit_map)
+  return ψ
 end
 
 """Construct the bond dim 2 representation of the cosh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
@@ -52,7 +50,7 @@ function cosh_itensornetwork(
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= 0.5
 
-  return ITensorNetworkFunction(ψ1 + ψ2, bit_map)
+  return ψ1 + ψ2
 end
 
 """Construct the bond dim 2 representation of the sinh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
@@ -70,7 +68,7 @@ function sinh_itensornetwork(
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= -0.5
 
-  return ITensorNetworkFunction(ψ1 + ψ2, bit_map)
+  return ψ1 + ψ2
 end
 
 """Construct the bond dim n representation of the tanh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
@@ -90,7 +88,7 @@ function tanh_itensornetwork(
     ψ = ψ + ψt
   end
 
-  return ITensorNetworkFunction(ψ, bit_map)
+  return ψ
 end
 
 """Construct the bond dim 2 representation of the cos(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
@@ -108,7 +106,7 @@ function cos_itensornetwork(
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= 0.5
 
-  return ITensorNetworkFunction(ψ1 + ψ2, bit_map)
+  return ψ1 + ψ2
 end
 
 """Construct the bond dim 2 representation of the sin(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
@@ -126,7 +124,7 @@ function sin_itensornetwork(
   ψ1[first(vertices(ψ1))] *= -0.5 * im
   ψ2[first(vertices(ψ1))] *= 0.5 * im
 
-  return ITensorNetworkFunction(ψ1 + ψ2, bit_map)
+  return ψ1 + ψ2
 end
 
 # #FUNCTIONS NEEDED TO IMPLEMENT POLYNOMIALS
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
new file mode 100644
index 0000000..6757885
--- /dev/null
+++ b/src/itensornetworks_elementary_operators.jl
@@ -0,0 +1,73 @@
+function plus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
+  @assert is_tree(s)
+  ttn_op = OpSum()
+  dim_vertices = vertices(bit_map, dimension)
+  L = length(dim_vertices)
+
+  string_site = [("S+", vertex(bit_map, dimension, L))]
+  add!(ttn_op, 1.0, "S+", vertex(bit_map, dimension, L))
+  for i in L:-1:2
+    pop!(string_site)
+    push!(string_site, ("S-", vertex(bit_map, dimension, i)))
+    push!(string_site, ("S+", vertex(bit_map, dimension, i - 1)))
+    add!(ttn_op, 1.0, (string_site...)...)
+  end
+
+  return TTN(ttn_op, s; algorithm="svd")
+end
+
+function minus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
+  @assert is_tree(s)
+  ttn_op = OpSum()
+  dim_vertices = vertices(bit_map, dimension)
+  L = length(dim_vertices)
+
+  string_site = [("S-", vertex(bit_map, dimension, L))]
+  add!(ttn_op, 1.0, "S-", vertex(bit_map, dimension, L))
+  for i in L:-1:2
+    pop!(string_site)
+    push!(string_site, ("S+", vertex(bit_map, dimension, i)))
+    push!(string_site, ("S-", vertex(bit_map, dimension, i - 1)))
+    add!(ttn_op, 1.0, (string_site...)...)
+  end
+
+  return TTN(ttn_op, s; algorithm="svd")
+end
+
+function no_shift_TTN(s::IndsNetwork)
+  ttn_op = OpSum()
+  string_site_full = [("I", v) for v in vertices(s)]
+  add!(ttn_op, 1.0, (string_site_full...)...)
+  return TTN(ttn_op, s; algorithm="svd")
+end
+
+function stencil(
+  s::IndsNetwork,
+  bit_map,
+  shifts::Vector{Float64},
+  delta_power::Int64;
+  dimension=default_dimension(),
+  truncate_kwargs...,
+)
+  @assert length(shifts) == 3
+  plus_shift = first(shifts) * plus_shift_TTN(s, bit_map; dimension)
+  minus_shift = last(shifts) * minus_shift_TTN(s, bit_map; dimension)
+  no_shift = shifts[2] * no_shift_TTN(s)
+
+  stencil_op = plus_shift + minus_shift + no_shift
+  stencil_op = truncate(stencil_op; truncate_kwargs...)
+
+  for v in vertices(stencil_op)
+    stencil_op[v] = (2^delta_power) * stencil_op[v]
+  end
+
+  return truncate(stencil_op; truncate_kwargs...)
+end
+
+function Laplacian_operator(s::IndsNetwork, bit_map; kwargs...)
+  return stencil(s, bit_map, [1.0, -2.0, 1.0], 2; kwargs...)
+end
+
+function derivative_operator(s::IndsNetwork, bit_map; kwargs...)
+  return 0.5 * stencil(s, bit_map, [1.0, 0.0, -1.0], 1; kwargs...)
+end
diff --git a/test/test_itensorfunction.jl b/test/test_itensorfunction.jl
index 7f19053..8018034 100644
--- a/test/test_itensorfunction.jl
+++ b/test/test_itensorfunction.jl
@@ -119,13 +119,12 @@ end
     vertex_to_bit_map = Dictionary(vertices(g), [v[2] for v in vertices(g)])
     bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
 
-
     ψ_fxyz = const_itn(s, bit_map; c)
 
-    x, y, z  = 0.5, 0.25, 0.0
-    vertex_to_bit_value_map = calculate_bit_values(ψ_fxyz, [x, y, z], [1,2, 3])
+    x, y, z = 0.5, 0.25, 0.0
+    vertex_to_bit_value_map = calculate_bit_values(ψ_fxyz, [x, y, z], [1, 2, 3])
 
-    fx_xyz = calculate_fxyz(ψ_fxyz, [x,y,z], [1,2,3])
+    fx_xyz = calculate_fxyz(ψ_fxyz, [x, y, z], [1, 2, 3])
     @test fx_xyz ≈ c
   end
 
@@ -140,7 +139,9 @@ end
   L = 10
   g = named_grid((L, 1))
   vertex_to_dimension_map = Dictionary(vertices(g), [(v[1] % 2) + 1 for v in vertices(g)])
-  vertex_to_bit_map = Dictionary(vertices(g), [ceil(Int64, v[1] * 0.5) for v in vertices(g)])
+  vertex_to_bit_map = Dictionary(
+    vertices(g), [ceil(Int64, v[1] * 0.5) for v in vertices(g)]
+  )
   bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
   x, y = 0.625, 0.25
 
@@ -150,11 +151,11 @@ end
       k = 0.125
       s = siteinds("S=1/2", g)
 
-      ψ_fx = net_func(s, bit_map; k, a, dimension = 1)
-      ψ_fy = net_func(s, bit_map; k, a, dimension = 2)
+      ψ_fx = net_func(s, bit_map; k, a, dimension=1)
+      ψ_fy = net_func(s, bit_map; k, a, dimension=2)
 
       ψ_fxy = ψ_fx + ψ_fy
-      fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1,2])
+      fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1, 2])
       @test func(k * x + a) + func(k * y + a) ≈ fxy_xy
     end
   end
@@ -173,12 +174,11 @@ end
     bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
 
     x, y = 0.625, 0.875
-    ψ_fx = tanh_itn(s, bit_map; k, a, nterms, dimension = 1)
-    ψ_fy = tanh_itn(s, bit_map; k, a, nterms, dimension = 2)
+    ψ_fx = tanh_itn(s, bit_map; k, a, nterms, dimension=1)
+    ψ_fy = tanh_itn(s, bit_map; k, a, nterms, dimension=2)
 
     ψ_fxy = ψ_fx + ψ_fy
-    fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1,2])
+    fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1, 2])
     @test tanh(k * x + a) + tanh(k * y + a) ≈ fxy_xy
   end
-
-end
\ No newline at end of file
+end
diff --git a/test/test_operators.jl b/test/test_operators.jl
new file mode 100644
index 0000000..2000029
--- /dev/null
+++ b/test/test_operators.jl
@@ -0,0 +1,46 @@
+using Test
+using TensorNetworkFunctionals
+using ITensors
+using ITensorNetworks
+using Random
+using Distributions
+using Graphs
+
+using TensorNetworkFunctionals: itensornetwork
+
+@testset "test laplacian on MPS" begin
+  g = named_grid((12, 1))
+  L = nv(g)
+  s = siteinds("S=1/2", g)
+  bit_map = BitMap(g)
+
+  ∇sq = Laplacian_operator(s, bit_map; cutoff=1e-10)
+  @test maxlinkdim(∇sq) == 3
+
+  ψ_fx = sin_itn(s, bit_map; k=Float64(pi))
+  ∂2x_ψ_fx = operate(∇sq, ψ_fx; truncate_kwargs=(; cutoff=1e-12))
+
+  xs = [0.00025, 0.09765625, 0.25, 0.625, 0.875]
+  for x in xs
+    ∂2x_ψ_fx_x = real(calculate_fx(∂2x_ψ_fx, x))
+    @test ∂2x_ψ_fx_x ≈ -pi * pi * sin(pi * x) atol = 1e-3
+  end
+end
+
+@testset "test derivative on tree" begin
+  g = named_comb_tree((3, 4))
+  L = nv(g)
+  s = siteinds("S=1/2", g)
+  bit_map = BitMap(g)
+
+  ∂_∂x = derivative_operator(s, bit_map; cutoff=1e-10)
+
+  ψ_fx = sin_itn(s, bit_map; k=Float64(pi))
+  ∂x_ψ_fx = operate(∂_∂x, ψ_fx; truncate_kwargs=(; cutoff=1e-12))
+
+  xs = [0.025, 0.1, 0.25, 0.625, 0.875]
+  for x in xs
+    ∂x_ψ_fx_x = real(calculate_fx(∂x_ψ_fx, x))
+    @test ∂x_ψ_fx_x ≈ pi * cos(pi * x) atol = 1e-3
+  end
+end

From 4caa5b2b64eb6d697dfd71ba905d673c54c3b9c8 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Thu, 28 Mar 2024 18:17:00 -0400
Subject: [PATCH 02/10] formating

---
 src/itensornetworksutils.jl | 36 +++++++++++++++---------------------
 1 file changed, 15 insertions(+), 21 deletions(-)

diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index 9ebddb1..d30f802 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -1,22 +1,3 @@
-#Imports
-using ITensors
-using Graphs
-using NamedGraphs
-using ITensorNetworks
-using EllipsisNotation
-using Dictionaries
-using DataGraphs
-
-using ITensorNetworks: delta_network, data_graph, data_graph_type
-using NamedGraphs: add_edges, random_bfs_tree, rem_edges
-
-using SplitApplyCombine
-
-using Random
-using Distributions
-
-include("itensornetworks_elementary_functions.jl")
-
 """Build the order L tensor corresponding to fx(x): x ∈ [0,1]."""
 function build_full_rank_tensor(L::Int64, fx::Function)
   inds = [Index(2, "$i") for i in 1:L]
@@ -47,11 +28,24 @@ function c_tensor(phys_ind::Index, virt_inds::Vector)
 end
 
 function copy_tensor_network(s::IndsNetwork, linkdim::Int64)
-  tn = randomITensorNetwork(s; link_space = linkdim)
+  tn = randomITensorNetwork(s; link_space=linkdim)
   for v in vertices(tn)
     virt_inds = setdiff(inds(tn[v]), Index[only(s[v])])
     tn[v] = c_tensor(only(s[v]), virt_inds)
   end
 
   return tn
-end
\ No newline at end of file
+end
+
+function operate(
+  operator::AbstractITensorNetwork,
+  ψ::ITensorNetworkFunction;
+  truncate_kwargs=(;),
+  kwargs...,
+)
+  ψ_tn = TTN(itensornetwork(ψ))
+  ψO_tn = noprime(ITensors.contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
+  ψO_tn = truncate(ψO_tn; truncate_kwargs...)
+
+  return ITensorNetworkFunction(ITensorNetwork(ψO_tn), bit_map(ψ))
+end

From 57020849c959bcbb91dbda914b2eef4cfe66ce16 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Fri, 29 Mar 2024 14:37:01 -0400
Subject: [PATCH 03/10] Cleared up NameSPace formatting in tests

---
 src/TensorNetworkFunctionals.jl             |   6 +-
 src/itensornetworkfunction.jl               |   4 +
 src/itensornetworks_elementary_operators.jl |  83 +++++++++++++--
 src/itensornetworksutils.jl                 |  13 ---
 test/test_bitmaps.jl                        |   8 +-
 test/test_itensorfunction.jl                |  17 +--
 test/test_operators.jl                      | 109 ++++++++++++++++++--
 7 files changed, 194 insertions(+), 46 deletions(-)

diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index 3793f60..d45cf42 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -11,11 +11,11 @@ using SplitApplyCombine: group
 using ITensorNetworks: delta_network
 using NamedGraphs: add_edges, random_bfs_tree, rem_edges
 
+include("itensornetworksutils.jl")
 include("bitmaps.jl")
 include("itensornetworkfunction.jl")
 include("itensornetworks_elementary_functions.jl")
 include("itensornetworks_elementary_operators.jl")
-include("itensornetworksutils.jl")
 
 export ITensorNetworkFunction
 export BitMap,
@@ -29,10 +29,10 @@ export const_itensornetwork,
   sin_itensornetwork,
   get_edge_toward_root,
   polynomial_itensornetwork,
-  Laplacian_operator,
+  laplacian_operator,
   derivative_operator
 export const_itn, poly_itn, cosh_itn, sinh_itn, tanh_itn, exp_itn, sin_itn, cos_itn
 export calculate_fx, calculate_fxyz
-export operate
+export operate, apply_gx_operator, multiply
 
 end
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index ac6e0b4..2a3c9a7 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -64,6 +64,10 @@ function calculate_fxyz(
   return ITensors.contract(fitn_xyz)[]
 end
 
+function calculate_fxyz(fitn::ITensorNetworkFunction, xs::Vector{Float64})
+  return calculate_fxyz(fitn, xs, [i for i in 1:length(xs)])
+end
+
 function calculate_fx(fitn::ITensorNetworkFunction, x::Float64)
   @assert dimension(fitn) == 1
   return calculate_fxyz(fitn, [x], [1])
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index 6757885..4482a99 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -1,4 +1,4 @@
-function plus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
+function plus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
   @assert is_tree(s)
   ttn_op = OpSum()
   dim_vertices = vertices(bit_map, dimension)
@@ -16,7 +16,7 @@ function plus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
   return TTN(ttn_op, s; algorithm="svd")
 end
 
-function minus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
+function minus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
   @assert is_tree(s)
   ttn_op = OpSum()
   dim_vertices = vertices(bit_map, dimension)
@@ -34,7 +34,7 @@ function minus_shift_TTN(s::IndsNetwork, bit_map; dimension=default_dimension())
   return TTN(ttn_op, s; algorithm="svd")
 end
 
-function no_shift_TTN(s::IndsNetwork)
+function no_shift_ttn(s::IndsNetwork)
   ttn_op = OpSum()
   string_site_full = [("I", v) for v in vertices(s)]
   add!(ttn_op, 1.0, (string_site_full...)...)
@@ -50,24 +50,91 @@ function stencil(
   truncate_kwargs...,
 )
   @assert length(shifts) == 3
-  plus_shift = first(shifts) * plus_shift_TTN(s, bit_map; dimension)
-  minus_shift = last(shifts) * minus_shift_TTN(s, bit_map; dimension)
-  no_shift = shifts[2] * no_shift_TTN(s)
+  plus_shift = first(shifts) * plus_shift_ttn(s, bit_map; dimension)
+  minus_shift = last(shifts) * minus_shift_ttn(s, bit_map; dimension)
+  no_shift = shifts[2] * no_shift_ttn(s)
 
   stencil_op = plus_shift + minus_shift + no_shift
   stencil_op = truncate(stencil_op; truncate_kwargs...)
 
-  for v in vertices(stencil_op)
+  for v in vertices(bit_map, dimension)
     stencil_op[v] = (2^delta_power) * stencil_op[v]
   end
 
   return truncate(stencil_op; truncate_kwargs...)
 end
 
-function Laplacian_operator(s::IndsNetwork, bit_map; kwargs...)
+function laplacian_operator(s::IndsNetwork, bit_map; kwargs...)
   return stencil(s, bit_map, [1.0, -2.0, 1.0], 2; kwargs...)
 end
 
 function derivative_operator(s::IndsNetwork, bit_map; kwargs...)
   return 0.5 * stencil(s, bit_map, [1.0, 0.0, -1.0], 1; kwargs...)
 end
+
+function apply_gx_operator(gx::ITensorNetworkFunction)
+  gx = copy(gx)
+  operator = itensornetwork(gx)
+  s = siteinds(operator)
+  for v in vertices(operator)
+    sind = s[v]
+    sindsim = sim(sind)
+    operator[v] = replaceinds!(operator[v], sind, sindsim)
+    operator[v] = operator[v] * delta(vcat(sind, sindsim, sind'))
+  end
+  return operator
+end
+
+function multiply(gx::ITensorNetworkFunction, fx::ITensorNetworkFunction)
+  @assert vertices(gx) == vertices(fx)
+  fx, fxgx = copy(fx), copy(gx)
+  s = siteinds(fxgx)
+  for v in vertices(fxgx)
+    ssim = sim(s[v])
+    temp_tensor = replaceinds!(fx[v], s[v], ssim)
+    fxgx[v] = noprime!(fxgx[v] * delta(s[v], s[v]', ssim) * temp_tensor)
+  end
+
+  return combine_linkinds(fxgx)
+end
+
+function multiply(
+  gx::ITensorNetworkFunction,
+  fx::ITensorNetworkFunction,
+  hx::ITensorNetworkFunction,
+  fs::ITensorNetworkFunction...,
+)
+  return multiply(multiply(gx, fx), hx, fs...)
+end
+
+Base.:*(fs::ITensorNetworkFunction...) = multiply(fs...)
+
+function operate(
+  operator::TreeTensorNetwork, ψ::ITensorNetworkFunction; truncate_kwargs=(;), kwargs...
+)
+  ψ_tn = TTN(itensornetwork(ψ))
+  ψO_tn = noprime(ITensors.contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
+  ψO_tn = truncate(ψO_tn; truncate_kwargs...)
+
+  return ITensorNetworkFunction(ITensorNetwork(ψO_tn), bit_map(ψ))
+end
+
+function operate(operator::ITensorNetwork, ψ::ITensorNetworkFunction; kwargs...)
+  return operate(TTN(operator), ψ; kwargs...)
+end
+
+function operate(
+  operators::Vector{TreeTensorNetwork{V}}, ψ::ITensorNetworkFunction; kwargs...
+) where {V}
+  ψ = copy(ψ)
+  for op in operators
+    ψ = operate(op, ψ; kwargs...)
+  end
+  return ψ
+end
+
+function operate(
+  operators::Vector{ITensorNetwork{V}}, ψ::ITensorNetworkFunction; kwargs...
+) where {V}
+  return operate(TTN.(operators), ψ; kwargs...)
+end
diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index d30f802..9624697 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -36,16 +36,3 @@ function copy_tensor_network(s::IndsNetwork, linkdim::Int64)
 
   return tn
 end
-
-function operate(
-  operator::AbstractITensorNetwork,
-  ψ::ITensorNetworkFunction;
-  truncate_kwargs=(;),
-  kwargs...,
-)
-  ψ_tn = TTN(itensornetwork(ψ))
-  ψO_tn = noprime(ITensors.contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
-  ψO_tn = truncate(ψO_tn; truncate_kwargs...)
-
-  return ITensorNetworkFunction(ITensorNetwork(ψO_tn), bit_map(ψ))
-end
diff --git a/test/test_bitmaps.jl b/test/test_bitmaps.jl
index f78e0f2..f431c75 100644
--- a/test/test_bitmaps.jl
+++ b/test/test_bitmaps.jl
@@ -1,11 +1,9 @@
 using Test
 using TensorNetworkFunctionals
-using ITensors
-using ITensorNetworks
-using SplitApplyCombine
-using Dictionaries
 
-#include("../src/itensornetworksutils.jl")
+using NamedGraphs: named_grid, vertices
+using ITensors: siteinds
+using Dictionaries: Dictionary
 
 @testset "test single dimensional bit map" begin
   L = 4
diff --git a/test/test_itensorfunction.jl b/test/test_itensorfunction.jl
index 8018034..b965caf 100644
--- a/test/test_itensorfunction.jl
+++ b/test/test_itensorfunction.jl
@@ -1,11 +1,14 @@
 using Test
 using TensorNetworkFunctionals
-using ITensorNetworks
-using Random
-using Distributions
-using Graphs
 
-#include("../src/itensornetworksutils.jl")
+using Graphs: SimpleGraph, uniform_tree
+using NamedGraphs: NamedGraph, named_grid, vertices, named_comb_tree, rename_vertices
+using ITensors: siteinds
+using ITensorNetworks: randomITensorNetwork
+using Dictionaries: Dictionary
+using SplitApplyCombine: group
+using Random: seed!
+using Distributions: Uniform
 
 @testset "test constructor from ITensorNetwork" begin
   L = 10
@@ -90,8 +93,8 @@ end
 
     ###Generate a series of random polynomials on random graphs. Evaluate them at random x values"""
     for deg in degrees
-      Random.seed!(1234 * deg)
-      g = ITensorNetworks.NamedGraph(Graphs.SimpleGraph(uniform_tree(L)))
+      seed!(1234 * deg)
+      g = NamedGraph(SimpleGraph(uniform_tree(L)))
       g = rename_vertices(g, Dict(zip(vertices(g), [(v, 1) for v in vertices(g)])))
       s = siteinds("S=1/2", g)
 
diff --git a/test/test_operators.jl b/test/test_operators.jl
index 2000029..bbe26e8 100644
--- a/test/test_operators.jl
+++ b/test/test_operators.jl
@@ -1,20 +1,20 @@
 using Test
 using TensorNetworkFunctionals
-using ITensors
-using ITensorNetworks
-using Random
-using Distributions
-using Graphs
 
+using ITensors: siteinds
+using ITensorNetworks: maxlinkdim
+using Graphs: SimpleGraph, uniform_tree
+using NamedGraphs: named_grid, named_comb_tree, NamedGraph, nv, vertices
 using TensorNetworkFunctionals: itensornetwork
+using Dictionaries: Dictionary
 
-@testset "test laplacian on MPS" begin
+@testset "test laplacian in 1D on MPS" begin
   g = named_grid((12, 1))
   L = nv(g)
   s = siteinds("S=1/2", g)
   bit_map = BitMap(g)
 
-  ∇sq = Laplacian_operator(s, bit_map; cutoff=1e-10)
+  ∇sq = laplacian_operator(s, bit_map; cutoff=1e-10)
   @test maxlinkdim(∇sq) == 3
 
   ψ_fx = sin_itn(s, bit_map; k=Float64(pi))
@@ -27,8 +27,8 @@ using TensorNetworkFunctionals: itensornetwork
   end
 end
 
-@testset "test derivative on tree" begin
-  g = named_comb_tree((3, 4))
+@testset "test derivative in 1D on tree" begin
+  g = named_comb_tree((4, 4))
   L = nv(g)
   s = siteinds("S=1/2", g)
   bit_map = BitMap(g)
@@ -41,6 +41,95 @@ end
   xs = [0.025, 0.1, 0.25, 0.625, 0.875]
   for x in xs
     ∂x_ψ_fx_x = real(calculate_fx(∂x_ψ_fx, x))
-    @test ∂x_ψ_fx_x ≈ pi * cos(pi * x) atol = 1e-3
+    @test ∂x_ψ_fx_x ≈ pi * cos(pi * x) atol = 1e-4
+  end
+end
+
+@testset "test multiplication_operator_in_1D" begin
+  g = named_comb_tree((4, 4))
+  L = nv(g)
+  s = siteinds("S=1/2", g)
+  bit_map = BitMap(g)
+
+  ψ_gx = sin_itn(s, bit_map; k=0.5 * Float64(pi))
+  ψ_fx = cos_itn(s, bit_map; k=0.25 * Float64(pi))
+
+  ψ_fxgx = ψ_gx * ψ_fx
+  xs = [0.025, 0.1, 0.25, 0.625, 0.875]
+  for x in xs
+    ψ_fxgx_x = real(calculate_fx(ψ_fxgx, x))
+    @test ψ_fxgx_x ≈ sin(0.5 * pi * x) * cos(0.25 * pi * x) atol = 1e-4
+  end
+end
+
+@testset "test multiplication_operator_in_2D" begin
+  L = 10
+  g = NamedGraph(SimpleGraph(uniform_tree(L)))
+  s = siteinds("S=1/2", g)
+
+  vertex_to_dimension_map = Dictionary(
+    vertices(g), [(first(v) % 2) + 1 for v in vertices(g)]
+  )
+  vertex_to_bit_map = Dictionary(
+    vertices(g), [ceil(Int64, first(v) * 0.5) for v in vertices(g)]
+  )
+  bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
+
+  ψ_fx = cos_itn(s, bit_map; k=0.25 * Float64(pi), dimension=1)
+  ψ_gy = sin_itn(s, bit_map; k=0.5 * Float64(pi), dimension=2)
+  @assert dimension(ψ_fx) == dimension(ψ_gy) == 2
+
+  ψ_fxgy = ψ_fx * ψ_gy
+
+  xs = [0.125, 0.25, 0.625, 0.875]
+  ys = [0.125, 0.25, 0.625, 0.875]
+  for x in xs
+    for y in ys
+      ψ_fx_x = real(calculate_fxyz(ψ_fx, [x, y]))
+      ψ_gy_y = real(calculate_fxyz(ψ_gy, [x, y]))
+      @test ψ_fx_x ≈ cos(0.25 * pi * x)
+      @test ψ_gy_y ≈ sin(0.5 * pi * y)
+      ψ_fxgy_xy = real(calculate_fxyz(ψ_fxgy, [x, y]))
+      @test ψ_fxgy_xy ≈ cos(0.25 * pi * x) * sin(0.5 * pi * y) atol = 1e-3
+    end
+  end
+end
+
+@testset "test differentiation_operator_on_3D_function" begin
+  L = 60
+  g = named_grid((L, 1))
+  s = siteinds("S=1/2", g)
+
+  vertex_to_dimension_map = Dictionary(
+    vertices(g), [(first(v) % 3) + 1 for v in vertices(g)]
+  )
+  vertex_to_bit_map = Dictionary(
+    vertices(g), [ceil(Int64, first(v) / 3) for v in vertices(g)]
+  )
+  bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
+  ψ_fx = sin_itn(s, bit_map; k=Float64(pi), dimension=1)
+  ψ_gy = sin_itn(s, bit_map; k=Float64(pi), dimension=2)
+  ψ_hz = sinh_itn(s, bit_map; k=Float64(pi), dimension=3)
+  @assert dimension(ψ_fx) == dimension(ψ_gy) == dimension(ψ_hz) == 3
+
+  ψ_fxgyhz = ψ_fx * ψ_gy * ψ_hz
+
+  ∂_∂y = derivative_operator(s, bit_map; dimension=2, cutoff=1e-10)
+
+  ∂_∂y_ψ_fxgyhz = operate([∂_∂y], ψ_fxgyhz; truncate_kwargs=(; cutoff=1e-10))
+
+  xs = [0.125, 0.25, 0.675]
+  ys = [0.125, 0.25, 0.675]
+  zs = [0.125, 0.25, 0.675]
+  for x in xs
+    for y in ys
+      for z in zs
+        ψ_fxgyhz_xyz = real(calculate_fxyz(ψ_fxgyhz, [x, y, z]))
+        @test ψ_fxgyhz_xyz ≈ sin(pi * x) * sin(pi * y) * sinh(pi * z) atol = 1e-3
+
+        ∂_∂y_ψ_fxgyhz_xyz = real(calculate_fxyz(∂_∂y_ψ_fxgyhz, [x, y, z]))
+        @test ∂_∂y_ψ_fxgyhz_xyz ≈ pi * sin(pi * x) * cos(pi * y) * sinh(pi * z) atol = 1e-3
+      end
+    end
   end
 end

From 207286c3b64ebf0ec880a61db1e8ee0f3c7bd41c Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Mon, 1 Apr 2024 13:53:00 -0400
Subject: [PATCH 04/10] Namespace fixing

---
 src/TensorNetworkFunctionals.jl             | 10 +--------
 src/bitmaps.jl                              |  2 ++
 src/itensornetworkfunction.jl               |  7 +++++--
 src/itensornetworks_elementary_functions.jl | 23 +++++++++++++--------
 src/itensornetworks_elementary_operators.jl | 18 +++++++++++++++-
 src/itensornetworksutils.jl                 | 11 ++++++----
 6 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index d45cf42..04189d2 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -1,15 +1,7 @@
 module TensorNetworkFunctionals
 
-using ITensors
-using ITensorNetworks
-using NamedGraphs
-using EllipsisNotation
-using Dictionaries
 using Graphs
-
-using SplitApplyCombine: group
-using ITensorNetworks: delta_network
-using NamedGraphs: add_edges, random_bfs_tree, rem_edges
+using ITensorNetworks
 
 include("itensornetworksutils.jl")
 include("bitmaps.jl")
diff --git a/src/bitmaps.jl b/src/bitmaps.jl
index 059ef24..eed4357 100644
--- a/src/bitmaps.jl
+++ b/src/bitmaps.jl
@@ -1,3 +1,5 @@
+using Dictionaries: Dictionary, set!
+
 struct BitMap{VB,VD}
   vertex_bit::VB
   vertex_dimension::VD
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index 2a3c9a7..13a6f5b 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -1,3 +1,6 @@
+using ITensorNetworks: data_graph_type, AbstractITensorNetwork
+using ITensors: ITensor, dim, contract, siteinds
+
 struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},BM<:BitMap} <:
        AbstractITensorNetwork{V}
   itensornetwork::TN
@@ -49,7 +52,7 @@ function project(fitn::ITensorNetworkFunction, vertex_to_bit_value_map)
   s = siteinds(fitn)
   for v in keys(vertex_to_bit_value_map)
     proj = ITensor(
-      [i != vertex_to_bit_value_map[v] ? 0 : 1 for i in 0:(ITensors.dim(s[v]) - 1)], s[v]
+      [i != vertex_to_bit_value_map[v] ? 0 : 1 for i in 0:(dim(s[v]) - 1)], s[v]
     )
     fitn[v] = fitn[v] * proj
   end
@@ -61,7 +64,7 @@ function calculate_fxyz(
 )
   vertex_to_bit_value_map = calculate_bit_values(fitn, xs, dimensions)
   fitn_xyz = project(fitn, vertex_to_bit_value_map)
-  return ITensors.contract(fitn_xyz)[]
+  return contract(fitn_xyz)[]
 end
 
 function calculate_fxyz(fitn::ITensorNetworkFunction, xs::Vector{Float64})
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index cd7dac3..b901339 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -1,3 +1,8 @@
+using NamedGraphs:
+  vertices, random_bfs_tree, rem_edges, add_edges, undirected_graph, NamedEdge
+using ITensors: dim, commoninds
+using ITensorNetworks: IndsNetwork
+
 default_c_value() = 1.0
 default_a_value() = 0.0
 default_k_value() = 1.0
@@ -8,7 +13,7 @@ default_dimension() = 1
 function const_itensornetwork(
   s::IndsNetwork, bit_map; c::Union{Float64,ComplexF64}=default_c_value()
 )
-  ψ = copy_tensor_network(s, 1)
+  ψ = copy_tensor_network(s)
   inv_L = Float64(1.0 / nv(s))
   for v in vertices(ψ)
     ψ[v] *= c^inv_L
@@ -148,12 +153,12 @@ function Q_N_tensor(
   N::Int64, siteind::Index, αind::Vector{Index}, betaind::Index, xivals::Vector{Float64}
 )
   @assert length(αind) == N - 1
-  @assert length(xivals) == ITensors.dim(siteind)
-  n = ITensors.dim(betaind) - 1
-  @assert all(x -> x == n + 1, ITensors.dim.(αind))
+  @assert length(xivals) == dim(siteind)
+  n = dim(betaind) - 1
+  @assert all(x -> x == n + 1, dim.(αind))
 
   link_dims = [n + 1 for i in 1:N]
-  dims = vcat([ITensors.dim(siteind)], link_dims)
+  dims = vcat([dim(siteind)], link_dims)
   Q_N_array = zeros(Tuple(dims))
   for (i, xi) in enumerate(xivals)
     for j in 0:((n + 1)^(N) - 1)
@@ -198,11 +203,11 @@ function polynomial_itensornetwork(
   root_vertices_dim = filter(v -> v ∈ dimension_vertices, leaf_vertices(s_tree))
   @assert !isempty(root_vertices_dim)
   root_vertex = first(root_vertices_dim)
-  ψ = copy_tensor_network(s_tree, n + 1)
+  ψ = copy_tensor_network(s_tree; linkdim=n + 1)
   #Place the Q_n tensors, making sure we get the right index pointing towards the root
-  for v in vertices(ψ)
+  for v in dimension_vertices
     siteindex = s_tree[v][]
-    if v != root_vertex && v ∈ dimension_vertices
+    if v != root_vertex
       e = get_edge_toward_root(g_tree, v, root_vertex)
       betaindex = first(commoninds(ψ, e))
       alphas = setdiff(inds(ψ[v]), Index[siteindex, betaindex])
@@ -213,7 +218,7 @@ function polynomial_itensornetwork(
         betaindex,
         [0.0, (1.0 / (2^bit(bit_map, v)))],
       )
-    elseif v == root_vertex && v ∈ dimension_vertices
+    elseif v == root_vertex
       betaindex = Index(n + 1, "DummyInd")
       alphas = setdiff(inds(ψ[v]), Index[siteindex])
       ψv = Q_N_tensor(2, siteindex, alphas, betaindex, [0.0, (1.0 / (2^bit(bit_map, v)))])
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index 4482a99..f1f3c96 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -1,3 +1,19 @@
+using Graphs: is_tree
+using NamedGraphs: undirected_graph
+using ITensors:
+  sim,
+  OpSum,
+  siteinds,
+  noprime,
+  truncate,
+  replaceinds!,
+  delta,
+  add!,
+  prime,
+  noprime!,
+  contract
+using ITensorNetworks: IndsNetwork, TTN, TreeTensorNetwork, combine_linkinds
+
 function plus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
   @assert is_tree(s)
   ttn_op = OpSum()
@@ -113,7 +129,7 @@ function operate(
   operator::TreeTensorNetwork, ψ::ITensorNetworkFunction; truncate_kwargs=(;), kwargs...
 )
   ψ_tn = TTN(itensornetwork(ψ))
-  ψO_tn = noprime(ITensors.contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
+  ψO_tn = noprime(contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
   ψO_tn = truncate(ψO_tn; truncate_kwargs...)
 
   return ITensorNetworkFunction(ITensorNetwork(ψO_tn), bit_map(ψ))
diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index 9624697..d4fa1c0 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -1,3 +1,6 @@
+using ITensors: Index, dim, inds
+using ITensorNetworks: randomITensorNetwork, IndsNetwork
+
 """Build the order L tensor corresponding to fx(x): x ∈ [0,1]."""
 function build_full_rank_tensor(L::Int64, fx::Function)
   inds = [Index(2, "$i") for i in 1:L]
@@ -14,11 +17,11 @@ end
 
 function c_tensor(phys_ind::Index, virt_inds::Vector)
   inds = vcat(phys_ind, virt_inds)
-  @assert allequal(ITensors.dim.(virt_inds))
+  @assert allequal(dim.(virt_inds))
   #Build tensor to be delta on inds and independent of phys_ind
   T = ITensor(0.0, inds...)
-  for i in 1:ITensors.dim(phys_ind)
-    for j in 1:ITensors.dim(first(virt_inds))
+  for i in 1:dim(phys_ind)
+    for j in 1:dim(first(virt_inds))
       ind_array = [v => j for v in virt_inds]
       T[phys_ind => i, ind_array...] = 1.0
     end
@@ -27,7 +30,7 @@ function c_tensor(phys_ind::Index, virt_inds::Vector)
   return T
 end
 
-function copy_tensor_network(s::IndsNetwork, linkdim::Int64)
+function copy_tensor_network(s::IndsNetwork; linkdim::Int64=1)
   tn = randomITensorNetwork(s; link_space=linkdim)
   for v in vertices(tn)
     virt_inds = setdiff(inds(tn[v]), Index[only(s[v])])

From 7ecd5a5827967cb3823abee0caacb14cbafff111 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Mon, 1 Apr 2024 16:15:31 -0400
Subject: [PATCH 05/10] NameSpace Continued Fixes

---
 src/TensorNetworkFunctionals.jl             | 3 ---
 src/bitmaps.jl                              | 1 +
 src/itensornetworkfunction.jl               | 3 ++-
 src/itensornetworks_elementary_functions.jl | 5 +++--
 src/itensornetworks_elementary_operators.jl | 2 +-
 5 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index 04189d2..dbd9480 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -1,8 +1,5 @@
 module TensorNetworkFunctionals
 
-using Graphs
-using ITensorNetworks
-
 include("itensornetworksutils.jl")
 include("bitmaps.jl")
 include("itensornetworkfunction.jl")
diff --git a/src/bitmaps.jl b/src/bitmaps.jl
index eed4357..3b513c3 100644
--- a/src/bitmaps.jl
+++ b/src/bitmaps.jl
@@ -1,4 +1,5 @@
 using Dictionaries: Dictionary, set!
+using Graphs: Graphs
 
 struct BitMap{VB,VD}
   vertex_bit::VB
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index 13a6f5b..75ca43b 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -1,5 +1,6 @@
-using ITensorNetworks: data_graph_type, AbstractITensorNetwork
+using ITensorNetworks: ITensorNetworks, AbstractITensorNetwork, data_graph
 using ITensors: ITensor, dim, contract, siteinds
+using Graphs: Graphs
 
 struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},BM<:BitMap} <:
        AbstractITensorNetwork{V}
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index b901339..0184e66 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -1,7 +1,8 @@
+using Graphs: nv, vertices, edges, neighbors
 using NamedGraphs:
-  vertices, random_bfs_tree, rem_edges, add_edges, undirected_graph, NamedEdge
+  random_bfs_tree, rem_edges, add_edges, undirected_graph, NamedEdge, AbstractGraph, leaf_vertices, a_star
 using ITensors: dim, commoninds
-using ITensorNetworks: IndsNetwork
+using ITensorNetworks: IndsNetwork, underlying_graph
 
 default_c_value() = 1.0
 default_a_value() = 0.0
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index f1f3c96..74e42e8 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -12,7 +12,7 @@ using ITensors:
   prime,
   noprime!,
   contract
-using ITensorNetworks: IndsNetwork, TTN, TreeTensorNetwork, combine_linkinds
+using ITensorNetworks: IndsNetwork, ITensorNetwork, TTN, TreeTensorNetwork, combine_linkinds
 
 function plus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
   @assert is_tree(s)

From a6f7217cdf7333805cf7729713417d9a8fcb79f1 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Mon, 1 Apr 2024 17:52:17 -0400
Subject: [PATCH 06/10] Added ability to do n ary decomposition. Binary,
 trinary etc

---
 src/TensorNetworkFunctionals.jl             |  8 +++-
 src/bitmaps.jl                              | 53 ++++++++++++++-------
 src/itensornetworkfunction.jl               |  1 +
 src/itensornetworks_elementary_functions.jl | 26 ++++++++--
 src/itensornetworks_elementary_operators.jl |  2 +-
 test/test_bitmaps.jl                        | 24 ++++++++++
 test/test_itensorfunction.jl                | 29 ++++++++++-
 7 files changed, 118 insertions(+), 25 deletions(-)

diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index dbd9480..1c3b382 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -8,7 +8,13 @@ include("itensornetworks_elementary_operators.jl")
 
 export ITensorNetworkFunction
 export BitMap,
-  default_dimension_map, vertex, calculate_xyz, calculate_x, calculate_bit_values, dimension
+  default_dimension_map,
+  vertex,
+  calculate_xyz,
+  calculate_x,
+  calculate_bit_values,
+  dimension,
+  base
 export const_itensornetwork,
   exp_itensornetwork,
   cosh_itensornetwork,
diff --git a/src/bitmaps.jl b/src/bitmaps.jl
index 3b513c3..8010581 100644
--- a/src/bitmaps.jl
+++ b/src/bitmaps.jl
@@ -2,39 +2,51 @@ using Dictionaries: Dictionary, set!
 using Graphs: Graphs
 
 struct BitMap{VB,VD}
-  vertex_bit::VB
+  vertex_digit::VB
   vertex_dimension::VD
+  base::Int64
 end
 
-vertex_bit(bm::BitMap) = bm.vertex_bit
+default_base() = 2
+
+vertex_digit(bm::BitMap) = bm.vertex_digit
 vertex_dimension(bm::BitMap) = bm.vertex_dimension
+base(bm::BitMap) = bm.base
 
 default_bit_map(vertices::Vector) = Dictionary(vertices, [i for i in 1:length(vertices)])
 function default_dimension_map(vertices::Vector)
   return Dictionary(vertices, [1 for i in 1:length(vertices)])
 end
 
-BitMap(g) = BitMap(default_bit_map(vertices(g)), default_dimension_map(vertices(g)))
-function BitMap(dimension_vertices::Vector{Vector{V}}) where {V}
-  vertex_bit = Dictionary()
+function BitMap(g; base::Int64=default_base())
+  return BitMap(default_bit_map(vertices(g)), default_dimension_map(vertices(g)), base)
+end
+function BitMap(vertex_digit, vertex_dimension; base::Int64=default_base())
+  return BitMap(vertex_digit, vertex_dimension, base)
+end
+function BitMap(dimension_vertices::Vector{Vector{V}}; base::Int64=default_base()) where {V}
+  vertex_digit = Dictionary()
   vertex_dimension = Dictionary()
   for (dimension, vertices) in enumerate(dimension_vertices)
     for (bit, v) in enumerate(vertices)
-      set!(vertex_bit, v, bit)
+      set!(vertex_digit, v, bit)
       set!(vertex_dimension, v, dimension)
     end
   end
-  return BitMap(vertex_bit, vertex_dimension)
+  return BitMap(vertex_digit, vertex_dimension, base)
 end
 
-Base.copy(bm::BitMap) = BitMap(copy(vertex_bit(bm)), copy(vertex_dimension(bm)))
+function Base.copy(bm::BitMap)
+  return BitMap(copy(vertex_digit(bm)), copy(vertex_dimension(bm)), copy(base(bm)))
+end
 
 dimension(bm::BitMap) = maximum(collect(values(vertex_dimension(bm))))
 dimension(bm::BitMap, vertex) = vertex_dimension(bm)[vertex]
-bit(bm::BitMap, vertex) = vertex_bit(bm)[vertex]
+digit(bm::BitMap, vertex) = vertex_digit(bm)[vertex]
+bit_value_to_scalar(bm::BitMap, vertex, value::Int64) = value / (base(bm)^digit(bm, vertex))
 
 function Graphs.vertices(bm::BitMap)
-  @assert keys(vertex_dimension(bm)) == keys(vertex_bit(bm))
+  @assert keys(vertex_dimension(bm)) == keys(vertex_digit(bm))
   return collect(keys(vertex_dimension(bm)))
 end
 function Graphs.vertices(bm::BitMap, dimension::Int64)
@@ -45,7 +57,7 @@ end
 function vertex(bm::BitMap, dimension::Int64, bit::Int64)
   return only(
     filter(
-      v -> vertex_dimension(bm)[v] == dimension && vertex_bit(bm)[v] == bit,
+      v -> vertex_dimension(bm)[v] == dimension && vertex_digit(bm)[v] == bit,
       keys(vertex_dimension(bm)),
     ),
   )
@@ -55,7 +67,7 @@ function calculate_xyz(bm::BitMap, vertex_to_bit_value_map, dimensions::Vector{I
   out = Float64[]
   for dimension in dimensions
     vs = vertices(bm, dimension)
-    push!(out, sum([vertex_to_bit_value_map[v] / (2^bit(bm, v)) for v in vs]))
+    push!(out, sum([bit_value_to_scalar(bm, v, vertex_to_bit_value_map[v]) for v in vs]))
   end
   return out
 end
@@ -79,13 +91,18 @@ function calculate_bit_values(
     dimension = dimensions[i]
     x_rn = copy(x)
     vs = vertices(bm, dimension)
-    sorted_vertices = sort(vs; by=vs -> bit(bm, vs))
+    sorted_vertices = sort(vs; by=vs -> digit(bm, vs))
     for v in sorted_vertices
-      if (x_rn >= 1.0 / (2^bit(bm, v)))
-        set!(vertex_to_bit_value_map, v, 1)
-        x_rn -= 1.0 / (2^bit(bm, v))
-      else
-        set!(vertex_to_bit_value_map, v, 0)
+      i = base(bm) - 1
+      vertex_set = false
+      while (!vertex_set)
+        if x_rn >= bit_value_to_scalar(bm, v, i)
+          set!(vertex_to_bit_value_map, v, i)
+          x_rn -= bit_value_to_scalar(bm, v, i)
+          vertex_set = true
+        else
+          i = i - 1
+        end
       end
     end
 
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index 75ca43b..b1d3b56 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -40,6 +40,7 @@ for f in [
   :calculate_bit_values,
   :calculate_x,
   :calculate_xyz,
+  :base,
 ]
   @eval begin
     function $f(fitn::ITensorNetworkFunction, args...; kwargs...)
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index 0184e66..e1b2f44 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -1,6 +1,13 @@
 using Graphs: nv, vertices, edges, neighbors
 using NamedGraphs:
-  random_bfs_tree, rem_edges, add_edges, undirected_graph, NamedEdge, AbstractGraph, leaf_vertices, a_star
+  random_bfs_tree,
+  rem_edges,
+  add_edges,
+  undirected_graph,
+  NamedEdge,
+  AbstractGraph,
+  leaf_vertices,
+  a_star
 using ITensors: dim, commoninds
 using ITensorNetworks: IndsNetwork, underlying_graph
 
@@ -35,7 +42,12 @@ function exp_itensornetwork(
   ψ = const_itensornetwork(s, bit_map)
   Lx = length(vertices(bit_map, dimension))
   for v in vertices(bit_map, dimension)
-    ψ[v] = ITensor([exp(a / Lx), exp(a / Lx) * exp(k / (2^bit(bit_map, v)))], inds(ψ[v]))
+    ψ[v] = ITensor(
+      [
+        exp(a / Lx) * exp(k * bit_value_to_scalar(bit_map, v, i)) for i in 0:(dim(s[v]) - 1)
+      ],
+      inds(ψ[v]),
+    )
   end
 
   return ψ
@@ -217,12 +229,18 @@ function polynomial_itensornetwork(
         siteindex,
         alphas,
         betaindex,
-        [0.0, (1.0 / (2^bit(bit_map, v)))],
+        [bit_value_to_scalar(bit_map, v, i) for i in 0:(dim(siteindex) - 1)],
       )
     elseif v == root_vertex
       betaindex = Index(n + 1, "DummyInd")
       alphas = setdiff(inds(ψ[v]), Index[siteindex])
-      ψv = Q_N_tensor(2, siteindex, alphas, betaindex, [0.0, (1.0 / (2^bit(bit_map, v)))])
+      ψv = Q_N_tensor(
+        2,
+        siteindex,
+        alphas,
+        betaindex,
+        [bit_value_to_scalar(bit_map, v, i) for i in 0:(dim(siteindex) - 1)],
+      )
       ψ[v] = ψv * ITensor(reverse(coeffs), betaindex)
     end
   end
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index 74e42e8..f898695 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -74,7 +74,7 @@ function stencil(
   stencil_op = truncate(stencil_op; truncate_kwargs...)
 
   for v in vertices(bit_map, dimension)
-    stencil_op[v] = (2^delta_power) * stencil_op[v]
+    stencil_op[v] = (base(bit_map)^delta_power) * stencil_op[v]
   end
 
   return truncate(stencil_op; truncate_kwargs...)
diff --git a/test/test_bitmaps.jl b/test/test_bitmaps.jl
index f431c75..42b3b83 100644
--- a/test/test_bitmaps.jl
+++ b/test/test_bitmaps.jl
@@ -14,6 +14,7 @@ using Dictionaries: Dictionary
 
   @test dimension(bit_map) == 1
   @test Set(vertices(bit_map)) == Set(vertices(g))
+  @test base(bit_map) == 2
 
   x = 0.625
   vertex_to_bit_value_map = calculate_bit_values(bit_map, x)
@@ -40,3 +41,26 @@ end
   xyzvals_approx = calculate_xyz(bit_map, vertex_to_bit_value_map)
   xyzvals ≈ xyzvals_approx
 end
+
+@testset "test multi dimensional trinary bit map" begin
+  L = 50
+  b = 3
+  g = named_grid((L, L))
+  vertex_to_dimension_map = Dictionary(vertices(g), [v[1] for v in vertices(g)])
+  vertex_to_bit_map = Dictionary(vertices(g), [v[2] for v in vertices(g)])
+  bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map, b)
+
+  @test base(bit_map) == b
+
+  x, y = (1.0 / 3.0), (5.0 / 9.0)
+  vertex_to_bit_value_map = calculate_bit_values(bit_map, [x, y], [1, 2])
+  xyz = calculate_xyz(bit_map, vertex_to_bit_value_map, [1, 2])
+  @test first(xyz) == x
+  @test last(xyz) == y
+
+  xyzvals = [rand() for i in 1:L]
+  vertex_to_bit_value_map = calculate_bit_values(bit_map, xyzvals)
+
+  xyzvals_approx = calculate_xyz(bit_map, vertex_to_bit_value_map)
+  xyzvals ≈ xyzvals_approx
+end
diff --git a/test/test_itensorfunction.jl b/test/test_itensorfunction.jl
index b965caf..e3989fa 100644
--- a/test/test_itensorfunction.jl
+++ b/test/test_itensorfunction.jl
@@ -22,6 +22,7 @@ using Distributions: Uniform
 
   @test vertices(fψ, 1) == vertices(fψ)
   @test dimension(fψ) == 1
+  @test base(fψ) == 2
 
   dimension_vertices = collect(values(group(v -> first(v) < Int64(0.5 * L), vertices(ψ))))
   fψ = ITensorNetworkFunction(ψ, dimension_vertices)
@@ -54,7 +55,7 @@ end
     ("sin", sin_itn, sin),
   ]
   for (name, net_func, func) in funcs
-    @testset "test $name" begin
+    @testset "test $name in binary" begin
       Lx, Ly = 2, 3
       g = named_comb_tree((2, 3))
       a = 1.2
@@ -70,6 +71,32 @@ end
     end
   end
 
+  funcs = [
+    ("cosh", cosh_itn, cosh),
+    ("sinh", sinh_itn, sinh),
+    ("exp", exp_itn, exp),
+    ("cos", cos_itn, cos),
+    ("sin", sin_itn, sin),
+  ]
+  for (name, net_func, func) in funcs
+    @testset "test $name in trinary" begin
+      Lx, Ly = 2, 3
+      g = named_comb_tree((2, 3))
+      a = 1.2
+      k = 0.125
+      b = 3
+      s = siteinds("S=1", g)
+
+      bit_map = BitMap(g; base=b)
+
+      x = (5.0 / 9.0)
+      ψ_fx = net_func(s, bit_map; k, a)
+      @test base(ψ_fx) == 3
+      fx_x = calculate_fx(ψ_fx, x)
+      @test func(k * x + a) ≈ fx_x
+    end
+  end
+
   @testset "test tanh" begin
     L = 10
     g = named_grid((L, 1))

From 60526c76b9d17a1ec383272b80d0f904d388f39a Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Tue, 2 Apr 2024 14:41:21 -0400
Subject: [PATCH 07/10] Onehot in favour of explicit projecting of local
 tensors

---
 src/itensornetworkfunction.jl | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index b1d3b56..fd46125 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -1,5 +1,5 @@
 using ITensorNetworks: ITensorNetworks, AbstractITensorNetwork, data_graph
-using ITensors: ITensor, dim, contract, siteinds
+using ITensors: ITensor, dim, contract, siteinds, onehot
 using Graphs: Graphs
 
 struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},BM<:BitMap} <:
@@ -53,10 +53,7 @@ function project(fitn::ITensorNetworkFunction, vertex_to_bit_value_map)
   fitn = copy(fitn)
   s = siteinds(fitn)
   for v in keys(vertex_to_bit_value_map)
-    proj = ITensor(
-      [i != vertex_to_bit_value_map[v] ? 0 : 1 for i in 0:(dim(s[v]) - 1)], s[v]
-    )
-    fitn[v] = fitn[v] * proj
+    fitn[v] = fitn[v] * onehot(only(s[v]) => vertex_to_bit_value_map[v] + 1)
   end
   return fitn
 end

From ae165fb65869e77985214eb04cf895235c8ad362 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Thu, 4 Apr 2024 16:44:45 -0400
Subject: [PATCH 08/10] Remove Weierstrass example for now. Add Laplace solver
 example

---
 examples/2d_laplace_solver.jl               |  58 ++++++++
 examples/weierstrass.jl                     | 150 --------------------
 src/TensorNetworkFunctionals.jl             |  14 +-
 src/bitmaps.jl                              |   8 ++
 src/itensornetworks_elementary_functions.jl |   5 +
 src/itensornetworks_elementary_operators.jl |  35 +++--
 src/itensornetworksutils.jl                 |  12 +-
 7 files changed, 112 insertions(+), 170 deletions(-)
 create mode 100644 examples/2d_laplace_solver.jl
 delete mode 100644 examples/weierstrass.jl

diff --git a/examples/2d_laplace_solver.jl b/examples/2d_laplace_solver.jl
new file mode 100644
index 0000000..46f4bf4
--- /dev/null
+++ b/examples/2d_laplace_solver.jl
@@ -0,0 +1,58 @@
+using Test
+using TensorNetworkFunctionals
+
+using Graphs: SimpleGraph, uniform_tree, binary_tree, random_regular_graph, is_tree
+using NamedGraphs:
+  NamedGraph,
+  named_grid,
+  vertices,
+  named_comb_tree,
+  rename_vertices,
+  random_bfs_tree,
+  undirected_graph
+using ITensors: ITensors, Index, siteinds, dim, tags, replaceprime!, MPO, MPS, inner
+using ITensorNetworks:
+  ITensorNetwork,
+  dmrg,
+  TTN,
+  maxlinkdim
+using Dictionaries: Dictionary
+using SplitApplyCombine: group
+using Random: seed!
+using Distributions: Uniform
+
+using UnicodePlots
+
+#Solve the 2D Laplace equation on a random tree
+seed!(1234)
+L = 14
+g = NamedGraph(SimpleGraph(uniform_tree(L)))
+s = siteinds("S=1/2", g)
+
+vertex_to_dimension_map = Dictionary(vertices(g), [(v[1] % 2) + 1 for v in vertices(g)])
+vertex_to_bit_map = Dictionary(vertices(g), [ceil(Int64, v[1] * 0.5) for v in vertices(g)])
+bit_map = BitMap(vertex_to_bit_map, vertex_to_dimension_map)
+
+ψ_fxy = 0.1 * rand_itn(s, bit_map; link_space=2)
+∇ = laplacian_operator(s, bit_map; scale=false)
+∇ = truncate(∇; cutoff=1e-12)
+@show maxlinkdim(∇)
+
+dmrg_kwargs = (nsweeps=25, normalize=true, maxdim=20, cutoff=1e-12, outputlevel=1, nsites=2)
+ϕ_fxy = dmrg(∇, TTN(itensornetwork(ψ_fxy)); dmrg_kwargs...)
+ϕ_fxy = ITensorNetworkFunction(ITensorNetwork(ϕ_fxy), bit_map)
+
+final_energy = inner(TTN(itensornetwork(ϕ_fxy))', ∇, TTN(itensornetwork(ϕ_fxy)))
+#Smallest eigenvalue in this case should be -8
+@show final_energy
+
+n_grid = 100
+x_vals, y_vals = grid_points(bit_map, n_grid, 1), grid_points(bit_map, n_grid, 2)
+vals = zeros((length(x_vals), length(y_vals)))
+for (i, x) in enumerate(x_vals)
+  for (j, y) in enumerate(y_vals)
+    vals[i, j] = real(calculate_fxyz(ϕ_fxy, [x, y]))
+  end
+end
+
+show(heatmap(vals))
diff --git a/examples/weierstrass.jl b/examples/weierstrass.jl
deleted file mode 100644
index cd6d113..0000000
--- a/examples/weierstrass.jl
+++ /dev/null
@@ -1,150 +0,0 @@
-using ITensors
-using ITensorNetworks
-using NamedGraphs
-using EllipsisNotation
-using Graphs
-using LaTeXStrings
-
-using ITensorNetworks: delta_network, contract_inner, distance_to_leaf
-using NamedGraphs: add_edges, random_bfs_tree, rename_vertices
-
-using Plots
-using Random
-
-using Statistics
-
-include("../src/QTT_utils.jl")
-
-plot_font = "Computer Modern"
-default(; fontfamily=plot_font, linewidth=1.5, framestyle=:box, label=nothing, grid=false)
-
-"""ITensorNetwork representation of the weierstrass function (first nterms) over [0,1]: see https://mathworld.wolfram.com/WeierstrassFunction.html"""
-function weierstrass_itn(
-  s::IndsNetwork, vertex_map::Dict, a::Int64, nterms::Int64; svd_kwargs...
-)
-  ψ = sin_itn(s, vertex_map; a=0.0, k=Float64(pi))
-  ψ[first(vertices(ψ))] *= (1.0 / pi)
-  for n in 2:nterms
-    omega = pi * (n^a)
-    ψt = sin_itn(s, vertex_map; a=0.0, k=omega)
-    ψt[first(vertices(ψt))] *= (1.0 / (omega))
-    ψ = ψ + ψt
-    if isa(ψ, TreeTensorNetwork)
-      ψ = truncate(ψ; svd_kwargs...)
-    end
-  end
-
-  return ψ
-end
-
-"""Evaluate Weierstrass(x) (first nterms) see: see https://mathworld.wolfram.com/WeierstrassFunction.html"""
-function weierstrass(a::Int64, nterms::Int64, x::Float64)
-  out = 0
-  for n in 1:nterms
-    omega = pi * (n^a)
-    out += sin(omega * x) / omega
-  end
-
-  return out
-end
-
-function main()
-  L = 31
-  g = named_grid((L, 1))
-  a = 2
-  nterms = 200
-  s = siteinds("S=1/2", g)
-
-  #Define a map which determines a canonical ordering of the vertices of the network
-  vertex_map = Dict(vertices(g) .=> [i for i in 1:length(vertices(g))])
-
-  no_xs = 1000
-  xs = [(1.0 / no_xs) * i for i in 1:no_xs]
-  xevals = [calculate_x(calculate_xis(x, vertex_map), vertex_map) for x in xs]
-  bond_dims = [1, 2, 3, 4, 5, 6]
-  f_actual = weierstrass.(a, nterms, xevals)
-
-  p1 = Plots.plot(xevals, f_actual; legend=false, titlefont=font(10, "Computer Modern"))
-  xlabel!("x")
-  title!(L"W(x) = \sum_{k}\sin(\pi k^{a} x) / (\pi k^{a})")
-
-  mld, vertex_sequence = optimal_vertex_ordering(a, nterms, 1000, 1e-5)
-  #vertex_sequence = [i for i in 1:length(vertices(g))]
-
-  ψ = weierstrass_itn(s, vertex_map, a, nterms; cutoff=1e-16)
-  p2 = Plots.plot(; titlefont=font(10, "Computer Modern"), yaxis=:log)
-  MPS_avg_error = zeros(length(bond_dims))
-  MPS_memory_costs = zeros(length(bond_dims))
-  for (χindex, χ) in enumerate(bond_dims)
-    println("Truncating Down to Bond Dimension $χ")
-    ψtrunc = truncate(ψ; maxdim=χ)
-    f_itn = zeros(no_xs)
-    for (index, x) in enumerate(xs)
-      xis = calculate_xis(x, vertex_map)
-      eval_point = calculate_x(xis, vertex_map)
-      ψproj = get_bitstring_network(ψtrunc, s, xis)
-      f_itn[index] = real(ITensors.contract(ψproj)[])
-    end
-    MPS_avg_error[χindex] = mean(abs.(f_itn - f_actual))
-    MPS_memory_costs[χindex] = Base.summarysize(ψtrunc)
-    Plots.plot!(xevals, abs.(f_itn - f_actual); label="χ = $χ")
-  end
-  xlabel!("x")
-  ylabel!("|W(x) - f(x, χ)|")
-  mld = maxlinkdim(ψ)
-  title!("MPS f(x, χ). χmax = $mld"; fontsize=10)
-
-  Random.seed!(999)
-  g = NamedGraph(Graphs.SimpleGraph(binary_tree(5)))
-  g = rename_vertices(g, Dict(zip(vertices(g), [(v,) for v in vertices(g)])))
-  s = siteinds("S=1/2", g)
-  vertex_map = Dict(vertices(g) .=> vertex_sequence)
-  ψ = weierstrass_itn(s, vertex_map, a, nterms; cutoff=1e-16)
-  p3 = Plots.plot(; titlefont=font(10, "Computer Modern"), yaxis=:log)
-  Tree_avg_error = zeros(length(bond_dims))
-  Tree_memory_costs = zeros(length(bond_dims))
-
-  @show vertex_map
-  d = Dict(vertices(g) => [distance_to_leaf(g, v) for v in vertices(g)])
-  @show d
-
-  for (χindex, χ) in enumerate(bond_dims)
-    println("Truncating Down to Bond Dimension $χ")
-    ψtrunc = truncate(ψ; maxdim=χ)
-    f_itn = zeros(no_xs)
-    for (index, x) in enumerate(xs)
-      xis = calculate_xis(x, vertex_map)
-      eval_point = calculate_x(xis, vertex_map)
-      ψproj = get_bitstring_network(ψtrunc, s, xis)
-      f_itn[index] = real(ITensors.contract(ψproj)[])
-    end
-
-    Tree_avg_error[χindex] = mean(abs.(f_itn - f_actual))
-    Tree_memory_costs[χindex] = Base.summarysize(ψtrunc)
-
-    Plots.plot!(xevals, abs.(f_itn - f_actual); label="χ = $χ")
-  end
-  xlabel!("x")
-  ylabel!("|W(x) - f(x, χ)|")
-  mld = maxlinkdim(ψ)
-  title!("Binary Tree f(x, χ). χmax = $mld")
-
-  p4 = Plots.plot(
-    MPS_memory_costs / (1e6),
-    MPS_avg_error;
-    labels="MPS",
-    yaxis=:log,
-    titlefont=font(10, "Computer Modern"),
-    xguidefontsize=10,
-  )
-  Plots.plot!(Tree_memory_costs / (1e6), Tree_avg_error; labels="Tree", yaxis=:log)
-
-  xlabel!("Memory Cost (Mb)")
-  title!("Mean(|W(x) - f(x, χ)|)")
-
-  l = @layout [a d; b c]
-  plot(p1, p4, p2, p3; layout=l)
-  return Plots.savefig("examples/weierstrass_example_output.pdf")
-end
-
-main()
diff --git a/src/TensorNetworkFunctionals.jl b/src/TensorNetworkFunctionals.jl
index 1c3b382..b21a0a1 100644
--- a/src/TensorNetworkFunctionals.jl
+++ b/src/TensorNetworkFunctionals.jl
@@ -6,7 +6,7 @@ include("itensornetworkfunction.jl")
 include("itensornetworks_elementary_functions.jl")
 include("itensornetworks_elementary_operators.jl")
 
-export ITensorNetworkFunction
+export ITensorNetworkFunction, itensornetwork
 export BitMap,
   default_dimension_map,
   vertex,
@@ -14,7 +14,8 @@ export BitMap,
   calculate_x,
   calculate_bit_values,
   dimension,
-  base
+  base,
+  grid_points
 export const_itensornetwork,
   exp_itensornetwork,
   cosh_itensornetwork,
@@ -24,10 +25,13 @@ export const_itensornetwork,
   sin_itensornetwork,
   get_edge_toward_root,
   polynomial_itensornetwork,
+  random_itensornetworkfunction,
   laplacian_operator,
-  derivative_operator
-export const_itn, poly_itn, cosh_itn, sinh_itn, tanh_itn, exp_itn, sin_itn, cos_itn
+  derivative_operator,
+  identity_operator
+export const_itn,
+  poly_itn, cosh_itn, sinh_itn, tanh_itn, exp_itn, sin_itn, cos_itn, rand_itn
 export calculate_fx, calculate_fxyz
-export operate, apply_gx_operator, multiply
+export operate, operator, multiply
 
 end
diff --git a/src/bitmaps.jl b/src/bitmaps.jl
index 8010581..cfedb67 100644
--- a/src/bitmaps.jl
+++ b/src/bitmaps.jl
@@ -125,3 +125,11 @@ end
 function calculate_bit_values(bm::BitMap, x::Float64; kwargs...)
   return calculate_bit_values(bm, [x], [1]; kwargs...)
 end
+
+function grid_points(bm::BitMap, N::Int64, dimension::Int64)
+  vals = Vector{Float64}
+  L = length(vertices(bm, dimension))
+  a = round(base(bm)^L / N)
+  grid_points = [i * (a / base(bm)^L) for i in 0:(N + 1)]
+  return filter(x -> x <= 1, grid_points)
+end
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index e1b2f44..fe7e24a 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -248,6 +248,10 @@ function polynomial_itensornetwork(
   return ITensorNetworkFunction(ψ, bit_map)
 end
 
+function random_itensornetwork(s::IndsNetwork, bit_map; kwargs...)
+  return ITensorNetworkFunction(randomITensorNetwork(s; kwargs...), bit_map)
+end
+
 const const_itn = const_itensornetwork
 const poly_itn = polynomial_itensornetwork
 const cosh_itn = cosh_itensornetwork
@@ -256,3 +260,4 @@ const tanh_itn = tanh_itensornetwork
 const exp_itn = exp_itensornetwork
 const sin_itn = sin_itensornetwork
 const cos_itn = cos_itensornetwork
+const rand_itn = random_itensornetwork
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index f898695..608da0b 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -14,7 +14,9 @@ using ITensors:
   contract
 using ITensorNetworks: IndsNetwork, ITensorNetwork, TTN, TreeTensorNetwork, combine_linkinds
 
-function plus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
+function plus_shift_ttn(
+  s::IndsNetwork, bit_map; dimension=default_dimension(), boundary_value=[0.0]
+)
   @assert is_tree(s)
   ttn_op = OpSum()
   dim_vertices = vertices(bit_map, dimension)
@@ -63,6 +65,7 @@ function stencil(
   shifts::Vector{Float64},
   delta_power::Int64;
   dimension=default_dimension(),
+  scale=true,
   truncate_kwargs...,
 )
   @assert length(shifts) == 3
@@ -73,24 +76,38 @@ function stencil(
   stencil_op = plus_shift + minus_shift + no_shift
   stencil_op = truncate(stencil_op; truncate_kwargs...)
 
-  for v in vertices(bit_map, dimension)
-    stencil_op[v] = (base(bit_map)^delta_power) * stencil_op[v]
+  if scale
+    for v in vertices(bit_map, dimension)
+      stencil_op[v] = (base(bit_map)^delta_power) * stencil_op[v]
+    end
   end
 
-  return truncate(stencil_op; truncate_kwargs...)
+  return stencil_op
 end
 
-function laplacian_operator(s::IndsNetwork, bit_map; kwargs...)
-  return stencil(s, bit_map, [1.0, -2.0, 1.0], 2; kwargs...)
+function laplacian_operator(
+  s::IndsNetwork, bit_map; dimensions=[i for i in 1:dimension(bit_map)], kwargs...
+)
+  remaining_dims = copy(dimensions)
+  ∇ = stencil(s, bit_map, [1.0, -2.0, 1.0], 2; dimension=first(remaining_dims), kwargs...)
+  popfirst!(remaining_dims)
+  for rd in remaining_dims
+    ∇ += stencil(s, bit_map, [1.0, -2.0, 1.0], 2; dimension=rd, kwargs...)
+  end
+  return ∇
 end
 
 function derivative_operator(s::IndsNetwork, bit_map; kwargs...)
   return 0.5 * stencil(s, bit_map, [1.0, 0.0, -1.0], 1; kwargs...)
 end
 
-function apply_gx_operator(gx::ITensorNetworkFunction)
-  gx = copy(gx)
-  operator = itensornetwork(gx)
+function identity_operator(s::IndsNetwork, bit_map; kwargs...)
+  return stencil(s, bit_map, [0.0, 1.0, 0.0], 0; kwargs...)
+end
+
+function operator(fx::ITensorNetworkFunction)
+  fx = copy(fx)
+  operator = itensornetwork(fx)
   s = siteinds(operator)
   for v in vertices(operator)
     sind = s[v]
diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index d4fa1c0..2b70f77 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -2,13 +2,13 @@ using ITensors: Index, dim, inds
 using ITensorNetworks: randomITensorNetwork, IndsNetwork
 
 """Build the order L tensor corresponding to fx(x): x ∈ [0,1]."""
-function build_full_rank_tensor(L::Int64, fx::Function)
-  inds = [Index(2, "$i") for i in 1:L]
-  dims = Tuple([2 for i in 1:L])
+function build_full_rank_tensor(L::Int64, fx::Function; base::Int64=2)
+  inds = [Index(base, "$i") for i in 1:L]
+  dims = Tuple([base for i in 1:L])
   array = zeros(dims)
-  for i in 0:(2^(L) - 1)
-    xis = digits(i; base=2, pad=L)
-    x = sum([xis[i] / (2^i) for i in 1:L])
+  for i in 0:(base^(L) - 1)
+    xis = digits(i; base, pad=L)
+    x = sum([xis[i] / (base^i) for i in 1:L])
     array[Tuple(xis + ones(Int64, (L)))...] = fx(x)
   end
 

From e92bc88dafaa6bd1eb49c1951c6d4b3b946a7094 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Tue, 9 Apr 2024 11:41:00 -0400
Subject: [PATCH 09/10] Bump compat to 0.6.0. Add eltype spec to onehot

---
 Project.toml                                |  2 +-
 examples/2d_laplace_solver.jl               |  6 +-----
 src/itensornetworkfunction.jl               |  3 ++-
 src/itensornetworks_elementary_functions.jl |  2 +-
 src/itensornetworks_elementary_operators.jl | 14 +++++++-------
 src/itensornetworksutils.jl                 |  4 ++--
 test/test_itensorfunction.jl                |  4 ++--
 7 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/Project.toml b/Project.toml
index 7535449..4e3cfef 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,4 +18,4 @@ SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ITensorNetworks = "0.4"
+ITensorNetworks = "0.6.0"
diff --git a/examples/2d_laplace_solver.jl b/examples/2d_laplace_solver.jl
index 46f4bf4..7167d2e 100644
--- a/examples/2d_laplace_solver.jl
+++ b/examples/2d_laplace_solver.jl
@@ -11,11 +11,7 @@ using NamedGraphs:
   random_bfs_tree,
   undirected_graph
 using ITensors: ITensors, Index, siteinds, dim, tags, replaceprime!, MPO, MPS, inner
-using ITensorNetworks:
-  ITensorNetwork,
-  dmrg,
-  TTN,
-  maxlinkdim
+using ITensorNetworks: ITensorNetwork, dmrg, TTN, maxlinkdim
 using Dictionaries: Dictionary
 using SplitApplyCombine: group
 using Random: seed!
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index fd46125..bf81398 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -53,7 +53,8 @@ function project(fitn::ITensorNetworkFunction, vertex_to_bit_value_map)
   fitn = copy(fitn)
   s = siteinds(fitn)
   for v in keys(vertex_to_bit_value_map)
-    fitn[v] = fitn[v] * onehot(only(s[v]) => vertex_to_bit_value_map[v] + 1)
+    fitn[v] =
+      fitn[v] * onehot(eltype(fitn[v]), only(s[v]) => vertex_to_bit_value_map[v] + 1)
   end
   return fitn
 end
diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index fe7e24a..ae6392a 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -249,7 +249,7 @@ function polynomial_itensornetwork(
 end
 
 function random_itensornetwork(s::IndsNetwork, bit_map; kwargs...)
-  return ITensorNetworkFunction(randomITensorNetwork(s; kwargs...), bit_map)
+  return ITensorNetworkFunction(random_tensornetwork(s; kwargs...), bit_map)
 end
 
 const const_itn = const_itensornetwork
diff --git a/src/itensornetworks_elementary_operators.jl b/src/itensornetworks_elementary_operators.jl
index 608da0b..b8a5845 100644
--- a/src/itensornetworks_elementary_operators.jl
+++ b/src/itensornetworks_elementary_operators.jl
@@ -12,7 +12,7 @@ using ITensors:
   prime,
   noprime!,
   contract
-using ITensorNetworks: IndsNetwork, ITensorNetwork, TTN, TreeTensorNetwork, combine_linkinds
+using ITensorNetworks: IndsNetwork, ITensorNetwork, TreeTensorNetwork, combine_linkinds, ttn
 
 function plus_shift_ttn(
   s::IndsNetwork, bit_map; dimension=default_dimension(), boundary_value=[0.0]
@@ -31,7 +31,7 @@ function plus_shift_ttn(
     add!(ttn_op, 1.0, (string_site...)...)
   end
 
-  return TTN(ttn_op, s; algorithm="svd")
+  return ttn(ttn_op, s; algorithm="svd")
 end
 
 function minus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
@@ -49,14 +49,14 @@ function minus_shift_ttn(s::IndsNetwork, bit_map; dimension=default_dimension())
     add!(ttn_op, 1.0, (string_site...)...)
   end
 
-  return TTN(ttn_op, s; algorithm="svd")
+  return ttn(ttn_op, s; algorithm="svd")
 end
 
 function no_shift_ttn(s::IndsNetwork)
   ttn_op = OpSum()
   string_site_full = [("I", v) for v in vertices(s)]
   add!(ttn_op, 1.0, (string_site_full...)...)
-  return TTN(ttn_op, s; algorithm="svd")
+  return ttn(ttn_op, s; algorithm="svd")
 end
 
 function stencil(
@@ -145,7 +145,7 @@ Base.:*(fs::ITensorNetworkFunction...) = multiply(fs...)
 function operate(
   operator::TreeTensorNetwork, ψ::ITensorNetworkFunction; truncate_kwargs=(;), kwargs...
 )
-  ψ_tn = TTN(itensornetwork(ψ))
+  ψ_tn = ttn(itensornetwork(ψ))
   ψO_tn = noprime(contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
   ψO_tn = truncate(ψO_tn; truncate_kwargs...)
 
@@ -153,7 +153,7 @@ function operate(
 end
 
 function operate(operator::ITensorNetwork, ψ::ITensorNetworkFunction; kwargs...)
-  return operate(TTN(operator), ψ; kwargs...)
+  return operate(ttn(operator), ψ; kwargs...)
 end
 
 function operate(
@@ -169,5 +169,5 @@ end
 function operate(
   operators::Vector{ITensorNetwork{V}}, ψ::ITensorNetworkFunction; kwargs...
 ) where {V}
-  return operate(TTN.(operators), ψ; kwargs...)
+  return operate(ttn.(operators), ψ; kwargs...)
 end
diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index 2b70f77..9d4fe15 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -1,5 +1,5 @@
 using ITensors: Index, dim, inds
-using ITensorNetworks: randomITensorNetwork, IndsNetwork
+using ITensorNetworks: random_tensornetwork, IndsNetwork
 
 """Build the order L tensor corresponding to fx(x): x ∈ [0,1]."""
 function build_full_rank_tensor(L::Int64, fx::Function; base::Int64=2)
@@ -31,7 +31,7 @@ function c_tensor(phys_ind::Index, virt_inds::Vector)
 end
 
 function copy_tensor_network(s::IndsNetwork; linkdim::Int64=1)
-  tn = randomITensorNetwork(s; link_space=linkdim)
+  tn = random_tensornetwork(s; link_space=linkdim)
   for v in vertices(tn)
     virt_inds = setdiff(inds(tn[v]), Index[only(s[v])])
     tn[v] = c_tensor(only(s[v]), virt_inds)
diff --git a/test/test_itensorfunction.jl b/test/test_itensorfunction.jl
index e3989fa..dc32f9f 100644
--- a/test/test_itensorfunction.jl
+++ b/test/test_itensorfunction.jl
@@ -4,7 +4,7 @@ using TensorNetworkFunctionals
 using Graphs: SimpleGraph, uniform_tree
 using NamedGraphs: NamedGraph, named_grid, vertices, named_comb_tree, rename_vertices
 using ITensors: siteinds
-using ITensorNetworks: randomITensorNetwork
+using ITensorNetworks: random_tensornetwork
 using Dictionaries: Dictionary
 using SplitApplyCombine: group
 using Random: seed!
@@ -16,7 +16,7 @@ using Distributions: Uniform
   g = named_grid((L, 1))
   s = siteinds("S=1/2", g)
 
-  ψ = randomITensorNetwork(s; link_space=2)
+  ψ = random_tensornetwork(s; link_space=2)
 
   fψ = ITensorNetworkFunction(ψ)
 

From f81cad5eb9478956b7441fc2ba5115e0404ee3ae Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Tue, 9 Apr 2024 11:55:52 -0400
Subject: [PATCH 10/10] Remove c_tensor_network in favour of
 const_tensornetwork

---
 src/itensornetworks_elementary_functions.jl | 13 +++++++------
 src/itensornetworksutils.jl                 | 14 ++------------
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/src/itensornetworks_elementary_functions.jl b/src/itensornetworks_elementary_functions.jl
index ae6392a..0fe1011 100644
--- a/src/itensornetworks_elementary_functions.jl
+++ b/src/itensornetworks_elementary_functions.jl
@@ -17,14 +17,15 @@ default_k_value() = 1.0
 default_nterms() = 20
 default_dimension() = 1
 
-"""Construct the product state representation of the function f(x) = const."""
+"""Build a representation of the function f(x,y,z,...) = c, with flexible choice of linkdim"""
 function const_itensornetwork(
-  s::IndsNetwork, bit_map; c::Union{Float64,ComplexF64}=default_c_value()
+  s::IndsNetwork, bit_map; c::Union{Float64,ComplexF64}=default_c_value(), linkdim::Int64=1
 )
-  ψ = copy_tensor_network(s)
+  ψ = random_tensornetwork(s; link_space=linkdim)
   inv_L = Float64(1.0 / nv(s))
   for v in vertices(ψ)
-    ψ[v] *= c^inv_L
+    virt_inds = setdiff(inds(ψ[v]), Index[only(s[v])])
+    ψ[v] = c^inv_L * c_tensor(only(s[v]), virt_inds)
   end
 
   return ITensorNetworkFunction(ψ, bit_map)
@@ -216,7 +217,7 @@ function polynomial_itensornetwork(
   root_vertices_dim = filter(v -> v ∈ dimension_vertices, leaf_vertices(s_tree))
   @assert !isempty(root_vertices_dim)
   root_vertex = first(root_vertices_dim)
-  ψ = copy_tensor_network(s_tree; linkdim=n + 1)
+  ψ = const_itensornetwork(s_tree, bit_map; linkdim=n + 1)
   #Place the Q_n tensors, making sure we get the right index pointing towards the root
   for v in dimension_vertices
     siteindex = s_tree[v][]
@@ -245,7 +246,7 @@ function polynomial_itensornetwork(
     end
   end
 
-  return ITensorNetworkFunction(ψ, bit_map)
+  return ψ
 end
 
 function random_itensornetwork(s::IndsNetwork, bit_map; kwargs...)
diff --git a/src/itensornetworksutils.jl b/src/itensornetworksutils.jl
index 9d4fe15..05ac429 100644
--- a/src/itensornetworksutils.jl
+++ b/src/itensornetworksutils.jl
@@ -1,7 +1,7 @@
 using ITensors: Index, dim, inds
 using ITensorNetworks: random_tensornetwork, IndsNetwork
 
-"""Build the order L tensor corresponding to fx(x): x ∈ [0,1]."""
+"""Build the order L tensor corresponding to fx(x): x ∈ [0,1], default decomposition is binary"""
 function build_full_rank_tensor(L::Int64, fx::Function; base::Int64=2)
   inds = [Index(base, "$i") for i in 1:L]
   dims = Tuple([base for i in 1:L])
@@ -15,10 +15,10 @@ function build_full_rank_tensor(L::Int64, fx::Function; base::Int64=2)
   return ITensor(array, inds)
 end
 
+"""Build the tensor C such that C_{phys_ind, virt_inds...} = delta_{virt_inds...}"""
 function c_tensor(phys_ind::Index, virt_inds::Vector)
   inds = vcat(phys_ind, virt_inds)
   @assert allequal(dim.(virt_inds))
-  #Build tensor to be delta on inds and independent of phys_ind
   T = ITensor(0.0, inds...)
   for i in 1:dim(phys_ind)
     for j in 1:dim(first(virt_inds))
@@ -29,13 +29,3 @@ function c_tensor(phys_ind::Index, virt_inds::Vector)
 
   return T
 end
-
-function copy_tensor_network(s::IndsNetwork; linkdim::Int64=1)
-  tn = random_tensornetwork(s; link_space=linkdim)
-  for v in vertices(tn)
-    virt_inds = setdiff(inds(tn[v]), Index[only(s[v])])
-    tn[v] = c_tensor(only(s[v]), virt_inds)
-  end
-
-  return tn
-end