87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
|
|
class ProximalDecoder:
|
|
"""Class implementing the Proximal Decoding algorithm. See "Proximal Decoding for LDPC Codes" by Tadashi
|
|
Wadayama, and Satoshi Takabe.
|
|
"""
|
|
def __init__(self, H: np.array, K: int = 10, step_size: float = 0.01, gamma: float = 0.05):
|
|
"""Construct a new ProximalDecoder Object.
|
|
|
|
:param H: Parity Check Matrix
|
|
:param K: Max number of iterations to perform when decoding
|
|
:param step_size: Step size for the gradient descent process
|
|
:param gamma: Positive constant. Arises in the approximation of the prior PDF
|
|
"""
|
|
self._H = H
|
|
self._K = K
|
|
self._step_size = step_size
|
|
self._gamma = gamma
|
|
|
|
@staticmethod
|
|
def _L_awgn(s: np.array, y: np.array) -> np.array:
|
|
"""Variation of the negative log-likelihood for the special case of AWGN noise. See 4.1, p. 4."""
|
|
return s - y
|
|
|
|
def _grad_h(self, x: np.array) -> np.array:
|
|
"""Gradient of the code-constraint polynomial. See 2.3, p. 2."""
|
|
# Calculate first term
|
|
|
|
result = 4 * (x**2 - 1) * x
|
|
|
|
# Calculate second term
|
|
|
|
for k, x_k in enumerate(x):
|
|
# TODO: Perform this operation for each row simultaneously
|
|
B_k = np.argwhere(self._H[:, k] == 1)
|
|
B_k = B_k[:, 0] # Get rid of one layer of arrays
|
|
|
|
# TODO: Perform the summation with np.sum()
|
|
sum_result = 0
|
|
for i in B_k:
|
|
# TODO: Perform this operation for each column simultaneously
|
|
A_i = np.argwhere(self._H[i] == 1)
|
|
A_i = A_i[:, 0] # Get rid of one layer of arrays
|
|
|
|
prod = 1
|
|
for j in A_i:
|
|
prod *= x[j]
|
|
|
|
sum_result += prod**2 - prod
|
|
|
|
term_2 = 2 / x_k * sum_result
|
|
|
|
result[k] += term_2
|
|
|
|
return np.array(result)
|
|
|
|
def _check_parity(self, y_hat: np.array) -> bool:
|
|
"""Perform a parity check for a given codeword.
|
|
|
|
:param y_hat: codeword to be checked
|
|
:return: True if the parity check passes, i.e. the codeword is valid. False otherwise
|
|
"""
|
|
syndrome = np.dot(self._H, y_hat) % 2
|
|
return not np.any(syndrome)
|
|
|
|
def decode(self, y: np.array) -> np.array:
|
|
"""Decode a received signal. The algorithm is detailed in 3.2, p.3.
|
|
|
|
This function assumes an AWGN channel.
|
|
|
|
:param y: Vector of received values
|
|
:return: Most probably sent symbol
|
|
"""
|
|
s = 0
|
|
x_hat = 0
|
|
for k in range(self._K):
|
|
r = s - self._step_size * self._L_awgn(s, y)
|
|
s = r - self._gamma * self._grad_h(r)
|
|
x_hat = (np.sign(s) == 1) * 1
|
|
|
|
if self._check_parity(x_hat):
|
|
break
|
|
|
|
return x_hat
|