Merge pull request #149 from sp-nitech/dtw_merge

Add dtw_merge

Author: takenori-y

diffsptk/functional.py

@@ -571,6 +571,38 @@ def dtw(
     )
 
 
+def dtw_merge(x: Tensor, y: Tensor, indices: Tensor) -> tuple[Tensor, Tensor]:
+    """Align two vector sequences according to the given path.
+
+    Parameters
+    ----------
+    x : Tensor [shape=(T1, ...)]
+        The query vector sequence.
+
+    y : Tensor [shape=(T2, ...)]
+        The reference vector sequence.
+
+    indices : Tensor [shape=(T, 2)]
+        The indices of the path.
+
+    Returns
+    -------
+    x_align : Tensor [shape=(T, ...)]
+        The aligned query vector sequence.
+
+    y_align : Tensor [shape=(T, ...)]
+        The aligned reference vector sequence.
+
+    """
+    if x.dim() != y.dim():
+        raise ValueError("x and y must have the same number of dimensions.")
+    if indices.dim() != 2 or indices.size(-1) != 2:
+        raise ValueError("The shape of indices must be (T, 2).")
+    x_align = x[indices[:, 0]]
+    y_align = y[indices[:, 1]]
+    return x_align, y_align
+
+
 def entropy(p: Tensor, out_format: str = "nat") -> Tensor:
     """Calculate the entropy of a probability distribution.
 

diffsptk/modules/dtw.py

@@ -32,8 +32,8 @@ def _soft_dtw_core(
     B, T1, T2 = D.shape
 
     R = torch.full_like(D, float("inf"))
-    R[:, 0, 0] = D[:, 0, 0]
     R_ = R.clone() if has_two_step_transition else None
+    R[:, 0, 0] = D[:, 0, 0]
 
     if return_indices:
         P = torch.full((B, T1, T2, 2), -1, device=D.device, dtype=torch.long)
@@ -300,14 +300,13 @@ def _forward(
         dist_func: Callable,
         dtw_func: Callable,
     ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
-        d = x.dim()
-        if d == 1:
+        if x.dim() == 1:
             x = x.view(1, -1, 1)
             y = y.view(1, -1, 1)
-        elif d == 2:
+        elif x.dim() == 2:
             x = x.unsqueeze(0)
             y = y.unsqueeze(0)
-        else:
+        if x.dim() != 3:
             raise ValueError("x and y must be 1D, 2D, or 3D tensor.")
         if x.dim() != y.dim():
             raise ValueError("x and y must have the same number of dimensions.")

docs/source/modules/dtw.rst

@@ -9,3 +9,7 @@ dtw
     :members:
 
 .. autofunction:: diffsptk.functional.dtw
+
+.. seealso::
+
+    :ref:`dtw_merge`

docs/source/modules/dtw_merge.rst

@@ -0,0 +1,10 @@
+.. _dtw_merge:
+
+dtw_merge
+=========
+
+.. autofunction:: diffsptk.functional.dtw_merge
+
+.. seealso::
+
+    :ref:`dtw`

tests/test_dtw.py

@@ -14,7 +14,9 @@
 # limitations under the License.                                           #
 # ------------------------------------------------------------------------ #
 
+import numpy as np
 import pytest
+import torch
 
 import diffsptk
 import tests.utils as U
@@ -23,7 +25,7 @@
 @pytest.mark.parametrize("module", [False, True])
 @pytest.mark.parametrize("metric", [0, 1, 2, 3])
 @pytest.mark.parametrize("p", [0, 1, 2, 3, 4, 5, 6])
-def test_compatibility(device, dtype, module, metric, p, M=0, T1=10, T2=10):
+def test_compatibility(device, dtype, module, metric, p, M=2, T1=10, T2=10):
     dtw = U.choice(
         module,
         diffsptk.DynamicTimeWarping,
@@ -35,40 +37,41 @@ def test_compatibility(device, dtype, module, metric, p, M=0, T1=10, T2=10):
     tmp1 = "dtw.tmp1"
     tmp2 = "dtw.tmp2"
     tmp3 = "dtw.tmp3"
+
+    def _dtw(x, y):
+        distance, indices = dtw(x, y, return_indices=True)
+        output_distance = distance.item()
+        target_distance = np.fromfile(tmp3, dtype=np.float64)
+        assert np.allclose(output_distance, target_distance), (
+            f"Output: {output_distance}\nTarget: {target_distance}"
+        )
+        return torch.cat(diffsptk.functional.dtw_merge(x, y, indices[0]), dim=-1)
+
     U.check_compatibility(
         device,
         dtype,
-        dtw,
+        _dtw,
         [
             f"nrand -s 1 -l {T1 * (M + 1)} {sopr} > {tmp1}",
             f"nrand -s 2 -l {T2 * (M + 1)} {sopr} > {tmp2}",
         ],
         [f"cat {tmp1}", f"cat {tmp2}"],
-        (
-            f"dtw -m {M} -d {metric} -p {p} < {tmp1} {tmp2} -S {tmp3} > /dev/null "
-            f"&& cat {tmp3}"
-        ),
+        (f"dtw -m {M} -d {metric} -p {p} < {tmp1} {tmp2} -S {tmp3}"),
         [f"rm {tmp1} {tmp2} {tmp3}"],
         dx=M + 1,
+        dy=2 * (M + 1),
     )
 
-    U.check_compatibility(
-        device,
-        dtype,
+    U.check_differentiability(device, dtype, dtw, [(T1, M + 1), (T2, M + 1)])
+
+
+def test_various_shape(T=10):
+    dtw = diffsptk.DynamicTimeWarping()
+    U.check_various_shape(
         dtw,
         [
-            f"nrand -s 1 -l {T1} {sopr} > {tmp1}",
-            f"nrand -s 2 -l {T2} {sopr} > {tmp2}",
+            [(T,), (T,)],
+            [(T, 1), (T, 1)],
+            [(1, T, 1), (1, T, 1)],
         ],
-        [f"cat {tmp1}", f"cat {tmp2}"],
-        (
-            f"dtw -m 0 -d {metric} -p {p} < {tmp1} {tmp2} -P {tmp3} > /dev/null "
-            f"&& x2x +id {tmp3}"
-        ),
-        [f"rm {tmp1} {tmp2} {tmp3}"],
-        dy=2,
-        opt={"return_indices": True},
-        get=[1, 0],
     )
-
-    U.check_differentiability(device, dtype, dtw, [(T1, M + 1), (T2, M + 1)])

tests/utils.py

@@ -148,9 +148,6 @@ def check_compatibility(
     if dy is not None:
         y = y.reshape(-1, dy)
 
-    for cmd in teardown:
-        call(cmd, get=False)
-
     module = compose(*modules)
     if len(key) == 0:
         y_hat = module(*x, **opt)
@@ -180,6 +177,9 @@ def check_compatibility(
     else:
         assert eq(y_hat, y, **kwargs), f"Output: {y_hat}\nTarget: {y}"
 
+    for cmd in teardown:
+        call(cmd, get=False)
+
 
 def check_confidence(
     device,
@@ -256,12 +256,18 @@ def check_differentiability(
 
 
 def check_various_shape(module, shapes, *, preprocess=None):
-    x = torch.randn(*shapes[0])
+    if is_array(shapes[0][0]):
+        xs = [torch.randn(*shape) for shape in shapes[0]]
+    else:
+        xs = [torch.randn(*shapes[0])]
     if preprocess is not None:
-        x = preprocess(x)
+        xs = [preprocess(x) for x in xs]
     for i, shape in enumerate(shapes):
-        x = x.view(shape)
-        y = module(x).view(-1)
+        if is_array(shapes[0][0]):
+            x = [x.view(*shape) for x, shape in zip(xs, shape)]
+        else:
+            x = [xs[0].view(shape)]
+        y = module(*x).view(-1)
         if i == 0:
             target = y
         else: