Refactor sites methods

bsc-quantic · May 20, 2024 · 0c67e56 · 0c67e56
1 parent 6d6e8e4
commit 0c67e56
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/Ansatz.jl b/src/Ansatz.jl
@@ -24,7 +24,6 @@ for f in [
  :noutputs,
  :inputs,
  :outputs,
- :sites,
  :nsites,
  :nlanes,
  :socket,
@@ -46,6 +45,8 @@ function Base.summary(io::IO, tn::A) where {A<:Ansatz}
 end
 Base.show(io::IO, tn::A) where {A<:Ansatz} = summary(io, tn)
 
+sites(tn::Ansatz; kwargs...) = sites(Quantum(tn); kwargs...)
+
 function Tenet.inds(tn::Ansatz; kwargs...)
  if keys(kwargs) === (:bond,)
  inds(tn, Val(:bond), kwargs[:bond]...)

diff --git a/src/Quantum.jl b/src/Quantum.jl
@@ -150,7 +150,15 @@ Base.show(io::IO, q::Quantum) = print(io, "Quantum (inputs=$(ninputs(q)), output
 
 Returns the sites of a [`Quantum`](@ref) Tensor Network.
 """
-sites(tn::Quantum) = collect(keys(tn.sites))
+function sites(tn::Quantum; kwargs...)
+ if isempty(kwargs)
+ collect(keys(tn.sites))
+ elseif keys(kwargs) === (:at,)
+ findfirst(i -> i === kwargs[:at], tn.sites)
+ else
+ throw(MethodError(sites, (Quantum,), kwargs))
+ end
+end
 
 """
  nsites(q::Quantum)