Refactor select method and add getindex function

to TensorNetwork

src/TensorNetwork.jl

@@ -265,10 +265,15 @@ end
 Return tensors whose indices match with the list of indices `i`.
 """
 select(tn::AbstractTensorNetwork, i::Symbol) = copy(tn.indexmap[i])
-select(tn::AbstractTensorNetwork, is::AbstractVecOrTuple{Symbol}) =
-    filter(tn.indexmap[first(is)]) do tensor
-        is ⊆ inds(tensor)
-    end
+select(tn::AbstractTensorNetwork, is::AbstractVecOrTuple{Symbol}) = select(⊆, tn, is)
+
+function select(selector, tn::TensorNetwork, is::AbstractVecOrTuple{Symbol})
+    filter(Base.Fix1(selector, is) ∘ inds, tn.indexmap[first(is)])
+end
+
+function Base.getindex(tn::TensorNetwork, is::Symbol...; mul::Int = 1)
+    first(Iterators.drop(Iterators.filter(Base.Fix1(issetequal, is) ∘ inds, tn.indexmap[first(is)]), mul - 1))
+end
 
 """
     in(tensor::Tensor, tn::AbstractTensorNetwork)

test/TensorNetwork_test.jl

@@ -193,6 +193,23 @@
         @test isempty(select(tn, (:j, :l)))
     end
 
+    @testset "getindex" begin
+        t_ij = Tensor(zeros(2, 2), (:i, :j))
+        t_ik = Tensor(zeros(2, 2), (:i, :k))
+        t_ilm = Tensor(zeros(2, 2, 2), (:i, :l, :m))
+        t_lm = Tensor(zeros(2, 2), (:l, :m))
+        tn = TensorNetwork([t_ij, t_ik, t_ilm, t_lm])
+
+        @test t_ij === tn[:i, :j]
+        @test t_ik === tn[:i, :k]
+        @test t_ilm === tn[:i, :l, :m]
+        @test t_lm === tn[:l, :m]
+
+        # NOTE although it should throw `KeyError`, it throws `ArgumentError` due to implementation 
+        @test_throws ArgumentError tn[:i, :x]
+        @test_throws ArgumentError tn[:i, :j, :k]
+    end
+
     # @testset "selectdim" begin
     #     tn = rand(TensorNetwork, 10, 3)
     #     label = first(inds(tn))