From d49f8a1c043ab639cedcc94b1f3547971f39eb7e Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 28 Mar 2019 20:15:03 -0400 Subject: [PATCH] Write and test the first set of measurements --- CDTools/tools/measurements.py | 78 ++++++++++++++++++++++++++++++++ setup.py | 2 +- tests/tools/test_measurements.py | 52 +++++++++++++++++++++ 3 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 CDTools/tools/measurements.py create mode 100644 tests/tools/test_measurements.py diff --git a/CDTools/tools/measurements.py b/CDTools/tools/measurements.py new file mode 100644 index 0000000..fdcbbf1 --- /dev/null +++ b/CDTools/tools/measurements.py @@ -0,0 +1,78 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools import cmath +import torch as t +import numpy as np + +# +# This file will host tools to turn a propagated wavefield into a measured +# intensity pattern on a detector +# + +__all__ = ['intensity', 'incoherent sum', 'quadratic_background'] + + +def intensity(wavefield, detector_slice=None): + """Returns the intensity of a wavefield + + The intensity is defined as the magnitude squared of the + wavefront. If a detector slice is given, the returned array + will only include that slice from the simulated wavefront. + + Args: + wavefield (torch.Tensor) : A complex MxNx2 wavefield + detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return + + Returns: + torch.Tensor : A real MxN array storing the wavefield's intensities + """ + if detector_slice is None: + return cmath.cabssq(wavefield) + else: + return cmath.cabssq(wavefield[detector_slice]) + + +def incoherent_sum(wavefields, detector_slice=None): + """Returns the incoherent sum of the intensities of the wavefields + + The intensity is defined as the sum of the magnitudes squared of + the wavefields. If a detector slice is given, the returned array + will only include that slice from the simulated wavefronts. + + Args: + wavefields (torch.Tensor) : A JxMxNx2 stack of complex wavefields + detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return + + Returns: + torch.Tensor : A real MxN array storing the incoherently summed intensities + """ + # This syntax just adds an axis to the slice to preserve the J direction + if detector_slice is None: + return t.sum(cmath.cabssq(wavefields),dim=-3) + else: + return t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=-3) + + + +def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity): + """Returns the intensity of a wavefield plus a background + + The intensity is calculated via the given measurment function + Then, the square of the given background is added. This kind + of background model is commonly used to enforce positivity of the + background model. + + Args: + wavefield (torch.Tensor) : A complex MxNx2 wavefield + background (torch.Tensor) : An tensor storing the square root of the detector background + detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return + measurement (function) : Optional, the measurement function to use. The default is measurements.intensity + + Returns: + torch.Tensor : A real MxN array storing the wavefield's intensities + """ + if detector_slice is None: + return measurement(wavefield) + background**2 + else: + return measurement(wavefield, detector_slice) \ + + background[detector_slice]**2 diff --git a/setup.py b/setup.py index de2f1ef..5f37bab 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,7 @@ setuptools.setup( "numpy", "scipy", "matplotlib", - "dateutil", + "python-dateutil", #"pytorch", "h5py"], packages=setuptools.find_packages(), diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py new file mode 100644 index 0000000..93e45dd --- /dev/null +++ b/tests/tools/test_measurements.py @@ -0,0 +1,52 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools import measurements +from CDTools.tools import cmath +import torch as t +import numpy as np +import pytest + + +def test_intensity(): + wavefield = t.rand((10,10,2)) + np_result = np.abs(cmath.torch_to_complex(wavefield))**2 + assert t.allclose(measurements.intensity(wavefield), + t.tensor(np_result)) + + det_slice = np.s_[3:,5:8] + assert t.allclose(measurements.intensity(wavefield,det_slice), + t.tensor(np_result[det_slice])) + + + +def test_incoherent_sum(): + wavefields = t.rand((4,10,10,2)) + np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + assert t.allclose(measurements.incoherent_sum(wavefields), + t.tensor(np_result)) + + det_slice = np.s_[3:,5:8] + assert t.allclose(measurements.incoherent_sum(wavefields,det_slice), + t.tensor(np_result[det_slice])) + + + +def test_quadratic_background(): + # test with intensity + wavefield = t.rand((10,10,2)) + background = t.rand((10,10)) + np_result = np.abs(cmath.torch_to_complex(wavefield))**2 + background.numpy()**2 + det_slice = np.s_[3:,5:8] + + result = measurements.quadratic_background(wavefield,background, + detector_slice=det_slice, + measurement=measurements.intensity) + assert t.allclose(result, t.tensor(np_result[det_slice])) + + # test with incoherent sum but no slice + wavefields = t.rand((4,10,10,2)) + np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + np_result += background.numpy()**2 + result = measurements.quadratic_background(wavefields, background, + measurement=measurements.incoherent_sum) + assert t.allclose(result, t.tensor(np_result))