Blahut-Arimoto algorithm Matlab源码

For a discrete memoryless channel , the capacity is defined as

where  and  denote the input and output variables of the channel respectively, and the maximization is taken over all input distributions .

Given a channel transition matrix whose -entry is the conditional probability , the Blahut-Arimoto algorithm computes the capacity of the discrete memoryless channel, and the input distribution  that attains the maximum.

Reference: Chapter 9 in the book Information Theory and Network Coding by Raymond Yeung.

function [C r] = BlahutArimoto(p)

disp('BlahutArimoto')

% Capacity of discrete memoryless channel
% Blahut-Arimoto algorithm

% Input
% p: m x n matrix
% p is the transition matrix for a channel with m inputs and n outputs
%
% The input matrix p should contain no zero row and no zero column.
%
% p(i,j) is the condition probability that the channel output
% is j given that the input is i
% (i=1,2,...,m and j = 1,2,...,n)
%
%
% Output
% capacity : capacity in bits
% r: channel input distribution which achieves capacity
%

% For example, the transition matrix for the erasure channel is
% can be calculated as
% e = 0.5;
% p = [1-e e 0; 0 e 1-e]; % conditional prob. for erasure channel
% The capacity can be calculated by BlahutArimoto(p), and is equal to 1-e
%

% Check that the entries of input matrix p are non-negative
if ~isempty(find(p < 0))
    disp('Error: some entry in the input matrix is negative')
    C = 0; return;
end

% Check that the input matrix p does not have zero column
column_sum = sum(p);
if ~isempty(find(column_sum == 0))
    disp('Error: there is a zero column in the input matrix');
    C = 0; return;
end

% Check that the input matrix p does not have zero row
row_sum = sum(p,2);
if ~isempty(find(row_sum == 0))
    disp('Error: there is a zero row in the input matrix');
    C = 0; return;
else
    p = diag(sum(p,2))^(-1) * p; % Make sure that the row sums are 1
end

[m n] = size(p);

r = ones(1,m)/m; % initial distribution for channel input
q = zeros(m,n);
error_tolerance = 1e-5/m;
r1 = [];
for i = 1:m
    p(i,:) = p(i,:)/sum(p(i,:));
end
for iter = 1:10000
    for j = 1:n
        q(:,j) = r'.*p(:,j);
        q(:,j) = q(:,j)/sum(q(:,j));
    end

    for i = 1:m
        r1(i) = prod(q(i,:).^p(i,:));
    end

    r1 = r1/sum(r1);
    if norm(r1 - r) < error_tolerance
        break
    else
        r = r1;
    end
end

C = 0;
for i = 1:m
    for j = 1:n
        if r(i) > 0 && q(i,j) > 0
            C = C+ r(i)*p(i,j)* log(q(i,j)/r(i));
        end
    end
end

C = C/log(2); % Capacity in bits