Add gradient of textbook fidelity and a test with FD (#573)

* Add gradient of textbook fidelity and a test with FD

* Fix preexisting tests specializes.jl

* Add fidelity2 and fidelity2'

Implement a squared version of the fidelity.
This improves numerical stability for algorithms that intentionally minimize the overlap,
since, in contrast to the existing gradient fidelity',
no division by the overlap occurs in the gradient.

Additionally:
-Implement tests for fidelity2 and fidelity2'
-Add a docstring for fidelity2 and fidelity2' and integrate them into the interface
-Add the internal function pure_state_fidelity2 which uses abs2 instead of abs

The docstrings are concise but reference the fidelity docstring and clarify the differences.

* Fix docstring and clean up tests

* Remove unused circuits in tests

* Let fidelity2 call into fidelity

---------

Co-authored-by: Yanick Kind <[email protected]>

Author: YanickKi

lib/YaoAPI/src/registers.jl

@@ -652,6 +652,17 @@ julia> fidelity(reg1, reg2)
 """
 @interface fidelity
 
+
+"""
+    fidelity2(register1, register2) -> Real/Vector{<:Real}
+    fidelity2'(pair_or_reg1, pair_or_reg2) -> (g1, g2)
+
+Return the [`fidelity`](@ref) squared.
+
+`fidelity2'` returns the corresponding gradient of `fidelity2` with respect to registers and circuit parameters, similar to `fidelity'`.
+"""
+@interface fidelity2
+
 """
     tracedist(register1, register2)
 

lib/YaoArrayRegister/src/YaoArrayRegister.jl

@@ -76,6 +76,7 @@ export AbstractRegister,
     von_neumann_entropy,
     mutual_information,
     fidelity,
+    fidelity2,
     focus!,
     focus,
     insert_qudits!,

lib/YaoArrayRegister/src/density_matrix.jl

@@ -103,6 +103,8 @@ function YaoAPI.fidelity(m::DensityMatrix, n::DensityMatrix)
     return density_matrix_fidelity(m.state, n.state)
 end
 
+YaoAPI.fidelity2(m::DensityMatrix, n::DensityMatrix) = YaoAPI.fidelity(m, n)^2
+
 function YaoAPI.purify(r::DensityMatrix{D}; num_env::Int = nactive(r)) where {D}
     Ne = D ^ num_env
     Ns = size(r.state, 1)

lib/YaoArrayRegister/src/register.jl

@@ -370,6 +370,8 @@ function YaoAPI.fidelity(r1::ArrayReg, r2::ArrayReg)
     end
 end
 
+YaoAPI.fidelity2(r1::AbstractArrayReg, r2::AbstractArrayReg) = YaoAPI.fidelity(r2, r1).^2
+
 YaoAPI.tracedist(r1::ArrayReg, r2::ArrayReg) = tracedist(density_matrix(r1), density_matrix(r2))
 YaoAPI.tracedist(r1::BatchedArrayReg, r2::BatchedArrayReg) = tracedist.(r1, r2)
 

lib/YaoArrayRegister/test/density_matrix.jl

@@ -37,6 +37,49 @@ using LinearAlgebra
     ) |> all
 end
 
+@testset "test fidelity2" begin
+    reg = rand_state(3)
+    reg_ = rand_state(3)
+    reg2 = clone(reg, 3)
+    @test fidelity2(reg, reg) ≈ 1
+    @test fidelity2(reg, reg_) < 1
+    @test fidelity2(reg2, reg2) ≈ [1, 1, 1]
+    @test isapprox(fidelity2(reg, reg_), fidelity(reg, reg_)^2)
+    # mix
+    reg4 = join(reg, reg)
+    reg5 = join(reg_, reg_)
+    focus!(reg4, 1:3)
+    focus!(reg5, 1:3)
+    @test isapprox(fidelity2(reg, reg_), fidelity2(reg4, reg5), atol = 1e-5)
+    @test isapprox(fidelity2(reg4, reg5), fidelity(reg4, reg5)^2)
+
+    @test isapprox.(
+        fidelity2(reg, reg_),
+        fidelity2(clone(reg4, 3), clone(reg5, 3)),
+        atol = 1e-5,
+    ) |> all
+
+    @test isapprox.(
+        fidelity2(clone(reg4, 3), clone(reg5, 3)),
+        fidelity(clone(reg4, 3), clone(reg5, 3)).^2,
+    ) |> all
+
+    # batch
+    st = rand(ComplexF64, 8, 2)
+    reg1 = BatchedArrayReg(st)
+    reg2 = rand_state(3)
+
+    @test fidelity2(reg1, reg2) ≈
+          [fidelity2(ArrayReg(st[:, 1]), reg2), fidelity2(ArrayReg(st[:, 2]), reg2)]
+    @test isapprox(fidelity2(reg1, reg2), fidelity(reg1, reg2).^2)
+
+    @test isapprox.(
+        fidelity2(reg, reg_),
+        fidelity2(clone(reg4, 3), clone(reg5, 3)),
+        atol = 1e-5,
+    ) |> all
+end
+
 @testset "test trace distance" begin
     reg = rand_state(3)
     reg_ = rand_state(3)
@@ -162,15 +205,18 @@ end
     r2 = density_matrix(reg2, (2, 1))
     expected = abs(tr(sqrt(sqrt(r1.state) * r2.state * sqrt(r1.state))))
     @test expected ≈ fidelity(r1, r2) atol=1e-5
-
+    @test expected^2 ≈ fidelity2(r1, r2) atol=1e-5
     # focused state is viewed as mixed state
     f1 = focus!(copy(reg1), (2, 1))
     f2 = focus!(copy(reg2), (2, 1))
     @test fidelity(r1, r2) ≈ fidelity(f1, f2) atol=1e-6
-
+    @test fidelity2(r1, r2) ≈ fidelity2(f1, f2) atol=1e-6
+    @test fidelity2(r1, r2) ≈ fidelity(f1, f2)^2 atol=1e-6
     # fidelity between focused and pure state
     f1 = rand_state(2)
     @test fidelity(density_matrix(f1, (1, 2)), r2) ≈ fidelity(f1, f2) atol=1e-6
+    @test fidelity2(density_matrix(f1, (1, 2)), r2) ≈ fidelity2(f1, f2) atol=1e-6
+    @test fidelity2(density_matrix(f1, (1, 2)), r2) ≈ fidelity(f1, f2)^2 atol=1e-6
 
     dm = rand_density_matrix(2)
     @test is_density_matrix(dm.state)

lib/YaoBlocks/src/YaoBlocks.jl

@@ -28,6 +28,7 @@ import YaoAPI:
     dispatch!,
     expect,
     fidelity,
+    fidelity2,
     focus!,
     getiparams,
     iparams_eltype,

lib/YaoBlocks/src/autodiff/specializes.jl

@@ -1,4 +1,4 @@
-for F in [:expect, :fidelity, :operator_fidelity]
+for F in [:expect, :fidelity, :fidelity2, :operator_fidelity]
     @eval Base.adjoint(::typeof($F)) = Adjoint($F)
     @eval Base.show(io::IO, ::Adjoint{Any,typeof($F)}) = print(io, "$($F)'")
     @eval Base.show(io::IO, ::MIME"text/plain", ::Adjoint{Any,typeof($F)}) =
@@ -24,6 +24,7 @@ end
 _eval(p::Pair{<:AbstractRegister,<:AbstractBlock}) = copy(p.first) |> p.second
 _eval(reg::AbstractRegister) = reg
 YaoAPI.fidelity(p1, p2) = fidelity(_eval(p1), _eval(p2))
+YaoAPI.fidelity2(p1, p2) = fidelity2(_eval(p1), _eval(p2))
 
 function (::Adjoint{Any,typeof(fidelity)})(
     reg1::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
@@ -32,6 +33,13 @@ function (::Adjoint{Any,typeof(fidelity)})(
     fidelity_g(reg1, reg2)
 end
 
+function (::Adjoint{Any,typeof(fidelity2)})(
+    reg1::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
+    reg2::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
+)
+    fidelity2_g(reg1, reg2)
+end
+
 function fidelity_g(
     reg1::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
     reg2::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
@@ -82,6 +90,56 @@ please file an issue if you really need this feature.",
     return res1, res2
 end
 
+function fidelity2_g(
+    reg1::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
+    reg2::Union{AbstractArrayReg,Pair{<:AbstractArrayReg,<:AbstractBlock}},
+)
+    if reg1 isa Pair
+        in1, c1 = reg1
+        out1 = copy(in1) |> c1
+    else
+        out1 = reg1
+    end
+
+    if reg2 isa Pair
+        in2, c2 = reg2
+        out2 = copy(in2) |> c2
+    else
+        out2 = reg2
+    end
+    if nremain(out1) != 0
+        throw(
+            ArgumentError(
+                "The gradient of registers with environment is not implemented yet.
+However, back propagating over a focused register is possible,
+please file an issue if you really need this feature.",
+            ),
+        )
+    end
+    overlap = out1' * out2
+
+    out1δ = copy(out2)
+    regscale!.(viewbatch.(Ref(out1δ), 1:length(overlap)), 2.0.*conj.(overlap))
+    out2δ = copy(out1)
+    regscale!.(viewbatch.(Ref(out2δ), 1:length(overlap)), 2.0.*overlap)
+
+    if reg1 isa Pair
+        (_, in1δ), params1δ = apply_back((out1, out1δ), c1)
+        res1 = in1δ => params1δ
+    else
+        res1 = out1δ
+    end
+
+    if reg2 isa Pair
+        (_, in2δ), params2δ = apply_back((out2, out2δ), c2)
+        res2 = in2δ => params2δ
+    else
+        res2 = out2δ
+    end
+
+    return res1, res2
+end
+
 function (::Adjoint{Any,typeof(operator_fidelity)})(b1::AbstractBlock, b2::AbstractBlock)
     operator_fidelity_g(b1, b2)
 end

lib/YaoBlocks/test/autodiff/specializes.jl

@@ -105,6 +105,57 @@ end
     #@test isapprox(vec(g2.state), state_numgrad(reg2->sum(fidelity(reg1=>c1, reg2=>c2)), reg2), atol=1e-5)
 end
 
+@testset "fidelity2 grad" begin
+    nbit = 4
+    Random.seed!(2)
+    for nbatch in [NoBatch(), 10]
+        reg1 = rand_state(nbit; nbatch = nbatch)
+        reg2 = rand_state(nbit; nbatch = nbatch)
+        c1 = qftcirc(nbit)
+        c2 = chain(put(nbit, 2 => Rx(0.5)), control(nbit, 1, 3 => Ry(0.5)))
+
+        g1, g2 = fidelity2'(reg1, reg2)
+        @test isapprox(
+            vec(g1.state),
+            state_numgrad(reg1 -> sum(fidelity2(reg1, reg2)), reg1),
+            atol = 1e-4,
+        )
+        @test isapprox(
+            vec(g2.state),
+            state_numgrad(reg2 -> sum(fidelity2(reg1, reg2)), reg2),
+            atol = 1e-4,
+        )
+
+        (g1, pg1), (g2, pg2) = fidelity2'(reg1 => c1, reg2 => c2)
+        npg1 = YaoBlocks.AD.ng(
+            x -> sum(fidelity2(reg1 => dispatch!(c1, x), reg2 => c2)),
+            parameters(c1),
+        )
+        npg2 = YaoBlocks.AD.ng(
+            x -> sum(fidelity2(reg1 => c1, reg2 => dispatch!(c2, x))),
+            parameters(c2),
+        )
+        @test isapprox(pg1, vec(npg1), atol = 1e-5)
+        @test isapprox(pg2, vec(npg2), atol = 1e-5)
+        @test isapprox(
+            vec(g1.state),
+            state_numgrad(reg1 -> sum(fidelity2(reg1 => c1, reg2 => c2)), reg1),
+            atol = 1e-4,
+        )
+        @test isapprox(
+            vec(g2.state),
+            state_numgrad(reg2 -> sum(fidelity2(reg1 => c1, reg2 => c2)), reg2),
+            atol = 1e-4,
+        )
+    end
+
+    nbatch = NoBatch()
+    reg1 = rand_state(nbit; nbatch = nbatch) |> focus!(2, 1, 4)
+    reg2 = rand_state(nbit; nbatch = nbatch) |> focus!(2, 1, 4)
+
+    @test_throws ArgumentError fidelity2'(reg1, reg2)
+end
+
 @testset "operator fideliy" begin
     nbit = 4
     c1 = qftcirc(nbit)