From c0cdda16fc3abd0c085e990f8dd303874d9e0e29 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Tue, 16 Apr 2024 16:44:24 -0400
Subject: [PATCH] IndsNetworkMap to wrap IndsNetwork and IndexMap

---
 examples/2d_laplace_solver.jl                 |   7 +-
 .../construct_multi_dimensional_function.jl   |   9 +-
 src/ITensorNumericalAnalysis.jl               |   5 +-
 src/elementary_functions.jl                   |  80 ++++----
 src/elementary_operators.jl                   | 105 +++++------
 src/indexmap.jl                               |  43 +----
 src/indsnetworkmap.jl                         | 105 +++++++++++
 src/itensornetworkfunction.jl                 |  33 ++--
 src/utils.jl                                  |  16 +-
 test/test_derivative_operators.jl             | 139 --------------
 test/test_indexmaps.jl                        |  19 +-
 test/test_itensorfunction.jl                  |  37 ++--
 ...t_shift_operators.jl => test_operators.jl} | 176 +++++++++++++++---
 13 files changed, 394 insertions(+), 380 deletions(-)
 create mode 100644 src/indsnetworkmap.jl
 delete mode 100644 test/test_derivative_operators.jl
 rename test/{test_shift_operators.jl => test_operators.jl} (56%)

diff --git a/examples/2d_laplace_solver.jl b/examples/2d_laplace_solver.jl
index 285debd..99ac27b 100644
--- a/examples/2d_laplace_solver.jl
+++ b/examples/2d_laplace_solver.jl
@@ -15,11 +15,10 @@ seed!(1234)
 L = 12
 g = NamedGraph(SimpleGraph(uniform_tree(L)))
 
-s = continuous_siteinds(g)
-index_map = IndexMap(s; map_dimension=2)
+s = continuous_siteinds(g; map_dimension=2)
 
-ψ_fxy = 0.1 * rand_itn(s, index_map; link_space=2)
-∇ = laplacian_operator(s, index_map; scale=false, cutoff=1e-8)
+ψ_fxy = 0.1 * rand_itn(s; link_space=2)
+∇ = laplacian_operator(s; scale=false, cutoff=1e-8)
 println("2D Laplacian constructed for this tree, bond dimension is $(maxlinkdim(∇))")
 
 init_energy =
diff --git a/examples/construct_multi_dimensional_function.jl b/examples/construct_multi_dimensional_function.jl
index 84dbc84..42e05c8 100644
--- a/examples/construct_multi_dimensional_function.jl
+++ b/examples/construct_multi_dimensional_function.jl
@@ -8,15 +8,14 @@ using Random: Random
 L = 12
 Random.seed!(1234)
 g = NamedGraph(SimpleGraph(uniform_tree(L)))
-s = continuous_siteinds(g)
-index_map = IndexMap(s; map_dimension=3)
+s = continuous_siteinds(g; map_dimension=3)
 
 println(
   "Constructing the 3D function f(x,y,z) = x³(y + y²) + cosh(πz) as a tensor network on a randomly chosen tree with $L vertices",
 )
-ψ_fx = poly_itn(s, index_map, [0.0, 0.0, 0.0, 1.0]; dimension=1)
-ψ_fy = poly_itn(s, index_map, [0.0, 1.0, 1.0, 0.0]; dimension=2)
-ψ_fz = cosh_itn(s, index_map; k=Float64(pi), dimension=3)
+ψ_fx = poly_itn(s, [0.0, 0.0, 0.0, 1.0]; dimension=1)
+ψ_fy = poly_itn(s, [0.0, 1.0, 1.0, 0.0]; dimension=2)
+ψ_fz = cosh_itn(s; k=Float64(pi), dimension=3)
 ψxyz = ψ_fx * ψ_fy + ψ_fz
 
 ψxyz = truncate(ψxyz; cutoff=1e-12)
diff --git a/src/ITensorNumericalAnalysis.jl b/src/ITensorNumericalAnalysis.jl
index baf8f32..fdbdd6f 100644
--- a/src/ITensorNumericalAnalysis.jl
+++ b/src/ITensorNumericalAnalysis.jl
@@ -2,6 +2,7 @@ module ITensorNumericalAnalysis
 
 include("utils.jl")
 include("indexmap.jl")
+include("indsnetworkmap.jl")
 include("polynomialutils.jl")
 include("itensornetworkfunction.jl")
 include("elementary_functions.jl")
@@ -11,6 +12,7 @@ export continuous_siteinds
 export ITensorNetworkFunction, itensornetwork, dimension_vertices
 export IndexMap,
   default_dimension_map,
+  dimension_inds,
   calculate_xyz,
   calculate_x,
   calculate_ind_values,
@@ -18,6 +20,7 @@ export IndexMap,
   dimensions,
   grid_points
 export IndsNetworkMap,
+  continuous_siteinds,
   indsnetwork,
   indexmap,
   vertex_dimension,
@@ -33,7 +36,7 @@ export const_itensornetwork,
   sin_itensornetwork,
   get_edge_toward_root,
   polynomial_itensornetwork,
-  random_itensornetworkfunction,
+  random_itensornetwork,
   laplacian_operator,
   first_derivative_operator,
   second_derivative_operator,
diff --git a/src/elementary_functions.jl b/src/elementary_functions.jl
index 2acb5bf..a04d799 100644
--- a/src/elementary_functions.jl
+++ b/src/elementary_functions.jl
@@ -19,12 +19,9 @@ default_dimension() = 1
 
 """Build a representation of the function f(x,y,z,...) = c, with flexible choice of linkdim"""
 function const_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
-  c::Union{Float64,ComplexF64}=default_c_value(),
-  linkdim::Int64=1,
+  s::IndsNetworkMap; c::Union{Float64,ComplexF64}=default_c_value(), linkdim::Int64=1
 )
-  ψ = random_tensornetwork(s; link_space=linkdim)
+  ψ = random_itensornetwork(s; link_space=linkdim)
   inv_L = Float64(1.0 / nv(s))
   for v in vertices(ψ)
     sinds = inds(s, v)
@@ -32,25 +29,22 @@ function const_itensornetwork(
     ψ[v] = c^inv_L * c_tensor(only(sinds), virt_inds)
   end
 
-  return ITensorNetworkFunction(ψ, indexmap)
+  return ψ
 end
 
 """Construct the product state representation of the exp(kx+a) 
 function for x ∈ [0,1] as an ITensorNetworkFunction, along the specified dim"""
 function exp_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ = const_itensornetwork(s, indexmap)
-  Lx = length(inds(indexmap, dimension))
+  ψ = const_itensornetwork(s)
+  Lx = length(dimension_vertices(ψ, dimension))
   for v in dimension_vertices(ψ, dimension)
     sind = only(inds(s, v))
-    ψ[v] = ITensor(
-      exp(a / Lx) * exp.(k * index_values_to_scalars(indexmap, sind)), inds(ψ[v])
-    )
+    ψ[v] = ITensor(exp(a / Lx) * exp.(k * index_values_to_scalars(s, sind)), inds(ψ[v]))
   end
 
   return ψ
@@ -59,14 +53,13 @@ end
 """Construct the bond dim 2 representation of the cosh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
 defines the network geometry. Vertex map provides the ordering of the sites as bits"""
 function cosh_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ1 = exp_itensornetwork(s, indexmap; a, k, dimension)
-  ψ2 = exp_itensornetwork(s, indexmap; a=-a, k=-k, dimension)
+  ψ1 = exp_itensornetwork(s; a, k, dimension)
+  ψ2 = exp_itensornetwork(s; a=-a, k=-k, dimension)
 
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= 0.5
@@ -77,14 +70,13 @@ end
 """Construct the bond dim 2 representation of the sinh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
 defines the network geometry. Vertex map provides the ordering of the sites as bits"""
 function sinh_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ1 = exp_itensornetwork(s, indexmap; a, k, dimension)
-  ψ2 = exp_itensornetwork(s, indexmap; a=-a, k=-k, dimension)
+  ψ1 = exp_itensornetwork(s; a, k, dimension)
+  ψ2 = exp_itensornetwork(s; a=-a, k=-k, dimension)
 
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= -0.5
@@ -95,16 +87,15 @@ end
 """Construct the bond dim n representation of the tanh(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
 defines the network geometry. Vertex map provides the ordering of the sites as bits"""
 function tanh_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   nterms::Int64=default_nterms(),
   dimension::Int64=default_dimension(),
 )
-  ψ = const_itensornetwork(s, indexmap)
+  ψ = const_itensornetwork(s)
   for n in 1:nterms
-    ψt = exp_itensornetwork(s, indexmap; a=-2 * n * a, k=-2 * k * n, dimension)
+    ψt = exp_itensornetwork(s; a=-2 * n * a, k=-2 * k * n, dimension)
     ψt[first(vertices(ψt))] *= 2 * ((-1)^n)
     ψ = ψ + ψt
   end
@@ -115,14 +106,13 @@ end
 """Construct the bond dim 2 representation of the cos(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
 defines the network geometry. Vertex map provides the ordering of the sites as bits"""
 function cos_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ1 = exp_itensornetwork(s, indexmap; a=a * im, k=k * im, dimension)
-  ψ2 = exp_itensornetwork(s, indexmap; a=-a * im, k=-k * im, dimension)
+  ψ1 = exp_itensornetwork(s; a=a * im, k=k * im, dimension)
+  ψ2 = exp_itensornetwork(s; a=-a * im, k=-k * im, dimension)
 
   ψ1[first(vertices(ψ1))] *= 0.5
   ψ2[first(vertices(ψ1))] *= 0.5
@@ -133,14 +123,13 @@ end
 """Construct the bond dim 2 representation of the sin(kx+a) function for x ∈ [0,1] as an ITensorNetwork, using an IndsNetwork which 
 defines the network geometry. Vertex map provides the ordering of the sites as bits"""
 function sin_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap;
+  s::IndsNetworkMap;
   k::Union{Float64,ComplexF64}=default_k_value(),
   a::Union{Float64,ComplexF64}=default_a_value(),
   dimension::Int64=default_dimension(),
 )
-  ψ1 = exp_itensornetwork(s, indexmap; a=a * im, k=k * im, dimension)
-  ψ2 = exp_itensornetwork(s, indexmap; a=-a * im, k=-k * im, dimension)
+  ψ1 = exp_itensornetwork(s; a=a * im, k=k * im, dimension)
+  ψ2 = exp_itensornetwork(s; a=-a * im, k=-k * im, dimension)
 
   ψ1[first(vertices(ψ1))] *= -0.5 * im
   ψ2[first(vertices(ψ1))] *= 0.5 * im
@@ -151,23 +140,22 @@ end
 """Build a representation of the function f(x) = sum_{i=0}^{n}coeffs[i+1]*(x)^{i} on the graph structure specified
 by indsnetwork"""
 function polynomial_itensornetwork(
-  s::IndsNetwork,
-  indexmap::IndexMap,
-  coeffs::Vector{Float64};
-  dimension::Int64=default_dimension(),
+  s::IndsNetworkMap, coeffs::Vector{Float64}; dimension::Int64=default_dimension()
 )
   n = length(coeffs)
   #First treeify the index network (ignore edges that form loops)
-  g = underlying_graph(s)
+  _s = indsnetwork(s)
+  g = underlying_graph(_s)
   g_tree = undirected_graph(random_bfs_tree(g, first(vertices(g))))
-  s_tree = add_edges(rem_edges(s, edges(g)), edges(g_tree))
+  s_tree = add_edges(rem_edges(_s, edges(g)), edges(g_tree))
+  s_tree = IndsNetworkMap(s_tree, indexmap(s))
 
-  ψ = const_itensornetwork(s_tree, indexmap; linkdim=n)
+  ψ = const_itensornetwork(s_tree; linkdim=n)
   dim_vertices = dimension_vertices(ψ, dimension)
   source_vertex = first(dim_vertices)
 
   for v in dim_vertices
-    siteindex = only(s_tree[v])
+    siteindex = only(inds(s_tree, v))
     if v != source_vertex
       e = get_edge_toward_vertex(g_tree, v, source_vertex)
       betaindex = only(commoninds(ψ, e))
@@ -177,7 +165,7 @@ function polynomial_itensornetwork(
         siteindex,
         alphas,
         betaindex,
-        index_values_to_scalars(indexmap, siteindex),
+        index_values_to_scalars(s_tree, siteindex),
       )
     elseif v == source_vertex
       betaindex = Index(n, "DummyInd")
@@ -187,7 +175,7 @@ function polynomial_itensornetwork(
         siteindex,
         alphas,
         betaindex,
-        index_values_to_scalars(indexmap, siteindex),
+        index_values_to_scalars(s_tree, siteindex),
       )
       ψ[v] = ψv * ITensor(coeffs, betaindex)
     end
@@ -196,7 +184,7 @@ function polynomial_itensornetwork(
   #Put the transfer tensors in, these are special tensors that
   # go on the digits (sites) that don't correspond to the desired dimension
   for v in setdiff(vertices(ψ), dim_vertices)
-    siteindex = only(s_tree[v])
+    siteindex = only(inds(s_tree, v))
     e = get_edge_toward_vertex(g_tree, v, source_vertex)
     betaindex = only(commoninds(ψ, e))
     alphas = setdiff(inds(ψ[v]), Index[siteindex, betaindex])
@@ -206,8 +194,8 @@ function polynomial_itensornetwork(
   return ψ
 end
 
-function random_itensornetwork(s::IndsNetwork, bit_map; kwargs...)
-  return ITensorNetworkFunction(random_tensornetwork(s; kwargs...), bit_map)
+function random_itensornetwork(s::IndsNetworkMap; kwargs...)
+  return ITensorNetworkFunction(random_tensornetwork(indsnetwork(s); kwargs...), s)
 end
 
 const const_itn = const_itensornetwork
diff --git a/src/elementary_operators.jl b/src/elementary_operators.jl
index 3d6453e..826cc3c 100644
--- a/src/elementary_operators.jl
+++ b/src/elementary_operators.jl
@@ -46,42 +46,38 @@ function ITensors.op(::OpName"Dup", ::SiteType"Digit", s::Index)
 end
 
 function forward_shift_opsum(
-  s::IndsNetwork,
-  index_map::IndexMap;
-  dimension=default_dimension(),
-  boundary=default_boundary(),
-  n::Int64=0,
+  s::IndsNetworkMap; dimension=default_dimension(), boundary=default_boundary(), n::Int64=0
 )
   @assert is_tree(s)
   @assert base(s) == 2
   ttn_op = OpSum()
-  dim_vertices = dimension_vertices(s, index_map, dimension)
+  dim_vertices = dimension_vertices(s, dimension)
   L = length(dim_vertices)
 
-  string_site = [("D+", vertex(s, index_map, dimension, L - n))]
-  add!(ttn_op, 1.0, "D+", vertex(s, index_map, dimension, L - n))
+  string_site = [("D+", vertex(s, dimension, L - n))]
+  add!(ttn_op, 1.0, "D+", vertex(s, dimension, L - n))
   for i in (L - n):-1:2
     pop!(string_site)
-    push!(string_site, ("D-", vertex(s, index_map, dimension, i)))
-    push!(string_site, ("D+", vertex(s, index_map, dimension, i - 1)))
+    push!(string_site, ("D-", vertex(s, dimension, i)))
+    push!(string_site, ("D+", vertex(s, dimension, i - 1)))
     add!(ttn_op, 1.0, (string_site...)...)
   end
 
   if boundary == "Neumann"
     string_site = [
       if j <= (L - n)
-        ("Dup", vertex(s, index_map, dimension, j))
+        ("Dup", vertex(s, dimension, j))
       else
-        ("I", vertex(s, index_map, dimension, j))
+        ("I", vertex(s, dimension, j))
       end for j in 1:L
     ]
     add!(ttn_op, 1.0, (string_site...)...)
   elseif boundary == "Periodic"
     string_site = [
       if j <= (L - n)
-        ("D-", vertex(s, index_map, dimension, j))
+        ("D-", vertex(s, dimension, j))
       else
-        ("I", vertex(s, index_map, dimension, j))
+        ("I", vertex(s, dimension, j))
       end for j in 1:L
     ]
     add!(ttn_op, 1.0, (string_site...)...)
@@ -91,42 +87,38 @@ function forward_shift_opsum(
 end
 
 function backward_shift_opsum(
-  s::IndsNetwork,
-  index_map::IndexMap;
-  dimension=default_dimension(),
-  boundary=default_boundary(),
-  n::Int64=0,
+  s::IndsNetworkMap; dimension=default_dimension(), boundary=default_boundary(), n::Int64=0
 )
   @assert is_tree(s)
   @assert base(s) == 2
   ttn_op = OpSum()
-  dim_vertices = dimension_vertices(s, index_map, dimension)
+  dim_vertices = dimension_vertices(s, dimension)
   L = length(dim_vertices)
 
-  string_site = [("D-", vertex(s, index_map, dimension, L - n))]
-  add!(ttn_op, 1.0, "D-", vertex(s, index_map, dimension, L - n))
+  string_site = [("D-", vertex(s, dimension, L - n))]
+  add!(ttn_op, 1.0, "D-", vertex(s, dimension, L - n))
   for i in (L - n):-1:2
     pop!(string_site)
-    push!(string_site, ("D+", vertex(s, index_map, dimension, i)))
-    push!(string_site, ("D-", vertex(s, index_map, dimension, i - 1)))
+    push!(string_site, ("D+", vertex(s, dimension, i)))
+    push!(string_site, ("D-", vertex(s, dimension, i - 1)))
     add!(ttn_op, 1.0, (string_site...)...)
   end
 
   if boundary == "Neumann"
     string_site = [
       if j <= (L - n)
-        ("Ddn", vertex(s, index_map, dimension, j))
+        ("Ddn", vertex(s, dimension, j))
       else
-        ("I", vertex(s, index_map, dimension, j))
+        ("I", vertex(s, dimension, j))
       end for j in 1:L
     ]
     add!(ttn_op, 1.0, (string_site...)...)
   elseif boundary == "Periodic"
     string_site = [
       if j <= (L - n)
-        ("D+", vertex(s, index_map, dimension, j))
+        ("D+", vertex(s, dimension, j))
       else
-        ("I", vertex(s, index_map, dimension, j))
+        ("I", vertex(s, dimension, j))
       end for j in 1:L
     ]
     add!(ttn_op, 1.0, (string_site...)...)
@@ -135,30 +127,25 @@ function backward_shift_opsum(
   return ttn_op
 end
 
-function no_shift_opsum(s::IndsNetwork)
+function no_shift_opsum(s::IndsNetworkMap)
   ttn_op = OpSum()
   string_site_full = [("I", v) for v in vertices(s)]
   add!(ttn_op, 1.0, (string_site_full...)...)
   return ttn_op
 end
 
-function backward_shift_op(
-  s::IndsNetwork, index_map::IndexMap; truncate_kwargs=(;), kwargs...
-)
-  ttn_opsum = backward_shift_opsum(s, index_map; kwargs...)
-  return ttn(ttn_opsum, s; algorithm="svd", truncate_kwargs...)
+function backward_shift_op(s::IndsNetworkMap; truncate_kwargs=(;), kwargs...)
+  ttn_opsum = backward_shift_opsum(s; kwargs...)
+  return ttn(ttn_opsum, indsnetwork(s); algorithm="svd", truncate_kwargs...)
 end
 
-function forward_shift_op(
-  s::IndsNetwork, index_map::IndexMap; truncate_kwargs=(;), kwargs...
-)
-  ttn_opsum = forward_shift_opsum(s, index_map; kwargs...)
-  return ttn(ttn_opsum, s; algorithm="svd", truncate_kwargs...)
+function forward_shift_op(s::IndsNetworkMap; truncate_kwargs=(;), kwargs...)
+  ttn_opsum = forward_shift_opsum(s; kwargs...)
+  return ttn(ttn_opsum, indsnetwork(s); algorithm="svd", truncate_kwargs...)
 end
 
 function stencil(
-  s::IndsNetwork,
-  index_map,
+  s::IndsNetworkMap,
   shifts::Vector{Float64},
   delta_power::Int64;
   dimension=default_dimension(),
@@ -175,7 +162,7 @@ function stencil(
     n = i == 1 ? 1 : 0
     if !iszero(shifts[i])
       stencil_opsum +=
-        shifts[i] * forward_shift_opsum(s, index_map; dimension, boundary=right_boundary, n)
+        shifts[i] * forward_shift_opsum(s; dimension, boundary=right_boundary, n)
     end
   end
 
@@ -183,14 +170,14 @@ function stencil(
     n = i == 5 ? 1 : 0
     if !iszero(shifts[i])
       stencil_opsum +=
-        shifts[i] * backward_shift_opsum(s, index_map; dimension, boundary=left_boundary, n)
+        shifts[i] * backward_shift_opsum(s; dimension, boundary=left_boundary, n)
     end
   end
 
-  stencil_op = ttn(stencil_opsum, s; algorithm="svd", kwargs...)
+  stencil_op = ttn(stencil_opsum, indsnetwork(s); algorithm="svd", kwargs...)
 
   if scale
-    for v in dimension_vertices(s, index_map, dimension)
+    for v in dimension_vertices(s, dimension)
       stencil_op[v] = (b^delta_power) * stencil_op[v]
     end
   end
@@ -198,36 +185,36 @@ function stencil(
   return stencil_op
 end
 
-function first_derivative_operator(s::IndsNetwork, index_map; kwargs...)
-  return stencil(s, index_map, [0.0, 0.5, 0.0, -0.5, 0.0], 1; kwargs...)
+function first_derivative_operator(s::IndsNetworkMap; kwargs...)
+  return stencil(s, [0.0, 0.5, 0.0, -0.5, 0.0], 1; kwargs...)
 end
 
-function second_derivative_operator(s::IndsNetwork, index_map; kwargs...)
-  return stencil(s, index_map, [0.0, 1.0, -2.0, 1.0, 0.0], 2; kwargs...)
+function second_derivative_operator(s::IndsNetworkMap; kwargs...)
+  return stencil(s, [0.0, 1.0, -2.0, 1.0, 0.0], 2; kwargs...)
 end
 
-function third_derivative_operator(s::IndsNetwork, index_map; kwargs...)
-  return stencil(s, index_map, [0.5, -1.0, 0.0, 1.0, -0.5], 3; kwargs...)
+function third_derivative_operator(s::IndsNetworkMap; kwargs...)
+  return stencil(s, [0.5, -1.0, 0.0, 1.0, -0.5], 3; kwargs...)
 end
 
-function fourth_derivative_operator(s::IndsNetwork, index_map; kwargs...)
-  return stencil(s, index_map, [1.0, -4.0, 6.0, -4.0, 1.0], 4; kwargs...)
+function fourth_derivative_operator(s::IndsNetworkMap; kwargs...)
+  return stencil(s, [1.0, -4.0, 6.0, -4.0, 1.0], 4; kwargs...)
 end
 
 function laplacian_operator(
-  s::IndsNetwork, index_map; dimensions=[i for i in 1:dimension(index_map)], kwargs...
+  s::IndsNetworkMap; dimensions=[i for i in 1:dimension(s)], kwargs...
 )
   remaining_dims = copy(dimensions)
-  ∇ = second_derivative_operator(s, index_map; dimension=first(remaining_dims), kwargs...)
+  ∇ = second_derivative_operator(s; dimension=first(remaining_dims), kwargs...)
   popfirst!(remaining_dims)
   for rd in remaining_dims
-    ∇ += second_derivative_operator(s, index_map; dimension=rd, kwargs...)
+    ∇ += second_derivative_operator(s; dimension=rd, kwargs...)
   end
   return ∇
 end
 
-function identity_operator(s::IndsNetwork, index_map; kwargs...)
-  return stencil(s, index_map, [0.0, 1.0, 0.0], 0; kwargs...)
+function identity_operator(s::IndsNetworkMap; kwargs...)
+  return stencil(s, [0.0, 1.0, 0.0], 0; kwargs...)
 end
 
 function operator(fx::ITensorNetworkFunction)
@@ -271,7 +258,7 @@ function operate(operator::TreeTensorNetwork, ψ::ITensorNetworkFunction; kwargs
   ψ_tn = ttn(itensornetwork(ψ))
   ψO_tn = noprime(contract(operator, ψ_tn; init=prime(copy(ψ_tn)), kwargs...))
 
-  return ITensorNetworkFunction(ITensorNetwork(ψO_tn), indexmap(ψ))
+  return ITensorNetworkFunction(ITensorNetwork(ψO_tn), indsnetworkmap(ψ))
 end
 
 function operate(operator::ITensorNetwork, ψ::ITensorNetworkFunction; kwargs...)
diff --git a/src/indexmap.jl b/src/indexmap.jl
index 0f4e913..e155f16 100644
--- a/src/indexmap.jl
+++ b/src/indexmap.jl
@@ -1,3 +1,4 @@
+using Base: Base
 using Dictionaries: Dictionary, set!
 using ITensors: ITensors, Index, dim
 using ITensorNetworks: IndsNetwork, vertex_data
@@ -67,7 +68,7 @@ function ITensors.inds(im::IndexMap)
   @assert keys(index_dimension(im)) == keys(index_digit(im))
   return collect(keys(index_dimension(im)))
 end
-function ITensors.inds(im::IndexMap, dimension::Int64)
+function dimension_inds(im::IndexMap, dimension::Int64)
   return collect(
     filter(i -> index_dimension(im)[i] == dimension, keys(index_dimension(im)))
   )
@@ -88,7 +89,7 @@ end
 function calculate_xyz(im::IndexMap, ind_to_ind_value_map, dimensions::Vector{Int64})
   out = Float64[]
   for dimension in dimensions
-    indices = inds(im, dimension)
+    indices = dimension_inds(im, dimension)
     push!(
       out,
       sum([index_value_to_scalar(im, ind, ind_to_ind_value_map[ind]) for ind in indices]),
@@ -114,8 +115,8 @@ function calculate_ind_values(
   ind_to_ind_value_map = Dictionary()
   for (i, x) in enumerate(xs)
     dimension = dimensions[i]
-    x_rn = copy(x)
-    indices = inds(im, dimension)
+    x_rn = x
+    indices = dimension_inds(im, dimension)
     sorted_inds = sort(indices; by=indices -> digit(im, indices))
     for ind in sorted_inds
       ind_val = dim(ind) - 1
@@ -152,42 +153,12 @@ function calculate_ind_values(im::IndexMap, x::Float64; kwargs...)
 end
 
 function grid_points(im::IndexMap, N::Int64, dimension::Int64)
-  dims = dim.(inds(im, dimension))
+  dims = dim.(dimension_inds(im, dimension))
   @assert all(y -> y == first(dims), dims)
   base = first(dims)
   vals = Vector{Float64}
-  L = length(inds(im, dimension))
+  L = length(dimension_inds(im, dimension))
   a = round(base^L / N)
   grid_points = [i * (a / base^L) for i in 0:(N + 1)]
   return filter(x -> x <= 1, grid_points)
 end
-
-#Functions for translating onto vertices through an IndsNetwork
-function vertices_dimensions(s::IndsNetwork, im::IndexMap, verts::Vector)
-  indices = inds(s, verts)
-  return [dimension(im, i) for i in indices]
-end
-
-function vertices_digits(s::IndsNetwork, im::IndexMap, verts::Vector)
-  indices = inds(s, verts)
-  return [digit(im, i) for i in indices]
-end
-
-function vertex_dimension(s::IndsNetwork, im::IndexMap, v)
-  ind = only(inds(s, v))
-  return dimension(im, ind)
-end
-
-function vertex_digit(s::IndsNetwork, im::IndexMap, v)
-  ind = only(inds(s, v))
-  return digit(im, ind)
-end
-
-function dimension_vertices(s::IndsNetwork, im::IndexMap, dimension::Int64)
-  return filter(v -> vertex_dimension(s, im, v) == dimension, vertices(s))
-end
-
-function vertex(s::IndsNetwork, im::IndexMap, dimension::Int64, digit::Int64)
-  index = ind(im, dimension, digit)
-  return only(filter(v -> index ∈ s[v], vertices(s)))
-end
diff --git a/src/indsnetworkmap.jl b/src/indsnetworkmap.jl
new file mode 100644
index 0000000..f6fd04d
--- /dev/null
+++ b/src/indsnetworkmap.jl
@@ -0,0 +1,105 @@
+using Base: Base
+using Graphs: Graphs
+using NamedGraphs: NamedGraphs
+using DataGraphs: DataGraphs
+using ITensors: ITensors
+using ITensorNetworks:
+  ITensorNetworks, AbstractIndsNetwork, IndsNetwork, data_graph, underlying_graph
+
+struct IndsNetworkMap{V,I,IN<:IndsNetwork{V,I},IM} <: AbstractIndsNetwork{V,I}
+  indsnetwork::IN
+  indexmap::IM
+end
+
+indsnetwork(inm::IndsNetworkMap) = inm.indsnetwork
+indexmap(inm::IndsNetworkMap) = inm.indexmap
+
+indtype(inm::IndsNetworkMap) = indtype(typeof(indsnetwork(inm)))
+indtype(::Type{<:IndsNetworkMap{V,I,IN,IM}}) where {V,I,IN,IM} = I
+ITensorNetworks.data_graph(inm::IndsNetworkMap) = data_graph(indsnetwork(inm))
+function DataGraphs.underlying_graph(inm::IndsNetworkMap)
+  return underlying_graph(data_graph(indsnetwork(inm)))
+end
+NamedGraphs.vertextype(::Type{<:IndsNetworkMap{V,I,IN,IM}}) where {V,I,IN,IM} = V
+DataGraphs.underlying_graph_type(G::Type{<:IndsNetworkMap}) = NamedGraph{vertextype(G)}
+Graphs.is_directed(::Type{<:IndsNetworkMap}) = false
+
+function Base.copy(inm::IndsNetworkMap)
+  return IndsNetworkMap(indsnetwork(inm), indexmap(inm))
+end
+
+#Constructors 
+function IndsNetworkMap(
+  s::IndsNetwork, dimension_vertices::Vector{Vector{V}}; kwargs...
+) where {V}
+  return IndsNetworkMap(s, IndexMap(s, dimension_vertices; kwargs...))
+end
+
+function IndsNetworkMap(s::IndsNetwork, dimension_indices::Vector{Vector{Index}})
+  return IndsNetworkMap(s, IndexMap(dimension_indices; kwargs...))
+end
+
+function IndsNetworkMap(s::IndsNetwork; kwargs...)
+  return IndsNetworkMap(s, IndexMap(s; kwargs...))
+end
+
+function IndsNetworkMap(g::AbstractGraph, args...; base::Int=2, kwargs...)
+  s = digit_siteinds(g; base)
+  return IndsNetworkMap(s, args...; kwargs...)
+end
+
+const continuous_siteinds = IndsNetworkMap
+
+#Forward functionality from indexmap
+for f in [
+  :ind,
+  :dimension,
+  :dimension_inds,
+  :dimensions,
+  :digit,
+  :digits,
+  :calculate_ind_values,
+  :calculate_x,
+  :calculate_xyz,
+  :grid_points,
+  :index_value_to_scalar,
+  :index_values_to_scalars,
+]
+  @eval begin
+    function $f(inm::IndsNetworkMap, args...; kwargs...)
+      return $f(indexmap(inm), args...; kwargs...)
+    end
+  end
+end
+
+#Functions on indsnetwork that don't seem to autoforward
+function ITensors.inds(inm::IndsNetworkMap, args...; kwargs...)
+  return inds(indsnetwork(inm), args...; kwargs...)
+end
+
+base(inm::IndsNetworkMap) = base(indsnetwork(inm))
+
+function vertices_dimensions(inm::IndsNetworkMap, verts::Vector)
+  return [dimension(inm, i) for i in inds(inm, verts)]
+end
+
+function vertices_digits(inm::IndsNetworkMap, verts::Vector)
+  return [digit(inm, i) for i in inds(inm, verts)]
+end
+
+function vertex_dimension(inm::IndsNetworkMap, v)
+  return dimension(inm, only(inds(inm, v)))
+end
+
+function vertex_digit(inm::IndsNetworkMap, v)
+  return digit(inm, only(inds(inm, v)))
+end
+
+function dimension_vertices(inm::IndsNetworkMap, dimension::Int64)
+  return filter(v -> vertex_dimension(inm, v) == dimension, vertices(inm))
+end
+
+function vertex(inm::IndsNetworkMap, dimension::Int64, digit::Int64)
+  index = ind(inm, dimension, digit)
+  return only(filter(v -> index ∈ inm[v], vertices(inm)))
+end
diff --git a/src/itensornetworkfunction.jl b/src/itensornetworkfunction.jl
index d5590d0..412977e 100644
--- a/src/itensornetworkfunction.jl
+++ b/src/itensornetworkfunction.jl
@@ -1,15 +1,16 @@
+using Base: Base
 using ITensorNetworks: ITensorNetworks, AbstractITensorNetwork, data_graph, data_graph_type
 using ITensors: ITensor, dim, contract, siteinds, onehot
 using Graphs: Graphs
 
-struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},IM<:IndexMap} <:
+struct ITensorNetworkFunction{V,TN<:AbstractITensorNetwork{V},INM<:IndsNetworkMap} <:
        AbstractITensorNetwork{V}
   itensornetwork::TN
-  indexmap::IM
+  indsnetworkmap::INM
 end
 
 itensornetwork(fitn::ITensorNetworkFunction) = fitn.itensornetwork
-indexmap(fitn::ITensorNetworkFunction) = fitn.indexmap
+indsnetworkmap(fitn::ITensorNetworkFunction) = fitn.indsnetworkmap
 
 #Needed for interface from AbstractITensorNetwork
 function ITensorNetworks.data_graph_type(TN::Type{<:ITensorNetworkFunction})
@@ -17,22 +18,21 @@ function ITensorNetworks.data_graph_type(TN::Type{<:ITensorNetworkFunction})
 end
 ITensorNetworks.data_graph(fitn::ITensorNetworkFunction) = data_graph(itensornetwork(fitn))
 function Base.copy(fitn::ITensorNetworkFunction)
-  return ITensorNetworkFunction(copy(itensornetwork(fitn)), copy(indexmap(fitn)))
+  return ITensorNetworkFunction(copy(itensornetwork(fitn)), copy(indsnetworkmap(fitn)))
 end
 
 function ITensorNetworkFunction(
   itn::AbstractITensorNetwork, dimension_vertices::Vector{Vector{V}}
 ) where {V}
   s = siteinds(itn)
-  return ITensorNetworkFunction(itn, IndexMap(s, dimension_vertices))
+  return ITensorNetworkFunction(itn, IndsNetworkMap(s, dimension_vertices))
 end
 
-#Constructor, assume one-dimensional and ordered as vertices of the itn
 function ITensorNetworkFunction(itn::AbstractITensorNetwork)
-  return ITensorNetworkFunction(itn, IndexMap(siteinds(itn)))
+  return ITensorNetworkFunction(itn, IndsNetworkMap(siteinds(itn)))
 end
 
-#Forward functionality from indexmap
+#Forward functionality from indsnetworkmap
 for f in [
   :ind,
   :dimension,
@@ -43,15 +43,6 @@ for f in [
   :calculate_x,
   :calculate_xyz,
   :grid_points,
-]
-  @eval begin
-    function $f(fitn::ITensorNetworkFunction, args...; kwargs...)
-      return $f(indexmap(fitn), args...; kwargs...)
-    end
-  end
-end
-
-for f in [
   :vertices_dimensions,
   :vertices_digits,
   :vertex_digit,
@@ -60,14 +51,16 @@ for f in [
 ]
   @eval begin
     function $f(fitn::ITensorNetworkFunction, args...; kwargs...)
-      return $f(siteinds(fitn), indexmap(fitn), args...; kwargs...)
+      return $f(indsnetworkmap(fitn), args...; kwargs...)
     end
   end
 end
 
+ITensors.siteinds(fitn::ITensorNetworkFunction) = indsnetwork(indsnetworkmap(fitn))
+
 function project(fitn::ITensorNetworkFunction, ind_to_ind_value_map)
   fitn = copy(fitn)
-  s = siteinds(fitn)
+  s = indsnetwork(indsnetworkmap(fitn))
   for v in vertices(fitn)
     indices = inds(s, v)
     for ind in indices
@@ -97,5 +90,5 @@ end
 function ITensorNetworks.truncate(fitn::ITensorNetworkFunction; kwargs...)
   @assert is_tree(fitn)
   ψ = truncate(ttn(itensornetwork(fitn)); kwargs...)
-  return ITensorNetworkFunction(ITensorNetwork(ψ), indexmap(fitn))
+  return ITensorNetworkFunction(ITensorNetwork(ψ), indsnetworkmap(fitn))
 end
diff --git a/src/utils.jl b/src/utils.jl
index 5ee6b2f..97096a9 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -43,17 +43,17 @@ function ITensors.inds(s::IndsNetwork)
   return inds(s, vertices(s))
 end
 
-function continuous_siteinds(g::AbstractGraph; base=2)
-  is = IndsNetwork(g)
-  for v in vertices(g)
-    is[v] = [Index(base, "Digit, V$(vertex_tag(v))")]
-  end
-  return is
-end
-
 function base(s::IndsNetwork)
   indices = inds(s)
   dims = dim.(indices)
   @assert all(d -> d == first(dims), dims)
   return first(dims)
 end
+
+function digit_siteinds(g::AbstractGraph; base=2)
+  is = IndsNetwork(g)
+  for v in vertices(g)
+    is[v] = [Index(base, "Digit, V$(vertex_tag(v))")]
+  end
+  return is
+end
diff --git a/test/test_derivative_operators.jl b/test/test_derivative_operators.jl
deleted file mode 100644
index 65b4b60..0000000
--- a/test/test_derivative_operators.jl
+++ /dev/null
@@ -1,139 +0,0 @@
-using Test
-using ITensorNumericalAnalysis
-
-using ITensors: siteinds
-using ITensorNetworks: maxlinkdim
-using Graphs: SimpleGraph, uniform_tree
-using NamedGraphs: named_grid, named_comb_tree, NamedGraph, nv, vertices
-using ITensorNumericalAnalysis: itensornetwork
-using Dictionaries: Dictionary
-
-@testset "test differentiation in 1D on MPS" begin
-  g = named_grid((9, 1))
-  L = nv(g)
-  delta = (2.0)^(-Float64(L))
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s)
-  left_boundary, right_boundary = "Periodic", "Periodic"
-
-  f1 = first_derivative_operator(s, index_map; cutoff=1e-12, left_boundary, right_boundary)
-  f2 = second_derivative_operator(s, index_map; cutoff=1e-12, left_boundary, right_boundary)
-  f3 = third_derivative_operator(s, index_map; cutoff=1e-12, left_boundary, right_boundary)
-  f4 = fourth_derivative_operator(s, index_map; cutoff=1e-12, left_boundary, right_boundary)
-
-  ψ_fx = sin_itn(s, index_map; k=2.0 * Float64(pi))
-
-  ψ_f1x = operate(f1, ψ_fx; cutoff=1e-8)
-  ψ_f2x = operate(f2, ψ_fx; cutoff=1e-8)
-  ψ_f3x = operate(f3, ψ_fx; cutoff=1e-8)
-  ψ_f4x = operate(f4, ψ_fx; cutoff=1e-8)
-
-  xs = [0.0, 0.25, 0.625, 0.875, 1.0 - delta]
-  for x in xs
-    @test 1.0 + calculate_fx(ψ_fx, x) ≈ 1.0 + sin(2.0 * pi * x) rtol = 1e-3
-    @test 1.0 + calculate_fx(ψ_f1x, x) ≈ 1.0 + 2.0 * pi * cos(2.0 * pi * x) rtol = 1e-3
-    @test 1.0 + calculate_fx(ψ_f2x, x) ≈ 1.0 + -1.0 * (2.0 * pi)^2 * sin(2.0 * pi * x) rtol =
-      1e-3
-    @test 1.0 + calculate_fx(ψ_f3x, x) ≈ 1.0 + -1.0 * (2.0 * pi)^3 * cos(2.0 * pi * x) rtol =
-      1e-3
-    @test 1.0 + calculate_fx(ψ_f4x, x) ≈ 1.0 + 1.0 * (2.0 * pi)^4 * sin(2.0 * pi * x) rtol =
-      1e-3
-  end
-end
-
-@testset "test derivative in 1D on tree" begin
-  g = named_comb_tree((4, 3))
-  L = nv(g)
-  delta = 2.0^(-Float64(L))
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s)
-
-  ∂_∂x = first_derivative_operator(s, index_map; cutoff=1e-10)
-
-  ψ_fx = sin_itn(s, index_map; k=Float64(pi))
-  ∂x_ψ_fx = operate(∂_∂x, ψ_fx; cutoff=1e-12)
-
-  xs = [delta, 0.125, 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
-
-@testset "test multiplication_operator_in_1D" begin
-  g = named_comb_tree((4, 3))
-  L = nv(g)
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s)
-
-  ψ_gx = sin_itn(s, index_map; k=0.5 * Float64(pi))
-  ψ_fx = cos_itn(s, index_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-3
-  end
-end
-
-@testset "test multiplication_operator_in_2D" begin
-  L = 8
-  g = NamedGraph(SimpleGraph(uniform_tree(L)))
-
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s; map_dimension=2)
-
-  ψ_fx = cos_itn(s, index_map; k=0.25 * Float64(pi), dimension=1)
-  ψ_gy = sin_itn(s, index_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 = 45
-  g = named_grid((L, 1))
-
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s; map_dimension=3)
-
-  ψ_fx = poly_itn(s, index_map, [0.0, -1.0, 1.0]; dimension=1)
-  ψ_gy = sin_itn(s, index_map; k=Float64(pi), dimension=2)
-  ψ_hz = sin_itn(s, index_map; k=Float64(pi), dimension=3)
-  @assert dimension(ψ_fx) == dimension(ψ_gy) == dimension(ψ_hz) == 3
-
-  ψ_fxgyhz = ψ_fx * ψ_gy * ψ_hz
-
-  ∂_∂y = first_derivative_operator(s, index_map; dimension=2, cutoff=1e-10)
-
-  ∂_∂y_ψ_fxgyhz = operate([∂_∂y], ψ_fxgyhz; 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 ≈ (x^2 - x) * sin(pi * y) * sin(pi * z) atol = 1e-3
-
-        ∂_∂y_ψ_fxgyhz_xyz = real(calculate_fxyz(∂_∂y_ψ_fxgyhz, [x, y, z]))
-        @test ∂_∂y_ψ_fxgyhz_xyz ≈ pi * (x^2 - x) * cos(pi * y) * sin(pi * z) atol = 1e-3
-      end
-    end
-  end
-end
diff --git a/test/test_indexmaps.jl b/test/test_indexmaps.jl
index 0dc0ddf..3f62f46 100644
--- a/test/test_indexmaps.jl
+++ b/test/test_indexmaps.jl
@@ -11,32 +11,29 @@ using Dictionaries: Dictionary
   g = named_grid((L, L))
 
   s = continuous_siteinds(g)
-  index_map = IndexMap(s)
 
-  @test dimension(index_map) == 1
-  @test Set(inds(index_map)) == Set(inds(s))
+  @test dimension(s) == 1
 
   x = 0.625
-  ind_to_ind_value_map = calculate_ind_values(index_map, x)
+  ind_to_ind_value_map = calculate_ind_values(s, x)
   @test Set(keys(ind_to_ind_value_map)) == Set(inds(s))
-  @test calculate_x(index_map, ind_to_ind_value_map) == x
+  @test calculate_x(s, ind_to_ind_value_map) == x
 end
 
 @testset "test multi dimensional index map" begin
   L = 10
   g = named_grid((L, L))
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s, [[(i, j) for i in 1:L] for j in 1:L])
+  s = continuous_siteinds(g, [[(i, j) for i in 1:L] for j in 1:L])
 
   x, y = 0.5, 0.75
-  ind_to_ind_value_map = calculate_ind_values(index_map, [x, y], [1, 2])
-  xyz = calculate_xyz(index_map, ind_to_ind_value_map, [1, 2])
+  ind_to_ind_value_map = calculate_ind_values(s, [x, y], [1, 2])
+  xyz = calculate_xyz(s, ind_to_ind_value_map, [1, 2])
   @test first(xyz) == x
   @test last(xyz) == y
 
   xyzvals = [rand() for i in 1:L]
-  ind_to_ind_value_map = calculate_ind_values(index_map, xyzvals)
+  ind_to_ind_value_map = calculate_ind_values(s, xyzvals)
 
-  xyzvals_approx = calculate_xyz(index_map, ind_to_ind_value_map)
+  xyzvals_approx = calculate_xyz(s, ind_to_ind_value_map)
   xyzvals ≈ xyzvals_approx
 end
diff --git a/test/test_itensorfunction.jl b/test/test_itensorfunction.jl
index 3ae35b0..c35978a 100644
--- a/test/test_itensorfunction.jl
+++ b/test/test_itensorfunction.jl
@@ -4,7 +4,6 @@ using ITensorNumericalAnalysis
 using Graphs: SimpleGraph, uniform_tree
 using NamedGraphs: NamedGraph, named_grid, vertices, named_comb_tree, rename_vertices
 using ITensors: siteinds
-using ITensorNetworks: random_tensornetwork
 using Dictionaries: Dictionary
 using SplitApplyCombine: group
 using Random: seed!
@@ -16,7 +15,7 @@ using Distributions: Uniform
   g = named_grid((L, 1))
   s = continuous_siteinds(g)
 
-  ψ = random_tensornetwork(s; link_space=2)
+  ψ = random_itensornetwork(s; link_space=2)
 
   fψ = ITensorNetworkFunction(ψ)
 
@@ -35,10 +34,9 @@ end
     L = 3
     g = named_grid((L, L))
     s = continuous_siteinds(g)
-    index_map = IndexMap(s)
     c = 1.5
 
-    ψ_fx = const_itn(s, index_map; c)
+    ψ_fx = const_itn(s; c)
 
     x = 0.5
     ind_to_ind_value_map = calculate_ind_values(ψ_fx, x)
@@ -62,10 +60,9 @@ end
       k = 0.125
 
       s = continuous_siteinds(g)
-      index_map = IndexMap(s)
 
       x = 0.625
-      ψ_fx = net_func(s, index_map; k, a)
+      ψ_fx = net_func(s; k, a)
       fx_x = calculate_fx(ψ_fx, x)
       @test func(k * x + a) ≈ fx_x
     end
@@ -87,10 +84,9 @@ end
       b = 3
 
       s = continuous_siteinds(g; base=3)
-      index_map = IndexMap(s)
 
       x = (5.0 / 9.0)
-      ψ_fx = net_func(s, index_map; k, a)
+      ψ_fx = net_func(s; k, a)
       fx_x = calculate_fx(ψ_fx, x)
       @test func(k * x + a) ≈ fx_x
     end
@@ -104,10 +100,9 @@ end
     nterms = 50
 
     s = continuous_siteinds(g)
-    index_map = IndexMap(s)
 
     x = 0.625
-    ψ_fx = tanh_itn(s, index_map; k, a, nterms)
+    ψ_fx = tanh_itn(s; k, a, nterms)
     fx_x = calculate_fx(ψ_fx, x)
 
     @test tanh(k * x + a) ≈ fx_x
@@ -123,12 +118,11 @@ end
       g = NamedGraph(SimpleGraph(uniform_tree(L)))
       g = rename_vertices(g, Dict(zip(vertices(g), [(v, 1) for v in vertices(g)])))
       s = continuous_siteinds(g)
-      index_map = IndexMap(s)
 
       coeffs = [rand(Uniform(-2, 2)) for i in 1:(deg + 1)]
 
       x = 0.875
-      ψ_fx = poly_itn(s, index_map, coeffs)
+      ψ_fx = poly_itn(s, coeffs)
       fx_x = calculate_fx(ψ_fx, x)
 
       fx_exact = sum([coeffs[i] * (x^(i - 1)) for i in 1:(deg + 1)])
@@ -141,12 +135,11 @@ end
   #Constant function but represented in three dimension
   @testset "test const" begin
     g = named_grid((3, 3))
-    s = continuous_siteinds(g)
-    index_map = IndexMap(s; map_dimension=3)
+    s = continuous_siteinds(g; map_dimension=3)
 
     c = 1.5
 
-    ψ_fxyz = const_itn(s, index_map; c)
+    ψ_fxyz = const_itn(s; c)
 
     x, y, z = 0.5, 0.25, 0.0
 
@@ -164,8 +157,7 @@ end
   ]
   L = 10
   g = named_grid((L, 1))
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s; map_dimension=2)
+  s = continuous_siteinds(g; map_dimension=2)
   x, y = 0.625, 0.25
 
   for (name, net_func, func) in funcs
@@ -173,8 +165,8 @@ end
       a = 1.2
       k = 0.125
 
-      ψ_fx = net_func(s, index_map; k, a, dimension=1)
-      ψ_fy = net_func(s, index_map; k, a, dimension=2)
+      ψ_fx = net_func(s; k, a, dimension=1)
+      ψ_fy = net_func(s; k, a, dimension=2)
 
       ψ_fxy = ψ_fx + ψ_fy
       fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1, 2])
@@ -189,12 +181,11 @@ end
     a = 1.3
     k = 0.15
     nterms = 10
-    s = continuous_siteinds(g)
-    index_map = IndexMap(s; map_dimension=2)
+    s = continuous_siteinds(g; map_dimension=2)
 
     x, y = 0.625, 0.875
-    ψ_fx = tanh_itn(s, index_map; k, a, nterms, dimension=1)
-    ψ_fy = tanh_itn(s, index_map; k, a, nterms, dimension=2)
+    ψ_fx = tanh_itn(s; k, a, nterms, dimension=1)
+    ψ_fy = tanh_itn(s; k, a, nterms, dimension=2)
 
     ψ_fxy = ψ_fx + ψ_fy
     fxy_xy = calculate_fxyz(ψ_fxy, [x, y], [1, 2])
diff --git a/test/test_shift_operators.jl b/test/test_operators.jl
similarity index 56%
rename from test/test_shift_operators.jl
rename to test/test_operators.jl
index 87bc35a..9308b35 100644
--- a/test/test_shift_operators.jl
+++ b/test/test_operators.jl
@@ -8,34 +8,156 @@ using NamedGraphs: named_grid, named_comb_tree, NamedGraph, nv, vertices
 using ITensorNumericalAnalysis: itensornetwork
 using Dictionaries: Dictionary
 
-using ITensorNumericalAnalysis: backward_shift_op, forward_shift_op
+@testset "test differentiation in 1D on MPS" begin
+  g = named_grid((9, 1))
+  L = nv(g)
+  delta = (2.0)^(-Float64(L))
+  s = continuous_siteinds(g)
+  left_boundary, right_boundary = "Periodic", "Periodic"
+
+  f1 = first_derivative_operator(s; cutoff=1e-12, left_boundary, right_boundary)
+  f2 = second_derivative_operator(s; cutoff=1e-12, left_boundary, right_boundary)
+  f3 = third_derivative_operator(s; cutoff=1e-12, left_boundary, right_boundary)
+  f4 = fourth_derivative_operator(s; cutoff=1e-12, left_boundary, right_boundary)
+
+  ψ_fx = sin_itn(s; k=2.0 * Float64(pi))
+
+  ψ_f1x = operate(f1, ψ_fx; cutoff=1e-8)
+  ψ_f2x = operate(f2, ψ_fx; cutoff=1e-8)
+  ψ_f3x = operate(f3, ψ_fx; cutoff=1e-8)
+  ψ_f4x = operate(f4, ψ_fx; cutoff=1e-8)
+
+  xs = [0.0, 0.25, 0.625, 0.875, 1.0 - delta]
+  for x in xs
+    @test 1.0 + calculate_fx(ψ_fx, x) ≈ 1.0 + sin(2.0 * pi * x) rtol = 1e-3
+    @test 1.0 + calculate_fx(ψ_f1x, x) ≈ 1.0 + 2.0 * pi * cos(2.0 * pi * x) rtol = 1e-3
+    @test 1.0 + calculate_fx(ψ_f2x, x) ≈ 1.0 + -1.0 * (2.0 * pi)^2 * sin(2.0 * pi * x) rtol =
+      1e-3
+    @test 1.0 + calculate_fx(ψ_f3x, x) ≈ 1.0 + -1.0 * (2.0 * pi)^3 * cos(2.0 * pi * x) rtol =
+      1e-3
+    @test 1.0 + calculate_fx(ψ_f4x, x) ≈ 1.0 + 1.0 * (2.0 * pi)^4 * sin(2.0 * pi * x) rtol =
+      1e-3
+  end
+end
+
+@testset "test differentiation in 1D on tree" begin
+  g = named_comb_tree((4, 3))
+  L = nv(g)
+  delta = 2.0^(-Float64(L))
+  s = continuous_siteinds(g)
+
+  ∂_∂x = first_derivative_operator(s; cutoff=1e-10)
+
+  ψ_fx = sin_itn(s; k=Float64(pi))
+  ∂x_ψ_fx = operate(∂_∂x, ψ_fx; cutoff=1e-12)
+
+  xs = [delta, 0.125, 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
+
+@testset "test differentiation_operator_on_3D_function" begin
+  L = 45
+  g = named_grid((L, 1))
+
+  s = continuous_siteinds(g; map_dimension=3)
+
+  ψ_fx = poly_itn(s, [0.0, -1.0, 1.0]; dimension=1)
+  ψ_gy = sin_itn(s; k=Float64(pi), dimension=2)
+  ψ_hz = sin_itn(s; k=Float64(pi), dimension=3)
+  @assert dimension(ψ_fx) == dimension(ψ_gy) == dimension(ψ_hz) == 3
+
+  ψ_fxgyhz = ψ_fx * ψ_gy * ψ_hz
+
+  ∂_∂y = first_derivative_operator(s; dimension=2, cutoff=1e-10)
+
+  ∂_∂y_ψ_fxgyhz = operate([∂_∂y], ψ_fxgyhz; 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 ≈ (x^2 - x) * sin(pi * y) * sin(pi * z) atol = 1e-3
+
+        ∂_∂y_ψ_fxgyhz_xyz = real(calculate_fxyz(∂_∂y_ψ_fxgyhz, [x, y, z]))
+        @test ∂_∂y_ψ_fxgyhz_xyz ≈ pi * (x^2 - x) * cos(pi * y) * sin(pi * z) atol = 1e-3
+      end
+    end
+  end
+end
+
+@testset "test multiplication_operator_in_1D" begin
+  g = named_comb_tree((4, 3))
+  L = nv(g)
+  s = continuous_siteinds(g)
+
+  ψ_gx = sin_itn(s; k=0.5 * Float64(pi))
+  ψ_fx = cos_itn(s; 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-3
+  end
+end
+
+@testset "test multiplication_operator_in_2D" begin
+  L = 8
+  g = NamedGraph(SimpleGraph(uniform_tree(L)))
+
+  s = continuous_siteinds(g; map_dimension=2)
+
+  ψ_fx = cos_itn(s; k=0.25 * Float64(pi), dimension=1)
+  ψ_gy = sin_itn(s; 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 shift operators in 1D on Tree" begin
   g = named_comb_tree((2, 3))
   L = nv(g)
   delta = 2.0^(-1.0 * L)
   s = continuous_siteinds(g)
-  index_map = IndexMap(s)
   xs = [0.0, delta, 0.25, 0.5, 0.625, 0.875, 1.0 - delta]
-  ψ_fx = poly_itn(s, index_map, [1.0, 0.5, 0.25])
+  ψ_fx = poly_itn(s, [1.0, 0.5, 0.25])
 
   forward_shift_dirichlet = forward_shift_op(
-    s, index_map; boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_dirichlet = backward_shift_op(
-    s, index_map; boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_pbc = forward_shift_op(
-    s, index_map; boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_pbc = backward_shift_op(
-    s, index_map; boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_neumann = forward_shift_op(
-    s, index_map; boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_neumann = backward_shift_op(
-    s, index_map; boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
   )
 
   ψ_fx_pshift_dirichlet = operate(forward_shift_dirichlet, ψ_fx; cutoff=1e-12)
@@ -76,28 +198,27 @@ end
   L = nv(g)
   delta = 2.0^(-1.0 * L)
   s = continuous_siteinds(g)
-  index_map = IndexMap(s)
   xs = [0.0, delta, 0.25, 0.5, 0.625, 0.875, 1.0 - delta]
-  ψ_fx = poly_itn(s, index_map, [1.0, 0.5, 0.25])
+  ψ_fx = poly_itn(s, [1.0, 0.5, 0.25])
   n = 1
 
   forward_shift_dirichlet = forward_shift_op(
-    s, index_map; n, boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_dirichlet = backward_shift_op(
-    s, index_map; n, boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Dirichlet", truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_pbc = forward_shift_op(
-    s, index_map; n, boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_pbc = backward_shift_op(
-    s, index_map; n, boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Periodic", truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_neumann = forward_shift_op(
-    s, index_map; n, boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_neumann = backward_shift_op(
-    s, index_map; n, boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
+    s; n, boundary="Neumann", truncate_kwargs=(; cutoff=1e-10)
   )
 
   ψ_fx_pshift_dirichlet = operate(forward_shift_dirichlet, ψ_fx; cutoff=1e-12)
@@ -136,33 +257,32 @@ end
 
 @testset "test shift operators in 2D on Tree" begin
   g = named_comb_tree((3, 3))
-  s = continuous_siteinds(g)
-  index_map = IndexMap(s; map_dimension=2)
-  L = length(inds(index_map, 2))
+  s = continuous_siteinds(g; map_dimension=2)
+  L = length(dimension_inds(s, 2))
   delta = 2.0^(-1.0 * L)
   x = 0.5
   ys = [0.0, delta, 0.25, 0.5, 0.625, 0.875, 1.0 - delta]
-  ψ_fx = poly_itn(s, index_map, [1.0, 0.5, 0.25]; dimension=1)
-  ψ_fy = cos_itn(s, index_map; dimension=2)
+  ψ_fx = poly_itn(s, [1.0, 0.5, 0.25]; dimension=1)
+  ψ_fy = cos_itn(s; dimension=2)
   ψ_fxy = ψ_fx + ψ_fx
 
   forward_shift_dirichlet = forward_shift_op(
-    s, index_map; boundary="Dirichlet", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Dirichlet", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_dirichlet = backward_shift_op(
-    s, index_map; boundary="Dirichlet", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Dirichlet", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_pbc = forward_shift_op(
-    s, index_map; boundary="Periodic", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Periodic", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_pbc = backward_shift_op(
-    s, index_map; boundary="Periodic", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Periodic", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
   forward_shift_neumann = forward_shift_op(
-    s, index_map; boundary="Neumann", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Neumann", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
   backward_shift_neumann = backward_shift_op(
-    s, index_map; boundary="Neumann", dimension=2, truncate_kwargs=(; cutoff=1e-10)
+    s; boundary="Neumann", dimension=2, truncate_kwargs=(; cutoff=1e-10)
   )
 
   ψ_fxy_pshift_dirichlet = operate(forward_shift_dirichlet, ψ_fxy; cutoff=1e-12)