From 2dec460453dad17fae162badf283143d4527bd0f Mon Sep 17 00:00:00 2001
From: jasonhan3 <98669315+jasonhan3@users.noreply.github.com>
Date: Sat, 3 Aug 2024 16:01:45 -0400
Subject: [PATCH] Added tests for the amplitude damping noise model

---
 src/BPGates.jl            |  33 +++--
 test/runtests.jl          |   1 +
 test/test_mixedstateop.jl | 293 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 319 insertions(+), 8 deletions(-)
 create mode 100644 test/test_mixedstateop.jl

diff --git a/src/BPGates.jl b/src/BPGates.jl
index 7a5ecf3..0be4b34 100644
--- a/src/BPGates.jl
+++ b/src/BPGates.jl
@@ -216,6 +216,7 @@ end
 
 """The permutations realized by [`BellPauliPermutation`](@ref)."""
 const pauli_perm_tuple = (
+    # 3 is actually 2; 2 is actually 3
     (1, 2, 3, 4),
     # Z gate?
     (3, 4, 1, 2), ## X flip
@@ -231,9 +232,16 @@ const pauli_perm_qc = (sId1,sX,sZ,sY)
 function QuantumClifford.apply!(state::BellState, op::BellPauliPermutation) # TODO abstract away the permutation application as it is used by other gates too
     phase = state.phases
     @inbounds phase_idx = bit_to_int(phase[op.sidx*2-1],phase[op.sidx*2])
+    # println("QuantumClifford.apply!, BellPauliPermutation, op.sidx: $(op.sidx)")
+    # println("QuantumClifford.apply!, BellPauliPermutation, phase_idx: $phase_idx")
+    # println("QuantumClifford.apply!, BellPauliPermutation, op.pidx: $(op.pidx)")
     @inbounds perm = pauli_perm_tuple[op.pidx]
     @inbounds permuted_idx = perm[phase_idx]
+    # println("QuantumClifford.apply!, BellPauliPermutation, permuted_idx: $permuted_idx")
+    # this makes no sense. Make bit 2 go before bit1.
     bit1, bit2 = int_to_bit(permuted_idx,Val(2))
+    # println("QuantumClifford.apply!, BellPauliPermutation, bit1: $bit1")
+    # println("QuantumClifford.apply!, BellPauliPermutation, bit2: $bit2")
     @inbounds phase[op.sidx*2-1] = bit1
     @inbounds phase[op.sidx*2] = bit2
     return state
@@ -567,9 +575,9 @@ function get_mixed_transition_probs(λ::Float64, cur_state_idx::Int)
     # 0 <= λ <= 1 (λ = 1 - e ^(-t / T1))
     mixed_state_tuple = (
         (0.5 * λ^2 - λ + 1, 0.5 * λ * (1 - λ), 0.5 * λ^2, 0.5 * λ * (1 - λ)),
-        (0.5 * λ, 1 - λ, 0.5 * λ, 0),
+        (0.5 * λ, 1 - λ, 0.5 * λ, 0.0),
         (0.5 * λ^2, 0.5 * λ * (1 - λ), 0.5 * λ^2 - λ + 1, 0.5 * λ * (1 - λ)),
-        (0.5 * λ, 0, 0.5 * λ, 1 - λ)
+        (0.5 * λ, 0.0, 0.5 * λ, 1 - λ)
     )
     return mixed_state_tuple[cur_state_idx]
 end
@@ -583,13 +591,9 @@ struct MixedStateOp <: AbstractNoiseBellOp
     lambda::Float64
 end
 
-function QuantumClifford.apply!(state::BellState, g::MixedStateOp)
-    state_idx = g.idx
-    phase = state.phases
-    @inbounds phase_idx = bit_to_int(phase[state_idx * 2 - 1],phase[state_idx * 2])
-    rand_prob_val = rand()
-    transition_probs = get_mixed_transition_probs(g.lambda, phase_idx)
+function select_rand_state(transition_probs::NTuple{4, Float64})
     # TODO: hardcoded for speed, BUT change this to a macro for readability.
+    rand_prob_val = rand()
     new_phase_idx = 0
     if rand_prob_val < transition_probs[1]
         new_phase_idx = 1
@@ -600,6 +604,15 @@ function QuantumClifford.apply!(state::BellState, g::MixedStateOp)
     else
         new_phase_idx = 4
     end
+    return new_phase_idx
+end
+
+function QuantumClifford.apply!(state::BellState, g::MixedStateOp)
+    state_idx = g.idx
+    phase = state.phases
+    @inbounds phase_idx = bit_to_int(phase[state_idx * 2 - 1],phase[state_idx * 2])
+    transition_probs = get_mixed_transition_probs(g.lambda, phase_idx)
+    new_phase_idx = select_rand_state(transition_probs)
     # this should never happen; can remove for speed
     @assert new_phase_idx != 0
     state_bit1, state_bit2 = int_to_bit(new_phase_idx, Val(2))
@@ -621,6 +634,10 @@ function QuantumClifford.apply!(state::BellState, g::AsymmDepOp)
     # ^ ?
     r = rand()
     # 2: Z, 3: X, 4: Y
+    # println("QuantumClifford.apply!, AsymmDepOp")
+    # print("QuantumClifford.apply!, AsymmDepOp, px_bell: $px_bell")
+    # print("QuantumClifford.apply!, AsymmDepOp, py_bell: $py_bell")
+    # print("QuantumClifford.apply!, AsymmDepOp, pz_bell: $pz_bell")
     if r<g.px
         apply!(state, BellPauliPermutation(3, i))
     elseif r<g.px+g.pz
diff --git a/test/runtests.jl b/test/runtests.jl
index 663fd04..c25bb30 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -28,6 +28,7 @@ println("Starting tests with $(Threads.nthreads()) threads out of `Sys.CPU_THREA
 get(ENV,"JET_TEST","")=="true" && @doset "jet"
 @doset "doctests"
 @doset "bpgates"
+@doset "mixedstateop"
 
 using Aqua
 doset("aqua") && begin
diff --git a/test/test_mixedstateop.jl b/test/test_mixedstateop.jl
new file mode 100644
index 0000000..ab6d5aa
--- /dev/null
+++ b/test/test_mixedstateop.jl
@@ -0,0 +1,293 @@
+using Test
+using Logging
+
+using BPGates: get_mixed_transition_probs, select_rand_state
+
+using QuantumClifford: apply!
+
+# TODO: have tolerance for numerical precision for noise calculations
+
+# Set up the logger to enable debug level
+function enable_debug_logging()
+    global_logger(ConsoleLogger(stderr, Logging.Info))
+end
+
+# Get the bell states used for testing
+function get_bell_states()
+    return (1 / sqrt(2)) * [
+        [1.0; 0.0; 0.0; 1.0],
+        [0.0; 1.0; 1.0; 0.0],
+        [1.0; 0.0; 0.0; -1.0],
+        [0.0; 1.0; -1.0; 0.0]
+    ]
+end
+
+# Define kraus operators to test
+function get_ampdamp_kraus_ops(λ::Float64)::Vector{Matrix{Complex}}
+    # Define Kraus operators
+    E1 = [1 0; 0 sqrt(1 - λ)]
+    E2 = [0 sqrt(λ); 0 0]
+    return [E1, E2]
+end
+
+# Returns an index selected from the indices of a list, weighted by the probability at that index
+function get_weighted_index(prob_dist::Vector{Float64})
+    if !isapprox(sum(prob_dist), 1.0, atol=1e-9)
+        @error "get_weighted_index: prob_dist sums to $(sum(prob_dist)) and not 1"
+    end
+    rand_num = rand()
+    cur_sum = 0.0
+    for (idx, prob_val) in enumerate(prob_dist)
+        if rand_num < cur_sum + prob_val
+            return idx
+        end
+        cur_sum += prob_val
+    end
+end
+
+# Converts a bit to an int from left to right (LSB is the rightmost bit)
+function bit_to_int_testing(bits)
+    reduce(⊻,(bit<<(length(bits) - index) for (index,bit) in enumerate(bits))) + 1 # +1 so that we use julia indexing convenctions
+end
+
+# Returns the tensor product of a set of Kraus operators
+function get_tensored_kraus_ops(kraus_ops::Vector{Matrix{Complex}})
+    tensored_kraus_ops = []
+    for op1 in kraus_ops
+        for op2 in kraus_ops
+            push!(tensored_kraus_ops, kron(op1, op2))
+        end
+    end
+    return tensored_kraus_ops
+end
+
+# Returns the diagonal of a matrix
+function get_diagonal(matrix::Matrix{ComplexF64})::Vector{Float64}
+    diagonal_length = min(size(matrix)[1], size(matrix)[2])
+    if size(matrix)[1] != size(matrix)[2]
+        @error("get_diagonal: matrix is not square, rows: $(size(matrix)[1]) != cols: $(size(matrix)[2])")
+    end
+    diag_elements = []
+    for idx in 1:diagonal_length
+        push!(diag_elements, matrix[idx, idx])
+    end
+    return diag_elements
+end
+
+# Returns the bell state corresponding to the input phases
+function get_bell_state(phases::BitVector)
+    state_idx = bit_to_int_testing(phases)
+    bell_states = get_bell_states()
+    return bell_states[state_idx]
+end
+
+# Numerically computes the transition probabilities of one state to another,
+# given a mixed probability distribution of states and a set of kraus operators
+function get_numerical_trans_probs(mixed_state_probs::Vector{Float64}, kraus_ops::Vector{Matrix{Complex}})
+    bell_states = get_bell_states()
+    ρ_bell = zeros(ComplexF64, length(bell_states), length(bell_states))
+    for (bell_idx, bell_prob) in enumerate(mixed_state_probs)
+        ρ_bell += bell_prob * bell_states[bell_idx] * bell_states[bell_idx]'
+    end
+    @debug("get_numerical_trans_probs: ρ_bell: $ρ_bell")
+    tensored_ops = get_tensored_kraus_ops(kraus_ops)
+    # Apply the Kraus operators to the Bell state
+    ρ_new = zeros(ComplexF64, size(ρ_bell)[1], size(ρ_bell)[2])
+    for K in tensored_ops
+        ρ_new += K * ρ_bell * adjoint(K)
+    end
+    U_bell = hcat(get_bell_states()...)
+    @debug("get_numerical_trans_probs: ρ_new: $ρ_new")
+    @debug("get_numerical_trans_probs: U_bell: $U_bell")
+    ρ_Bell_basis = adjoint(U_bell) * ρ_new * U_bell
+    @debug("get_numerical_trans_probs: ρ_Bell_basis: $ρ_Bell_basis")
+    @debug("get_numerical_trans_probs: size(ρ_Bell_basis): $(size(ρ_Bell_basis))")
+    return get_diagonal(ρ_Bell_basis)
+end
+
+# Returns the simulated transition probabilities from the mixed state probabilities,
+# the noise parameter, simulated for num_samples times
+function get_simulated_trans_probs(mixed_state_probs::Vector{Float64}, noise_param::Float64, num_samples::Integer)
+    sampled_state_count = fill(0, length(get_bell_states()))
+    for _ in 1:num_samples
+        state_idx = get_weighted_index(mixed_state_probs)
+        transition_probs = get_mixed_transition_probs(noise_param, state_idx)
+        new_state_idx = select_rand_state(transition_probs)
+        # TODO: ensure that the index is inbounds for my list
+        sampled_state_count[new_state_idx] += 1
+    end
+    return sampled_state_count / num_samples
+end
+
+# Testing functions
+
+# Test our simulated noise for the pure states in the bell basis
+function test_bell_basis_probs(kraus_ops_factory::Function)
+    statedist_list = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+    lambda = 1 - exp(-500e-9 / 260e-6)
+    num_samples = 1000000
+    kraus_ops = kraus_ops_factory(lambda)
+    for state_dist in statedist_list
+        @info "test_bell_basis_probs: state_dist: $state_dist"
+        numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+        simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+        # @info("test_A_state_probs: numerical_trans_probs: $numerical_trans_probs")
+        # @info("test_A_state_probs: simulated_trans_probs: $simulated_trans_probs")
+        @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+    end
+end
+
+# Test our simulated noise when there should be no noise applied for states in the bell basis
+function test_nonoise_bell_basis(kraus_ops_factory::Function)
+    statedist_list = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+    lambda = 0.0
+    num_samples = 1000000
+    kraus_ops = kraus_ops_factory(lambda)
+    for state_dist in statedist_list
+        @info "test_nonoise_bell_basis: state_dist: $state_dist"
+        numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+        simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+        # @info("test_A_state_probs: numerical_trans_probs: $numerical_trans_probs")
+        # @info("test_A_state_probs: simulated_trans_probs: $simulated_trans_probs")
+        @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+    end
+end
+
+# Test our simulated noise with a noise factor close to 1 for states in the bell basis
+function test_almost_all_noise(kraus_ops_factory::Function)
+    statedist_list = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+    lambda = 0.999999999
+    num_samples = 1000000
+    kraus_ops = kraus_ops_factory(lambda)
+    for state_dist in statedist_list
+        @info "test_almost_all_noise: state_dist: $state_dist"
+        numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+        simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+        @info("test_almost_all_noise: numerical_trans_probs: $numerical_trans_probs")
+        @info("test_almost_all_noise: simulated_trans_probs: $simulated_trans_probs")
+        @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+    end
+end
+
+# Test random noise factors for states in the bell basis
+function test_random_lambdas(kraus_ops_factory::Function)
+    statedist_list = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+    num_samples = 1000000
+    num_random_lambdas = 5
+    for _ in 1:num_random_lambdas
+        lambda = rand()
+        num_samples = 1000000
+        kraus_ops = kraus_ops_factory(lambda)
+        for state_dist in statedist_list
+            @info "test_random_lambdas: state_dist: $state_dist"
+            @info "test_random_lambdas: lambda: $lambda"
+            numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+            simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+            @info("test_random_lambdas: numerical_trans_probs: $numerical_trans_probs")
+            @info("test_random_lambdas: simulated_trans_probs: $simulated_trans_probs")
+            @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+        end
+    end
+end
+
+# Test all mixed states composed of two states in the Bell basis with equal probability
+function test_equal_mixed_probs_2(kraus_ops_factory::Function)
+    bell_states = get_bell_states()
+    lambda = 1 - exp(-500e-9 / 260e-6)
+    kraus_ops = kraus_ops_factory(lambda)
+    num_samples = 1000000
+    for first_phases_idx in 1:length(bell_states)
+        for second_phases_idx in 1:length(bell_states)
+            if first_phases_idx == second_phases_idx
+                continue
+            end
+            state_dist = [0.0 for _ in 1:length(bell_states)]
+            for basis_idx in 1:length(bell_states)
+                if basis_idx == first_phases_idx || basis_idx == second_phases_idx
+                    state_dist[basis_idx] = 0.5
+                end
+            end
+            numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+            simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+            @info("test_equal_mixed_probs_2: numerical_trans_probs: $numerical_trans_probs")
+            @info("test_equal_mixed_probs_2: simulated_trans_probs: $simulated_trans_probs")
+            @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+        end
+    end
+end
+
+# Test a mixed state composed of all 4 bell basis states with equal probability
+function test_equal_mixed_probs_4(kraus_ops_factory::Function)
+    state_dist = [0.25 for _ in 1:length(get_bell_states())]
+    lambda = 1 - exp(-500e-9 / 260e-6)
+    num_samples = 1000000
+    kraus_ops = kraus_ops_factory(lambda)
+    @info "test_equal_mixed_probs_4: state_dist: $state_dist"
+    numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+    simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+    @info("test_equal_mixed_probs_4: numerical_trans_probs: $numerical_trans_probs")
+    @info("test_equal_mixed_probs_4: simulated_trans_probs: $simulated_trans_probs")
+    @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+end
+
+# Test random mixed states in the bell basis
+function test_random_mixed_states(kraus_ops_factory::Function)
+    num_random_states = 20
+    statedist_list = []
+    for _ in 1:num_random_states
+        state_dist = [rand() for _ in 1:length(get_bell_states())]
+        state_dist = state_dist / sum(state_dist)
+        push!(statedist_list, state_dist)
+    end
+    @debug "length(statedist_list): $(length(statedist_list))"
+    lambda = 1 - exp(-500e-9 / 260e-6)
+    num_samples = 1000000
+    kraus_ops = kraus_ops_factory(lambda)
+    for state_dist in statedist_list
+        @info "test_random_mixed_states: state_dist: $state_dist"
+        numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+        simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+        @info("test_random_mixed_states: numerical_trans_probs: $numerical_trans_probs")
+        @info("test_random_mixed_states: simulated_trans_probs: $simulated_trans_probs")
+        @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+    end
+end
+
+# Test random mixed states with random noise parameters in the Bell basis
+function test_random_mixed_states_random_lambdas(kraus_ops_factory::Function)
+    num_random_states = 20
+    statedist_list = []
+    for _ in 1:num_random_states
+        state_dist = [rand() for _ in 1:length(get_bell_states())]
+        state_dist = state_dist / sum(state_dist)
+        push!(statedist_list, state_dist)
+    end
+    num_random_lambdas = 20
+    num_samples = 1000000
+    for _ in 1:num_random_lambdas
+        lambda = rand()
+        @info "test_random_mixed_states_random_lambdas: lambda: $lambda"
+        kraus_ops = kraus_ops_factory(lambda)
+        for state_dist in statedist_list
+            @info "test_random_mixed_states_random_lambdas: state_dist: $state_dist"
+            numerical_trans_probs = get_numerical_trans_probs(state_dist, kraus_ops)
+            simulated_trans_probs = get_simulated_trans_probs(state_dist, lambda, num_samples)
+            @info("test_random_mixed_states_random_lambdas: numerical_trans_probs: $numerical_trans_probs")
+            @info("test_random_mixed_states_random_lambdas: simulated_trans_probs: $simulated_trans_probs")
+            @test all(isapprox.(numerical_trans_probs, simulated_trans_probs, atol=(2 / sqrt(num_samples))))
+        end
+    end
+end
+
+# Run all tests
+@testset "BPGates.jl, get_mixed_transition_probs tests" begin
+    enable_debug_logging()
+    test_bell_basis_probs(get_ampdamp_kraus_ops)
+    test_nonoise_bell_basis(get_ampdamp_kraus_ops)
+    test_almost_all_noise(get_ampdamp_kraus_ops)
+    test_random_lambdas(get_ampdamp_kraus_ops)
+    test_equal_mixed_probs_2(get_ampdamp_kraus_ops)
+    test_equal_mixed_probs_4(get_ampdamp_kraus_ops)
+    test_random_mixed_states(get_ampdamp_kraus_ops)
+    test_random_mixed_states_random_lambdas(get_ampdamp_kraus_ops)
+end