Add native complex conversion

Author: mpariente

asteroid_filterbanks/transforms.py

@@ -324,6 +324,36 @@ def from_torchaudio(tensor, dim: int = -2):
     return torch.cat(torch.chunk(tensor, 2, dim=-1), dim=dim).squeeze(-1)
 
 
+def to_torch_complex(tensor, dim: int = -2):
+    """Converts complex-like torch tensor to native PyTorch complex tensor.
+
+    Args:
+        tensor (torch.tensor): asteroid-style complex-like torch tensor.
+        dim(int, optional): the frequency (or equivalent) dimension along which
+            real and imaginary values are concatenated.
+
+    Returns:
+        :class:`torch.Tensor`:
+            Pytorch native complex-like torch tensor.
+    """
+    return torch.view_as_complex(to_torchaudio(tensor, dim=dim))
+
+
+def from_torch_complex(tensor, dim: int = -2):
+    """Converts Pytorch native complex tensor to complex-like torch tensor.
+
+    Args:
+        tensor (torch.tensor): PyTorch native complex-like torch tensor.
+        dim(int, optional): the frequency (or equivalent) dimension along which
+            real and imaginary values are concatenated.
+
+    Returns:
+        :class:`torch.Tensor`:
+            asteroid-style complex-like torch tensor.
+    """
+    return torch.cat([tensor.real, tensor.imag], dim=dim)
+
+
 def angle(tensor, dim: int = -2):
     """Return the angle of the complex-like torch tensor.
 

tests/transforms_test.py

@@ -190,6 +190,19 @@ def test_torchaudio_format(dim, max_tested_ndim):
     assert ta_tensor.shape[-1] == 2
 
 
+@pytest.mark.parametrize("dim", [0, 1, 2, 3, -1, -2, -3])
+@pytest.mark.parametrize("max_tested_ndim", [4, 5])
+def test_torch_complex_format(dim, max_tested_ndim):
+    # Random tensor shape
+    tensor_shape = [random.randint(1, 10) for _ in range(max_tested_ndim)]
+    # Make sure complex dimension has even shape
+    tensor_shape[dim] = 2 * tensor_shape[dim]
+    complex_tensor = torch.randn(tensor_shape)
+    ta_tensor = transforms.to_torch_complex(complex_tensor, dim=dim)
+    tensor_back = transforms.from_torch_complex(ta_tensor, dim=dim)
+    assert_allclose(complex_tensor, tensor_back)
+
+
 def test_magphase():
     spec_shape = [2, 514, 100]
     spec = torch.randn(*spec_shape)