Apply prime, noprime and sim to selected indices

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.16.2"
+version = "0.16.3"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/index.jl

@@ -236,15 +236,30 @@ decoration(n) = (;)
 decoration(n::IndexName) = (; tags = tags(n), plev = plev(n))
 
 """
-    prime(i)
-    prime(t::AbstractNamedTensor)
+    prime(i, plinc = 1)
+    prime(a::AbstractNamedTensor, plinc = 1)
+    prime(a::AbstractNamedTensor, is)
+    prime(a::AbstractNamedTensor, plinc, is)
+    prime(predicate, a::AbstractNamedTensor, plinc = 1)
 
-Increment the prime level of an index or index name by one, returning a new index that
+Increment the prime level of an index or index name by `plinc`, returning a new index that
 is distinct from `i`. Priming is the usual way to make a second copy of an index that
 carries the same tags but is not contracted against the original. The inverse is
 [`noprime`](@ref), which resets the prime level to zero. Given a tensor, prime all of its
 indices.
 
+Given a tensor and `is`, prime only those indices, leaving the rest alone. `is` is a
+collection, or a single index or index name, and `plinc` goes before it. Given a predicate
+instead, prime the indices of `a` for which `predicate` is true, where `plinc` trails the
+predicate because the predicate has to come first. An index of `a` is selected
+by its full name, prime level included, so `noprime(prime(a, i), i)` leaves `a` unchanged:
+the tensor holds `i'`, not `i`. Name the level you mean with `prime(i, 2)`. An index that `a`
+does not have is ignored.
+
+A negative `plinc` lowers the prime level, so `prime(a, -1)` undoes one level of priming.
+Nothing clamps at zero: going below it yields a negative prime level, which is legal but
+prints a warning.
+
 # Examples
 
 ```jldoctest
@@ -255,42 +270,90 @@ false
 
 julia> noprime(prime(i)) == i
 true
+
+julia> prime(i, 2) == prime(prime(i))
+true
+
+julia> j = Index(3);
+
+julia> a = NamedTensor(zeros(2, 3), (i, j));
+
+julia> inds(prime(a, i)) == [prime(i), j]
+true
+
+julia> inds(prime(a, 2, i)) == [prime(i, 2), j]
+true
+
+julia> prime(prime(a, 2), -2) == a
+true
+
+julia> inds(prime(n -> ITensorBase.plev(n) == 0, prime(a, i))) == [prime(i), prime(j)]
+true
 ```
 
-See also [`noprime`](@ref), [`Index`](@ref).
+See also [`noprime`](@ref), [`sim`](@ref), [`Index`](@ref).
 """
 function prime end
 
 """
     noprime(i)
-    noprime(t::AbstractNamedTensor)
+    noprime(a::AbstractNamedTensor)
+    noprime(a::AbstractNamedTensor, is)
+    noprime(predicate, a::AbstractNamedTensor)
 
 Reset the prime level of an index or index name to zero, returning a new index. This
 undoes any number of [`prime`](@ref) calls. Given a tensor, reset the prime level of all of
 its indices.
 
+Given a tensor and `is`, reset only those indices, leaving the rest alone. `is` is a
+collection, or a single index or index name. Given a predicate instead, reset the indices of
+`a` for which `predicate` is true, which is the way to select on something other than the
+exact name, such as a tag. Selection is by full name, so the index passed must carry the
+prime level it has on `a`, and `noprime(a, i)` does nothing to a tensor holding `i'`.
+
 # Examples
 
 ```jldoctest
 julia> i = Index(2);
 
 julia> noprime(prime(i)) == i
 true
+
+julia> j = Index(3);
+
+julia> a = prime(NamedTensor(zeros(2, 3), (i, j)));
+
+julia> inds(noprime(a, prime(j))) == [prime(i), j]
+true
+
+julia> s = ITensorBase.settags(Index(2), "Site");
+
+julia> b = prime(NamedTensor(zeros(2, 3), (s, j)));
+
+julia> inds(noprime(x -> ITensorBase.hastag(x, "Site"), b)) == [s, prime(j)]
+true
 ```
 
-See also [`prime`](@ref), [`Index`](@ref).
+See also [`prime`](@ref), [`sim`](@ref), [`Index`](@ref).
 """
 function noprime end
 
 """
     sim(i)
-    sim(t::AbstractNamedTensor)
+    sim(a::AbstractNamedTensor)
+    sim(a::AbstractNamedTensor, is)
+    sim(predicate, a::AbstractNamedTensor)
 
 Return a "similar" index: a new index (or, given a tensor, a tensor with all of its indices
 replaced) carrying the same tags and prime level as `i` but a fresh unique identifier, so it
 is distinct from `i` and will not contract against it. This is the index-manipulation
 spelling of [`uniquename`](@ref) on an index.
 
+Given a tensor and `is`, replace only those indices, leaving the rest alone. `is` is a
+collection, or a single index or index name. Given a predicate instead, replace the indices
+of `a` for which `predicate` is true. Selection is by full name, prime level included, and an
+index that `a` does not have is ignored.
+
 # Examples
 
 ```jldoctest
@@ -301,13 +364,23 @@ false
 
 julia> length(sim(i))
 2
+
+julia> j = Index(3);
+
+julia> a = NamedTensor(zeros(2, 3), (i, j));
+
+julia> inds(sim(a, i))[1] == i
+false
+
+julia> inds(sim(a, i))[2] == j
+true
 ```
 
-See also [`uniquename`](@ref), [`prime`](@ref).
+See also [`uniquename`](@ref), [`prime`](@ref), [`noprime`](@ref).
 """
 function sim end
 
-prime(n::IndexName) = setplev(n, plev(n) + 1)
+prime(n::IndexName, plinc::Integer = 1) = setplev(n, plev(n) + plinc)
 noprime(n::IndexName) = setplev(n, 0)
 sim(n::IndexName) = uniquename(n)
 
@@ -401,15 +474,47 @@ emptytags(i::Index) = setname(i, emptytags(name(i)))
 decoration(i::Index) = decoration(name(i))
 
 setplev(i::Index, plev) = setname(i, setplev(name(i), plev))
-prime(i::Index) = setname(i, prime(name(i)))
+prime(i::Index, plinc::Integer = 1) = setname(i, prime(name(i), plinc))
 noprime(i::Index) = setname(i, noprime(name(i)))
 sim(i::Index) = setname(i, sim(name(i)))
 
 # Whole-tensor index manipulation: relabel every index name-only via `rename`, leaving the
 # data and spaces untouched.
-prime(a::AbstractNamedTensor) = rename(prime, a)
-noprime(a::AbstractNamedTensor) = rename(noprime, a)
-sim(a::AbstractNamedTensor) = rename(sim, a)
+#
+# The selected forms take a collection of indices, so a lone index needs methods of its own:
+# an `Index` is a `NamedUnitRange`, which iterates its own elements, and would otherwise be
+# read as a collection of integers rather than as one index.
+for f in [:prime, :noprime, :sim]
+    @eval begin
+        $f(a::AbstractNamedTensor) = rename($f, a)
+        $f(a::AbstractNamedTensor, i::AbstractName) = $f(a, (i,))
+        $f(a::AbstractNamedTensor, i::AbstractNamedArray) = $f(a, (i,))
+        function $f(a::AbstractNamedTensor, is)
+            return rename(a, (name(i) => $f(name(i)) for i in is)...)
+        end
+        # `predicate` is typed, since an untyped first argument would be ambiguous against the
+        # collection form above whenever both arguments are tensors.
+        function $f(predicate::Function, a::AbstractNamedTensor)
+            return $f(a, Iterators.filter(predicate, inds(a)))
+        end
+    end
+end
+
+# Only `prime` counts levels, so only `prime` takes an increment. It precedes the selection,
+# as in `prime(i, plinc)` and in both earlier ITensor generations. `Integer` is disjoint from
+# the name and index types above, so this adds no ambiguity. A negative `plinc` lowers the
+# level, which is what an `unprime` would do.
+prime(a::AbstractNamedTensor, plinc::Integer) = rename(n -> prime(n, plinc), a)
+prime(a::AbstractNamedTensor, plinc::Integer, i::AbstractName) = prime(a, plinc, (i,))
+prime(a::AbstractNamedTensor, plinc::Integer, i::AbstractNamedArray) = prime(a, plinc, (i,))
+function prime(a::AbstractNamedTensor, plinc::Integer, is)
+    return rename(a, (name(i) => prime(name(i), plinc) for i in is)...)
+end
+# `plinc` trails the predicate but precedes a selection. The predicate has to come first, so
+# there is nowhere else for it to go, and this is the order legacy settled on too.
+function prime(predicate::Function, a::AbstractNamedTensor, plinc::Integer)
+    return prime(a, plinc, Iterators.filter(predicate, inds(a)))
+end
 
 function primestring(plev)
     if plev < 0

test/test_basics.jl

@@ -1,7 +1,7 @@
 using ITensorBase: ITensorBase, AbstractNamedTensor, ITensor, Index, IndexName, NamedTensor,
     commonind, commoninds, gettag, hascommoninds, hastag, id, inds, name, names, nametype,
     noncommonind, noncommoninds, noprime, operator, plev, prime, rename, setplev, settag,
-    sim, tags, trycommonind, trynoncommonind, tryuniqueind, unioninds, uniqueind,
+    settags, sim, tags, trycommonind, trynoncommonind, tryuniqueind, unioninds, uniqueind,
     uniqueinds, uniquename, unname, unnamed, unsettag, uuid
 using Test: @test, @test_broken, @test_throws, @testset
 using UUIDs: UUID
@@ -114,6 +114,13 @@ using UUIDs: UUID
         @test unnamed(a) == x
         @test plev(i) == 0
         @test plev(prime(i)) == 1
+        # `plinc` lives on the index, never on a tensor, so a tensor's second argument is
+        # always a selection. `prime(i, n)` is how a level-`n` index gets named.
+        @test plev(prime(i, 2)) == 2
+        @test plev(prime(name(i), 3)) == 3
+        @test prime(i, 1) == prime(i)
+        @test prime(i, 0) == i
+        @test prime(prime(i), 2) == prime(i, 3)
         @test length(tags(i)) == 0
         a′ = rename(prime, a)
         @test unnamed(a′) == x
@@ -202,6 +209,83 @@ using UUIDs: UUID
         @test sprint(show, "text/plain", i) ==
             "Index(2|id=$(first(string(uuid(i)), 8))|X=>Y)"
     end
+    @testset "selected-index manipulation" begin
+        elt = Float64
+        i, j, k = Index.((2, 3, 4))
+        a = randn(elt, i, j, k)
+
+        # `sim` is absent here because it mints a fresh id on every call, so two separate calls
+        # never compare equal; it is checked on its own below. Each function is given a tensor
+        # it actually changes, so the comparisons are not trivially true.
+        for (f, b) in ((prime, a), (noprime, prime(a)))
+            bi, _, bk = inds(b)
+            # A lone index or index name stands for the one-element collection. An `Index` is a
+            # `NamedUnitRange`, so without its own method it would be read as a range of integers.
+            @test f(b, bi) == f(b, (bi,))
+            @test f(b, name(bi)) == f(b, (bi,))
+            # Any collection of indices works.
+            @test f(b, [bi, bk]) == f(b, (bi, bk))
+            # Selecting every index agrees with the whole-tensor form.
+            @test f(b, inds(b)) == f(b)
+        end
+
+        # Properties shared by all three, each read off a single call.
+        for f in (prime, noprime, sim)
+            # The unselected indices are left alone.
+            @test inds(f(a, (i, k)))[2] == j
+            # Selecting none is a no-op, as is naming an index the tensor does not have.
+            @test f(a, ()) == a
+            @test f(a, Index(5)) == a
+            # Relabeling is name-only, so the data is untouched.
+            @test unnamed(f(a, i)) == unnamed(a)
+        end
+
+        # `prime` and `noprime` select on the full name, prime level included, so a primed index
+        # has to be named as such. The predicate form is the way to avoid spelling it out.
+        a′ = prime(a, (i, j))
+        @test issetequal(inds(a′), (prime(i), prime(j), k))
+        @test noprime(a′, i) == a′
+        @test noprime(a′, prime(i)) == prime(a, j)
+        @test prime(n -> plev(n) == 0, a′) == prime(a)
+        # A predicate that selects every primed index is just `noprime(a)`, so test one that
+        # leaves a primed index behind.
+        a′′ = prime(a′, (prime(i),))
+        @test issetequal(inds(a′′), (prime(prime(i)), prime(j), k))
+        @test noprime(n -> plev(n) == 2, a′′) == prime(a, j)
+        # Naming the level-2 index directly is the alternative to the predicate.
+        @test noprime(a′′, prime(i, 2)) == prime(a, j)
+
+        # `prime` alone takes a level increment, before the selection. `noprime` and `sim`
+        # have no level to count, so they take no such argument.
+        @test prime(a, 2) == prime(prime(a))
+        @test inds(prime(a, 2, i)) == [prime(i, 2), j, k]
+        @test inds(prime(a, 2, (i, k))) == [prime(i, 2), j, prime(k, 2)]
+        @test prime(a, 2, name(i)) == prime(a, 2, i)
+        @test prime(a, 1) == prime(a)
+        @test prime(a, 0) == a
+        # Stepping up twice matches one increment of two only when the second step names the
+        # index at the level it has reached, since selection includes the prime level.
+        @test prime(a, 2, i) == prime(prime(a, i), prime(i))
+        # A negative increment lowers the level, which is what an `unprime` would do.
+        @test prime(prime(a, 2), -2) == a
+        @test prime(prime(a, 2, i), -1, prime(i, 2)) == prime(a, i)
+        # The predicate form takes the increment too, trailing the predicate.
+        @test prime(n -> plev(n) == 0, a, 2) == prime(a, 2)
+        @test prime(x -> x == i, a, 2) == prime(a, 2, i)
+        @test prime(n -> plev(n) == 0, a, 1) == prime(n -> plev(n) == 0, a)
+
+        # The predicate sees the indices, not the names, so index-level accessors work.
+        i_t = settags(i, "Site")
+        a_t = randn(elt, i_t, j, k)
+        @test prime(x -> hastag(x, "Site"), a_t) == prime(a_t, i_t)
+
+        # `sim` mints a fresh id for the selected index only.
+        a_s = sim(a, i)
+        @test inds(a_s)[1] != i
+        @test issetequal(inds(a_s)[2:3], (j, k))
+        @test unnamed(a_s) == unnamed(a)
+    end
+
     @testset "whole-tensor index manipulation" begin
         elt = Float64
         i, j = Index.((2, 3))