Allow custom axes when constructing operators, SymmetrySectors.jl and ITensorBase.jl extensions (#17)

Author: mtfishman

Project.toml

@@ -1,20 +1,31 @@
 name = "QuantumOperatorDefinitions"
 uuid = "826dd319-6fd5-459a-a990-3a4f214664bf"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [weakdeps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
+LabelledNumbers = "f856a3a6-4152-4ec4-b2a7-02c1a55d7993"
 ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
+NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
+SymmetrySectors = "f8a8ad64-adbc-4fce-92f7-ffe2bb36a86e"
 
 [extensions]
-QuantumOperatorDefinitionsITensorBaseExt = "ITensorBase"
+QuantumOperatorDefinitionsITensorBaseExt = ["ITensorBase", "NamedDimsArrays"]
+QuantumOperatorDefinitionsSymmetrySectorsExt = ["BlockArrays", "GradedUnitRanges", "LabelledNumbers", "SymmetrySectors"]
 
 [compat]
+BlockArrays = "1.3.0"
+GradedUnitRanges = "0.1.2"
 ITensorBase = "0.1.10"
+LabelledNumbers = "0.1.0"
 LinearAlgebra = "1.10"
+NamedDimsArrays = "0.4.0"
 Random = "1.10"
+SymmetrySectors = "0.1.3"
 julia = "1.10"

README.md

@@ -33,7 +33,6 @@ julia> Pkg.add("QuantumOperatorDefinitions")
 
 ````julia
 using QuantumOperatorDefinitions: OpName, SiteType, StateName, ⊗, controlled, op, state
-using LinearAlgebra: Diagonal
 using SparseArrays: SparseMatrixCSC, SparseVector
 using Test: @test
 
@@ -64,8 +63,6 @@ using Test: @test
 @test op("Y") == [0 -im; im 0]
 @test op("Z") == [1 0; 0 -1]
 
-@test op("Z") isa Diagonal
-
 @test op(Float32, "X") == [0 1; 1 0]
 @test eltype(op(Float32, "X")) === Float32
 @test op(SparseMatrixCSC, "X") == [0 1; 1 0]

examples/README.jl

@@ -38,7 +38,6 @@ julia> Pkg.add("QuantumOperatorDefinitions")
 # ## Examples
 
 using QuantumOperatorDefinitions: OpName, SiteType, StateName, ⊗, controlled, op, state
-using LinearAlgebra: Diagonal
 using SparseArrays: SparseMatrixCSC, SparseVector
 using Test: @test
 
@@ -69,8 +68,6 @@ using Test: @test
 @test op("Y") == [0 -im; im 0]
 @test op("Z") == [1 0; 0 -1]
 
-@test op("Z") isa Diagonal
-
 @test op(Float32, "X") == [0 1; 1 0]
 @test eltype(op(Float32, "X")) === Float32
 @test op(SparseMatrixCSC, "X") == [0 1; 1 0]

ext/QuantumOperatorDefinitionsITensorBaseExt/QuantumOperatorDefinitionsITensorBaseExt.jl

@@ -1,24 +1,58 @@
 module QuantumOperatorDefinitionsITensorBaseExt
 
-using ITensorBase: ITensor, Index, dag, gettag, prime
+using ITensorBase: ITensorBase, ITensor, Index, dag, gettag, prime, settag
+using NamedDimsArrays: dename
 using QuantumOperatorDefinitions:
-  QuantumOperatorDefinitions, OpName, SiteType, StateName, has_fermion_string
+  QuantumOperatorDefinitions,
+  @OpName_str,
+  OpName,
+  SiteType,
+  StateName,
+  has_fermion_string,
+  name
 
 function QuantumOperatorDefinitions.SiteType(r::Index)
-  return SiteType(gettag(r, "sitetype", "Qudit"); dim=Int(length(r)))
+  # We pass the axis of the (unnamed) Index because
+  # the Index may have originated from a slice, in which
+  # case the start may not be 1 (for NonContiguousIndex,
+  # which we need to add support for, it may not even
+  # be a unit range).
+  return SiteType(
+    gettag(r, "sitetype", "Qudit"); dim=Int.(length(r)), range=only(axes(dename(r)))
+  )
 end
 
+function (rangetype::Type{<:Index})(t::SiteType)
+  return settag(rangetype(AbstractUnitRange(t)), "sitetype", String(name(t)))
+end
+
+# TODO: Define in terms of `OpName` directly, and define a generic
+# forwarding method `has_fermion_string(n::String, t) = has_fermion_string(OpName(n), t)`.
 function QuantumOperatorDefinitions.has_fermion_string(n::String, r::Index)
   return has_fermion_string(OpName(n), SiteType(r))
 end
 
-function Base.AbstractArray(n::OpName, r::Index)
-  # TODO: Define this with mapped dimnames.
-  return ITensor(AbstractArray(n, SiteType(r)), (prime(r), dag(r)))
+function Base.axes(::OpName, domain::Tuple{Vararg{Index}})
+  return (prime.(domain)..., dag.(domain)...)
 end
+## function Base.axes(::OpName"SWAP", domain::Tuple{Vararg{Index}})
+##   return (prime.(reverse(domain))..., dag.(domain)...)
+## end
 
-function Base.AbstractArray(n::StateName, r::Index)
-  return ITensor(AbstractArray(n, SiteType(r)), (r,))
+# Fix ambiguity error with generic `AbstractArray` version.
+function ITensorBase.ITensor(n::Union{OpName,StateName}, domain::Index...)
+  return ITensor(n, domain)
+end
+# Fix ambiguity error with generic `AbstractArray` version.
+function ITensorBase.ITensor(n::Union{OpName,StateName}, domain::Tuple{Vararg{Index}})
+  return ITensor(AbstractArray(n, domain), axes(n, domain))
+end
+function (arrtype::Type{<:AbstractArray})(
+  n::Union{OpName,StateName}, domain::Tuple{Vararg{Index}}
+)
+  # Convert to `SiteType` in case the Index specifies a `"sitetype"` tag.
+  # TODO: Try to build this into the generic codepath.
+  return ITensor(arrtype(n, SiteType.(domain)), axes(n, domain))
 end
 
 end

ext/QuantumOperatorDefinitionsSymmetrySectorsExt/QuantumOperatorDefinitionsSymmetrySectorsExt.jl

@@ -0,0 +1,96 @@
+module QuantumOperatorDefinitionsSymmetrySectorsExt
+
+using BlockArrays: blocklasts, blocklengths
+using GradedUnitRanges: AbstractGradedUnitRange, GradedOneTo, gradedrange
+using LabelledNumbers: label, labelled, unlabel
+using QuantumOperatorDefinitions:
+  QuantumOperatorDefinitions,
+  @SiteType_str,
+  @GradingType_str,
+  SiteType,
+  GradingType,
+  OpName,
+  name
+using SymmetrySectors: ×, dual, SectorProduct, U1, Z
+
+function Base.axes(::OpName, domain::Tuple{Vararg{AbstractGradedUnitRange}})
+  return (domain..., dual.(domain)...)
+end
+
+sortedunion(a, b) = sort(union(a, b))
+function QuantumOperatorDefinitions.combine_axes(a1::GradedOneTo, a2::GradedOneTo)
+  return gradedrange(
+    map(blocklengths(a1), blocklengths(a2)) do s1, s2
+      l1 = unlabel(s1)
+      l2 = unlabel(s2)
+      @assert l1 == l2
+      labelled(l1, label(s1) × label(s2))
+    end,
+  )
+end
+QuantumOperatorDefinitions.combine_axes(a::GradedOneTo, b::Base.OneTo) = a
+QuantumOperatorDefinitions.combine_axes(a::Base.OneTo, b::GradedOneTo) = b
+
+function Base.AbstractUnitRange(::GradingType"N", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; N=U1(i - 1))) => 1, 1:length(t)))
+end
+function Base.AbstractUnitRange(::GradingType"Sz", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; Sz=U1(i - 1))) => 1, 1:length(t)))
+end
+function Base.AbstractUnitRange(::GradingType"Sz↑", t::SiteType)
+  return AbstractUnitRange(GradingType"Sz"(), t)
+end
+function Base.AbstractUnitRange(::GradingType"Sz↓", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; Sz=U1(-(i - 1)))) => 1, 1:length(t)))
+end
+
+function sector(gradingtype::GradingType, sec)
+  sectorname = Symbol(get(gradingtype, :name, name(gradingtype)))
+  return SectorProduct(NamedTuple{(sectorname,)}((sec,)))
+end
+
+function Base.AbstractUnitRange(s::GradingType"Nf", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+# TODO: Write in terms of `GradingType"Nf"` definition.
+function Base.AbstractUnitRange(s::GradingType"NfParity", t::SiteType"Fermion")
+  return gradedrange([sector(s, Z{2}(0)) => 1, sector(s, Z{2}(1)) => 1])
+end
+function Base.AbstractUnitRange(s::GradingType"Sz", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+function Base.AbstractUnitRange(s::GradingType"Sz↑", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+function Base.AbstractUnitRange(s::GradingType"Sz↓", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(-1)) => 1])
+end
+
+# TODO: Write in terms of `SiteType"Fermion"` definitions.
+function Base.AbstractUnitRange(s::GradingType"Nf", t::SiteType"Electron")
+  return gradedrange([
+    sector(s, U1(0)) => 1,
+    sector(s, U1(1)) => 1,
+    sector(s, U1(1)) => 1,
+    sector(s, U1(2)) => 1,
+  ])
+end
+# TODO: Write in terms of `GradingType"Nf"` definition.
+function Base.AbstractUnitRange(s::GradingType"NfParity", t::SiteType"Electron")
+  return gradedrange([
+    sector(s, Z{2}(0)) => 1,
+    sector(s, Z{2}(1)) => 1,
+    sector(s, Z{2}(1)) => 1,
+    sector(s, Z{2}(0)) => 1,
+  ])
+end
+function Base.AbstractUnitRange(s::GradingType"Sz", t::SiteType"Electron")
+  return gradedrange([
+    sector(s, U1(0)) => 1,
+    sector(s, U1(1)) => 1,
+    sector(s, U1(-1)) => 1,
+    sector(s, U1(0)) => 1,
+  ])
+end
+
+end

src/op.jl

@@ -8,8 +8,8 @@ struct OpName{Name,Params}
 end
 name(::OpName{Name}) where {Name} = Name
 params(n::OpName) = getfield(n, :params)
-
 Base.getproperty(n::OpName, name::Symbol) = getfield(params(n), name)
+Base.get(t::OpName, name::Symbol, default) = get(params(t), name, default)
 
 OpName{N}(; kwargs...) where {N} = OpName{N}((; kwargs...))
 
@@ -54,9 +54,14 @@ end
 # Generic to `StateName` or `OpName`.
 const StateOrOpName = Union{StateName,OpName}
 alias(n::StateOrOpName) = n
-function (arrtype::Type{<:AbstractArray})(n::StateOrOpName, domain::Integer...)
+function (arrtype::Type{<:AbstractArray})(
+  n::StateOrOpName, domain::Union{Integer,AbstractUnitRange}...
+)
   return arrtype(n, domain)
 end
+function (arrtype::Type{<:AbstractArray})(n::StateOrOpName, domain::Tuple{Vararg{Integer}})
+  return arrtype(n, Base.oneto.(domain))
+end
 (arrtype::Type{<:AbstractArray})(n::StateOrOpName, ts::SiteType...) = arrtype(n, ts)
 function (n::StateOrOpName)(domain...)
   # TODO: Try one alias at a time?
@@ -87,32 +92,58 @@ function nsites(n::StateOrOpName)
   return nsites(n′)
 end
 
-function op_convert(
-  arrtype::Type{<:AbstractArray{<:Any,N}},
-  domain::Tuple{Vararg{Integer}},
-  a::AbstractArray{<:Any,N},
-) where {N}
-  # TODO: Check the dimensions.
-  return convert(arrtype, a)
+# TODO: This does some unwanted conversions, like turning
+# `Diagonal` dense.
+function array(a::AbstractArray, ax::Tuple{Vararg{AbstractUnitRange}})
+  return a[ax...]
 end
-function op_convert(
-  arrtype::Type{<:AbstractArray}, domain::Tuple{Vararg{Integer}}, a::AbstractArray
-)
-  # TODO: Check the dimensions.
-  return convert(arrtype, a)
+
+function Base.axes(::OpName, domain::Tuple{Vararg{AbstractUnitRange}})
+  return (domain..., domain...)
+end
+function Base.axes(n::StateOrOpName, domain::Tuple{Vararg{Integer}})
+  return axes(n, Base.OneTo.(domain))
+end
+function Base.axes(n::StateOrOpName, domain::Tuple{Vararg{SiteType}})
+  return axes(n, AbstractUnitRange.(domain))
+end
+
+## function Base.axes(::OpName"SWAP", domain::Tuple{Vararg{AbstractUnitRange}})
+##   return (reverse(domain)..., domain...)
+## end
+
+function reversed_sites(n::StateOrOpName, domain)
+  return reverse_sites(n, reshape(n(domain...), length.(axes(n, reverse(domain)))))
+end
+function reverse_sites(n::OpName, a::AbstractArray)
+  ndomain = Int(ndims(a)//2)
+  perm1 = reverse(ntuple(identity, ndomain))
+  perm2 = perm1 .+ ndomain
+  perm = (perm1..., perm2...)
+  return permutedims(a, perm)
 end
-function op_convert(
-  arrtype::Type{<:AbstractArray{<:Any,N}}, domain::Tuple{Vararg{Integer}}, a::AbstractArray
-) where {N}
-  size = (domain..., domain...)
-  @assert length(size) == N
-  return convert(arrtype, reshape(a, size))
+
+function state_or_op_convert(
+  n::StateOrOpName,
+  arrtype::Type{<:AbstractArray},
+  domain::Tuple{Vararg{AbstractUnitRange}},
+  a::AbstractArray,
+)
+  ax = axes(n, domain)
+  a′ = reshape(a, length.(ax))
+  a′′ = array(a′, ax)
+  return convert(arrtype, a′′)
 end
-function (arrtype::Type{<:AbstractArray})(n::OpName, domain::Tuple{Vararg{SiteType}})
-  return op_convert(arrtype, length.(domain), n(domain...))
+
+function (arrtype::Type{<:AbstractArray})(n::StateOrOpName, domain::Tuple{Vararg{SiteType}})
+  domain′ = AbstractUnitRange.(domain)
+  return state_or_op_convert(n, arrtype, domain′, reversed_sites(n, domain))
 end
-function (arrtype::Type{<:AbstractArray})(n::OpName, domain::Tuple{Vararg{Integer}})
-  return op_convert(arrtype, domain, n(Int.(domain)...))
+function (arrtype::Type{<:AbstractArray})(
+  n::StateOrOpName, domain::Tuple{Vararg{AbstractUnitRange}}
+)
+  # TODO: Make `(::OpName)(domain...)` constructor process more general inputs.
+  return state_or_op_convert(n, arrtype, domain, reversed_sites(n, Int.(length.(domain))))
 end
 
 function op(arrtype::Type{<:AbstractArray}, n::String, domain...; kwargs...)
@@ -475,13 +506,13 @@ function (n::OpName"Controlled")(domain...)
   # Number of control sites.
   nc = get(params(n), :ncontrol, length(domain) - nt)
   @assert length(domain) == nc + nt
-  d_control = prod(to_dim.(domain[1:nc]))
+  d_control = prod(to_dim.(domain)) - prod(to_dim.(domain[(nc + 1):end]))
   return cat(I(d_control), n.arg(domain[(nc + 1):end]...); dims=(1, 2))
 end
-@op_alias "CNOT" "Controlled" op = OpName"X"()
-@op_alias "CX" "Controlled" op = OpName"X"()
-@op_alias "CY" "Controlled" op = OpName"Y"()
-@op_alias "CZ" "Controlled" op = OpName"Z"()
+@op_alias "CNOT" "Controlled" arg = OpName"X"()
+@op_alias "CX" "Controlled" arg = OpName"X"()
+@op_alias "CY" "Controlled" arg = OpName"Y"()
+@op_alias "CZ" "Controlled" arg = OpName"Z"()
 function alias(n::OpName"CPhase")
   return controlled(OpName"Phase"(; params(n)...))
 end
@@ -504,17 +535,17 @@ function alias(::OpName"CRn")
 end
 @op_alias "CRn̂" "CRn"
 
-@op_alias "CCNOT" "Controlled" ncontrol = 2 op = OpName"X"()
+@op_alias "CCNOT" "Controlled" ncontrol = 2 arg = OpName"X"()
 @op_alias "Toffoli" "CCNOT"
 @op_alias "CCX" "CCNOT"
 @op_alias "TOFF" "CCNOT"
 
-@op_alias "CSWAP" "Controlled" ncontrol = 2 op = OpName"SWAP"()
+@op_alias "CSWAP" "Controlled" ncontrol = 2 arg = OpName"SWAP"()
 @op_alias "Fredkin" "CSWAP"
 @op_alias "CSwap" "CSWAP"
 @op_alias "CS" "CSWAP"
 
-@op_alias "CCCNOT" "Controlled" ncontrol = 3 op = OpName"X"()
+@op_alias "CCCNOT" "Controlled" ncontrol = 3 arg = OpName"X"()
 
 ## # 1-qudit rotation around generic axis n̂.
 ## # exp(-im * α / 2 * n̂ ⋅ σ⃗)