Allow custom axes when constructing operators, ITensor extension

Author: mtfishman

Project.toml

@@ -9,12 +9,14 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [weakdeps]
 ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
+NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 
 [extensions]
-QuantumOperatorDefinitionsITensorBaseExt = "ITensorBase"
+QuantumOperatorDefinitionsITensorBaseExt = ["ITensorBase", "NamedDimsArrays"]
 
 [compat]
 ITensorBase = "0.1.10"
 LinearAlgebra = "1.10"
+NamedDimsArrays = "0.4.0"
 Random = "1.10"
 julia = "1.10"

ext/QuantumOperatorDefinitionsITensorBaseExt/QuantumOperatorDefinitionsITensorBaseExt.jl

@@ -1,11 +1,18 @@
 module QuantumOperatorDefinitionsITensorBaseExt
 
 using ITensorBase: ITensor, Index, dag, gettag, prime
+using NamedDimsArrays: dename
 using QuantumOperatorDefinitions:
   QuantumOperatorDefinitions, OpName, SiteType, StateName, has_fermion_string
 
 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 (and 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 QuantumOperatorDefinitions.has_fermion_string(n::String, r::Index)
@@ -14,6 +21,7 @@ end
 
 function Base.AbstractArray(n::OpName, r::Index)
   # TODO: Define this with mapped dimnames.
+  # Generalize beyond prime levels with codomain and domain indices.
   return ITensor(AbstractArray(n, SiteType(r)), (prime(r), dag(r)))
 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,7 +54,9 @@ 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
 (arrtype::Type{<:AbstractArray})(n::StateOrOpName, ts::SiteType...) = arrtype(n, ts)
@@ -87,32 +89,50 @@ function nsites(n::StateOrOpName)
   return nsites(n′)
 end
 
+function array(a::AbstractArray, ax::Tuple{Vararg{AbstractUnitRange}})
+  return a[ax...]
+end
+
 function op_convert(
   arrtype::Type{<:AbstractArray{<:Any,N}},
-  domain::Tuple{Vararg{Integer}},
+  domain::Tuple{Vararg{AbstractUnitRange}},
   a::AbstractArray{<:Any,N},
 ) where {N}
-  # TODO: Check the dimensions.
-  return convert(arrtype, a)
+  ax = (domain..., domain...)
+  a′ = array(a, ax)
+  return convert(arrtype, a′)
 end
 function op_convert(
-  arrtype::Type{<:AbstractArray}, domain::Tuple{Vararg{Integer}}, a::AbstractArray
+  arrtype::Type{<:AbstractArray}, domain::Tuple{Vararg{AbstractUnitRange}}, a::AbstractArray
 )
-  # TODO: Check the dimensions.
-  return convert(arrtype, a)
+  ax = (domain..., domain...)
+  a′ = array(a, ax)
+  return convert(arrtype, a′)
 end
 function op_convert(
-  arrtype::Type{<:AbstractArray{<:Any,N}}, domain::Tuple{Vararg{Integer}}, a::AbstractArray
+  arrtype::Type{<:AbstractArray{<:Any,N}},
+  domain::Tuple{Vararg{AbstractUnitRange}},
+  a::AbstractArray,
 ) where {N}
-  size = (domain..., domain...)
-  @assert length(size) == N
-  return convert(arrtype, reshape(a, size))
+  ax = (domain..., domain...)
+  @assert length(ax) == N
+  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...))
+  domain′ = AbstractUnitRange.(domain)
+  return op_convert(arrtype, domain′, n(domain...))
+end
+
+function (arrtype::Type{<:AbstractArray})(
+  n::OpName, domain::Tuple{Vararg{AbstractUnitRange}}
+)
+  # TODO: Make `(::OpName)(domain...)` constructor process more general inputs.
+  return op_convert(arrtype, domain, n(Int.(length.(domain))...))
 end
 function (arrtype::Type{<:AbstractArray})(n::OpName, domain::Tuple{Vararg{Integer}})
-  return op_convert(arrtype, domain, n(Int.(domain)...))
+  return arrtype(n, Base.oneto.(domain))
 end
 
 function op(arrtype::Type{<:AbstractArray}, n::String, domain...; kwargs...)

src/sitetype.jl

@@ -4,14 +4,15 @@ struct SiteType{T,Params}
     return new{N,typeof(params)}(params)
   end
 end
+value(::SiteType{T}) where {T} = T
 params(t::SiteType) = getfield(t, :params)
 Base.getproperty(t::SiteType, name::Symbol) = getfield(params(t), name)
+Base.get(t::SiteType, name::Symbol, default) = get(params(t), name, default)
 
 SiteType{N}(; kwargs...) where {N} = SiteType{N}((; kwargs...))
 
 SiteType(s::AbstractString; kwargs...) = SiteType{Symbol(s)}(; kwargs...)
 SiteType(i::Integer; kwargs...) = SiteType{Symbol(i)}(; kwargs...)
-value(::SiteType{T}) where {T} = T
 macro SiteType_str(s)
   return SiteType{Symbol(s)}
 end
@@ -26,7 +27,11 @@ function Base.length(t::SiteType)
   end
   return length(t′)
 end
-Base.AbstractUnitRange(t::SiteType) = Base.OneTo(length(t))
+function Base.AbstractUnitRange(t::SiteType)
+  # This logic allows specifying a range with extra properties,
+  # like ones with symmetry sectors.
+  return get(t, :range, Base.OneTo(length(t)))
+end
 Base.size(t::SiteType) = (length(t),)
 Base.size(t::SiteType, dim::Integer) = size(t)[dim]
 Base.axes(t::SiteType) = (AbstractUnitRange(t),)

src/state.jl

@@ -6,14 +6,15 @@ struct StateName{Name,Params}
     return new{N,typeof(params)}(params)
   end
 end
+name(::StateName{N}) where {N} = N
 params(n::StateName) = getfield(n, :params)
 Base.getproperty(n::StateName, name::Symbol) = getfield(params(n), name)
+Base.get(t::StateName, name::Symbol, default) = get(params(t), name, default)
 
 StateName{N}(; kwargs...) where {N} = StateName{N}((; kwargs...))
 
 StateName(s::AbstractString; kwargs...) = StateName{Symbol(s)}(; kwargs...)
 StateName(s::Symbol; kwargs...) = StateName{s}(; kwargs...)
-name(::StateName{N}) where {N} = N
 macro StateName_str(s)
   return StateName{Symbol(s)}
 end