Use Mathematica to Implement Blahut–Arimoto Algorithm (Part: Algorithm for Channel Capacity)

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

BlahutArimoto[pYx_, logBase_: 2, thresh_: 1*^-12, maxIter_: 1000] := Module[
  {
    m = Length[pYx],
    n = Length[pYx[[1]]],
    q, r1, tolerance, iteration, c
  },

  (* Input test *)
  If[Abs[Mean[Total[pYx, {2}]] - 1] >= 1*^-6, Return[$Failed]];
  If[m <= 1, Return[$Failed]];

  r = ConstantArray[1/m, m];

  (* Compute the r(x) that maximizes the capacity *)
  For[iteration = 1, iteration <= maxIter, iteration++,
    q = Transpose[r]pYx;
    q = Transpose[Transpose[q]/Total[q, {1}]];
    r1 = Table[Product[(q[[i, j]] + 1*^-12)^pYx[[i, j]], {j, n}], {i, m}];
    r1 = r1 / Total[r1];
    tolerance = Norm[r1 - r];
    r = r1;
    If[tolerance < thresh, Break[]]
  ];

  (* Calculate the capacity *)
  c = 0;
  For[i = 1, i <= m, i++,
    If[r[[i]] > 0,
      c += Total[
        r[[i]] * pYx[[i]] * Log[q[[i]] / r[[i]] + 1*^-16]
      ]
    ]
  ];
  c = c / Log[logBase];
  {c, r}
]

It has passed these test cases.

(*Binary symmetric channel*)
e = 0.2;
{Capacity, r} = BlahutArimoto[{{1 - e, e}, {e, 1 - e}}];
Print["Capacity: ", Capacity];
Print["The prior: ", r];

HPe = -(e*Log[2, e] + (1 - e)*Log[2, 1 - e]);
Print["Anatliyic capacity: ", 1 - HPe];

(*Symmetric channel*)
vec = {0.233, 0.267, 0.233, 0.267};
symbolsSend = Length[vec];
On[Assert];
Assert[Total[vec] == 1]
pYx = Table[RotateRight[vec, offset], {offset, 1, symbolsRecv}];
{Capacity, r} = BlahutArimoto[pYx];
Print["Capacity: ", Capacity // N];
Print["The prior: ", r];

Hvec = -Total[Map[#*Log[2, #] &, vec]];
Print["Anatliyic capacity: ", Log[2, Length[vec]] - Hvec // N];

(* The q-ary symmetric channel (q-SC) with error probability p *)
m = 4;
q = 2^m;
p = 1/15;
vec = {1 - p}~Join~ConstantArray[p/(q - 1), {q - 1}];
pYx = Table[RotateRight[vec, offset], {offset, 1, Length[vec]}];
{Capacity, r} = BlahutArimoto[pYx];
Print["Capacity: ", Capacity // N];
Print["The prior: ", r];

Hp = -p*Log[2, p] - (1 - p)*Log[2, 1 - p];
Print["Anatliyic capacity: ", m - Hp - p*Log[2, q - 1] // N];

(*Erasure channel*)
e = 0.1;
{Capacity, r} = BlahutArimoto[{{1 - e, e, 0}, {0, e, 1 - e}}];
Print["Capacity: ", Capacity];
Print["The prior: ", r];

Print["Anatliyic capacity: ", 1 - e];

(*Z channel*)
p = 0.1234;
{Capacity, r} = BlahutArimoto[{{1, 0}, {p, 1 - p}}];
Print["Capacity: ", Capacity];
Print["The prior: ", r];

Hp = -p*Log[2, p] - (1 - p)*Log[2, 1 - p];
sp = Hp/(1 - p);
Print["Anatliyic capacity: ", Log[2, 1 + 2^(-sp)]];

(* Quasi-symmetric channel *)
e = 0.12;
p = 0.23;
vec1 = {1 - p - e, p - e};
vec2 = {e, e, 0, 0};
symbolsSend = Max[Length@vec1, Length@vec2];
(A1 = Table[RotateRight[vec1, offset], {offset, 1, symbolsSend}]) // 
  MatrixForm;
(A2 = Table[RotateRight[vec2, offset], {offset, 1, symbolsSend}]) // 
  MatrixForm;
On[Assert];
Assert[Length@*Union@Total[A1, {1}] == 1]
Assert[Length@*Union@Total[A1, {2}] == 1]
Assert[Length@*Union@Total[A2, {1}] == 1]
Assert[Length@*Union@Total[A2, {2}] == 1]
(pYx = ArrayFlatten[{{A1, A2}}]) // MatrixForm
{Capacity, r} = BlahutArimoto[pYx];
Print["Capacity: ", Capacity // N];
Print["The prior: ", r];


n1 = Total@A1[[1]];
m1 = Total[Transpose[A1][[1]]];
n2 = Total@A2[[1]];
m2 = Total[Transpose[A2][[1]]];
AnatliyicCapacity = 
  Log[2, symbolsSend] + 
   Total[Map[#*Log[2, (# + 1*^-16)] &, vec1~Join~vec2]] - 
   n1*Log[2, (m1 + 1*^-16)] - n2*Log[2, (m2 + 1*^-16)];

Print["Anatliyic capacity: ", AnatliyicCapacity // N];