site[s] convenient constructors for unit ranges/axes (#21)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumOperatorDefinitions"
 uuid = "826dd319-6fd5-459a-a990-3a4f214664bf"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.5"
+version = "0.1.6"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

ext/QuantumOperatorDefinitionsITensorBaseExt/QuantumOperatorDefinitionsITensorBaseExt.jl

@@ -23,7 +23,12 @@ function QuantumOperatorDefinitions.SiteType(r::Index)
 end
 
 function (rangetype::Type{<:Index})(t::SiteType)
-  return settag(rangetype(AbstractUnitRange(t)), "sitetype", String(name(t)))
+  i = rangetype(AbstractUnitRange(t))
+  i = settag(i, "sitetype", String(name(t)))
+  if haskey(t, :site)
+    i = settag(i, "site", string(t.site))
+  end
+  return i
 end
 
 # TODO: Define in terms of `OpName` directly, and define a generic

src/sitetype.jl

@@ -95,3 +95,26 @@ to_dim(d::Integer) = d
 # TODO: Decide on this.
 # TODO: Move to `sitetype.jl`.
 default_sitetype() = SiteType"Qubit"()
+
+# TODO: Do we want to define this?
+# (t::SiteType)() = AbstractUnitRange(t)
+
+function site(rangetype::Type{<:AbstractUnitRange}, name::String; kwargs...)
+  return rangetype(SiteType(name; kwargs...))
+end
+site(name::String; kwargs...) = site(AbstractUnitRange, name; kwargs...)
+
+function sites(rangetype::Type{<:AbstractUnitRange}, name::String, positions; kwargs...)
+  return map(position -> site(rangetype, name; site=position, kwargs...), positions)
+end
+function sites(
+  rangetype::Type{<:AbstractUnitRange}, name::String, npositions::Integer; kwargs...
+)
+  return sites(rangetype, name, Base.OneTo(npositions); kwargs...)
+end
+function sites(name::String, positions; kwargs...)
+  return sites(AbstractUnitRange, name, positions; kwargs...)
+end
+function sites(name::String, npositions::Integer; kwargs...)
+  return sites(AbstractUnitRange, name, Base.OneTo(npositions); kwargs...)
+end

test/test_basics.jl

@@ -1,4 +1,5 @@
-using QuantumOperatorDefinitions: OpName, SiteType, ⊗, expand, op, opexpr, state, nsites
+using QuantumOperatorDefinitions:
+  OpName, SiteType, ⊗, expand, nsites, op, opexpr, site, sites, state
 using LinearAlgebra: Diagonal, I
 using Test: @test, @testset
 
@@ -217,6 +218,35 @@ const elts = (real_elts..., complex_elts...)
     @test op("X", 2) == op("X", Base.OneTo(2))
     @test op("X", 3) == op("X", Base.OneTo(3))
   end
+  @testset "site, sites" begin
+    s = site("Qudit"; dim=3)
+    @test s === Base.OneTo(3)
+
+    s = site(AbstractUnitRange{Int32}, "Qudit"; dim=3)
+    @test s === Base.OneTo(Int32(3))
+
+    for ss in (sites("Qudit", 4; dim=3), sites("Qudit", 2:5; dim=3))
+      @test length(ss) == 4
+      @test ss == map(Returns(Base.OneTo(3)), 1:4)
+      for s in ss
+        @test s === Base.OneTo(3)
+      end
+    end
+
+    for ss in (
+      sites(AbstractUnitRange{Int32}, "Qudit", 4; dim=3),
+      sites(AbstractUnitRange{Int32}, "Qudit", 2:5; dim=3),
+    )
+      @test length(ss) == 4
+      @test ss == map(Returns(Base.OneTo(3)), 1:4)
+      for s in ss
+        @test s === Base.OneTo(Int32(3))
+      end
+    end
+
+    ss = sites("Qudit", (1, 2, 3); dim=3)
+    @test ss === (Base.OneTo(3), Base.OneTo(3), Base.OneTo(3))
+  end
   @testset "Electron/tJ" begin
     for (ns, x) in (
       (("0", "Emp"), [1, 0, 0, 0]),

test/test_itensorbaseext.jl

@@ -1,5 +1,6 @@
-using ITensorBase: ITensor, Index, gettag, prime, settag
-using QuantumOperatorDefinitions: OpName, SiteType, StateName, op, state
+using ITensorBase: ITensor, Index, gettag, hastag, prime, settag
+using NamedDimsArrays: dename
+using QuantumOperatorDefinitions: OpName, SiteType, StateName, op, site, sites, state
 using Test: @test, @testset
 
 @testset "ITensorBaseExt" begin
@@ -61,4 +62,30 @@ using Test: @test, @testset
   @test a[i1′[2], i2′, i1[1], i2] == zeros(i2′, i2)
   @test a[i1′[1], i2′, i1[2], i2] == zeros(i2′, i2)
   @test a[i1′[2], i2′, i1[2], i2] == op("X", i2)
+
+  i = site(Index, "Qudit"; dim=3)
+  @test dename(i) == Base.OneTo(3)
+  @test gettag(i, "sitetype") == "Qudit"
+  @test !hastag(i, "site")
+
+  for is in (
+    sites(Index, "Qudit", 3; dim=3),
+    sites(Index, "Qudit", 1:3; dim=3),
+    sites(Index, "Qudit", (1, 2, 3); dim=3),
+  )
+    @test length(is) == 3
+    for (pos, i) in pairs(is)
+      @test dename(i) == Base.OneTo(3)
+      @test gettag(i, "sitetype") == "Qudit"
+      @test gettag(i, "site") == "$pos"
+    end
+  end
+
+  is = sites(Index, "Qudit", 2:4; dim=3)
+  @test dename(is[1]) == Base.OneTo(3)
+  @test gettag(is[1], "site") == "2"
+  @test dename(is[2]) == Base.OneTo(3)
+  @test gettag(is[2], "site") == "3"
+  @test dename(is[3]) == Base.OneTo(3)
+  @test gettag(is[3], "site") == "4"
 end