modify dtw_merge

Author: takenori-y

diffsptk/functional.py

@@ -572,36 +572,27 @@ def dtw(
     )
 
 
-def dtw_merge(x: Tensor, y: Tensor, indices: Tensor) -> tuple[Tensor, Tensor]:
-    """Align two vector sequences according to the given path.
+def dtw_merge(x: Tensor, y: Tensor, indices: Tensor) -> Tensor:
+    """Merge two sequences according to the given indices.
 
     Parameters
     ----------
-    x : Tensor [shape=(T1, ...)]
+    x : Tensor [shape=(T1, D) or (T1,)]
         The query vector sequence.
 
-    y : Tensor [shape=(T2, ...)]
+    y : Tensor [shape=(T2, D) or (T2,)]
         The reference vector sequence.
 
     indices : Tensor [shape=(T, 2)]
-        The indices of the path.
+        The indices of the viterbi path.
 
     Returns
     -------
-    x_align : Tensor [shape=(T, ...)]
-        The aligned query vector sequence.
-
-    y_align : Tensor [shape=(T, ...)]
-        The aligned reference vector sequence.
+    z : Tensor [shape=(T, 2D) or (T, 2)]
+        The merged 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
+    return nn.DynamicTimeWarping.merge(x=x, y=y, indices=indices)
 
 
 def entropy(p: Tensor, out_format: str = "nat") -> Tensor:

diffsptk/modules/dtw.py

@@ -327,3 +327,55 @@ def _forward(
         if return_indices:
             return distance, indices
         return distance
+
+    @staticmethod
+    def merge(
+        x: torch.Tensor,
+        y: torch.Tensor,
+        indices: torch.Tensor,
+    ) -> torch.Tensor:
+        """Merge two sequences according to the given indices.
+
+        Parameters
+        ----------
+        x : Tensor [shape=(T1, D) or (T1,)]
+            The query vector sequence.
+
+        y : Tensor [shape=(T2, D) or (T2,)]
+            The reference vector sequence.
+
+        indices : Tensor [shape=(T, 2)]
+            The indices of the viterbi path.
+
+        Returns
+        -------
+        z : Tensor [shape=(T, 2D) or (T, 2)]
+            The merged vector sequence.
+
+        Examples
+        --------
+        >>> import diffsptk
+        >>> dtw = diffsptk.DynamicTimeWarping(p=1)
+        >>> x = torch.tensor([1., 3., 6., 9.])
+        >>> y = torch.tensor([2., 3., 8., 8.])
+        >>> _, indices = dtw(x, y, return_indices=True)
+        >>> z = dtw.merge(x, y, indices[0])
+        >>> z
+        tensor([[1., 2.],
+                [3., 3.],
+                [6., 8.],
+                [9., 8.],
+                [9., 8.]])
+
+        """
+        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_expanded = x[indices[:, 0]]
+        y_expanded = y[indices[:, 1]]
+        if x.dim() == 1:
+            z = torch.stack([x_expanded, y_expanded], dim=-1)
+        else:
+            z = torch.cat([x_expanded, y_expanded], dim=-1)
+        return z

tests/test_dtw.py

@@ -45,7 +45,7 @@ def _dtw(x, y):
         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)
+        return diffsptk.functional.dtw_merge(x, y, indices[0])
 
     U.check_compatibility(
         device,
@@ -75,3 +75,13 @@ def test_various_shape(T=10):
             [(1, T, 1), (1, T, 1)],
         ],
     )
+
+    indices = torch.stack([torch.arange(T), torch.arange(T)], dim=-1)
+    dtw_merge = lambda x, y: diffsptk.functional.dtw_merge(x, y, indices)
+    U.check_various_shape(
+        dtw_merge,
+        [
+            [(T,), (T,)],
+            [(T, 1), (T, 1)],
+        ],
+    )