Spaces:
Sleeping
Sleeping
| """Beamformer module.""" | |
| from distutils.version import LooseVersion | |
| from typing import Tuple | |
| from typing import Union | |
| import torch | |
| from torch_complex import functional as FC | |
| from torch_complex.tensor import ComplexTensor | |
| EPS = torch.finfo(torch.double).eps | |
| is_torch_1_8_plus = LooseVersion(torch.__version__) >= LooseVersion("1.8.0") | |
| is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0") | |
| def new_complex_like( | |
| ref: Union[torch.Tensor, ComplexTensor], | |
| real_imag: Tuple[torch.Tensor, torch.Tensor], | |
| ): | |
| if isinstance(ref, ComplexTensor): | |
| return ComplexTensor(*real_imag) | |
| elif is_torch_complex_tensor(ref): | |
| return torch.complex(*real_imag) | |
| else: | |
| raise ValueError( | |
| "Please update your PyTorch version to 1.9+ for complex support." | |
| ) | |
| def is_torch_complex_tensor(c): | |
| return is_torch_1_9_plus and torch.is_complex(c) | |
| def is_complex(c): | |
| return isinstance(c, ComplexTensor) or is_torch_complex_tensor(c) | |
| def to_double(c): | |
| if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c): | |
| return c.to(dtype=torch.complex128) | |
| else: | |
| return c.double() | |
| def to_float(c): | |
| if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c): | |
| return c.to(dtype=torch.complex64) | |
| else: | |
| return c.float() | |
| def complex_norm(c: Union[torch.Tensor, ComplexTensor]) -> torch.Tensor: | |
| if not is_complex(c): | |
| raise TypeError("Input is not a complex tensor.") | |
| return torch.sqrt((c.real ** 2 + c.imag ** 2).sum(dim=-1, keepdim=True) + EPS) | |
| def einsum(equation, *operands): | |
| # NOTE: Do not mix ComplexTensor and torch.complex in the input! | |
| # NOTE (wangyou): Until PyTorch 1.9.0, torch.einsum does not support | |
| # mixed input with complex and real tensors. | |
| if len(operands) == 1: | |
| if isinstance(operands[0], (tuple, list)): | |
| operands = operands[0] | |
| complex_module = FC if isinstance(operands[0], ComplexTensor) else torch | |
| return complex_module.einsum(equation, *operands) | |
| elif len(operands) != 2: | |
| raise ValueError("0 or More than 2 operands are not supported.") | |
| a, b = operands | |
| if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): | |
| return FC.einsum(equation, a, b) | |
| elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): | |
| if not torch.is_complex(a): | |
| o_real = torch.einsum(equation, a, b.real) | |
| o_imag = torch.einsum(equation, a, b.imag) | |
| return torch.complex(o_real, o_imag) | |
| elif not torch.is_complex(b): | |
| o_real = torch.einsum(equation, a.real, b) | |
| o_imag = torch.einsum(equation, a.imag, b) | |
| return torch.complex(o_real, o_imag) | |
| else: | |
| return torch.einsum(equation, a, b) | |
| else: | |
| return torch.einsum(equation, a, b) | |
| def inverse( | |
| c: Union[torch.Tensor, ComplexTensor] | |
| ) -> Union[torch.Tensor, ComplexTensor]: | |
| if isinstance(c, ComplexTensor): | |
| return c.inverse2() | |
| else: | |
| return c.inverse() | |
| def matmul( | |
| a: Union[torch.Tensor, ComplexTensor], b: Union[torch.Tensor, ComplexTensor] | |
| ) -> Union[torch.Tensor, ComplexTensor]: | |
| # NOTE: Do not mix ComplexTensor and torch.complex in the input! | |
| # NOTE (wangyou): Until PyTorch 1.9.0, torch.matmul does not support | |
| # multiplication between complex and real tensors. | |
| if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): | |
| return FC.matmul(a, b) | |
| elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): | |
| if not torch.is_complex(a): | |
| o_real = torch.matmul(a, b.real) | |
| o_imag = torch.matmul(a, b.imag) | |
| return torch.complex(o_real, o_imag) | |
| elif not torch.is_complex(b): | |
| o_real = torch.matmul(a.real, b) | |
| o_imag = torch.matmul(a.imag, b) | |
| return torch.complex(o_real, o_imag) | |
| else: | |
| return torch.matmul(a, b) | |
| else: | |
| return torch.matmul(a, b) | |
| def reverse(a: Union[torch.Tensor, ComplexTensor], dim=0): | |
| if isinstance(a, ComplexTensor): | |
| return FC.reverse(a, dim=dim) | |
| else: | |
| return torch.flip(a, dims=(dim,)) | |
| def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]): | |
| """Solve the linear equation ax = b.""" | |
| # NOTE: Do not mix ComplexTensor and torch.complex in the input! | |
| # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support | |
| # mixed input with complex and real tensors. | |
| if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): | |
| if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor): | |
| return FC.solve(b, a)[0] | |
| else: | |
| return matmul(inverse(a), b) | |
| elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): | |
| if torch.is_complex(a) and torch.is_complex(b): | |
| return torch.linalg.solve(a, b) | |
| else: | |
| return matmul(inverse(a), b) | |
| else: | |
| if is_torch_1_8_plus: | |
| return torch.linalg.solve(a, b) | |
| else: | |
| return torch.solve(b, a)[0] | |