Added some docstrings, did some minor code refactoring

tupak/gw/utils.py

@@ -12,50 +12,70 @@ from tupak.core.utils import gps_time_to_gmst, ra_dec_to_theta_phi, speed_of_lig
 def asd_from_freq_series(freq_data, df):
     """
     Calculate the ASD from the frequency domain output of gaussian_noise()
-    Input:
-    freq_data - array of complex frequency domain data
-    df - spacing of freq_data, 1/(segment length) used to generate the gaussian noise
-    Output:
-    asd = array of real-valued normalized frequency domain ASD data
+
+    Parameters
+    -------
+    freq_data: array_like
+        Array of complex frequency domain data
+    df: float
+        Spacing of freq_data, 1/(segment length) used to generate the gaussian noise
+
+    Returns
+    -------
+    array_like: array of real-valued normalized frequency domain ASD data
+
     """
-    asd = np.absolute(freq_data) * (2 * df)**0.5
-    return asd
+    return np.absolute(freq_data) * (2 * df)**0.5
 
 
 def psd_from_freq_series(freq_data, df):
     """
     Calculate the PSD from the frequency domain output of gaussian_noise()
     Calls asd_from_freq_series() and squares the output
-    Input:
-    freq_data - array of complex frequency domain data
-    df - spacing of freq_data, 1/(segment length) used to generate the gaussian noise
-    Output:
-    psd - array of real-valued normalized frequency domain PSD data
+
+    Parameters
+    -------
+    freq_data: array_like
+        Array of complex frequency domain data
+    df: float
+        Spacing of freq_data, 1/(segment length) used to generate the gaussian noise
+
+    Returns
+    -------
+    array_like: Real-valued normalized frequency domain PSD data
+
     """
-    psd = np.power(asd_from_freq_series(freq_data, df), 2)
-    return psd
+    return np.power(asd_from_freq_series(freq_data, df), 2)
 
 
 def time_delay_geocentric(detector1, detector2, ra, dec, time):
     """
     Calculate time delay between two detectors in geocentric coordinates based on XLALArrivaTimeDiff in TimeDelay.c
-    Input:
-    detector1 - cartesian coordinate vector for the first detector in the geocentric frame
-                generated by the Interferometer class as self.vertex
-    detector2 - cartesian coordinate vector for the second detector in the geocentric frame
-    To get time delay from Earth center, use detector2 = np.array([0,0,0])
-    ra - right ascension of the source in radians
-    dec - declination of the source in radians
-    time - GPS time in the geocentric frame
-    Output:
-    delta_t - time delay between the two detectors in the geocentric frame
+    Parameters
+    -------
+    detector1: array_like
+        Cartesian coordinate vector for the first detector in the geocentric frame
+        generated by the Interferometer class as self.vertex.
+    detector2: array_like
+        Cartesian coordinate vector for the second detector in the geocentric frame.
+        To get time delay from Earth center, use detector2 = np.array([0,0,0])
+    ra: float
+        Right ascension of the source in radians
+    dec: float
+        Declination of the source in radians
+    time: float
+        GPS time in the geocentric frame
+
+    Returns
+    -------
+    float: Time delay between the two detectors in the geocentric frame
+
     """
     gmst = gps_time_to_gmst(time)
     theta, phi = ra_dec_to_theta_phi(ra, dec, gmst)
     omega = np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)])
     delta_d = detector2 - detector1
-    delta_t = np.dot(omega, delta_d) / speed_of_light
-    return delta_t
+    return np.dot(omega, delta_d) / speed_of_light
 
 
 def get_polarization_tensor(ra, dec, time, psi, mode):
@@ -66,13 +86,23 @@ def get_polarization_tensor(ra, dec, time, psi, mode):
     [u, v, w] represent the Earth-frame
     [m, n, omega] represent the wave-frame
     Note: there is a typo in the definition of the wave-frame in Nishizawa et al.
+    Parameters
+    -------
+    ra: float
+        right ascension in radians
+    dec: float
+        declination in radians
+    time: float
+        geocentric GPS time
+    psi: float
+        binary polarisation angle counter-clockwise about the direction of propagation
+    mode: str
+        polarisation mode
+
+    Returns
+    -------
+    array_like: A 3x3 representation of the polarization_tensor for the specified mode.
 
-    :param ra: right ascension in radians
-    :param dec: declination in radians
-    :param time: geocentric GPS time
-    :param psi: binary polarisation angle counter-clockwise about the direction of propagation
-    :param mode: polarisation mode
-    :return: polarization_tensor(ra, dec, time, psi, mode): polarization tensor for the specified mode.
     """
     greenwich_mean_sidereal_time = gps_time_to_gmst(time)
     theta, phi = ra_dec_to_theta_phi(ra, dec, greenwich_mean_sidereal_time)
@@ -106,6 +136,20 @@ def get_vertex_position_geocentric(latitude, longitude, elevation):
 
     Based on arXiv:gr-qc/0008066 Eqs. B11-B13 except for the typo in the definition of the local radius.
     See Section 2.1 of LIGO-T980044-10 for the correct expression
+
+    Parameters
+    -------
+    latitude: float
+        Latitude in radians
+    longitude:
+        Longitude in radians
+    elevation:
+        Elevation in meters
+
+    Returns
+    -------
+    array_like: A 3D representation of the geocentric vertex position
+
     """
     semi_major_axis = 6378137  # for ellipsoid model of Earth, in m
     semi_minor_axis = 6356752.314  # in m
@@ -121,25 +165,24 @@ def inner_product(aa, bb, frequency, PSD):
     """
     Calculate the inner product defined in the matched filter statistic
 
-    arguments:
-    aai, bb: single-sided Fourier transform, created, e.g., by the nfft function above
-    frequency: an array of frequencies associated with aa, bb, also returned by nfft
-    PSD: PSD object
+    Parameters
+    -------
+    aa, bb: array_like
+        Single-sided Fourier transform, created, e.g., by the nfft function above
+    frequency: array_like
+        An array of frequencies associated with aa, bb, also returned by nfft
+    PSD: tupak.gw.detector.PowerSpectralDensity
 
-    Returns:
+    Returns
+    -------
     The matched filter inner product for aa and bb
+
     """
     PSD_interp = PSD.power_spectral_density_interpolated(frequency)
-
-    # calculate the inner product
     integrand = np.conj(aa) * bb / PSD_interp
-
     df = frequency[1] - frequency[0]
     integral = np.sum(integrand) * df
-
-    product = 4. * np.real(integral)
-
-    return product
+    return 4. * np.real(integral)
 
 
 def noise_weighted_inner_product(aa, bb, power_spectral_density, time_duration):
@@ -148,32 +191,61 @@ def noise_weighted_inner_product(aa, bb, power_spectral_density, time_duration):
 
     Parameters
     ----------
-    aa: array
+    aa: array_like
         Array to be complex conjugated
-    bb: array
+    bb: array_like
         Array not to be complex conjugated
-    power_spectral_density: array
-        Power spectral density
+    power_spectral_density: array_like
+        Power spectral density of the noise
     time_duration: float
         time_duration of the data
 
-    Return
+    Returns
     ------
     Noise-weighted inner product.
     """
 
-    # caluclate the inner product
     integrand = np.conj(aa) * bb / power_spectral_density
-    product = 4 / time_duration * np.sum(integrand)
-    return product
+    return 4 / time_duration * np.sum(integrand)
 
 
 def matched_filter_snr_squared(signal, interferometer, time_duration):
+    """
+
+    Parameters
+    ----------
+    signal: array_like
+        Array containing the signal
+    interferometer: tupak.gw.detector.Interferometer
+        Interferometer which we want to have the data and noise from
+    time_duration: float
+        Time duration of the signal
+
+    Returns
+    -------
+    float: The matched filter signal to noise ratio squared
+
+    """
     return noise_weighted_inner_product(signal, interferometer.data, interferometer.power_spectral_density_array,
                                         time_duration)
 
 
 def optimal_snr_squared(signal, interferometer, time_duration):
+    """
+
+    Parameters
+    ----------
+    signal: array_like
+        Array containing the signal
+    interferometer: tupak.gw.detector.Interferometer
+        Interferometer which we want to have the data and noise from
+    time_duration: float
+        Time duration of the signal
+
+    Returns
+    -------
+    float: The optimal signal to noise ratio possible squared
+    """
     return noise_weighted_inner_product(signal, signal, interferometer.power_spectral_density_array, time_duration)
 
 
@@ -195,7 +267,7 @@ def get_event_time(event):
     event: str
         Event descriptor, this can deal with some prefixes, e.g., '150914', 'GW150914', 'LVT151012'
 
-    Return
+    Returns
     ------
     event_time: float
         Merger time
@@ -284,6 +356,7 @@ def read_frame_file(file_name, t1, t2, channel=None, **kwargs):
 
     """
     loaded = False
+    strain = None
     if channel is not None:
         try:
             strain = TimeSeries.read(source=file_name, channel=channel, start=t1, end=t2, **kwargs)
@@ -301,7 +374,7 @@ def read_frame_file(file_name, t1, t2, channel=None, **kwargs):
                 loaded = True
                 logging.info('Successfully loaded {}.'.format(channel))
             except RuntimeError:
-                None
+                pass
 
     if loaded:
         return strain
@@ -310,40 +383,23 @@ def read_frame_file(file_name, t1, t2, channel=None, **kwargs):
         return None
 
 
-def process_strain_data(
-        strain, alpha=0.25, filter_freq=1024, **kwargs):
+def process_strain_data(strain, alpha=0.25, filter_freq=1024, **kwargs):
     """
     Helper function to obtain an Interferometer instance with appropriate
     PSD and data, given an center_time.
 
     Parameters
     ----------
-    name: str
-        Detector name, e.g., 'H1'.
-    center_time: float
-        GPS time of the center_time about which to perform the analysis.
-        Note: the analysis data is from `center_time-T/2` to `center_time+T/2`.
-    T: float
-        The total time (in seconds) to analyse. Defaults to 4s.
+    strain: array_like
+        Strain data to be processed
     alpha: float
         The tukey window shape parameter passed to `scipy.signal.tukey`.
-    psd_offset, psd_duration: float
-        The power spectral density (psd) is estimated using data from
-        `center_time+psd_offset` to `center_time+psd_offset + psd_duration`.
-    outdir: str
-        Directory where the psd files are saved
-    plot: bool
-        If true, create an ASD + strain plot
     filter_freq: float
         Low pass filter frequency
-    **kwargs:
-        All keyword arguments are passed to
-        `gwpy.timeseries.TimeSeries.fetch_open_data()`.
 
     Returns
     -------
-    interferometer: `tupak.detector.Interferometer`
-        An Interferometer instance with a PSD and frequency-domain strain data.
+    tupak.detector.Interferometer: An Interferometer instance with a PSD and frequency-domain strain data.
 
     """
 
@@ -355,12 +411,8 @@ def process_strain_data(
     strain = strain.crop(*strain.span.contract(1))
 
     time_series = strain.times.value
-    time_duration = time_series[-1] - time_series[0]
 
     # Apply Tukey window
-    N = len(time_series)
-    strain = strain * signal.windows.tukey(N, alpha=alpha)
-
+    strain = strain * signal.windows.tukey(len(time_series), alpha=alpha)
     frequency_domain_strain, frequencies = nfft(strain.value, sampling_frequency)
-
-    return frequency_domain_strain, frequencies
\ No newline at end of file
+    return frequency_domain_strain, frequencies