Merge pull request #30 from GillesOrban/encircled_energy

added phase structure function and encircled energy calculations

aotools/functions/_functions.py

@@ -54,9 +54,63 @@ def aziAvg(data):
     """
 
     size = data.shape[0]
-    avg = numpy.empty(size / 2, dtype="float")
-    for i in range(size / 2):
+    avg = numpy.empty(int(size / 2), dtype="float")
+    for i in range(int(size / 2)):
         ring = pupil.circle(i + 1, size) - pupil.circle(i, size)
         avg[i] = (ring * data).sum() / (ring.sum())
 
     return avg
+
+
+def encircledEnergy(data,
+                    fraction=0.5, center=None,
+                    eeDiameter=True):
+    """
+        Return the encircled energy diameter for a given fraction
+        (default is ee50d).
+        Can also return the encircled energy function.
+        Translated and extended from YAO.
+
+        Parameters:
+            data : 2-d array
+            fraction : energy fraction for diameter calculation
+                default = 0.5
+            center : default = center of image
+            eeDiameter : toggle option for return.
+                If False returns two vectors: (x, ee(x))
+                Default = True
+        Returns:
+            Encircled energy diameter
+            or
+            2 vectors: diameters and encircled energies
+
+    """
+    dim = data.shape[0] // 2
+    center = [dim, dim]
+    xc = center[0]
+    yc = center[1]
+    e = 1.9
+    npt = 20
+    rad = numpy.linspace(0, dim**(1. / e), npt)**e
+    ee = numpy.empty(rad.shape)
+
+    for i in range(npt):
+        pup = pupil.circle(rad[i],
+                           int(dim) * 2,
+                           circle_centre=(xc + 0.5,
+                                          yc + 0.5),
+                           origin='corner')
+        rad[i] = numpy.sqrt(numpy.sum(pup) * 4 / numpy.pi)  # diameter
+        ee[i] = numpy.sum(pup * data)
+
+    rad = numpy.append(0, rad)
+    ee = numpy.append(0, ee)
+    ee /= numpy.sum(data)
+    xi = numpy.linspace(0, dim, int(2 * dim))
+    yi = numpy.interp(xi, rad, ee)
+
+    if eeDiameter is False:
+        return xi, yi
+    else:
+        ee50d = float(xi[numpy.argmin(numpy.abs(yi - fraction))])
+        return ee50d

aotools/turbulence/slopecovariance.py

@@ -415,6 +415,53 @@ def structure_function_vk(seperation, r0, L0):
 
     return D_vk
 
+
+def structure_function_kolmogorov(separation, r0):
+    '''
+        Compute the Kolmogorov phase structure function
+
+        Parameters:
+            separation (ndarray, float): float or array of data representing
+                separations between points
+            r0 (float): Fried parameter for atmosphere
+
+        Returns:
+            ndarray, float: Structure function for separation(s)
+    '''
+    return 6.88 * (separation / r0)**(5. / 3.)
+
+
+def calculate_structure_function(phase, nbOfPoint=None, step=None):
+    '''
+        Compute the structure function of an 2D array, along the first
+         dimension.
+        SF defined as sf[j]= < (phase - phase_shifted_by_j)^2 >
+        Translated from a YAO function.
+
+        Parameters:
+            phase (ndarray, 2d): 2d-array
+            nbOfPoint (int): final size of the structure function vector
+                Default is phase.shape[1] / 4
+            step (int): step in pixel when computing the sf.
+                (step * sampling_phase) gives the sf sampling in meters.
+                Default is 1
+
+        Returns:
+            ndarray, float: values for the structure function of the data.
+    '''
+    if nbOfPoint is None:
+        nbOfPoint = phase.shape[1] / 4
+    if step is None:
+        step = 1
+    step = int(step)
+    xm = int(numpy.min([nbOfPoint, phase.shape[1] / step - 1]))
+    sf_x = numpy.empty(xm)
+    for i in range(step, xm * step, step):
+        sf_x[int(i / step)] = numpy.mean((phase[0:-i, :] - phase[i:, :])**2)
+
+    return sf_x
+
+
 def mirror_covariance_matrix(cov_mat, n_subaps):
     """
     Mirrors a covariance matrix around the axis of the diagonal.

test/test_functions.py

@@ -11,3 +11,19 @@ def test_gaussian2d_2d():
     gaussian = functions.gaussian2d((10, 8), (3, 2), 10., (4, 3))
     print(gaussian.shape)
     assert gaussian.shape == (10, 8)
+
+
+def test_encircledEnergy():
+    data = numpy.random.rand(32, 32)
+    ee50d = functions.encircledEnergy(data)
+    print(ee50d)
+    assert type(ee50d) == float
+
+
+def test_encircledEnergy_func():
+        data = numpy.random.rand(32, 32)
+        x, y = functions.encircledEnergy(data, eeDiameter=False)
+        print(y.min(), y.max())
+        assert len(x) == len(y)
+        assert numpy.max(y) <= 1
+        assert numpy.min(y) >= 0

test/test_slopecovariance.py

@@ -200,4 +200,20 @@ def test_slopecovmat_makecovmat_multithreaded():
 
     # from matplotlib import pyplot
     # pyplot.imshow(cov_mat.covariance_matrix)
-    # pyplot.show()
+    # pyplot.show()
+
+
+def test_structure_function_kolmogorov():
+    seps = numpy.arange(0, 10, 0.1)
+    r0 = 0.16
+    sf = aotools.structure_function_kolmogorov(seps, r0)
+
+    assert len(sf) == len(seps)
+    assert all(numpy.isnan(sf)) is False
+
+
+def test_calculate_structure_function():
+    phase = numpy.random.randn(32, 32)
+    sf = aotools.calculate_structure_function(phase)
+    assert all(numpy.isnan(sf)) is False
+    assert len(sf) == int(phase.shape[1] / 4)