More missing functionality

Author: mtfishman

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.2"
 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.11.0"
 MapBroadcast = "0.1.5"
 NamedDimsArrays = "0.3.0"
 UnallocatedArrays = "0.1.1"

src/ITensorBase.jl

@@ -18,6 +18,7 @@ using NamedDimsArrays:
   named,
   nameddimsindices,
   setname,
+  setnameddimsindices,
   unname
 
 const Tag = String
@@ -79,7 +80,9 @@ struct Index{T,Value<:AbstractUnitRange{T}} <: AbstractNamedUnitRange{T,Value,In
   name::IndexName
 end
 
-Index(length::Int; kwargs...) = Index(Base.OneTo(length), IndexName(; kwargs...))
+function Index(length::Int; tags, 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
@@ -177,17 +180,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
@@ -216,6 +232,25 @@ 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
+
+prime(a::AbstractITensor) = setinds(a, prime.(inds(a)))
+
 include("quirks.jl")
 
 end

src/quirks.jl

@@ -4,11 +4,28 @@ dag(i::Index) = i
 dim(i::Index) = dename(length(i))
 # TODO: Define this properly.
 hasqns(i::Index) = false
-inds(a::ITensor) = nameddimsindices(a)
 # 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