Support for specifying symmetries of axes

Author: mtfishman

Project.toml

@@ -8,15 +8,24 @@ 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", "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"

ext/QuantumOperatorDefinitionsITensorBaseExt/QuantumOperatorDefinitionsITensorBaseExt.jl

@@ -1,6 +1,6 @@
 module QuantumOperatorDefinitionsITensorBaseExt
 
-using ITensorBase: ITensor, Index, dag, gettag, prime
+using ITensorBase: ITensorBase, ITensor, Index, dag, gettag, prime
 using NamedDimsArrays: dename
 using QuantumOperatorDefinitions:
   QuantumOperatorDefinitions, OpName, SiteType, StateName, has_fermion_string
@@ -15,6 +15,10 @@ function QuantumOperatorDefinitions.SiteType(r::Index)
   )
 end
 
+function ITensorBase.Index(t::SiteType; kwargs...)
+  return Index(AbstractUnitRange(t); kwargs...)
+end
+
 function QuantumOperatorDefinitions.has_fermion_string(n::String, r::Index)
   return has_fermion_string(OpName(n), SiteType(r))
 end

ext/QuantumOperatorDefinitionsSymmetrySectorsExt/QuantumOperatorDefinitionsSymmetrySectorsExt.jl

@@ -0,0 +1,86 @@
+module QuantumOperatorDefinitionsSymmetrySectorsExt
+
+using BlockArrays: blocklasts, blocklengths
+using GradedUnitRanges: GradedOneTo, gradedrange
+using LabelledNumbers: label, labelled, unlabel
+using QuantumOperatorDefinitions:
+  QuantumOperatorDefinitions, @SiteType_str, @SymmetryType_str, SiteType, SymmetryType, name
+using SymmetrySectors: ×, SectorProduct, U1, Z
+
+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(::SymmetryType"N", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; N=U1(i - 1))) => 1, 1:length(t)))
+end
+function Base.AbstractUnitRange(::SymmetryType"Sz", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; Sz=U1(i - 1))) => 1, 1:length(t)))
+end
+function Base.AbstractUnitRange(::SymmetryType"Sz↑", t::SiteType)
+  return AbstractUnitRange(SymmetryType"Sz"(), t)
+end
+function Base.AbstractUnitRange(::SymmetryType"Sz↓", t::SiteType)
+  return gradedrange(map(i -> SectorProduct((; Sz=U1(-(i - 1)))) => 1, 1:length(t)))
+end
+
+function sector(symmetrytype::SymmetryType, sec)
+  sectorname = Symbol(get(symmetrytype, :name, name(symmetrytype)))
+  return SectorProduct(NamedTuple{(sectorname,)}((sec,)))
+end
+
+function Base.AbstractUnitRange(s::SymmetryType"Nf", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+# TODO: Write in terms of `SymmetryType"Nf"` definition.
+function Base.AbstractUnitRange(s::SymmetryType"NfParity", t::SiteType"Fermion")
+  return gradedrange([sector(s, Z{2}(0)) => 1, sector(s, Z{2}(1)) => 1])
+end
+function Base.AbstractUnitRange(s::SymmetryType"Sz", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+function Base.AbstractUnitRange(s::SymmetryType"Sz↑", t::SiteType"Fermion")
+  return gradedrange([sector(s, U1(0)) => 1, sector(s, U1(1)) => 1])
+end
+function Base.AbstractUnitRange(s::SymmetryType"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::SymmetryType"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 `SymmetryType"Nf"` definition.
+function Base.AbstractUnitRange(s::SymmetryType"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::SymmetryType"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/sitetype.jl

@@ -4,33 +4,77 @@ struct SiteType{T,Params}
     return new{N,typeof(params)}(params)
   end
 end
-value(::SiteType{T}) where {T} = T
+name(::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)
-
+Base.haskey(t::SiteType, name::Symbol) = haskey(params(t), name)
 SiteType{N}(; kwargs...) where {N} = SiteType{N}((; kwargs...))
-
 SiteType(s::AbstractString; kwargs...) = SiteType{Symbol(s)}(; kwargs...)
 SiteType(i::Integer; kwargs...) = SiteType{Symbol(i)}(; kwargs...)
 macro SiteType_str(s)
-  return SiteType{Symbol(s)}
+  return :(SiteType{$(Expr(:quote, Symbol(s)))})
 end
 
 alias(t::SiteType) = t
 alias(i::Integer) = i
 
+# Like `Base.Broadcast.axistype` (https://github.com/JuliaLang/julia/blob/v1.11.3/base/broadcast.jl#L536-L538)
+# and `BlockArrays.combine_blockaxes` (https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.3.0/src/blockbroadcast.jl#L37-L38).
+combine_axes(a::T, b::T) where {T} = a
+combine_axes(a::Base.OneTo, b::Base.OneTo) = Base.OneTo{Int}(a)
+function combine_axes(a, b)
+  return UnitRange{Int}(a)
+end
+combine_axes(a) = a
+combine_axes(a, b, rest...) = combine_axes(combine_axes(a, b), rest...)
+
+struct SymmetryType{T,Params}
+  params::Params
+  function SymmetryType{N}(params::NamedTuple) where {N}
+    return new{N,typeof(params)}(params)
+  end
+end
+name(::SymmetryType{T}) where {T} = T
+params(t::SymmetryType) = getfield(t, :params)
+Base.getproperty(t::SymmetryType, name::Symbol) = getfield(params(t), name)
+Base.get(t::SymmetryType, name::Symbol, default) = get(params(t), name, default)
+Base.haskey(t::SymmetryType, name::Symbol) = haskey(params(t), name)
+SymmetryType{N}(; kwargs...) where {N} = SymmetryType{N}((; kwargs...))
+SymmetryType(s::AbstractString; kwargs...) = SymmetryType{Symbol(s)}(; kwargs...)
+function SymmetryType(s::Pair{<:AbstractString,<:AbstractString}; kwargs...)
+  return SymmetryType(first(s); kwargs..., name=last(s))
+end
+function SymmetryType(s::Pair{<:AbstractString,<:NamedTuple}; kwargs...)
+  return SymmetryType(first(s); kwargs..., last(s)...)
+end
+macro SymmetryType_str(s)
+  return :(SymmetryType{$(Expr(:quote, Symbol(s)))})
+end
+
+function Base.AbstractUnitRange(symmetry::SymmetryType, t::SiteType)
+  return error("Not implemented.")
+end
+function Base.AbstractUnitRange(symmetry::SymmetryType"Trivial", t::SiteType)
+  return Base.OneTo(length(t))
+end
+
 function Base.length(t::SiteType)
   t′ = alias(t)
   if t == t′
-    return t.length
+    return t.dim
   end
   return length(t′)
 end
 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)))
+  haskey(t, :range) && return t.range
+  if haskey(t, :symmetries)
+    rs = map(symmetry -> AbstractUnitRange(SymmetryType(symmetry), t), t.symmetries)
+    return combine_axes(Base.OneTo(length(t)), rs...)
+  end
+  return Base.OneTo(length(t))
 end
 Base.size(t::SiteType) = (length(t),)
 Base.size(t::SiteType, dim::Integer) = size(t)[dim]