refactor: improve quantization classes and documentation

Author: rwnobrega

src/komm/_quantization/LloydMaxQuantizer.py

@@ -147,30 +147,26 @@ def lloyd_max_quantizer(
 ) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
     # See [Say06, eqs. (9.27) and (9.28)].
     x_min, x_max = input_range
-    delta = (x_max - x_min) / num_levels
 
     # Initial guess
-    levels = np.linspace(x_min + delta / 2, x_max - delta / 2, num=num_levels)
-    thresholds = np.empty(num_levels + 1, dtype=float)
-    new_levels = np.empty_like(levels)
+    delta = (x_max - x_min) / num_levels
+    y = np.linspace(x_min + delta / 2, x_max - delta / 2, num=num_levels)
 
-    for _ in range(max_iter):
-        thresholds[0] = x_min
-        thresholds[1:-1] = 0.5 * (levels[:-1] + levels[1:])
-        thresholds[-1] = x_max
+    λ = np.concatenate([[x_min], np.empty(num_levels - 1), [x_max]])
+    y_new = np.empty_like(y)
 
+    for _ in range(max_iter):
+        λ[1:-1] = 0.5 * (y[:-1] + y[1:])
         for i in range(num_levels):
-            left, right = thresholds[i], thresholds[i + 1]
-            x = np.linspace(left, right, num=points_per_interval, dtype=float)
+            x = np.linspace(λ[i], λ[i + 1], num=points_per_interval, dtype=float)
             pdf = input_pdf(x)
-            numerator = np.trapezoid(x * pdf, x)
             denominator = np.trapezoid(pdf, x)
             if denominator != 0:
-                new_levels[i] = numerator / denominator
+                y_new[i] = np.trapezoid(x * pdf, x) / denominator
             else:  # Keep old level
-                new_levels[i] = levels[i]
-        if np.allclose(levels, new_levels):
+                y_new[i] = y[i]
+        if np.allclose(y, y_new):
             break
-        levels = new_levels.copy()
+        y = y_new.copy()
 
-    return new_levels, thresholds[1:-1]
+    return y_new, λ[1:-1]

src/komm/_quantization/ScalarQuantizer.py

@@ -9,25 +9,25 @@
 
 class ScalarQuantizer(abc.ScalarQuantizer):
     r"""
-    General scalar quantizer. It is defined by a list of *levels*, $v_0, v_1, \ldots, v_{L-1}$, and a list of *thresholds*, $t_0, t_1, \ldots, t_L$, satisfying
+    General scalar quantizer. It is defined by a list of *levels*, $y_0, y_1, \ldots, y_{L-1}$, and a list of *thresholds*, $\lambda_0, \lambda_1, \ldots, \lambda_L$, satisfying
     $$
-        -\infty = t_0 < v_0 < t_1 < v_1 < \cdots < t_{L - 1} < v_{L - 1} < t_L = +\infty.
+        -\infty = \lambda_0 < y_0 < \lambda_1 < y_1 < \cdots < \lambda_{L - 1} < y_{L - 1} < \lambda_L = +\infty.
     $$
-    Given an input $x \in \mathbb{R}$, the output of the quantizer is given by $y = v_i$ if and only if $t_i \leq x < t_{i+1}$, where $i \in [0:L)$. For more details, see <cite>Say06, Ch. 9</cite>.
+    Given an input $x \in \mathbb{R}$, the output of the quantizer is given by $y = y_i$ if and only if $\lambda_i \leq x < \lambda_{i+1}$, where $i \in [0:L)$. For more details, see <cite>Say06, Ch. 9</cite>.
 
     Parameters:
-        levels: The quantizer levels $v_0, v_1, \ldots, v_{L-1}$. It should be a list floats of length $L$.
+        levels: The quantizer levels $y_0, y_1, \ldots, y_{L-1}$. It should be a list floats of length $L$.
 
-        thresholds: The quantizer finite thresholds $t_1, t_2, \ldots, t_{L-1}$. It should be a list of floats of length $L - 1$.
+        thresholds: The quantizer finite thresholds $\lambda_1, \lambda_2, \ldots, \lambda_{L-1}$. It should be a list of floats of length $L - 1$.
 
     Examples:
         The $5$-level scalar quantizer whose characteristic (input × output) curve is depicted in the figure below has levels
         $$
-            v_0 = -2, ~ v_1 = -1, ~ v_2 = 0, ~ v_3 = 1, ~ v_4 = 2,
+            y_0 = -2, ~ y_1 = -1, ~ y_2 = 0, ~ y_3 = 1, ~ y_4 = 2,
         $$
         and thresholds
         $$
-            t_0 = -\infty, ~ t_1 = -1.5, ~ t_2 = -0.3, ~ t_3 = 0.8, ~ t_4 = 1.4, ~ t_5 = \infty.
+            \lambda_0 = -\infty, ~ \lambda_1 = -1.5, ~ \lambda_2 = -0.3, ~ \lambda_3 = 0.8, ~ \lambda_4 = 1.4, ~ \lambda_5 = \infty.
         $$
 
         <figure markdown>

src/komm/_quantization/UniformQuantizer.py

@@ -11,11 +11,11 @@ class UniformQuantizer(abc.ScalarQuantizer):
     r"""
     Uniform scalar quantizer. It is a [scalar quantizer](/ref/ScalarQuantizer) in which the separation between levels is a constant $\Delta$, called the *quantization step*, and the thresholds are the mid-point between adjacent levels. More precisely, the levels are given by
     $$
-        v_i = (i - (L - 1)/2 + \theta) \Delta, \qquad i \in [0 : L),
+        y_i = (i - (L - 1)/2 + \theta) \Delta, \qquad i \in [0 : L),
     $$
     where $\theta \in \mathbb{R}$ is an arbitrary *offset* (normalized by $\Delta$), and the finite thresholds are given by
     $$
-        t_i = \frac{v_{i-1} + v_i}{2}, \qquad i \in [1 : L).
+        \lambda_i = \frac{y_{i-1} + y_i}{2}, \qquad i \in [1 : L).
     $$
     For more details, see <cite>Say06, Sec. 9.4</cite>.
 

src/komm/_quantization/base.py

@@ -11,15 +11,15 @@ class ScalarQuantizer(ABC):
     @abstractmethod
     def levels(self) -> npt.NDArray[np.floating]:
         r"""
-        The quantizer levels $v_0, v_1, \ldots, v_{L-1}$.
+        The quantizer levels $y_0, y_1, \ldots, y_{L-1}$.
         """
         raise NotImplementedError
 
     @cached_property
     @abstractmethod
     def thresholds(self) -> npt.NDArray[np.floating]:
         r"""
-        The quantizer finite thresholds $t_1, t_2, \ldots, t_{L-1}$.
+        The quantizer finite thresholds $\lambda_1, \lambda_2, \ldots, \lambda_{L-1}$.
         """
         raise NotImplementedError
 
@@ -68,20 +68,17 @@ def mean_squared_error(
         Parameters:
             input_pdf: The pdf $f_X(x)$ of the input signal.
             input_range: The range $(x_\mathrm{min}, x_\mathrm{max})$ of the input signal.
-            points_per_interval: The number of points per interval for numerical integration (default: 4096).
+            points_per_interval: The number of points per interval for numerical integration (default: `4096`).
 
         Returns:
             mse: The mean square quantization error.
         """
         # See [Say06, eq. (9.3)].
         x_min, x_max = input_range
-        thresholds = np.concatenate(([x_min], self.thresholds, [x_max]))
+        λ = np.concatenate(([x_min], self.thresholds, [x_max]))
         mse = 0.0
         for i, level in enumerate(self.levels):
-            left, right = thresholds[i], thresholds[i + 1]
-            x = np.linspace(left, right, num=points_per_interval, dtype=float)
-            pdf = input_pdf(x)
-            integrand: npt.NDArray[np.floating] = (level - x) ** 2 * pdf
-            integral = np.trapezoid(integrand, x)
-            mse += float(integral)
-        return mse
+            x = np.linspace(λ[i], λ[i + 1], num=points_per_interval, dtype=float)
+            integrand = (level - x) ** 2 * input_pdf(x)
+            mse += np.trapezoid(integrand, x)
+        return float(mse)