working conditional entropy

docs/src/api.md

@@ -22,6 +22,7 @@ entropy
 binary_entropy
 relative_entropy
 binary_relative_entropy
+conditional_entropy
 ```
 
 ## Measurements

src/entropy.jl

@@ -1,13 +1,13 @@
-_log(b::Real, x::Real) = x > 0 ? log(b, x) : zero(x)
+_log(base::Real, x::Real) = x > 0 ? log(base, x) : zero(x)
 
 """
-    relative_entropy([b=2,] ρ::AbstractMatrix, σ::AbstractMatrix)
+    relative_entropy([base=2,] ρ::AbstractMatrix, σ::AbstractMatrix)
 
-Computes the (quantum) relative entropy tr(`ρ` (log `ρ` - log `σ`)) between positive semidefinite matrices `ρ` and `σ` using a base `b` logarithm. Note that the support of `ρ` must be contained in the support of `σ` but for efficiency this is not checked.
+Computes the (quantum) relative entropy tr(`ρ` (log `ρ` - log `σ`)) between positive semidefinite matrices `ρ` and `σ` using a base `base` logarithm. Note that the support of `ρ` must be contained in the support of `σ` but for efficiency this is not checked.
 
 Reference: [Quantum relative entropy](https://en.wikipedia.org/wiki/Quantum_relative_entropy).
 """
-function relative_entropy(b::Real, ρ::AbstractMatrix{T}, σ::AbstractMatrix{S}) where {T,S<:Number}
+function relative_entropy(base::Real, ρ::AbstractMatrix{T}, σ::AbstractMatrix{S}) where {T,S<:Number}
     R = real(promote_type(T, S))
     if size(ρ) != size(σ)
         throw(ArgumentError("ρ and σ have the same size."))
@@ -21,8 +21,8 @@ function relative_entropy(b::Real, ρ::AbstractMatrix{T}, σ::AbstractMatrix{S})
         throw(ArgumentError("ρ and σ must be positive semidefinite."))
     end
     m = abs2.(ρ_U' * σ_U)
-    logρ_λ = _log.(Ref(b), ρ_λ)
-    logσ_λ = _log.(Ref(b), σ_λ)
+    logρ_λ = _log.(Ref(base), ρ_λ)
+    logσ_λ = _log.(Ref(base), σ_λ)
     d = size(ρ, 1)
     h = R(0)
     @inbounds for j = 1:d, i = 1:d
@@ -34,95 +34,102 @@ relative_entropy(ρ::AbstractMatrix, σ::AbstractMatrix) = relative_entropy(2, 
 export relative_entropy
 
 """
-    relative_entropy([b=2,] p::AbstractVector, q::AbstractVector)
+    relative_entropy([base=2,] p::AbstractVector, q::AbstractVector)
 
-Computes the relative entropy D(`p`||`q`) = Σᵢpᵢlog(pᵢ/qᵢ) between two non-negative vectors `p` and `q` using a base `b` logarithm. Note that the support of `p` must be contained in the support of `q` but for efficiency this is not checked.
+Computes the relative entropy D(`p`||`q`) = Σᵢpᵢlog(pᵢ/qᵢ) between two non-negative vectors `p` and `q` using a base `base` logarithm. Note that the support of `p` must be contained in the support of `q` but for efficiency this is not checked.
 
 Reference: [Relative entropy](https://en.wikipedia.org/wiki/Relative_entropy).
 """
-function relative_entropy(b::Real, p::AbstractVector{T}, q::AbstractVector{S}) where {T,S<:Real}
+function relative_entropy(base::Real, p::AbstractVector{T}, q::AbstractVector{S}) where {T,S<:Real}
     R = promote_type(T, S)
     if length(p) != length(q)
         throw(ArgumentError("`p` and q must have the same length."))
     end
     if any(p .< -Base.rtoldefault(R)) || any(q .< -Base.rtoldefault(R))
         throw(ArgumentError("p and q must be non-negative."))
     end
-    logp = _log.(Ref(b), p)
-    logq = _log.(Ref(b), q)
+    logp = _log.(Ref(base), p)
+    logq = _log.(Ref(base), q)
     h = sum(p[i] * (logp[i] - logq[i]) for i = 1:length(p))
     return h
 end
 relative_entropy(p::AbstractVector, q::AbstractVector) = relative_entropy(2, p, q)
 
 """
-    binary_relative_entropy([b=2,] p::Real, q::Real)
+    binary_relative_entropy([base=2,] p::Real, q::Real)
 
-Computes the binary relative entropy D(`p`||`q`) = p log(p/q) + (1-p) log((1-p)/(1-q)) between two probabilities `p` and `q` using a base `b` logarithm.
+Computes the binary relative entropy D(`p`||`q`) = p log(p/q) + (1-p) log((1-p)/(1-q)) between two probabilities `p` and `q` using a base `base` logarithm.
 
 Reference: [Relative entropy](https://en.wikipedia.org/wiki/Relative_entropy).
 """
-binary_relative_entropy(b::Real, p::Real, q::Real) = relative_entropy(b, [p, 1 - p], [q, 1 - q])
+binary_relative_entropy(base::Real, p::Real, q::Real) = relative_entropy(base, [p, 1 - p], [q, 1 - q])
 binary_relative_entropy(p::Real, q::Real) = binary_relative_entropy(2, p, q)
 export binary_relative_entropy
 
 """
-    entropy([b=2,] ρ::AbstractMatrix)
+    entropy([base=2,] ρ::AbstractMatrix)
 
-Computes the von Neumann entropy -tr(ρ log ρ) of a positive semidefinite operator `ρ` using a base `b` logarithm.
+Computes the von Neumann entropy -tr(ρ log ρ) of a positive semidefinite operator `ρ` using a base `base` logarithm.
 
 Reference: [von Neumann entropy](https://en.wikipedia.org/wiki/Von_Neumann_entropy).
 """
-function entropy(b::Real, ρ::AbstractMatrix)
+function entropy(base::Real, ρ::AbstractMatrix)
     if size(ρ, 1) != size(ρ, 2)
         throw(ArgumentError("ρ must be square."))
     end
     λ = LA.eigvals(ρ)
     if any(λ .< -Base.rtoldefault(eltype(λ)))
         throw(ArgumentError("ρ must be positive semidefinite."))
     end
-    h = -sum(λ[i] * _log(b, λ[i]) for i = 1:size(ρ, 1))
+    h = -sum(λ[i] * _log(base, λ[i]) for i = 1:size(ρ, 1))
     return h
 end
 entropy(ρ::AbstractMatrix) = entropy(2, ρ)
 export entropy
 
 """
-    entropy([b=2,] p::AbstractVector)
+    entropy([base=2,] p::AbstractVector)
 
-Computes the Shannon entropy -Σᵢpᵢlog(pᵢ) of a non-negative vector `p` using a base `b` logarithm.
+Computes the Shannon entropy -Σᵢpᵢlog(pᵢ) of a non-negative vector `p` using a base `base` logarithm.
 
 Reference: [Entropy (information theory)](https://en.wikipedia.org/wiki/Entropy_(information_theory)).
 """
-function entropy(b::Real, p::AbstractVector{T}) where {T<:Real}
+function entropy(base::Real, p::AbstractVector{T}) where {T<:Real}
     if any(p .< -Base.rtoldefault(T))
         throw(ArgumentError("p must be non-negative."))
     end
-    h = -sum(p[i] * _log(b, p[i]) for i = 1:length(p))
+    h = -sum(p[i] * _log(base, p[i]) for i = 1:length(p))
     return h
 end
 entropy(p::AbstractVector) = entropy(2, p)
 export entropy
 
 """
-    binary_entropy([b=2,] p::Real)
+    binary_entropy([base=2,] p::Real)
 
-Computes the Shannon entropy -p log(p) - (1-p)log(1-p) of a probability `p` using a base `b` logarithm.
+Computes the Shannon entropy -p log(p) - (1-p)log(1-p) of a probability `p` using a base `base` logarithm.
 
 Reference: [Entropy (information theory)](https://en.wikipedia.org/wiki/Entropy_(information_theory)).
 """
-binary_entropy(b::Real, p::Real) = p == 0 || p == 1 ? zero(p) : -p * log(b, p) - (1 - p) * log(b, 1 - p)
+binary_entropy(base::Real, p::Real) = p == 0 || p == 1 ? zero(p) : -p * log(base, p) - (1 - p) * log(base, 1 - p)
 binary_entropy(p::Real) = binary_entropy(2, p)
 export binary_entropy
 
-function conditional_entropy(b::Real, p::AbstractMatrix{T}) where {T<:Real}
-    nA, nB = size(p)
+"""
+    conditional_entropy([base=2,] pAB::AbstractMatrix)
+
+Computes the conditional (Shannon) entropy H(A|B) of the joint probability distribution `pAB` using a base `base` logarithm.
+
+Reference: [Conditional entropy](https://en.wikipedia.org/wiki/Conditional_entropy).
+"""
+function conditional_entropy(base::Real, pAB::AbstractMatrix{T}) where {T<:Real}
+    nA, nB = size(pAB)
     h = T(0)
-    pB = sum(p; dims = 2)
+    pB = sum(pAB; dims = 1)
     for a = 1:nA, b = 1:nB
-        h -= p[a, b] * _log(b, p[a, b] / pB[b])
+        h -= pAB[a, b] * _log(base, pAB[a, b] / pB[b])
     end
     return h
 end
-conditional_entropy(b::Real, p::AbstractMatrix) = conditional_entropy(2, p)
+conditional_entropy(pAB::AbstractMatrix) = conditional_entropy(2, pAB)
 export conditional_entropy

test/entropy.jl

@@ -36,4 +36,16 @@
             @test binary_relative_entropy(1, q) ≈ -log(2, q)
         end
     end
+
+    @testset "Conditional" begin
+        @test conditional_entropy(Diagonal(ones(2) / 2)) == 0.0
+        @test conditional_entropy(ones(2, 2) / 4) == 1.0
+        for R in [Float64, Double64, Float128, BigFloat]
+            pAB = ones(R, 3, 2) / 6
+            @test isa(conditional_entropy(pAB), R)
+            @test conditional_entropy(pAB) ≈ log2(R(6)) - 1
+            pAB = reshape(random_probability(R, 6), 2, 3)
+            @test conditional_entropy(pAB) ≈ conditional_entropy(ℯ, pAB) / log(R(2))
+        end
+    end
 end