Port functions from ITensorInfiniteMPS

Author: ryanlevy

src/itensor.jl

@@ -840,6 +840,7 @@ dirs(A::ITensor, is) = dirs(inds(A), is)
 
 # TODO: add isdiag(::Tensor) to NDTensors
 isdiag(T::ITensor)::Bool = (storage(T) isa Diag || storage(T) isa DiagBlockSparse)
+LinearAlgebra.isdiag(T::ITensor) = isdiag(T)
 
 diaglength(T::ITensor) = diaglength(tensor(T))
 

src/lib/ITensorMPS/src/mps.jl

@@ -1032,3 +1032,5 @@ end
 function expect(psi::MPS, op1::Matrix{<:Number}, ops::Matrix{<:Number}...; kwargs...)
   return expect(psi, (op1, ops...); kwargs...)
 end
+
+Base.getindex(ψ::MPS, r::UnitRange{Int}) = MPS([ψ[n] for n in r])

src/lib/SmallStrings/src/smallstring.jl

@@ -102,6 +102,33 @@ end
 Base.:(==)(s1::SmallString, s2::SmallString) = (s1.data == s2.data)
 Base.isless(s1::SmallString, s2::SmallString) = isless(s1.data, s2.data)
 
+maxlength(s::SmallString) = length(s.data)
+
+function Base.length(s::SmallString)
+  n = 1
+  while n <= maxlength(s) && s[n] != zero(eltype(s))
+    n += 1
+  end
+  return n - 1
+end
+
+Base.lastindex(s::SmallString) = length(s)
+Base.getindex(s::SmallString, r::UnitRange) = SmallString([s[n] for n in r])
+
+
+# TODO: make this work directly on a Tag, without converting
+# to String
+function Base.parse(::Type{T}, s::SmallString) where {T<:Integer}
+  return parse(T, string(s))
+end
+
+function Base.startswith(s::SmallString, subtag::SmallString)
+  for n in 1:length(subtag)
+    s[n] ≠ subtag[n] && return false
+  end
+  return true
+end
+
 ########################################################
 # Here are alternative SmallString comparison implementations
 #

src/lib/TagSets/src/TagSets.jl

@@ -215,6 +215,7 @@ data(T::TagSet) = T.data
 Base.length(T::TagSet) = T.length
 Base.@propagate_inbounds Base.getindex(T::TagSet, n::Integer) = SmallString(data(T)[n])
 Base.copy(ts::TagSet) = TagSet(data(ts), length(ts))
+Base.keys(ts::TagSet) = Base.OneTo(length(ts))
 
 function Base.:(==)(ts1::TagSet, ts2::TagSet)
   l1 = length(ts1)

src/tensor_operations/matrix_decomposition.jl

@@ -590,6 +590,27 @@ function sqrt_decomp(D::ITensor, u::Index, v::Index)
   return sqrtDL, prime(δᵤᵥ), sqrtDR
 end
 
+# Take the square root of T assuming it is Hermitian
+# TODO: add more general index structures
+function Base.sqrt(T::ITensor; ishermitian=true, atol=1e-15)
+  @assert ishermitian
+  # TODO diagonal version
+  #if isdiag(T) && order(T) == 2
+  #  return itensor(sqrt(tensor(T)))
+  #end
+  D, U = eigen(T; ishermitian=ishermitian)
+  sqrtD = D
+  for n in 1:mindim(D)
+    Dnn = D[n, n]
+    if Dnn < 0 && abs(Dnn) < atol
+      sqrtD[n, n] = 0
+    else
+      sqrtD[n, n] = sqrt(Dnn)
+    end
+  end
+  return U' * sqrtD * dag(U)
+end
+
 function factorize_svd(
   A::ITensor,
   Linds...;