Write and test the first set of measurements

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

setup.py

@@ -16,7 +16,7 @@ setuptools.setup(
         "numpy",
         "scipy",
         "matplotlib",
-        "dateutil",
+        "python-dateutil",
         #"pytorch",
         "h5py"],
     packages=setuptools.find_packages(),

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))