Merge pull request #18 from nicochunger/ge_uncert

Added functions to calculate the geometric elements with uncertainties

pystrometry/utils/du437_tools.py

@@ -17,6 +17,9 @@
 
 from .. import gaia_astrometry, pystrometry
 
+from uncertainties import unumpy as unp
+from uncertainties import correlated_values_norm
+
 try:
     from helpers.table_helpers import plot_columns_simple
 except ImportError:
@@ -697,3 +700,158 @@ def show_best_solution(file, out_dir):
 
     plot_columns_simple(data, os.path.join(out_dir, 'best_solution_simple_plots'),
                         name_seed='best_solution', units=units)
+
+
+def geometric_elements_with_uncertainties(thiele_innes_parameters, thiele_innes_parameters_errors=None, correlation_matrix=None):
+    """
+    Return geometrical orbit elements a, omega, OMEGA, i. If errors are not given
+    they are assumed to be 0 and correlation matrix is set to identity.
+    Complement to the pystrometry.geometric_elements function that allows to 
+    compute parameter uncertainties as well.
+
+    Parameters
+    ----------
+    thiele_innes_parameters : array
+        Array of Thiele Innes parameters [A,B,F,G] in milli-arcsecond
+    thiele_innes_parameters_errors : array, optional
+        Array of the errors of the Thiele Innes parameters [A,B,F,G] in milli-arcsecond
+    correlation_matrix : (4, 4) array, optional
+        Correlation matrix for the Thiele Innes parameters [A,B,F,G]
+
+    Returns
+    -------
+    geometric_parameters : array
+        Orbital elements [a_mas, omega_deg, OMEGA_deg, i_deg]
+    geometric_parameters_errors : array
+        Errors of the orbital elements [a_mas, omega_deg, OMEGA_deg, i_deg]
+
+    """
+
+    # Checks on the errors and correlation matrix
+    if (thiele_innes_parameters_errors is None) and (correlation_matrix is None):
+        # Define errors to 0 and correlation matrix to identity
+        thiele_innes_parameters_errors = [0,0,0,0]
+        correlation_matrix = np.identity(4)
+    elif (thiele_innes_parameters_errors is not None) and (correlation_matrix is not None):
+        # If both are given continue to the calculation
+        pass
+    else:
+        # If either one of them is provided but not the other raise an error
+        raise ValueError("thieles_innes_parameters_erros and correlation_matrix must be" \
+                          "specified together.")
+
+
+    # Define uncorrelated (value, uncertainty) pairs
+    A_u = (thiele_innes_parameters[0], thiele_innes_parameters_errors[0])
+    B_u = (thiele_innes_parameters[1], thiele_innes_parameters_errors[1])
+    F_u = (thiele_innes_parameters[2], thiele_innes_parameters_errors[2])
+    G_u = (thiele_innes_parameters[3], thiele_innes_parameters_errors[3])
+
+    # Create correlated quantities
+    A, B, F, G = correlated_values_norm([A_u, B_u, F_u, G_u], correlation_matrix)
+
+    p = (A ** 2 + B ** 2 + G ** 2 + F ** 2) / 2.
+    q = A * G - B * F
+
+    a_mas = unp.sqrt(p + unp.sqrt(p ** 2 - q ** 2))
+    i_rad = unp.arccos(q / (a_mas ** 2.))
+    omega_rad = (unp.arctan2(B - F, A + G) + unp.arctan2(-B - F, A - G)) / 2.
+    OMEGA_rad = (unp.arctan2(B - F, A + G) - unp.arctan2(-B - F, A - G)) / 2.
+
+    # Convert radians to degrees
+    i_deg = i_rad * 180 / np.pi
+    omega_deg = omega_rad * 180 / np.pi
+    OMEGA_deg = OMEGA_rad * 180 / np.pi
+
+    # Extract nominal values and standard deviations
+    geometric_parameters = np.array([unp.nominal_values(a_mas), 
+                                     unp.nominal_values(omega_deg), 
+                                     unp.nominal_values(OMEGA_deg), 
+                                     unp.nominal_values(i_deg)])
+
+    geometric_parameters_errors = np.array([unp.std_devs(a_mas), 
+                                            unp.std_devs(omega_deg), 
+                                            unp.std_devs(OMEGA_deg), 
+                                            unp.std_devs(i_deg)])
+
+    return geometric_parameters, geometric_parameters_errors
+
+
+def correlation_matrix(corr_vec):
+    """ 
+    This function reads the corr_vec from the nss_two_body_orbit table
+    and converts it to a numpy array with the full correlation matrix.
+    
+    Parameters
+    ----------
+    input_table : pandas Series or dict
+        A single row from the nss_two_body_orbit table for the desired target.
+
+    Returns
+    -------
+    corr_mat : ndarray
+        Correlation matrix for the specified parameters.
+    """
+
+    size = int((1+np.sqrt(8*len(corr_vec)+1))/2)
+
+    corr_mat = np.ones([size, size], dtype=float)
+
+    # Fill in lower triangle with corr_vec
+    corr_mat[np.tril_indices(size, k=-1)] = corr_vec
+    # Fill in upper triangle with mirror from lower triangle
+    upper_triangle = np.triu_indices(size, k=1)
+    corr_mat[upper_triangle] = corr_mat.T[upper_triangle]
+
+    return corr_mat
+
+
+def covariance_matrix(input_table, parameters=None):
+    """ 
+    Creates the covariance matrix from the corr_vec from the nss_two_body_orbit 
+    table. By default it works for the 12  'Orbital' parameter solution but it can
+    easily be adjusted to other solutions by providing the relevant parameters
+    (in the same order that they are listed in the Gaia documentation). 
+    
+    Parameters
+    ----------
+    input_table : pandas Series or dict
+        A single row from the nss_two_body_orbit table for the desired target.
+    parameters : array-like, optional
+        List of parameters for the corresponding solution of the desired
+        target. They have to be in the same order that they appear in the
+        Gaia documentation.
+
+    Returns
+    -------
+    cov_mat : ndarray
+        Covariance matrix for the specified parameters.
+    """
+
+    corr_vec = input_table['corr_vec']
+
+    if parameters is None:
+        # Use the 'Orbital' solution by default
+        parameters = ['ra', 'dec', 'parallax', 'pmra', 'pmdec', 'a_thiele_innes',
+                      'b_thiele_innes', 'f_thiele_innes', 'g_thiele_innes',
+                      'eccentricity', 'period', 't_periastron']
+    else:
+        # Check that the number of parameters provided matches the length of corr_vec
+        assert len(parameters) == int((1+np.sqrt(8*len(corr_vec)+1))/2), \
+               'Number of parameters does not match with the expected length of corr_vec'
+
+    size = len(parameters)
+
+    # First create correlation matrix
+    corr_mat = correlation_matrix(corr_vec)
+
+    # Copy of correlation matrix to construct covariance matrix
+    covar_mat = corr_mat.copy()
+
+    # Get uncertainties and apply variances to correlation matrix
+    pars_err = [par+'_error' for par in parameters]
+    error = np.array([input_table[err_col] for err_col in pars_err])
+    err_mat = np.dot(error.reshape(size, 1), error.reshape(1, size))
+    covar_mat = covar_mat * err_mat
+
+    return covar_mat