# Use Mathematica to Implement Blahut–Arimoto Algorithm (Part: Algorithm for Rate-Distortion)

Try to use Mathematica to implement Blahut–Arimoto algorithm (here we just focus on the part of computation of Rate-Distortion).

Here are some reference code snippets:

The algorithm looks like:

I have ported the python code in mathematica.

My question is:

Can mathematica do matirx manipulation efficiently like Python numpy or MATLAB? Is it possible?

• So you've got your code? Is there any problem with it? Is it not working? Commented May 14 at 8:34
• I'm pretty sure all three use LAPACK/BLAS for machine real dense matrix operations. One can call (at least some) BLAS functions, if that is necessary. Commented May 14 at 16:58

I don't find that, so I implement one by myself.

The results are completely consistent:

Try the Mathematica code online!

(* Define the Blahut-Arimoto function *)
BlahutArimoto[distMat_, oldpX_, beta_, maxIt_: 40, eps_: 1*^-4] :=
Module[{l, lHat, pCond, pHat, Iu, Du, DuPrev, it, pX=oldpX},
l = Length[distMat];
lHat = Length[distMat[[1]]];
pCond = Transpose@KroneckerProduct[ConstantArray[1, {lHat, 1}], pX];

pX = pX / Total[pX]; (* Normalize *)
pCond = pCond/Total[pCond, {2}];

it = 0;
DuPrev = 0;
Du = 2*eps;
While[it < maxIt && Abs[Du - DuPrev] > eps,
it++;
DuPrev = Du;
If[MatrixQ[Part[pCond,1]],pCond=Flatten[pCond,1]];
pHat = pX.pCond;

matrixA=Exp[-beta*distMat];
pCond = Table[matrixA[[i, j]] * pHat[[j]], {i, 1, Length[matrixA]}, {j, 1, Length[pHat]}];
If[it>=1,pCond =Transpose@@pCond;];
pCond = pCond/Total[pCond, {2}];
pCond = {pCond};

If[it==1, Iu = Total[pX . Flatten[(pCond * (Log[pCond / ConstantArray[pHat, Length[pCond]]])),1]];];
If[it>=2, pCond =Flatten[pCond,1]; Iu = Total[pX . (pCond *Log[E, Table[pCond[[i, j]]/pHat[[j]], {i, 1, Length[Part[pCond, 1]]}, {j, 1, Length[pHat]}]])] ];
If[it==1, Du = Total[pX . Flatten[(pCond * distMat),1]];];
If[it>=2, pCond ={pCond}; Du = Total[pX . Flatten[(pCond * distMat),1]];];
];

{Iu/Log[2], Du}
]

(* Define auxiliary functions *)
hammingDist[x_, y_] := If[x != y, 1.0, 0.0]

quadDist[x_, y_] := (x - y)^2

binEnt[x_] := -x*Log2[x] - (1 - x)*Log2[1 - x]

gaussPDF[x_] := (1/(Sqrt[2 Pi])) Exp[-x^2/2]

(* Blahut-Arimoto example *)
BlahutArimotoExample[] := Module[{beta, xx, xxHat, al, pX, X, XHat, distMat, R, D},
beta = 0.3;

(* Example 1: Bernoulli input with Hamming distortion *)
xx = {0, 1}; (* binary input *)
xxHat = {0, 1}; (* binary reconstruction *)
al = 0.4; (* P(X=1) = al *)
pX = {1 - al, al};

(*Create meshgrid equivalent*)
XHat = Flatten[Table[xx[[i]], {i, Length[xx]}, {j, Length[xxHat]}]];
X = Flatten[Table[xxHat[[j]], {i, Length[xx]}, {j, Length[xxHat]}]];
(*Reshape to match meshgrid output format*)
X = ArrayReshape[X, {Length[xx], Length[xxHat]}];
XHat = ArrayReshape[XHat, {Length[xx], Length[xxHat]}];
distMat = Table[hammingDist[X[[i,j]], XHat[[i,j]]], {i, Length[X]}, {j,Length[Part[X,1]]}];
distMat = Partition[distMat, Length[xxHat]];

{R, D} = BlahutArimoto[distMat, pX, beta]; (* evaluate at beta = 0.3 *)

Print["Hamming Binary:"];
Print["at beta = ", beta, ": D = ", D, ", R = ", R];
Print["Difference between true R(D) (binary): ",
Abs[binEnt[al] - binEnt[D] - R]]; (* difference between true and estimated *)

(* Example 2: (truncated) Gaussian input with quadratic distortion *)
xx = Range[-5, 5, 0.01]; (* source alphabet *)
xxHat = Range[-5, 5, 0.01]; (* reconstruction alphabet *)
pX = gaussPDF /@ xx; (* source pdf *)

(*Create meshgrid equivalent*)
XHat = Flatten[Table[xx[[i]], {i, Length[xx]}, {j, Length[xxHat]}]];
X = Flatten[Table[xxHat[[j]], {i, Length[xx]}, {j, Length[xxHat]}]];
(*Reshape to match meshgrid output format*)
X = ArrayReshape[X, {Length[xx], Length[xxHat]}];
XHat = ArrayReshape[XHat, {Length[xx], Length[xxHat]}];
distMat = Table[quadDist[X[[i,j]], XHat[[i,j]]], {i, Length[X]}, {j,Length[Part[X,1]]}];
distMat = Partition[distMat, Length[xxHat]];

{R, D} = BlahutArimoto[distMat, pX, beta]; (* evaluate at beta = 0.3 *)

Print["at beta = ", beta, ": D = ", D, ", R = ", R];
Print["Difference between true R(D) (quadratic Gaussian): ",
Abs[D - 2^(-2 R)]]; (* difference between true and estimated *)
]

(* Run the example *)
Print["Starting Blahut-Arimoto example..."];
BlahutArimotoExample[];


Actually using Mathematica to deal with matrix manipulation is so tedious and slow. It's better to use Python numpy or MATLAB to do that.

Correspondent matlab code (also written by me) is as follows:

clear all;close all;clc;

% Helper functions
hamming_dist = @(x, y) (x ~= y) + 0.0;
quad_dist = @(x, y) (x - y) .^ 2;
bin_ent = @(x) -x .* log2(x) - (1 - x) .* log2(1 - x);
Gauss_pdf = @(x) 1 / (2 * pi) * exp(-x .^ 2 / 2);

% Parameters
beta = 0.3;

% Example 1: Bernoulli input with Hamming distortion
xx = [0, 1]; % Binary input
xx_hat = [0, 1]; % Binary reconstruction
al = 0.4;  % P(X=1) = al
p_x = [1 - al, al];
[X, X_hat] = meshgrid(xx, xx_hat);  % Create distortion matrix
dist_mat = hamming_dist(X, X_hat);
% Evaluate Blahut-Arimoto at beta = 0.3
[R, D] = BlahutArimoto(dist_mat, p_x, beta, 40, 1e-4);
% Check against true R(D)
fprintf('Hamming Binary:\n');
fprintf('At beta = %.2f: D = %.4f, R = %.4f\n', beta, D, R);
fprintf('Difference between true R(D) (binary): %.4f\n', ...
abs(bin_ent(al) - bin_ent(D) - R));

% Example 2: (truncated) Gaussian input with quadratic distortion
xx = linspace(-5, 5, 1000); % Source alphabet
xx_hat = linspace(-5, 5, 1000); % Reconstruction alphabet
p_x = Gauss_pdf(xx); % Source PDF
[X, X_hat] = meshgrid(xx, xx_hat);  % Create distortion matrix
% Evaluate Blahut-Arimoto at beta = 0.3
[R, D] = BlahutArimoto(dist_mat, p_x, beta, 40, 1e-6);
% Check against true R(D)
fprintf('At beta = %.2f: D = %.4f, R = %.4f\n', beta, D, R);
fprintf('Difference between true R(D) (quadratic Gaussian): %.4f\n', ...
abs(D - 2 ^ (-2 * R)));

function [Iu, Du] = BlahutArimoto(dist_mat, p_x, beta, max_it, eps)
% Blahut-Arimoto Algorithm for Rate-Distortion Function
%
% Inputs:
%   dist_mat - Distortion matrix (matrix)
%   p_x - Probability mass function of the source (row vector)
%   beta - Slope of the rate-distortion function (scalar)
%   max_it - Maximal number of iterations (int)
%   eps - Accuracy threshold (float)
%
% Outputs:
%   Iu - Rate (in bits)
%   Du - Distortion

l = size(dist_mat, 1);
l_hat = size(dist_mat, 2);

p_cond = repmat(p_x, l_hat, 1).';
p_x = p_x ./ sum(p_x);
p_cond = p_cond ./ sum(p_cond, 2);

it = 0;
Du_prev = 0;
Du = 2 * eps;

while (it < max_it) && (abs(Du - Du_prev) > eps)
it = it + 1;
Du_prev = Du;

p_hat = p_x * p_cond;
p_cond = exp(-beta * dist_mat) .* p_hat;
p_cond = p_cond ./ sum(p_cond, 2);

Iu = sum(p_x * (p_cond .* log(p_cond ./ p_hat)) );
Du = sum(p_x * (p_cond .* dist_mat));
%     fprintf('it=%d lu = %.10f Du = %.10f\n',it, Iu,Du);
end

Iu = Iu / log(2); % Convert to bits
end

• I'd recommend to use Subtract[X, XHat]^2 instead of Table[quadDist[X[[i, j]], XHat[[i, j]]], {i, Length[X]}, {j, Length[Part[X, 1]]}]. And Table[matrixA[[i, j]]*pHat[[j]], {i, 1, Length[matrixA]}, {j, 1, Length[pHat]}] is probably just a matrixA.pHat. If you did these things this way in Matlab or Python similarly, they would be equally slow. (Well, plain Python will probably be even slower...) Commented May 14 at 20:42
• XHat = Flatten[Table[xx[[i]], {i, Length[xx]}, {j, Length[xxHat]}]]; XHat = ArrayReshape[XHat, {Length[xx], Length[xxHat]}]; would probably become faster and better to read with Transpose[ConstantArray[xx, Length[xxHat]]]. Commented May 14 at 20:44