Add missing functionality like tagging and priming (#4)

Author: mtfishman

Project.toml

@@ -1,12 +1,13 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 UnallocatedArrays = "43c9e47c-e622-40fb-bf18-a09fc8c466b6"
@@ -16,6 +17,7 @@ UnspecifiedTypes = "42b3faec-625b-4613-8ddc-352bf9672b8d"
 Accessors = "0.1.39"
 DerivableInterfaces = "0.3.7"
 FillArrays = "1.13.0"
+LinearAlgebra = "1.10"
 MapBroadcast = "0.1.5"
 NamedDimsArrays = "0.3.0"
 UnallocatedArrays = "0.1.1"

README.md

@@ -19,13 +19,8 @@ julia> Pkg.add("ITensorBase")
 
 ````julia
 using ITensorBase: ITensorBase, ITensor, Index
-````
-
-TODO: This should be `TensorAlgebra.qr`.
-
-````julia
 using LinearAlgebra: qr
-using NamedDimsArrays: NamedDimsArray, aligndims, dimnames, name, unname
+using NamedDimsArrays: aligndims, unname
 using Test: @test
 i = Index(2)
 j = Index(2)

examples/README.jl

@@ -20,9 +20,8 @@ julia> Pkg.add("ITensorBase")
 # ## Examples
 
 using ITensorBase: ITensorBase, ITensor, Index
-# TODO: This should be `TensorAlgebra.qr`.
 using LinearAlgebra: qr
-using NamedDimsArrays: NamedDimsArray, aligndims, dimnames, name, unname
+using NamedDimsArrays: aligndims, unname
 using Test: @test
 i = Index(2)
 j = Index(2)

src/ITensorBase.jl

@@ -1,5 +1,8 @@
 module ITensorBase
 
+export ITensor, Index
+
+using Accessors: @set
 using MapBroadcast: Mapped
 using NamedDimsArrays:
   NamedDimsArrays,
@@ -14,16 +17,53 @@ using NamedDimsArrays:
   name,
   named,
   nameddimsindices,
+  setname,
+  setnameddimsindices,
   unname
 
+const Tag = String
+const TagSet = Set{Tag}
+
+tagset(tags::String) = Set(filter(!isempty, String.(strip.(split(tags, ",")))))
+tagset(tags::TagSet) = tags
+
+function tagsstring(tags::TagSet)
+  str = ""
+  length(tags) == 0 && return str
+  tags_vec = collect(tags)
+  for n in 1:(length(tags_vec) - 1)
+    str *= "$(tags_vec[n]),"
+  end
+  str *= "$(tags_vec[end])"
+  return str
+end
+
 @kwdef struct IndexName <: AbstractName
   id::UInt64 = rand(UInt64)
+  tags::TagSet = TagSet()
   plev::Int = 0
-  tags::Set{String} = Set{String}()
-  namedtags::Dict{Symbol,String} = Dict{Symbol,String}()
 end
 NamedDimsArrays.randname(n::IndexName) = IndexName()
 
+id(n::IndexName) = n.id
+tags(n::IndexName) = n.tags
+plev(n::IndexName) = n.plev
+
+settags(n::IndexName, tags) = @set n.tags = tags
+addtags(n::IndexName, ts) = settags(n, tags(n) ∪ tagset(ts))
+
+setprime(n::IndexName, plev) = @set n.plev = plev
+prime(n::IndexName) = setprime(n, plev(n) + 1)
+
+function Base.show(io::IO, i::IndexName)
+  idstr = "id=$(id(i) % 1000)"
+  tagsstr = !isempty(tags(i)) ? "|\"$(tagsstring(tags(i)))\"" : ""
+  primestr = primestring(plev(i))
+  str = "IndexName($(idstr)$(tagsstr))$(primestr)"
+  print(io, str)
+  return nothing
+end
+
 struct IndexVal{Value<:Integer} <: AbstractNamedInteger{Value,IndexName}
   value::Value
   name::IndexName
@@ -41,7 +81,22 @@ struct Index{T,Value<:AbstractUnitRange{T}} <: AbstractNamedUnitRange{T,Value,In
   name::IndexName
 end
 
-Index(length::Int) = Index(Base.OneTo(length), IndexName())
+function Index(length::Int; tags=TagSet(), kwargs...)
+  return Index(Base.OneTo(length), IndexName(; tags=tagset(tags), kwargs...))
+end
+function Index(length::Int, tags::String; kwargs...)
+  return Index(Base.OneTo(length), IndexName(; kwargs..., tags=tagset(tags)))
+end
+
+# TODO: Define for `NamedDimsArrays.NamedViewIndex`.
+id(i::Index) = id(name(i))
+tags(i::Index) = tags(name(i))
+plev(i::Index) = plev(name(i))
+
+# TODO: Define for `NamedDimsArrays.NamedViewIndex`.
+addtags(i::Index, tags) = setname(i, addtags(name(i), tags))
+prime(i::Index) = setname(i, prime(name(i)))
+Base.adjoint(i::Index) = prime(i)
 
 # Interface
 # TODO: Overload `Base.parent` instead.
@@ -51,6 +106,29 @@ NamedDimsArrays.name(i::Index) = i.name
 # Constructor
 NamedDimsArrays.named(i::AbstractUnitRange, name::IndexName) = Index(i, name)
 
+function primestring(plev)
+  if plev < 0
+    return " (warning: prime level $plev is less than 0)"
+  end
+  if plev == 0
+    return ""
+  elseif plev > 3
+    return "'$plev"
+  else
+    return "'"^plev
+  end
+end
+
+function Base.show(io::IO, i::Index)
+  lenstr = "length=$(dename(length(i)))"
+  idstr = "|id=$(id(i) % 1000)"
+  tagsstr = !isempty(tags(i)) ? "|\"$(tagsstring(tags(i)))\"" : ""
+  primestr = primestring(plev(i))
+  str = "Index($(lenstr)$(idstr)$(tagsstr))$(primestr)"
+  print(io, str)
+  return nothing
+end
+
 struct NoncontiguousIndex{T,Value<:AbstractVector{T}} <:
        AbstractNamedVector{T,Value,IndexName}
   value::Value
@@ -103,17 +181,30 @@ struct AllocatableArrayInterface <: AbstractAllocatableArrayInterface end
 
 unallocatable(a::AbstractITensor) = NamedDimsArray(a)
 
-@interface ::AbstractAllocatableArrayInterface function Base.setindex!(
-  a::AbstractArray, value, I::Int...
-)
+function setindex_allocatable!(a::AbstractArray, value, I...)
   allocate!(specify_eltype!(a, typeof(value)))
   # TODO: Maybe use `@interface interface(a) a[I...] = value`?
   unallocatable(a)[I...] = value
   return a
 end
 
+# TODO: Combine these by using `Base.to_indices`.
+@interface ::AbstractAllocatableArrayInterface function Base.setindex!(
+  a::AbstractArray, value, I::Int...
+)
+  setindex_allocatable!(a, value, I...)
+  return a
+end
+@interface ::AbstractAllocatableArrayInterface function Base.setindex!(
+  a::AbstractArray, value, I::AbstractNamedInteger...
+)
+  setindex_allocatable!(a, value, I...)
+  return a
+end
+
 @derive AllocatableArrayInterface() (T=AbstractITensor,) begin
   Base.setindex!(::T, ::Any, ::Int...)
+  Base.setindex!(::T, ::Any, ::AbstractNamedInteger...)
 end
 
 mutable struct ITensor <: AbstractITensor
@@ -127,9 +218,42 @@ using Accessors: @set
 setdenamed(a::ITensor, denamed) = (@set a.parent = denamed)
 setdenamed!(a::ITensor, denamed) = (a.parent = denamed)
 
+function ITensor(elt::Type, I1::Index, I_rest::Index...)
+  I = (I1, I_rest...)
+  # TODO: Use `FillArrays.Zeros`.
+  return ITensor(zeros(elt, length.(dename.(I))...), I)
+end
+
 function ITensor(I1::Index, I_rest::Index...)
   I = (I1, I_rest...)
   return ITensor(Zeros{UnspecifiedZero}(length.(dename.(I))...), I)
 end
 
+function ITensor()
+  return ITensor(Zeros{UnspecifiedZero}(), ())
+end
+
+inds(a::AbstractITensor) = nameddimsindices(a)
+setinds(a::AbstractITensor, inds) = setnameddimsindices(a, inds)
+
+function uniqueinds(a1::AbstractITensor, a_rest::AbstractITensor...)
+  return setdiff(inds(a1), inds.(a_rest)...)
+end
+function uniqueind(a1::AbstractITensor, a_rest::AbstractITensor...)
+  return only(uniqueinds(a1, a_rest...))
+end
+
+function commoninds(a1::AbstractITensor, a_rest::AbstractITensor...)
+  return intersect(inds(a1), inds.(a_rest)...)
+end
+function commonind(a1::AbstractITensor, a_rest::AbstractITensor...)
+  return only(commoninds(a1, a_rest...))
+end
+
+# TODO: Use `replaceinds`/`mapinds`, based on
+# `replacenameddimsindices`/`mapnameddimsindices`.
+prime(a::AbstractITensor) = setinds(a, prime.(inds(a)))
+
+include("quirks.jl")
+
 end

src/quirks.jl

@@ -0,0 +1,31 @@
+# TODO: Define this properly.
+dag(i::Index) = i
+# TODO: Deprecate.
+dim(i::Index) = dename(length(i))
+# TODO: Define this properly.
+hasqns(i::Index) = false
+# TODO: Deprecate.
+itensor(parent::AbstractArray, nameddimsindices) = ITensor(parent, nameddimsindices)
+function itensor(parent::AbstractArray, i1::Index, i_rest::Index...)
+  return ITensor(parent, (i1, i_rest...))
+end
+
+# This seems to be needed to get broadcasting working.
+# TODO: Investigate this and see if we can get rid of it.
+Base.Broadcast.extrude(a::AbstractITensor) = a
+
+# TODO: Generalize this.
+# Maybe define it as `oneelement`, and base it on
+# `FillArrays.OneElement` (https://juliaarrays.github.io/FillArrays.jl/stable/#FillArrays.OneElement).
+function onehot(iv::Pair{<:Index,<:Int})
+  a = ITensor(first(iv))
+  a[last(iv)] = one(Bool)
+  return a
+end
+
+using LinearAlgebra: svd
+# TODO: Define this in `MatrixAlgebra.jl`/`TensorAlgebra.jl`.
+function factorize(a::AbstractITensor, args...; kwargs...)
+  U, S, V = svd(a, args...; kwargs...)
+  return U, S * V
+end