From 83673dfecd48475edd67adf8eee0825598fcfcca Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 20 Mar 2019 11:44:16 -0400 Subject: [PATCH] Add tests for the cmath functions --- .gitignore | 1 + CDTools/tools/cmath.py | 2 +- setup.py | 2 +- tests/tools/test_cmath.py | 126 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 tests/tools/test_cmath.py diff --git a/.gitignore b/.gitignore index 3c91904..ebb220f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,5 @@ \#*# *.pyc *.egg-info +.pytest_cache docs/build \ No newline at end of file diff --git a/CDTools/tools/cmath.py b/CDTools/tools/cmath.py index 9646627..34ab6a4 100644 --- a/CDTools/tools/cmath.py +++ b/CDTools/tools/cmath.py @@ -168,7 +168,7 @@ def cdiv(a,b): torch.Tensor : A tensor storing the elementwise complex quotient """ - return cmult(a, cconj(b)) / cabssq(b) + return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1) diff --git a/setup.py b/setup.py index 82c3bbc..c543375 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,7 @@ setuptools.setup( "numpy", "scipy", "matplotlib", - "pytorch", + #"pytorch", "h5py"], packages=setuptools.find_packages(), classifiers=[ diff --git a/tests/tools/test_cmath.py b/tests/tools/test_cmath.py new file mode 100644 index 0000000..8a8cdbe --- /dev/null +++ b/tests/tools/test_cmath.py @@ -0,0 +1,126 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools import cmath +import numpy as np +import torch as t +from scipy.fftpack import fftshift, ifftshift + +def test_complex_to_torch(): + arr = np.random.rand(100,4) + 1j * np.random.rand(100,4) + tensor = cmath.complex_to_torch(arr) + assert np.allclose(tensor[:,:,0].numpy(),np.real(arr)) + assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr)) + + +def test_torch_to_complex(): + tensor = t.rand(100,4,2) + arr = cmath.torch_to_complex(tensor) + assert np.allclose(tensor[:,:,0].numpy(),np.real(arr)) + assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr)) + + +def test_cabssq(): + arr = np.random.rand(56,23,2) + 1j * np.random.rand(56,23,2) + cabssq = cmath.cabssq(cmath.complex_to_torch(arr)) + assert np.allclose(cabssq.numpy(),np.abs(arr)**2) + + +def test_cabs(): + arr = np.random.rand(2,3,4,5) + 1j * np.random.rand(2,3,4,5) + cabs = cmath.cabs(cmath.complex_to_torch(arr)) + assert np.allclose(cabs.numpy(),np.abs(arr)) + + +def test_cconj(): + arr = np.random.rand(50) + 1j * np.random.rand(50) + cconj = cmath.cconj(cmath.complex_to_torch(arr)) + assert np.allclose(cmath.torch_to_complex(cconj), np.conj(arr)) + + +def test_cmult(): + arr1 = np.random.rand(50) + 1j * np.random.rand(50) + arr2 = np.random.rand(50) + 1j * np.random.rand(50) + mult = cmath.cmult(cmath.complex_to_torch(arr1), + cmath.complex_to_torch(arr2)) + assert np.allclose(cmath.torch_to_complex(mult),arr1*arr2) + + +def test_cdiv(): + arr1 = np.random.rand(50) + 1j * np.random.rand(50) + arr2 = np.random.rand(50) + 1j * np.random.rand(50) + div = cmath.cdiv(cmath.complex_to_torch(arr1), + cmath.complex_to_torch(arr2)) + assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2) + + +def test_cphase(): + arr = np.random.rand(50) + 1j * np.random.rand(50) + cphase = cmath.cphase(cmath.complex_to_torch(arr)) + assert np.allclose(cphase.numpy(), np.angle(arr)) + + +def test_scalars(): + arr1 = np.random.rand(1) + 1j * np.random.rand(1) + arr2 = np.random.rand(1) + 1j * np.random.rand(1) + div = cmath.cdiv(cmath.complex_to_torch(arr1), + cmath.complex_to_torch(arr2)) + assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2) + +def test_fftshift(): + #1D, even + arr = np.random.rand(300) + 1j * np.random.rand(300) + shifted = cmath.fftshift(cmath.complex_to_torch(arr)) + assert np.allclose(fftshift(arr), + cmath.torch_to_complex(shifted)) + #1D, odd + arr = np.random.rand(301) + 1j * np.random.rand(301) + shifted = cmath.fftshift(cmath.complex_to_torch(arr)) + assert np.allclose(fftshift(arr), + cmath.torch_to_complex(shifted)) + + #2D + arr = np.random.rand(20,21) + 1j * np.random.rand(20,21) + shifted = cmath.fftshift(cmath.complex_to_torch(arr)) + assert np.allclose(fftshift(arr), + cmath.torch_to_complex(shifted)) + #3D + arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) + shifted = cmath.fftshift(cmath.complex_to_torch(arr)) + assert np.allclose(fftshift(arr), + cmath.torch_to_complex(shifted)) + + #3D, choosing specific axes + arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) + shifted = cmath.fftshift(cmath.complex_to_torch(arr),dims=(0,1)) + assert np.allclose(fftshift(arr,axes=(0,1)), + cmath.torch_to_complex(shifted)) + + +def test_ifftshift(): + #1D, even + arr = np.random.rand(300) + 1j * np.random.rand(300) + shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) + assert np.allclose(ifftshift(arr), + cmath.torch_to_complex(shifted)) + #1D, odd + arr = np.random.rand(301) + 1j * np.random.rand(301) + shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) + assert np.allclose(ifftshift(arr), + cmath.torch_to_complex(shifted)) + + #2D + arr = np.random.rand(20,21) + 1j * np.random.rand(20,21) + shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) + assert np.allclose(ifftshift(arr), + cmath.torch_to_complex(shifted)) + #3D + arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) + shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) + assert np.allclose(ifftshift(arr), + cmath.torch_to_complex(shifted)) + + #3D, choosing specific axes + arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) + shifted = cmath.ifftshift(cmath.complex_to_torch(arr),dims=(0,1)) + assert np.allclose(ifftshift(arr,axes=(0,1)), + cmath.torch_to_complex(shifted))