add lbg

.coveragerc

@@ -1,6 +1,7 @@
 [report]
 exclude_lines =
     pragma: no cover
+    if self.verbose:
     raise NotImplementedError
     raise RuntimeError
     raise ValueError

diffsptk/core/__init__.py

@@ -37,6 +37,8 @@
 from .iulaw import MuLawExpansion
 from .ivq import InverseVectorQuantization
 from .lar2par import LogAreaRatioToParcorCoefficients
+from .lbg import LindeBuzoGrayAlgorithm
+from .lbg import LindeBuzoGrayAlgorithm as KMeans
 from .levdur import PseudoLevinsonDurbinRecursion
 from .levdur import PseudoLevinsonDurbinRecursion as LevinsonDurbinRecursion
 from .linear_intpl import LinearInterpolation

diffsptk/core/lbg.py

@@ -0,0 +1,186 @@
+# ------------------------------------------------------------------------ #
+# Copyright 2022 SPTK Working Group                                        #
+#                                                                          #
+# Licensed under the Apache License, Version 2.0 (the "License");          #
+# you may not use this file except in compliance with the License.         #
+# You may obtain a copy of the License at                                  #
+#                                                                          #
+#     http://www.apache.org/licenses/LICENSE-2.0                           #
+#                                                                          #
+# Unless required by applicable law or agreed to in writing, software      #
+# distributed under the License is distributed on an "AS IS" BASIS,        #
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. #
+# See the License for the specific language governing permissions and      #
+# limitations under the License.                                           #
+# ------------------------------------------------------------------------ #
+
+import logging
+
+import torch
+import torch.nn as nn
+
+from ..misc.utils import check_size
+from ..misc.utils import is_power_of_two
+from .vq import VectorQuantization
+
+
+class LindeBuzoGrayAlgorithm(nn.Module):
+    """See `this page <https://sp-nitech.github.io/sptk/latest/main/lbg.html>`_
+    for details. This module is not differentiable.
+
+    order : int >= 0 [scalar]
+        Order of vector.
+
+    codebook_size : int >= 1 [scalar]
+        Target codebook size, must be power of two.
+
+    min_data_per_cluster : int >= 1 [scalar]
+        Minimum number of data points in a cluster.
+
+    n_iter : int >= 1 [scalar]
+        Number of iterations.
+
+    eps : float >= 0 [scalar]
+        Convergence threshold.
+
+    perturb_factor : float > 0 [scalar]
+        Perturbation factor.
+
+    verbose : bool [scalar]
+        If True, print progress.
+
+    """
+
+    def __init__(
+        self,
+        order,
+        codebook_size,
+        min_data_per_cluster=1,
+        n_iter=100,
+        eps=1e-5,
+        perturb_factor=1e-5,
+        verbose=False,
+    ):
+        super(LindeBuzoGrayAlgorithm, self).__init__()
+
+        self.order = order
+        self.codebook_size = codebook_size
+        self.min_data_per_cluster = min_data_per_cluster
+        self.n_iter = n_iter
+        self.eps = eps
+        self.perturb_factor = perturb_factor
+        self.verbose = verbose
+
+        assert 0 <= self.order
+        assert is_power_of_two(self.codebook_size)
+        assert 1 <= self.min_data_per_cluster
+        assert 1 <= self.n_iter
+        assert 0 <= self.eps
+        assert 0 < self.perturb_factor
+
+        self.vq = VectorQuantization(order, codebook_size).eval()
+        self.vq.codebook[:] = 1e10
+
+        if self.verbose:
+            self.logger = logging.getLogger("lbg")
+            self.logger.setLevel(logging.INFO)
+            formatter = logging.Formatter(
+                "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
+            )
+            handler = logging.StreamHandler()
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+
+    def forward(self, x):
+        """Design a codebook.
+
+        Parameters
+        ----------
+        x : Tensor [shape=(..., M+1)]
+            Input vectors.
+
+        Returns
+        -------
+        codebook : Tensor [shape=(K, M+1)]
+            Codebook.
+
+        distance : Tensor [scalar]
+            Distance.
+
+        Examples
+        --------
+        >>> x = torch.randn(10, 1)
+        >>> lbg = diffsptk.KMeans(0, 2)
+        >>> codebook, distance = lbg(x)
+        >>> codebook
+        tensor([[-1.3889],
+                [ 0.6275]])
+        >>> distance
+        tensor(0.2489)
+
+        """
+        check_size(x.size(-1), self.order + 1, "dimension of input")
+
+        # Initalize codebook.
+        x = x.view(-1, x.size(-1))
+        mean = x.mean(0)
+        self.vq.codebook[0] = mean
+        distance = torch.inf
+
+        curr_codebook_size = 1
+        next_codebook_size = 2
+        while next_codebook_size <= self.codebook_size:
+            # Double codebook.
+            codebook = self.vq.codebook[:curr_codebook_size]
+            r = torch.randn_like(codebook) * self.perturb_factor
+            self.vq.codebook[curr_codebook_size:next_codebook_size] = codebook - r
+            self.vq.codebook[:curr_codebook_size] += r
+            curr_codebook_size = next_codebook_size
+            next_codebook_size *= 2
+
+            prev_distance = distance  # Suppress flake8 warnings.
+            for n in range(self.n_iter):
+                # E-step: evaluate model.
+                xq, indices, _ = self.vq(x)
+                distance = (x - xq).square().sum()
+                distance /= x.size(0)
+                if self.verbose:
+                    self.logger.info(f"K={curr_codebook_size} {n:5d}: {distance:g}")
+
+                # Check convergence.
+                if distance == 0:
+                    break
+                if 0 < n and (prev_distance - distance).abs() / distance < self.eps:
+                    break
+                prev_distance = distance
+
+                # Get number of data points for each cluster.
+                n_data = torch.histc(
+                    indices.float(),
+                    bins=curr_codebook_size,
+                    min=0,
+                    max=curr_codebook_size - 1,
+                )
+                mask = self.min_data_per_cluster <= n_data
+
+                # M-step: update centroids.
+                centroids = torch.zeros(
+                    (curr_codebook_size, self.order + 1), dtype=x.dtype, device=x.device
+                )
+                idx = indices.unsqueeze(1).repeat_interleave(self.order + 1, dim=1)
+                centroids.scatter_add_(0, idx, x)
+                centroids[mask] /= n_data[mask].unsqueeze(1)
+
+                if torch.any(~mask):
+                    # Get index of largest cluster.
+                    _, m = n_data.max(0)
+                    copied_centroids = centroids[m : m + 1].repeat_interleave(
+                        (~mask).sum(), dim=0
+                    )
+                    r = torch.randn_like(copied_centroids) * self.perturb_factor
+                    centroids[~mask] = copied_centroids - r
+                    centroids[m] += r.mean(0)
+
+                self.vq.codebook[:curr_codebook_size] = centroids
+
+        return self.vq.codebook, distance

docs/core/ivq.rst

@@ -6,4 +6,4 @@ ivq
 .. autoclass:: diffsptk.InverseVectorQuantization
    :members:
 
-.. seealso:: :ref:`vq`
+.. seealso:: :ref:`vq`  :ref:`lbg`

docs/core/lbg.rst

@@ -0,0 +1,11 @@
+.. _lbg:
+
+lbg
+---
+
+.. autoclass:: diffsptk.KMeans
+
+.. autoclass:: diffsptk.LindeBuzoGrayAlgorithm
+   :members:
+
+.. seealso:: :ref:`vq`  :ref:`ivq`

docs/core/vq.rst

@@ -6,4 +6,4 @@ vq
 .. autoclass:: diffsptk.VectorQuantization
    :members:
 
-.. seealso:: :ref:`ivq`
+.. seealso:: :ref:`ivq`  :ref:`lbg`

tests/test_lbg.py

@@ -0,0 +1,54 @@
+# ------------------------------------------------------------------------ #
+# Copyright 2022 SPTK Working Group                                        #
+#                                                                          #
+# Licensed under the Apache License, Version 2.0 (the "License");          #
+# you may not use this file except in compliance with the License.         #
+# You may obtain a copy of the License at                                  #
+#                                                                          #
+#     http://www.apache.org/licenses/LICENSE-2.0                           #
+#                                                                          #
+# Unless required by applicable law or agreed to in writing, software      #
+# distributed under the License is distributed on an "AS IS" BASIS,        #
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. #
+# See the License for the specific language governing permissions and      #
+# limitations under the License.                                           #
+# ------------------------------------------------------------------------ #
+
+import pytest
+import torch
+
+import diffsptk
+import tests.utils as U
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_compatibility(device, M=1, K=4, B=10, n_iter=10):
+    lbg = diffsptk.LindeBuzoGrayAlgorithm(M, K, n_iter=n_iter)
+
+    torch.manual_seed(1234)
+    torch.cuda.manual_seed(1234)
+
+    tmp1 = "lbg.tmp1"
+    tmp2 = "lbg.tmp2"
+    tmp3 = "lbg.tmp3"
+    tmp4 = "lbg.tmp4"
+    tmp5 = "lbg.tmp5"
+    U.check_compatibility(
+        device,
+        [lambda x: x[1], lbg],
+        [
+            f"nrand -u +2 -l {B*(M+1)} -s 1 > {tmp1}",
+            f"nrand -u -2 -l {B*(M+1)} -s 2 > {tmp2}",
+            f"nrand -u +4 -l {B*(M+1)} -s 3 > {tmp3}",
+            f"nrand -u -4 -l {B*(M+1)} -s 4 > {tmp4}",
+        ],
+        f"cat {tmp1} {tmp2} {tmp3} {tmp4}",
+        (
+            f"cat {tmp1} {tmp2} {tmp3} {tmp4} | "
+            f"lbg -m {M} -e {K} -i {n_iter} -s 1234 -S {tmp5} > /dev/null; "
+            f"cat {tmp5}"
+        ),
+        [f"rm {tmp1} {tmp2} {tmp3} {tmp4} {tmp5}"],
+        dx=M + 1,
+        rtol=0.1,
+    )

tests/utils.py

@@ -80,6 +80,7 @@ def check_compatibility(
     opt={},
     sr=None,
     verbose=False,
+    **kwargs,
 ):
     if device == "cuda" and not torch.cuda.is_available():
         return
@@ -131,9 +132,9 @@ def check_compatibility(
         print(f"Target: {y}")
 
     if eq is None:
-        assert allclose(y_hat, y), f"Output: {y_hat}\nTarget: {y}"
+        assert allclose(y_hat, y, **kwargs), f"Output: {y_hat}\nTarget: {y}"
     else:
-        assert eq(y_hat, y), f"Output: {y_hat}\nTarget: {y}"
+        assert eq(y_hat, y, **kwargs), f"Output: {y_hat}\nTarget: {y}"
 
 
 def check_differentiable(device, modules, shapes, opt={}, load=1):