added type annotations

aotools/astronomy/_astronomy.py

@@ -1,4 +1,5 @@
 import numpy
+from numpy import ndarray
 
 # Dictionary of flux values at the top of the atmosphere
 #                 band, lamda, dLamda, m=0 flux (Jy)
@@ -16,7 +17,7 @@
                 'z': [0.91, 0.13, 4810]}
 
 
-def photons_per_mag(mag, mask, pixel_scale, wvlBand, exposure_time):
+def photons_per_mag(mag: float, mask: ndarray, pixel_scale: float, wvlBand: float, exposure_time: int) -> float:
     """
     Calculates the photon flux for a given aperture, star magnitude and wavelength band
 
@@ -45,7 +46,7 @@ def photons_per_mag(mag, mask, pixel_scale, wvlBand, exposure_time):
     return photons
 
 
-def photons_per_band(mag, mask, pxlScale, expTime, waveband='V'):
+def photons_per_band(mag: float, mask: ndarray, pxlScale: float, expTime: float, waveband: str='V') -> float:
         '''
         Calculates the photon flux for a given aperture, star magnitude and wavelength band
 
@@ -74,7 +75,7 @@ def photons_per_band(mag, mask, pxlScale, expTime, waveband='V'):
         return photons
 
 
-def magnitude_to_flux(magnitude, waveband='V'):
+def magnitude_to_flux(magnitude: float, waveband: str='V') -> float:
     """
     Converts apparent magnitude to a flux of photons
 
@@ -92,7 +93,7 @@ def magnitude_to_flux(magnitude, waveband='V'):
     return flux_photons
 
 
-def flux_to_magnitude(flux, waveband='V'):
+def flux_to_magnitude(flux: float, waveband: str='V') -> float:
     """
     Converts incident flux of photons to the apparent magnitude
 

aotools/fouriertransform.py

@@ -2,9 +2,10 @@
 Module containing useful FFT based function and classes
 """
 import numpy
+from numpy import ndarray
 
 
-def ft(data, delta):
+def ft(data: ndarray, delta: float) -> ndarray:
     """
     A properly scaled 1-D FFT
 
@@ -21,7 +22,7 @@ def ft(data, delta):
             axes=(-1)) * delta
     return DATA
 
-def ift(DATA, delta_f):
+def ift(DATA: ndarray, delta_f: float) -> ndarray:
     """
     Scaled inverse 1-D FFT
 
@@ -41,7 +42,7 @@ def ift(DATA, delta_f):
     return data
 
 
-def ft2(data, delta):
+def ft2(data: ndarray, delta: float) -> ndarray:
     """
     A properly scaled 2-D FFT
 
@@ -61,7 +62,7 @@ def ft2(data, delta):
 
     return DATA
 
-def ift2(DATA, delta_f):
+def ift2(DATA: ndarray, delta_f: float) -> ndarray:
     """
     Scaled inverse 2-D FFT
 
@@ -79,7 +80,7 @@ def ift2(DATA, delta_f):
 
     return g
 
-def rft(data, delta):
+def rft(data: ndarray, delta: float) -> ndarray:
     """
     A properly scaled real 1-D FFT
 

aotools/functions/_functions.py

@@ -1,7 +1,9 @@
 import numpy
+from numpy import ndarray
+from typing import Optional, Tuple, Union
 
 
-def gaussian2d(size, width, amplitude=1., cent=None):
+def gaussian2d(size: Union[Tuple[int, int], int], width: Union[Tuple[int, int], int], amplitude: float=1., cent: Optional[Union[ndarray, Tuple[int, int]]]=None) -> ndarray:
     '''
     Generates 2D gaussian distribution
 

aotools/functions/karhunenLoeve.py

@@ -28,9 +28,11 @@
 import numpy as np
 import scipy
 from scipy.ndimage.interpolation import map_coordinates
+from numpy import int64, ndarray
+from typing import Dict, Optional, Tuple, Union
 
 
-def rebin(a, newshape):
+def rebin(a: ndarray, newshape: Union[Tuple[int, int], Tuple[int, int, int]]) -> ndarray:
     '''Rebin an array to a new shape.
     See scipy cookbook.
     It is intended to be similar to 'rebin' of IDL
@@ -45,7 +47,7 @@ def rebin(a, newshape):
     return a[tuple(indices)]
 
 
-def stf_kolmogorov(r):
+def stf_kolmogorov(r: ndarray) -> ndarray:
     '''
     Kolmogorov structure function with r = (D/r0)
     '''
@@ -71,7 +73,7 @@ def stf_vonKarman(r, L0):
     return D_vk
 
 
-def gkl_radii(ri, nr):
+def gkl_radii(ri: float, nr: int) -> ndarray:
     '''
     Generate n points evenly spaced in r^2 between r_i^2 and 1
 
@@ -95,7 +97,7 @@ def gkl_radii(ri, nr):
     return np.sqrt(r2)
 
 
-def gkl_kernel(ri, nr, rad, stfunc='kolmogorov', outerscale=None):
+def gkl_kernel(ri: float, nr: int, rad: ndarray, stfunc: str='kolmogorov', outerscale: None=None) -> ndarray:
     '''
     Calculation of the kernel L^p
 
@@ -158,7 +160,7 @@ def gkl_kernel(ri, nr, rad, stfunc='kolmogorov', outerscale=None):
     return kernel
 
 
-def piston_orth(nr):
+def piston_orth(nr: int) -> ndarray:
     '''
     Unitary matrix used to filter out piston term.
     Eq. 19 in Cannon 1996.
@@ -184,7 +186,7 @@ def piston_orth(nr):
     return s
 
 
-def gkl_fcom(ri, kernels, nfunc, verbose=False):
+def gkl_fcom(ri: float, kernels: ndarray, nfunc: int, verbose: bool=False) -> Tuple[ndarray, int64, ndarray, ndarray, ndarray]:
     '''
     Computation of the radial eigenvectors of the KL basis.
 
@@ -300,7 +302,7 @@ def gkl_fcom(ri, kernels, nfunc, verbose=False):
     return evals, nord, npo, oord, rabas
 
 
-def gkl_azimuthal(nord, npp):
+def gkl_azimuthal(nord: int64, npp: int) -> ndarray:
     '''
     Compute the azimuthal function of the KL basis.
 
@@ -327,8 +329,8 @@ def gkl_azimuthal(nord, npp):
     return gklazi
 
 
-def gkl_basis(ri=0.25, nr=40, npp=None, nfunc=500,
-              stf='kolstf', outerscale=None):
+def gkl_basis(ri: float=0.25, nr: int=40, npp: Optional[int]=None, nfunc: int=500,
+              stf: str='kolstf', outerscale: None=None) -> Dict[str, Union[int, float, str, ndarray, int64]]:
     '''
     Wrapper to create the radial and azimuthal K-L functions.
 
@@ -377,7 +379,7 @@ def gkl_basis(ri=0.25, nr=40, npp=None, nfunc=500,
     return gklbasis
 
 
-def gkl_sfi(kl_basis, i):
+def gkl_sfi(kl_basis: Dict[str, Union[int, float, str, ndarray, int64]], i: int) -> ndarray:
     '''
     return the i'th function from the generalized KL basis 'bas'.
     'bas' must be generated by 'gkl_basis'
@@ -397,7 +399,7 @@ def gkl_sfi(kl_basis, i):
     return sf
 
 
-def radii(nr, npp, ri):
+def radii(nr: int, npp: int, ri: float) -> ndarray:
     '''
     Use to generate a polar coordinate system.
 
@@ -422,7 +424,7 @@ def radii(nr, npp, ri):
     return ra
 
 
-def polang(r):
+def polang(r: ndarray) -> ndarray:
     '''
         Generate an array with the same dimensions as r, but containing the
         azimuthal values for a polar coordinate system.
@@ -437,7 +439,7 @@ def polang(r):
     return phi
 
 
-def set_pctr(bas, ncp=None, ncmar=None):
+def set_pctr(bas: Dict[str, Union[int, float, str, ndarray, int64]], ncp: Optional[int]=None, ncmar: Optional[int]=None) -> Dict[str, Union[ndarray, int]]:
     '''
     call pcgeom to build the dic geom_struct with the
     right initializtions.
@@ -455,7 +457,7 @@ def set_pctr(bas, ncp=None, ncmar=None):
     return pcgeom(bas['nr'], bas['np'], ncp, bas['ri'], ncmar)
 
 
-def setpincs(ax, ay, px, py, ri):
+def setpincs(ax: ndarray, ay: ndarray, px: ndarray, py: ndarray, ri: float) -> Tuple[ndarray, ndarray, ndarray]:
     '''
     determine a set of squares for interpolating from cartesian
     to polar coordinates,
@@ -496,7 +498,7 @@ def setpincs(ax, ay, px, py, ri):
     return pincx, pincy, pincw
 
 
-def pcgeom(nr, npp, ncp, ri, ncmar):
+def pcgeom(nr: int, npp: int, ncp: int, ri: float, ncmar: int) -> Dict[str, Union[ndarray, int]]:
     '''
     This routine builds a geom dic.
 
@@ -549,7 +551,7 @@ def pcgeom(nr, npp, ncp, ri, ncmar):
     return geom
 
 
-def pol2car(cpgeom, pol, mask=False):
+def pol2car(cpgeom: Dict[str, Union[ndarray, int]], pol: ndarray, mask: bool=False) -> ndarray:
     '''
     Polar to cartesian conversion.
 
@@ -565,8 +567,8 @@ def pol2car(cpgeom, pol, mask=False):
     return cd
 
 
-def make_kl(nmax, dim, ri=0.0, nr=40,
-            stf='kolmogorov', outerscale=None, mask=True):
+def make_kl(nmax: int, dim: int, ri: float=0.0, nr: int=40,
+            stf: str='kolmogorov', outerscale: None=None, mask: bool=True) -> Tuple[ndarray, ndarray, ndarray, Dict[str, Union[int, float, str, ndarray, int64]]]:
     '''
     Main routine to generatre a KL basis of dimension [nmax, dim, dim].
 

aotools/functions/pupil.py

@@ -7,9 +7,11 @@
 """
 
 import numpy
+from numpy import float64, ndarray
+from typing import Tuple, Union
 
 
-def circle(radius, size, circle_centre=(0, 0), origin="middle"):
+def circle(radius: Union[float, float64, int], size: int, circle_centre: Tuple[int, int]=(0, 0), origin: str="middle") -> ndarray:
     """
     Create a 2-D array: elements equal 1 within a circle and 0 outside.
 

aotools/functions/zernike.py

@@ -1,5 +1,7 @@
 import numpy
 from . import circle
+from numpy import ndarray
+from typing import List, Union
 
 # xrange just "range" in python3.
 # This code means fastest implementation used in 2 and 3
@@ -8,7 +10,7 @@
 except NameError:
     xrange = range
 
-def phaseFromZernikes(zCoeffs, size, norm="noll"):
+def phaseFromZernikes(zCoeffs: List[int], size: int, norm: str="noll") -> ndarray:
     """
     Creates an array of the sum of zernike polynomials with specified coefficeints
 
@@ -28,7 +30,7 @@ def phaseFromZernikes(zCoeffs, size, norm="noll"):
     return phase
 
 
-def zernike_noll(j, N):
+def zernike_noll(j: int, N: int) -> ndarray:
     """
      Creates the Zernike polynomial with mode index j,
      where j = 1 corresponds to piston.
@@ -44,7 +46,7 @@ def zernike_noll(j, N):
     return zernike_nm(n, m, N)
 
 
-def zernike_nm(n, m, N):
+def zernike_nm(n: int, m: int, N: int) -> ndarray:
     """
      Creates the Zernike polynomial with radial index, n, and azimuthal index, m.
 
@@ -75,7 +77,7 @@ def zernike_nm(n, m, N):
     return Z*circle(N/2., N)
 
 
-def zernikeRadialFunc(n, m, r):
+def zernikeRadialFunc(n: int, m: int, r: ndarray) -> ndarray:
     """
     Fucntion to calculate the Zernike radial function
 
@@ -100,7 +102,7 @@ def zernikeRadialFunc(n, m, r):
     return R
 
 
-def zernIndex(j):
+def zernIndex(j: int) -> List[int]:
     """
     Find the [n,m] list giving the radial order n and azimuthal order
     of the Zernike polynomial of Noll index j.
@@ -126,7 +128,7 @@ def zernIndex(j):
     return [n, m]
 
 
-def zernikeArray(J, N, norm="noll"):
+def zernikeArray(J: Union[List[int], int], N: int, norm: str="noll") -> ndarray:
     """
     Creates an array of Zernike Polynomials
 
@@ -172,7 +174,7 @@ def zernikeArray(J, N, norm="noll"):
     return Zs
 
 
-def makegammas(nzrad):
+def makegammas(nzrad: int) -> ndarray:
     """
     Make "Gamma" matrices which can be used to determine first derivative
     of Zernike matrices (Noll 1976).

aotools/image_processing/centroiders.py

@@ -7,9 +7,11 @@
 """
 
 import numpy
+from numpy import ndarray
+from typing import Union
 
 
-def correlation_centroid(im, ref, threshold=0., padding=1):
+def correlation_centroid(im: ndarray, ref: ndarray, threshold: float=0., padding: int=1) -> ndarray:
     """
     Correlation Centroider, currently only works for 3d im shape.
     Performs a simple thresholded COM on the correlation.
@@ -51,7 +53,8 @@ def correlation_centroid(im, ref, threshold=0., padding=1):
     return centroids
 
 
-def centre_of_gravity(img, threshold=0, min_threshold=0, **kwargs):
+def centre_of_gravity(img: ndarray, threshold: Union[float, int]=0, min_threshold: int=0, **kwargs
+) -> ndarray:
     """
     Centroids an image, or an array of images.
     Centroids over the last 2 dimensions.
@@ -90,7 +93,7 @@ def centre_of_gravity(img, threshold=0, min_threshold=0, **kwargs):
     return numpy.array([x_centroid, y_centroid])
 
 
-def brightest_pixel(img, threshold, **kwargs):
+def brightest_pixel(img: ndarray, threshold: float, **kwargs) -> ndarray:
     """
     Centroids using brightest Pixel Algorithm
     (A. G. Basden et al,  MNRAS, 2011)
@@ -123,7 +126,7 @@ def brightest_pixel(img, threshold, **kwargs):
     return centre_of_gravity(img)
 
 
-def cross_correlate(x, y, padding=1):
+def cross_correlate(x: ndarray, y: ndarray, padding: int=1) -> ndarray:
     """
     2D convolution using FFT, use to generate cross-correlations.
 
@@ -143,7 +146,7 @@ def cross_correlate(x, y, padding=1):
     return cross_correlation
 
 
-def quadCell(img, **kwargs):
+def quadCell(img: ndarray, **kwargs) -> ndarray:
     """
     Centroider to be used for 2x2 images.