Improved Grid_points and rand_p

src/ITensorNumericalAnalysis.jl

@@ -28,7 +28,8 @@ export AbstractIndexMap,
   calculate_ind_values,
   dimension,
   dimensions,
-  grid_points
+  grid_points,
+  rand_p
 export IndsNetworkMap,
   continuous_siteinds,
   complex_continuous_siteinds,

src/IndexMaps/abstractindexmap.jl

@@ -2,7 +2,7 @@ using Base: Base
 using Dictionaries: Dictionary, set!
 using ITensors: ITensors, Index, dim
 using ITensorNetworks: IndsNetwork, vertex_data
-
+using Random: AbstractRNG, default_rng
 abstract type AbstractIndexMap{VB,VD} end
 
 #These functions need to be defined on the concrete type for implementation
@@ -153,3 +153,9 @@ function grid_points(imap::AbstractIndexMap, d::Int)
   L = length(dimension_inds(imap, d))
   return grid_points(imap, base^L, d)
 end
+
+function rand_p(rng::AbstractRNG, imap::AbstractIndexMap)
+  return [rand_p(rng, imap, i) for i in 1:dimension(imap)]
+end
+
+rand_p(imap::AbstractIndexMap, args...) = rand_p(default_rng(), imap, args...)

src/IndexMaps/complexindexmap.jl

@@ -2,6 +2,7 @@ using Base: Base
 using Dictionaries: Dictionaries, Dictionary, set!
 using ITensors: ITensors, Index, dim, hastags
 using ITensorNetworks: IndsNetwork, vertex_data
+using Random: AbstractRNG
 
 struct ComplexIndexMap{VB,VD,VR} <: AbstractIndexMap{VB,VD}
   index_digit::VB
@@ -142,3 +143,15 @@ function grid_points(imap::ComplexIndexMap, N::Int, d::Int)
   ]
   return filter(x -> real(x) < 1 && imag(x) < 1, grid_points)
 end
+
+""" 
+  Picks a random grid point from `imap` given a dimension
+"""
+function rand_p(rng::AbstractRNG, imap::ComplexIndexMap, d::Integer)
+  dims = dim.(dimension_inds(imap, d))
+  @assert all(y -> y == first(dims), dims)
+  base = float(first(dims))
+  Lre, Lim = length(real_indices(imap, d)), length(imag_indices(imap, d))
+  return rand(rng, 0:(big(base)^Lre - 1)) / big(base)^Lre +
+         im * rand(rng, 0:(big(base)^Lim - 1)) / big(base)^Lim
+end

src/IndexMaps/realindexmap.jl

@@ -2,6 +2,7 @@ using Base: Base
 using Dictionaries: Dictionaries, Dictionary, set!
 using ITensors: ITensors, Index, dim
 using ITensorNetworks: IndsNetwork, vertex_data
+using Random: AbstractRNG
 
 struct RealIndexMap{VB,VD} <: AbstractIndexMap{VB,VD}
   index_digit::VB
@@ -75,12 +76,98 @@ function calculate_ind_values(imap::RealIndexMap, xs::Vector, dims::Vector{Int})
   return ind_to_ind_value_map
 end
 
-function grid_points(imap::RealIndexMap, N::Int, d::Int)
+""" 
+    grid_points(imap, N, d, span; exact_grid=true, enforced=[])
+
+  Gives `N` grid points from a given dimension of `imap` within a specified range.
+
+  # Arguments
+  - `imap::RealIndexMap`: An IndexMap specifying the structure of the TN being used
+  - `N::Int`: The number of grid points requested.
+  - `d::Int` The index of the dimension of `imap` requested
+  - `span::AbstractVector{<:Number}`: A two element number vector [a,b] with 0≤a<b≤1. The right endpoint of this span is not included as a gridpoint in the output.
+  - `exact_grid::Bool`: Flag specifying whether the function should give exact grid points 
+  (note: using `exact_grid=true` may cause less than `N` grid points to be returned)
+  - `enforced::AbstractVector{<:Number}`: A list of points that we want to enforce to show up in the grid.
+"""
+function grid_points(
+  imap::RealIndexMap,
+  N::Int,
+  d::Int;
+  span::AbstractVector{<:Number}=[0, 1],
+  exact_grid::Bool=true,
+  enforced=[],
+)
+  if length(span) != 2 || span[1] >= span[2] || span[1] < 0 || span[2] > 1
+    throw(
+      "expected a two-element vector [a,b] with 0≤a<b≤1 as input span, instead found $span"
+    )
+  end
+  dims = dim.(dimension_inds(imap, d))
+  @assert all(y -> y == first(dims), dims)
+  base = first(dims)
+  L = length(dimension_inds(imap, d))
+
+  # generate grid_points within the span (exclusive of right endpoint)
+  grid_points = collect(range(span[1], span[2] - (span[2] - span[1]) / N; length=N))
+  if exact_grid
+    for (i, point) in enumerate(grid_points)
+      grid_points[i] = round_to_nearest_exact_point(point, L)
+    end
+    # define lambda function
+    # round_near = point -> round_to_nearest_exact_point(point, L)
+    # grid_points = map(round_near, grid_points)
+    grid_points = unique(grid_points)
+  end
+
+  # add enforced points
+  if !isempty(enforced)
+    grid_points = sort(unique(vcat(grid_points, enforced_points)))
+  end
+
+  return grid_points
+end
+
+function round_to_nearest_exact_point(point::Number, L::Int)
+  if point < 0 || point >= 1
+    throw("Input point must be between 0 and 1")
+  end
+  return round(point * 2.0^L) / 2.0^L
+end
+
+function grid_points(imap::RealIndexMap, d::Int; kwargs...)
+  dims = dim.(dimension_inds(imap, d))
+  @assert all(y -> y == first(dims), dims)
+  base = dims[d]
+  L = length(dimension_inds(imap, d))
+  return grid_points(imap, base^L, d; kwargs...)
+end
+
+grid_points(imap::RealIndexMap; kwargs...) = grid_points(imap, 1; kwargs...)
+
+#multi-dimensional grid_points
+function grid_points(imap::RealIndexMap, Ns::Vector{Int}, dims::Vector{Int}; kwargs...)
+  if length(Ns) != length(dims)
+    throw("length of Ns and dims do not match!")
+  end
+  coords = [grid_points(imap, pair[1], pair[2]; kwargs...) for pair in zip(Ns, dims)]
+  gp = Base.Iterators.product(coords...)
+  return [collect(point) for point in gp]
+end
+
+function grid_points(imap::RealIndexMap, dims::Vector{Int}; kwargs...)
+  coords = [grid_points(imap, d; kwargs...) for d in dims]
+  gp = Base.Iterators.product(coords...)
+  return [collect(point) for point in gp]
+end
+
+""" 
+  Picks a random grid point from `imap` given a dimension
+"""
+function rand_p(rng::AbstractRNG, imap::RealIndexMap, d::Integer)
   dims = dim.(dimension_inds(imap, d))
   @assert all(y -> y == first(dims), dims)
-  base = float(first(dims))
+  base = dims[d]
   L = length(dimension_inds(imap, d))
-  a = round(base^L / N)
-  grid_points = [i * (a / base^L) for i in 0:(N + 1)]
-  return filter(x -> x < 1, grid_points)
+  return rand(rng, 0:(big(base)^L - 1)) / big(base)^L
 end

src/indsnetworkmap.jl

@@ -5,6 +5,7 @@ using NamedGraphs.GraphsExtensions: rem_vertex
 using ITensors: ITensors
 using ITensorNetworks:
   ITensorNetworks, AbstractIndsNetwork, IndsNetwork, data_graph, underlying_graph
+using Random: AbstractRNG
 
 struct IndsNetworkMap{V,I,IN<:IndsNetwork{V,I},IM<:AbstractIndexMap} <:
        AbstractIndsNetwork{V,I}
@@ -95,6 +96,7 @@ for f in [
   :calculate_ind_values,
   :calculate_p,
   :grid_points,
+  :rand_p,
   :index_value_to_scalar,
   :index_values_to_scalars,
 ]
@@ -104,6 +106,14 @@ for f in [
     end
   end
 end
+# Forward RNG functionality
+for f in [:rand_p]
+  @eval begin
+    function $f(rng::AbstractRNG, inm::IndsNetworkMap, args...; kwargs...)
+      return $f(rng, indexmap(inm), args...; kwargs...)
+    end
+  end
+end
 
 #Functions on indsnetwork that don't seem to autoforward
 function ITensors.inds(inm::IndsNetworkMap, args...; kwargs...)

test/test_indexmaps.jl

@@ -2,7 +2,7 @@ using Test
 using ITensorNumericalAnalysis
 
 using NamedGraphs: vertices
-using NamedGraphs.NamedGraphGenerators: named_grid
+using NamedGraphs.NamedGraphGenerators: named_grid, named_comb_tree
 using NamedGraphs.GraphsExtensions: is_tree
 using ITensors: siteinds, inds
 using Dictionaries: Dictionary
@@ -70,3 +70,90 @@ Random.seed!(1234)
     xyzvals ≈ xyzvals_approx
   end
 end
+
+@testset "grid_points tests" begin
+  @testset "test grid_points irregular span" begin
+    L = 16
+    base = 2
+    g = named_comb_tree((2, L ÷ 2))
+    s = continuous_siteinds(g; map_dimension=2)
+
+    #first set -- irregular span
+    N = 125
+    a = 0.12
+    b = 0.95
+
+    test_gridpoints = grid_points(s, N, 1; span=[a, b])
+    @test length(test_gridpoints) == N
+    @test test_gridpoints[1] >= a
+    @test test_gridpoints[end] < b
+  end
+
+  @testset "test grid_points spacing" begin
+    L = 16
+    g = named_comb_tree((2, L ÷ 2))
+    s = continuous_siteinds(g; map_dimension=2)
+
+    #second set
+    N = 32
+    a = 0.25
+    b = 0.5
+    test2 = grid_points(s, N, 1; span=[a, b])
+    @test length(test2) == N
+    @test test2[1] >= a
+    @test test2[end] < b
+    internal = (test2[2] - test2[1])
+    left = (test2[1] - a)
+    right = (b - test2[end])
+    @test internal >= left && internal >= right
+  end
+
+  @testset "test grid_points spacing 2" begin
+    L = 16
+    g = named_comb_tree((2, L ÷ 2))
+    s = continuous_siteinds(g; map_dimension=2)
+    a = 0.25
+    b = 0.5
+
+    #third set
+    test3 = grid_points(s, 1)
+    @test length(test3) == 256
+    @test test3[1] >= 0
+    @test test3[end] < 1
+    internal = (test3[2] - test3[1])
+    left = (test3[1] - a)
+    right = (b - test3[end])
+    @test internal >= left && internal >= right
+  end
+
+  @testset "test grid_points high precision" begin
+    #fourth set -- very large L
+    L = 140
+    base = 2
+    g = named_comb_tree((2, L ÷ 2))
+    s = continuous_siteinds(g; map_dimension=2)
+    n_grid = 16
+    @test length(grid_points(s, n_grid, 1)) == n_grid
+  end
+end
+
+@testset "test rand_p" begin
+  #test the rand_p() function to see if it succeeds on large L
+  L = 140
+  g = named_comb_tree((2, L ÷ 2))
+  s = continuous_siteinds(g; map_dimension=2)
+  ψ = cos_itn(s; dim=1) * cos_itn(s; dim=2)
+
+  #test the use of seedeed randomness
+  rng = Random.Xoshiro(42)
+  rand_gridpoint1 = rand_p(rng, s)
+  rand_gridpoint2 = rand_p(rng, s, 1)
+  x = real(ITensorNumericalAnalysis.evaluate(ψ, rand_gridpoint1))
+  @test x >= -1 && x <= 1 # check to make sure ψ can be evaluated at these points
+
+  #test the use of default rng
+  default_rng_gridpoint1 = rand_p(s)
+  default_rng_gridpoint2 = rand_p(s, 1)
+  y = real(ITensorNumericalAnalysis.evaluate(ψ, default_rng_gridpoint1))
+  @test y >= -1 && y <= 1 # check to make sure ψ can be evaluated at these points
+end