More general implementation

[SamuelBadr]

As discussed with Marc, here's a complete rewrite of the DiscretizedGrid type, tentatively called NewDiscretizedGrid.

Description & Interface

Everything not described in the docstring, which I just pasted below, (e.g. the conversion functions) works exactly as in the existing implementation.

NewDiscretizedGrid{D}

A discretized grid structure for D-dimensional grids with variable resolution,
supporting efficient conversion between quantics, grid indices, and original coordinates.
A NewDiscretizedGrid instance is intended to undergird a quantics tensor train
with a specific index structure, as defined in the indextable field.
For example, say indextable is [[(:a, 1), (:b, 2)], [(:a, 2)], [(:b, 1), (:a, 3)]],
then the corresponding tensor train has 3 tensor cores:

  a_1 b_2          a_2          b_1 a_3
   |   |            |            |   |
┌──┴───┴──┐    ┌────┴────┐    ┌──┴───┴──┐
│         │    │         │    │         │
│         │────│         │────│         │
│         │    │         │    │         │
└─────────┘    └─────────┘    └─────────┘

This object may be constructed with

julia> grid = NewDiscretizedGrid((:a, :b), [[(:a, 1), (:b, 2)], [(:a, 2)], [(:b, 1), (:a, 3)]])
NewDiscretizedGrid{2} with 8×4 = 32 grid points
├─ Variables: (a, b)
├─ Resolutions: (a: 3, b: 2)
├─ Domain: unit square [0, 1)²
├─ Grid spacing: (Δa = 0.125, Δb = 0.25)
└─ Tensor train: 3 sites (dimensions: 4-2-4)

and represents a 2^3 x 2^2 discretization of the unit square in the 2D plane (the x mark grid points):

   1.0  ┌───────────────────────────────┐
        │                               │
        │                               │
   0.75 x   x   x   x   x   x   x   x   │
        │                               │
        │                               │
b  0.5  x   x   x   x   x   x   x   x   │
        │                               │
        │                               │
   0.25 x   x   x   x   x   x   x   x   │
        │                               │
        │                               │
   0.0  x───x───x───x───x───x───x───x───┘
       0.0     0.25    0.5     0.75    1.0

                        a

If something other than a unit square is desired, lower_bound and upper_bound
can be specified. Also, bases different than the default base 2 can be used,
e.g. base=3 for a ternary grid.

In addition to the plain methods, there is a convenience layer for conversion
from the original coordinates

julia> origcoord_to_grididx(grid; a=0.5, b=0.25)
(5, 2)

julia> origcoord_to_quantics(grid; a=0.5, b=0.25)
3-element Vector{Int64}:
 4
 1
 1

and also from grid indices

julia> grididx_to_origcoord(grid; a=5, b=2)
(0.5, 0.25)

julia> grididx_to_quantics(grid; a=5, b=2)
3-element Vector{Int64}:
 4
 1
 1

For a simpler grid, we can just supply the resolution in each dimension:

julia> boring_grid = NewDiscretizedGrid((3, 9))
NewDiscretizedGrid{2} with 8×512 = 4096 grid points
├─ Resolutions: (1: 3, 2: 9)
├─ Domain: unit square [0, 1)²
├─ Grid spacing: (Δ1 = 0.125, Δ2 = 0.001953125)
└─ Tensor train: 12 sites (uniform dimension 2)

In this case, variable names are automatically generated as 1, 2, etc.

It is intended to be usable as a drop-in replacement for the existing DiscretizedGrid type except for some small differences in behavior, e.g. regarding upper_bound (which is shifted depending on includeendpoint so we don't need to store includeendpoint explicitly). It also doesn't subtype Grid{D} at the moment.

Performance

The new implementation includes a special path for the common base = 2 case using bit operations, so it is slightly faster.
Here's the results of a benchmark script, measuring the performance of quantics <-> origcoord conversion and comparing with DiscretizedGrid:

================================================================================
BENCHMARK SUMMARY
================================================================================
QUANTICS → ORIGCOORD
Config              Old (ns)     New (ns)      Speedup
--------------------------------------------------------------------------------
1D Small                87.3         45.4         1.92x
1D Medium              125.4         69.2         1.81x
2D Fused                84.7         61.2         1.38x
2D Interleaved         299.4         69.2         4.33x
2D Medium               96.5         75.7         1.27x
3D                      93.6         73.3         1.28x
Base-3 2D               82.9         77.3         1.07x
1D Large               137.5         80.9         1.70x

ORIGCOORD → QUANTICS
Config              Old (ns)     New (ns)      Speedup
--------------------------------------------------------------------------------
1D Small                75.4         47.9         1.57x
1D Medium              143.6         75.1         1.91x
2D Fused               106.1         62.3         1.70x
2D Interleaved         309.3         55.1         5.61x
2D Medium              165.3         64.2         2.57x
3D                     119.4         67.0         1.78x
Base-3 2D              101.6         83.3         1.22x
1D Large               163.7         61.9         2.64x

================================================================================
PERFORMANCE ANALYSIS
================================================================================
Average speedup (quantics→origcoord): 1.85x
Average speedup (origcoord→quantics): 2.38x

NewDiscretizedGrid faster in 7/8 forward tests
NewDiscretizedGrid faster in 8/8 reverse tests

Benchmark script

This needs Chairmarks.jl installed.

#!/usr/bin/env julia

using Pkg
Pkg.activate(@__DIR__)

using QuanticsGrids
using Chairmarks
using Random
using Printf

Random.seed!(42)

function create_equivalent_grids(::Val{D}, R, lower_bound, upper_bound; base=2, unfoldingscheme=:fused) where D
  old_grid = DiscretizedGrid{D}(R, lower_bound, upper_bound; base, unfoldingscheme)
  new_grid = NewDiscretizedGrid{D}(R; lower_bound, upper_bound, base, unfoldingscheme)
  return old_grid, new_grid
end

function generate_test_coordinates(grid, n_samples=1000)
  coord_samples = []
  D = typeof(grid) <: DiscretizedGrid ? length(grid.lower_bound) : ndims(grid)

  for _ in 1:n_samples
      coords = ntuple(D) do d
          lower = grid.lower_bound[d]
          upper = grid.upper_bound[d]
          lower + rand() * (upper - lower)
      end
      coords = D == 1 ? coords[1] : coords
      push!(coord_samples, coords)
  end
  return coord_samples
end

function generate_test_quantics(grid, n_samples=1000)
  quantics_samples = []

  for _ in 1:n_samples
      if typeof(grid) <: DiscretizedGrid
          D = length(grid.lower_bound)
          quantics = if grid.unfoldingscheme === :fused
              rand(1:grid.base^D, grid.R)
          else
              rand(1:grid.base, grid.R * D)
          end
      else
          local_dims = QuanticsGrids.localdimensions(grid)
          quantics = [rand(1:local_dim) for local_dim in local_dims]
      end
      push!(quantics_samples, quantics)
  end
  return quantics_samples
end

function benchmark_grid_config(config_name, D, R, lower_bound, upper_bound; base=2, unfoldingscheme=:fused, n_samples=1000)
  println("\n" * "="^60)
  println("$config_name")
  println("Dimensions: $D, R: $R, Base: $base, Unfolding: $unfoldingscheme")
  println("="^60)

  old_grid, new_grid = create_equivalent_grids(Val(D), R, lower_bound, upper_bound; base, unfoldingscheme)
  quantics_samples = generate_test_quantics(old_grid, n_samples)
  coord_samples = generate_test_coordinates(old_grid, n_samples)

  for i in 1:min(5, length(quantics_samples))
      quantics, coords = quantics_samples[i], coord_samples[i]

      old_result = quantics_to_origcoord(old_grid, quantics)
      new_result = quantics_to_origcoord(new_grid, quantics)

      if old_result isa Tuple && new_result isa Tuple
          all(isapprox.(old_result, new_result, rtol=1e-12)) || @warn "Results differ (quantics->origcoord)"
      elseif old_result isa Number && new_result isa Number
          isapprox(old_result, new_result, rtol=1e-12) || @warn "Results differ (quantics->origcoord)"
      end

      old_result_rev = origcoord_to_quantics(old_grid, coords)
      new_result_rev = origcoord_to_quantics(new_grid, coords)
      old_result_rev == new_result_rev || @warn "Results differ (origcoord->quantics)"
  end

  println("QUANTICS → ORIGCOORD")

  old_benchmark = @be begin
      for quantics in quantics_samples
          quantics_to_origcoord(old_grid, quantics)
      end
  end

  new_benchmark = @be begin
      for quantics in quantics_samples
          quantics_to_origcoord(new_grid, quantics)
      end
  end

  old_time = minimum([s.time for s in old_benchmark.samples]) / n_samples
  new_time = minimum([s.time for s in new_benchmark.samples]) / n_samples
  ratio = old_time / new_time

  @printf("  DiscretizedGrid:    %.1f ns/op\n", old_time * 1e9)
  @printf("  NewDiscretizedGrid: %.1f ns/op\n", new_time * 1e9)
  @printf("  Speedup: %.2fx %s\n", ratio, ratio > 1.1 ? "✅" : ratio > 0.9 ? "⚪" : "⚠️")

  println("\nORIGCOORD → QUANTICS")

  old_benchmark_rev = @be begin
      for coords in coord_samples
          origcoord_to_quantics(old_grid, coords)
      end
  end

  new_benchmark_rev = @be begin
      for coords in coord_samples
          origcoord_to_quantics(new_grid, coords)
      end
  end

  old_time_rev = minimum([s.time for s in old_benchmark_rev.samples]) / n_samples
  new_time_rev = minimum([s.time for s in new_benchmark_rev.samples]) / n_samples
  ratio_rev = old_time_rev / new_time_rev

  @printf("  DiscretizedGrid:    %.1f ns/op\n", old_time_rev * 1e9)
  @printf("  NewDiscretizedGrid: %.1f ns/op\n", new_time_rev * 1e9)
  @printf("  Speedup: %.2fx %s\n", ratio_rev, ratio_rev > 1.1 ? "✅" : ratio_rev > 0.9 ? "⚪" : "⚠️")

  return (
      old_time=old_time, new_time=new_time, ratio=ratio,
      old_time_rev=old_time_rev, new_time_rev=new_time_rev, ratio_rev=ratio_rev
  )
end
function main()
  println("QuanticsGrids Performance Benchmark")
  println("Julia $(VERSION), $(Threads.nthreads()) threads")

  configs = [
      ("1D Small", 1, 8, 0.0, 1.0, 2, :fused, 10000),
      ("1D Medium", 1, 16, -5.0, 5.0, 2, :fused, 5000),
      ("2D Fused", 2, 6, (0.0, 0.0), (1.0, 1.0), 2, :fused, 5000),
      ("2D Interleaved", 2, 6, (0.0, 0.0), (1.0, 1.0), 2, :interleaved, 5000),
      ("2D Medium", 2, 10, (-1.0, -1.0), (1.0, 1.0), 2, :fused, 2000),
      ("3D", 3, 5, (0.0, 0.0, 0.0), (1.0, 1.0, 1.0), 2, :fused, 1000),
      ("Base-3 2D", 2, 6, (0.0, 0.0), (1.0, 1.0), 3, :fused, 2000),
      ("1D Large", 1, 20, 0.0, 1.0, 2, :fused, 1000)
  ]

  results = []
  for (name, D, R, lower, upper, base, scheme, n_samples) in configs
      result = benchmark_grid_config(name, D, R, lower, upper;
          base=base, unfoldingscheme=scheme, n_samples=n_samples)
      push!(results, (name, result))
  end

  print_summary(results)
  return results
end

function print_summary(results)
  println("\n" * "="^80)
  println("BENCHMARK SUMMARY")
  println("="^80)

  println("QUANTICS → ORIGCOORD")
  println(@sprintf("%-15s %12s %12s %12s", "Config", "Old (ns)", "New (ns)", "Speedup"))
  println("-"^80)

  for (name, result) in results
      @printf("%-15s %12.1f %12.1f %12.2fx\n",
          name, result.old_time * 1e9, result.new_time * 1e9, result.ratio)
  end

  println("\nORIGCOORD → QUANTICS")
  println(@sprintf("%-15s %12s %12s %12s", "Config", "Old (ns)", "New (ns)", "Speedup"))
  println("-"^80)

  for (name, result) in results
      @printf("%-15s %12.1f %12.1f %12.2fx\n",
          name, result.old_time_rev * 1e9, result.new_time_rev * 1e9, result.ratio_rev)
  end

  ratios = [r[2].ratio for r in results]
  ratios_rev = [r[2].ratio_rev for r in results]

  println("\n" * "="^80)
  println("PERFORMANCE ANALYSIS")
  println("="^80)

  @printf("Average speedup (quantics→origcoord): %.2fx\n", sum(ratios) / length(ratios))
  @printf("Average speedup (origcoord→quantics): %.2fx\n", sum(ratios_rev) / length(ratios_rev))

  faster_count = count(r -> r[2].ratio > 1.1, results)
  faster_count_rev = count(r -> r[2].ratio_rev > 1.1, results)

  println("\nNewDiscretizedGrid faster in $faster_count/$(length(results)) forward tests")
  println("NewDiscretizedGrid faster in $faster_count_rev/$(length(results)) reverse tests")

  overall_improvement = (sum(ratios) + sum(ratios_rev)) / (2 * length(results))
  if overall_improvement > 1.1
      println("\n🎉 NewDiscretizedGrid shows significant performance improvements!")
  elseif overall_improvement > 0.9
      println("\n✅ NewDiscretizedGrid maintains equivalent performance")
  else
      println("\n⚠️  NewDiscretizedGrid shows performance regression")
  end
end

main()

[SamuelBadr]

In my view, this is now ready. I'm happy about any and all feedback/comments!

There's an accompanying PR to update QuanticsTCI.jl at tensor4all/QuanticsTCI.jl#14.

[rittermarc]

Nice work!

One question: Wouldn't it be better to infer the variablenames::NTuple{D,Symbol} argument in the constructor from the index table? Or is there a reason against that?

I have some small cleanup I'll propose in a separate commit

[shinaoka]

We are traveling now. Let us take some time.

[SamuelBadr]

Nice work!

Thanks!

One question: Wouldn't it be better to infer the variablenames::NTuple{D,Symbol} argument in the constructor from the index table? Or is there a reason against that?

The problem then is that you don't know the correct order of the dimensions, which is fine if you only call functions specifying the variable names (like origcoord_to_quantics(grid; x = 0.2, y = 0.0)) but not when you want to call stuff like origcoords_to_quantics(grid; ::NTuple), I think.

I have some small cleanup I'll propose in a separate commit

Perfect, looking forward to it!

[HironeIshida]

Thank you for the useful feature!

I tried this branch, and the tests passed locally.
Currently, I’m collaborating with researchers in the spintronics field, and this kind of discretization with different values of R is exactly what we need. I’m planning to apply the new DiscretizedGrid to this problem.

I noticed that the documentation for DiscretizedGrid{D} is listed under API Reference.
Is there any plan to integrate it into Home and revise the description into a more general overview?

[rittermarc]

OK, I finally got around to looking at everything, and it's great!

I think one place where this can be improved is the handling of invalid variable / index tables, see

QuanticsGrids.jl/test/newgrid/inherentdiscretegrid.jl

Line 726 in 6d7f303

invalid_indextable_unknown = [

Currently only one of these actually throws, with a non-specific AssertionError.

I'm also a bit afraid that it may be possible to bypass some of these parameter consistency checks, as they are only defined in some constructors, not others. In my opinion, the correct place to define them is in the place where the argument is actually used. For example, the consistency of the index table should probably be checked in _build_lookup_table, e.g. like this

QuanticsGrids.jl/src/inherent_discrete_grid.jl

Line 55 in 6d7f303

if isnothing(var_idx)

Other than that, the code is pretty good, and we're close to convergence!

PS: I think for the file structure you can just put your tests in test/, since your implementation will be the default, so there's no need for a subfolder.