Reorganize and simplify code (#2)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorQuantumOperatorDefinitions"
 uuid = "fd9b415b-e710-4e2a-b407-cba023081494"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/ITensorQuantumOperatorDefinitions.jl

@@ -1,7 +1,12 @@
 module ITensorQuantumOperatorDefinitions
 
 include("sitetype.jl")
-include("ITensorQuantumOperatorDefinitionsChainRulesCoreExt.jl")
+include("space.jl")
+include("val.jl")
+include("state.jl")
+include("op.jl")
+include("has_fermion_string.jl")
+
 include("sitetypes/aliases.jl")
 include("sitetypes/generic_sites.jl")
 include("sitetypes/qubit.jl")
@@ -13,4 +18,12 @@ include("sitetypes/tj.jl")
 include("sitetypes/qudit.jl")
 include("sitetypes/boson.jl")
 
+include("ITensorQuantumOperatorDefinitionsChainRulesCoreExt.jl")
+
+include("itensor/siteinds.jl")
+include("itensor/val.jl")
+include("itensor/state.jl")
+include("itensor/op.jl")
+include("itensor/has_fermion_string.jl")
+
 end

src/has_fermion_string.jl

@@ -0,0 +1 @@
+has_fermion_string(::OpName, ::SiteType) = nothing

src/itensor/has_fermion_string.jl

@@ -0,0 +1,25 @@
+using ITensorBase: Index
+
+has_fermion_string(operator::AbstractArray{<:Number}, s::Index; kwargs...)::Bool = false
+
+function has_fermion_string(opname::AbstractString, s::Index; kwargs...)::Bool
+  opname = strip(opname)
+
+  # Interpret operator names joined by *
+  # as acting sequentially on the same site
+  starpos = findfirst(isequal('*'), opname)
+  if !isnothing(starpos)
+    op1 = opname[1:prevind(opname, starpos)]
+    op2 = opname[nextind(opname, starpos):end]
+    return xor(has_fermion_string(op1, s; kwargs...), has_fermion_string(op2, s; kwargs...))
+  end
+
+  Ntags = length(tags(s))
+  stypes = _sitetypes(s)
+  opn = OpName(opname)
+  for st in stypes
+    res = has_fermion_string(opn, st)
+    !isnothing(res) && return res
+  end
+  return false
+end

src/itensor/op.jl

@@ -0,0 +1,138 @@
+using ITensorBase: Index, ITensor, dag, prime, tags
+
+op(::OpName, ::SiteType, ::Index...; kwargs...) = nothing
+function op(
+  ::OpName, ::SiteType, ::SiteType, sitetypes_inds::Union{SiteType,Index}...; kwargs...
+)
+  return nothing
+end
+
+_sitetypes(i::Index) = _sitetypes(tags(i))
+
+function commontags(i::Index...)
+  return union(tags.(i)...)
+end
+
+op(X::AbstractArray, s::Index...) = ITensor(X, (prime.(s)..., dag.(s)...))
+
+# TODO: Delete these.
+op(opname, s::Vector{<:Index}; kwargs...) = op(opname, s...; kwargs...)
+op(s::Vector{<:Index}, opname; kwargs...) = op(opname, s...; kwargs...)
+
+function op(name::AbstractString, s::Index...; adjoint::Bool=false, kwargs...)
+  name = strip(name)
+  # TODO: filter out only commons tags
+  # if there are multiple indices
+  commontags_s = commontags(s...)
+  # first we handle the + and - algebra, which requires a space between ops to avoid clashing
+  name_split = nothing
+  @ignore_derivatives name_split = String.(split(name, " "))
+  oplocs = findall(x -> x ∈ ("+", "-"), name_split)
+  if !isempty(oplocs)
+    @ignore_derivatives !isempty(kwargs) &&
+      error("Lazy algebra on parametric gates not allowed")
+    # the string representation of algebra ops: ex ["+", "-", "+"]
+    labels = name_split[oplocs]
+    # assign coefficients to each term: ex [+1, -1, +1]
+    coeffs = [1, [(-1)^Int(label == "-") for label in labels]...]
+    # grad the name of each operator block separated by an algebra op, and do so by
+    # making sure blank spaces between opnames are kept when building the new block.
+    start, opnames = 0, String[]
+    for oploc in oplocs
+      finish = oploc
+      opnames = vcat(
+        opnames, [prod([name_split[k] * " " for k in (start + 1):(finish - 1)])]
+      )
+      start = oploc
+    end
+    opnames = vcat(
+      opnames, [prod([name_split[k] * " " for k in (start + 1):length(name_split)])]
+    )
+    # build the vector of blocks and sum
+    op_list = [
+      coeff * (op(opname, s...; kwargs...)) for (coeff, opname) in zip(coeffs, opnames)
+    ]
+    return sum(op_list)
+  end
+  # the the multiplication come after
+  oploc = findfirst("*", name)
+  if !isnothing(oploc)
+    op1, op2 = nothing, nothing
+    @ignore_derivatives begin
+      op1 = name[1:prevind(name, oploc.start)]
+      op2 = name[nextind(name, oploc.start):end]
+      if !(op1[end] == ' ' && op2[1] == ' ')
+        @warn "($op1*$op2) composite op definition `A*B` deprecated: please use `A * B` instead (with spaces)"
+      end
+    end
+    return product(op(op1, s...; kwargs...), op(op2, s...; kwargs...))
+  end
+  common_stypes = _sitetypes(commontags_s)
+  @ignore_derivatives push!(common_stypes, SiteType("Generic"))
+  opn = OpName(name)
+  for st in common_stypes
+    op_mat = op(opn, st; kwargs...)
+    if !isnothing(op_mat)
+      rs = reverse(s)
+      res = ITensor(op_mat, (prime.(rs)..., dag.(rs)...))
+      adjoint && return swapprime(dag(res), 0 => 1)
+      return res
+    end
+  end
+  return throw(
+    ArgumentError(
+      "Overload of \"op\" or \"op!\" functions not found for operator name \"$name\" and Index tags: $(tags.(s)).",
+    ),
+  )
+end
+
+function op(opname, s::Vector{<:Index}, ns::NTuple{N,Integer}; kwargs...) where {N}
+  return op(opname, ntuple(n -> s[ns[n]], Val(N))...; kwargs...)
+end
+
+function op(opname, s::Vector{<:Index}, ns::Vararg{Integer}; kwargs...)
+  return op(opname, s, ns; kwargs...)
+end
+
+function op(s::Vector{<:Index}, opname, ns::Tuple{Vararg{Integer}}; kwargs...)
+  return op(opname, s, ns...; kwargs...)
+end
+
+function op(s::Vector{<:Index}, opname, ns::Integer...; kwargs...)
+  return op(opname, s, ns; kwargs...)
+end
+
+function op(s::Vector{<:Index}, opname, ns::Tuple{Vararg{Integer}}, kwargs::NamedTuple)
+  return op(opname, s, ns; kwargs...)
+end
+
+function op(s::Vector{<:Index}, opname, ns::Integer, kwargs::NamedTuple)
+  return op(opname, s, (ns,); kwargs...)
+end
+
+op(s::Vector{<:Index}, o::Tuple) = op(s, o...)
+
+op(o::Tuple, s::Vector{<:Index}) = op(s, o...)
+
+function op(
+  s::Vector{<:Index},
+  f::Function,
+  opname::AbstractString,
+  ns::Tuple{Vararg{Integer}};
+  kwargs...,
+)
+  return f(op(opname, s, ns...; kwargs...))
+end
+
+function op(
+  s::Vector{<:Index}, f::Function, opname::AbstractString, ns::Integer...; kwargs...
+)
+  return f(op(opname, s, ns; kwargs...))
+end
+
+# Here, Ref is used to not broadcast over the vector of indices
+# TODO: consider overloading broadcast for `op` with the example
+# here: https://discourse.julialang.org/t/how-to-broadcast-over-only-certain-function-arguments/19274/5
+# so that `Ref` isn't needed.
+ops(s::Vector{<:Index}, os::AbstractArray) = [op(oₙ, s) for oₙ in os]
+ops(os::AbstractVector, s::Vector{<:Index}) = [op(oₙ, s) for oₙ in os]

src/itensor/siteinds.jl

@@ -0,0 +1,78 @@
+using ITensorBase: Index, addtags
+
+function siteind(st::SiteType; addtags="", kwargs...)
+  sp = space(st; kwargs...)
+  isnothing(sp) && return nothing
+  return Index(sp, "Site, $(value(st)), $addtags")
+end
+
+function siteind(st::SiteType, n; kwargs...)
+  s = siteind(st; kwargs...)
+  !isnothing(s) && return addtags(s, "n=$n")
+  sp = space(st, n; kwargs...)
+  isnothing(sp) && error(space_error_message(st))
+  return Index(sp, "Site, $(value(st)), n=$n")
+end
+
+siteind(tag::String; kwargs...) = siteind(SiteType(tag); kwargs...)
+
+siteind(tag::String, n; kwargs...) = siteind(SiteType(tag), n; kwargs...)
+
+# Special case of `siteind` where integer (dim) provided
+# instead of a tag string
+#siteind(d::Integer, n::Integer; kwargs...) = Index(d, "Site,n=$n")
+function siteind(d::Integer, n::Integer; addtags="", kwargs...)
+  return Index(d, "Site,n=$n, $addtags")
+end
+
+siteinds(::SiteType, N; kwargs...) = nothing
+
+"""
+    siteinds(tag::String, N::Integer; kwargs...)
+
+Create an array of `N` physical site indices of type `tag`.
+Keyword arguments can be used to specify quantum number conservation,
+see the `space` function corresponding to the site type `tag` for
+supported keyword arguments.
+
+# Example
+
+```julia
+N = 10
+s = siteinds("S=1/2", N; conserve_qns=true)
+```
+"""
+function siteinds(tag::String, N::Integer; kwargs...)
+  st = SiteType(tag)
+
+  si = siteinds(st, N; kwargs...)
+  if !isnothing(si)
+    return si
+  end
+
+  return [siteind(st, j; kwargs...) for j in 1:N]
+end
+
+"""
+    siteinds(f::Function, N::Integer; kwargs...)
+
+Create an array of `N` physical site indices where the site type at site `n` is given
+by `f(n)` (`f` should return a string).
+"""
+function siteinds(f::Function, N::Integer; kwargs...)
+  return [siteind(f(n), n; kwargs...) for n in 1:N]
+end
+
+# Special case of `siteinds` where integer (dim)
+# provided instead of a tag string
+"""
+    siteinds(d::Integer, N::Integer; kwargs...)
+
+Create an array of `N` site indices, each of dimension `d`.
+
+# Keywords
+- `addtags::String`: additional tags to be added to all indices
+"""
+function siteinds(d::Integer, N::Integer; kwargs...)
+  return [siteind(d, n; kwargs...) for n in 1:N]
+end

src/itensor/state.jl

@@ -0,0 +1,24 @@
+using ITensorBase: ITensor, Index, onehot
+
+state(::StateName, ::SiteType, ::Index; kwargs...) = nothing
+
+# Syntax `state("Up", Index(2, "S=1/2"))`
+state(sn::String, i::Index; kwargs...) = state(i, sn; kwargs...)
+
+function state(s::Index, name::AbstractString; kwargs...)::ITensor
+  stypes = _sitetypes(s)
+  sname = StateName(name)
+  for st in stypes
+    v = state(sname, st; kwargs...)
+    !isnothing(v) && return ITensor(v, (s,))
+  end
+  return throw(
+    ArgumentError(
+      "Overload of \"state\" functions not found for state name \"$name\" and Index tags $(tags(s))",
+    ),
+  )
+end
+
+state(s::Index, n::Integer) = onehot(s => n)
+
+state(sset::Vector{<:Index}, j::Integer, st; kwargs...) = state(sset[j], st; kwargs...)

src/itensor/val.jl

@@ -0,0 +1,20 @@
+using ITensorBase: Index
+
+function val(s::Index, name::AbstractString)::Int
+  stypes = _sitetypes(s)
+  sname = ValName(name)
+
+  # Try calling val(::StateName"Name",::SiteType"Tag",)
+  for st in stypes
+    res = val(sname, st)
+    !isnothing(res) && return res
+  end
+
+  return throw(
+    ArgumentError("Overload of \"val\" function not found for Index tags $(tags(s))")
+  )
+end
+
+val(s::Index, n::Integer) = n
+
+val(sset::Vector{<:Index}, j::Integer, st) = val(sset[j], st)

src/op.jl

@@ -0,0 +1,24 @@
+# TODO: Add `params<:NamedTuple` field.
+struct OpName{Name} end
+OpName(s::AbstractString) = OpName{Symbol(s)}()
+OpName(s::Symbol) = OpName{s}()
+name(::OpName{N}) where {N} = N
+macro OpName_str(s)
+  return OpName{Symbol(s)}
+end
+
+# Default implementations of op
+op(::OpName; kwargs...) = nothing
+op(::OpName, ::SiteType; kwargs...) = nothing
+
+function _sitetypes(ts::Set)
+  return collect(SiteType, SiteType.(ts))
+end
+
+op(name::AbstractString; kwargs...) = error("Must input indices when creating an `op`.")
+
+# To ease calling of other op overloads,
+# allow passing a string as the op name
+op(opname::AbstractString, t::SiteType; kwargs...) = op(OpName(opname), t; kwargs...)
+
+op(f::Function, args...; kwargs...) = f(op(args...; kwargs...))