Source code for orbitdeterminator.kep_determination.ellipse_fit

"""Finds out the ellipse that best fits to a set of data points and calculates
   its keplerian elements.
"""

import math
import argparse
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize
from functools import partial

def __read_args():
    """Reads command line arguments.

       Returns:
           object: Parsed arguments.
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--file', type=str, help='path to .csv file', default='../example_data/orbit.csv')
    parser.add_argument('-u', '--units', type=str, help='units of distance (m or km)', default='km')
    return parser.parse_args()


def __cross_sum(data):
    """Returns the normalized sum of the cross products between consecutive vectors.

    Args:
        data(nx3 numpy array): A matrix where each column represents the x,y,z coordinates of each position vector.

    Returns:
        float: The normalized sum of the cross products between consecutive vectors.
    """

    cross_sum = 0
    for i in range(len(data)-1):
        v1 = data[i]
        v2 = data[i+1]
        cross_sum = cross_sum + np.cross(v1,v2)

    return cross_sum/np.linalg.norm(cross_sum)


def __plane_err(data,coeffs):
    """Calculates the total squared error of the data wrt a plane.

       The data should be a list of points. coeffs is an array of
       3 elements - the coefficients a,b,c in the plane equation
       ax+by+c = 0.

       Args:
           data(nx3 numpy array): A numpy array of points.
           coeffs(1x3 array): The coefficients of the plane ax+by+c=0.

       Returns:
           float: The total squared error wrt the plane defined by ax+by+cz = 0.
    """

    a,b,c = coeffs
    return np.sum((a*data[:,0]+b*data[:,1]+c*data[:,2])**2)/(a**2+b**2+c**2)


def __project_to_plane(points,coeffs):
    """Projects points onto a plane.

       Projects a list of points onto the plane ax+by+c=0,
       where a,b,c are elements of coeffs.

       Args:
           points(nx3 numpy array): A numpy array of points.
           coeffs(1x3 array): The coefficients of the plane ax+by+c=0.

       Returns:
           nx3 numpy array: A list of projected points.
    """

    a,b,c = coeffs

    proj_mat =  [[b**2+c**2,  -a*b   ,   -a*c  ],
                 [   -a*b  ,a**2+c**2,   -b*c  ],
                 [   -a*c  ,  -b*c   ,a**2+b**2]]

    return np.matmul(points,proj_mat)/(a**2+b**2+c**2)


def __conv_to_2D(points,x,y):
    """Finds coordinates of points in a plane wrt a basis.

       Given a list of points in a plane, and a basis of the plane,
       this function returns the coordinates of those points
       wrt this basis.

       Args:
           points(numpy array): A numpy array of points.
           x(3x1 numpy array): One vector of the basis.
           y(3x1 numpy array): Another vector of the basis.

       Returns:
           nx2 numpy array: Coordinates of the points wrt the basis [x,y].
    """

    mat = [x[0:2],y[0:2]]
    mat_inv = np.linalg.inv(mat)
    coords = np.matmul(points[:,0:2],mat_inv)

    return coords

def __cart_to_pol(points):
    """Converts a list of cartesian coordinates into polar ones.

       Args:
           points(nx2 numpy array): The list of points in the format [x,y].

       Returns:
           nx2 numpy array: A list of polar coordinates in the format [radius,angle].
    """

    pol = np.empty(points.shape)
    pol[:,0] = np.sqrt(points[:,0]**2+points[:,1]**2)
    pol[:,1] = np.arctan2(points[:,1],points[:,0])

    return pol

def __ellipse_err(polar_coords,params):
    """Calculates the total squared error of the data wrt an ellipse.

       params is a 3 element array used to define an ellipse.
       It contains 3 elements a,e, and t0.

       a is the semi-major axis
       e is the eccentricity
       t0 is the angle of the major axis wrt the x-axis.

       These 3 elements define an ellipse with one focus at origin.
       Equation of the ellipse is r = a(1-e^2)/(1+ecos(t-t0))

       The function calculates r for every theta in the data.
       It then takes the square of the difference and sums it.

       Args:
           polar_coords(nx2 numpy array): A list of polar coordinates in the format [radius,angle].
           params(1x3 numpy array): The array [a,e,t0].

       Returns:
           float: The total squared error of the data wrt the ellipse.
    """

    a,e,t0 = params
    dem = 1+e*np.cos(polar_coords[:,1]-t0)
    num = a*(1-e**2)
    r = np.divide(num,dem)
    err = np.sum((r - polar_coords[:,0])**2)
    return err


def __residuals(data,params,polar_coords,basis):
    """Calculates the residuals after fitting the ellipse.

       Residuals are the difference between the fitted points and
       the actual points.

       res_x = fitted_x - initial_x
       res_y = fitted_y - initial_y
       res_z = fitted_z - initial_z

       where fitted_x,y,z is the closest point on the ellipse to initial_x,y,z.

       However, it is computationally expensive to find the true nearest point.
       So we take an approximation. We consider the point on the ellipse with
       the same true anomaly as the initial point to be the nearest point to it.
       Since the eccentricities of the orbits involved are small, this approximation
       holds.

       Args:
           data(nx3 numpy array): The list of original points.
           params(1x3 numpy array): The array [semi-major axis, eccentricity, argument of periapsis]
                                    of the fitted ellipse.
           polar_coords(nx2 numpy array): The list of 2D polar coordinates of the original points after
                                          projecting them onto the best-fit plane.
           basis(3x2 numpy array): The basis of the best-fit plane.

       Returns:
            nx3 numpy array: Returns the residuals
    """

    a,e,t0 = params
    dem = 1+e*np.cos(polar_coords[:,1]-t0)
    num = a*(1-e**2)
    r = np.divide(num,dem)

    # convert to cartesian
    x_s = np.multiply(r,np.cos(polar_coords[:,1]))
    y_s = np.multiply(r,np.sin(polar_coords[:,1]))

    # convert to 3D
    filtered_coords = np.transpose(np.matmul(basis,[x_s,y_s]))

    residuals = filtered_coords - data

    return residuals

def __read_file(file_name):
    """Reads a space separated csv file with 4 columns in the format t x y z.

       Args:
           file_name(string): the path to the file

       Returns:
           nx3 numpy array: A numpy array with the columns [x y z]. Note that the t coloumn is discarded.
    """

    data = np.loadtxt(file_name,skiprows=1,usecols=(1,2,3))

    return data

[docs]def determine_kep(data): """Determines keplerian elements that fit a set of points. Args: data(nx3 numpy array): A numpy array of points in the format [x y z]. Returns: (kep,res) - The keplerian elements and the residuals as a tuple. kep: 1x6 numpy array res: nx3 numpy array For the keplerian elements: kep[0] - semi-major axis (in whatever units the data was provided in) kep[1] - eccentricity kep[2] - inclination (in degrees) kep[3] - argument of periapsis (in degrees) kep[4] - right ascension of ascending node (in degrees) kep[5] - true anomaly of the first row in the data (in degrees) For the residuals: (in whatever units the data was provided in) res[0] - residuals in x axis res[1] - residuals in y axis res[2] - residuals in z axis """ # try to fit a plane to the data first. # make a partial function of plane_err by supplying the data plane_err_data = partial(__plane_err,data) # plane is defined by ax+by+cz=0. p0 = __cross_sum(data) # make an initial guess # minimize the error p = minimize(plane_err_data,p0,method='nelder-mead',options={'maxiter':1000}).x p = p/np.linalg.norm(p) # normalize p # now p is the normal vector of the best-fit plane. # lan_vec is a vector along the line of intersection of the plane # and the x-y plane. lan_vec = np.cross([0,0,1],p) # if lan_vec is [0,0,0] it means that it is undefined and can take on # any value. So we set it to [1,0,0] so that the rest of the # calculation can proceed. if (np.array_equal(lan_vec,[0,0,0])): lan_vec = [1,0,0] # inclination is the angle between p and the z axis. inc = math.acos(np.clip(p[2]/np.linalg.norm(p),-1,1)) # lan is the angle between the lan_vec and the x axis. lan = math.atan2(lan_vec[1],lan_vec[0])%(2*math.pi) # now we try to convert the problem into a 2D problem. # project all the points onto the plane. proj_data = __project_to_plane(data,p) # p_x and p_y are 2 orthogonal unit vectors on the plane. p_x,p_y = lan_vec, np.cross(p,lan_vec) p_x,p_y = p_x/np.linalg.norm(p_x), p_y/np.linalg.norm(p_y) # find coordinates of the points wrt the basis [p_x,p_y]. coords_2D = __conv_to_2D(proj_data,p_x,p_y) # now try to fit an ellipse to these points. # convert them into polar coordinates polar_coords = __cart_to_pol(coords_2D) # make an initial guess for the parametres r_m = np.min(polar_coords[:,0]) r_M = np.max(polar_coords[:,0]) a0 = (r_m+r_M)/2 e0 = (r_M-r_m)/(r_M+r_m) t00 = polar_coords[np.argmin(polar_coords[:,0]),1] params0 = [a0,e0,t00] # initial guess # make a partial function of ellipse_err with the data ellipse_err_data = partial(__ellipse_err,polar_coords) # minimize the error params = minimize(ellipse_err_data,params0,method='nelder-mead',options={'maxiter':1000}).x params[2] = params[2]%(2*math.pi) # bring argp between 0-360 degrees # calculate the true anomaly of the first entry in the dataset true_anom = (polar_coords[0][1]-params[2])%(2*math.pi) # calculation of residuals res = __residuals(data,params,polar_coords,np.column_stack((p_x,p_y))) kep = np.empty((6,1)) kep[0] = params[0] kep[1] = params[1] kep[2] = math.degrees(inc) kep[3] = math.degrees(params[2]) kep[4] = math.degrees(lan) kep[5] = math.degrees(true_anom) return kep,res
def __print_kep(kep,res,unit): """Prints the keplerian elements and some information on residuals. Args: kep(1x6 numpy array): keplerian elements res(nx3 numpy array): residuals unit(string): units of distance used Returns: NIL """ # output the parameters print("Semi-major axis: ",kep[0][0],unit) print("Eccentricity: ",kep[1][0]) print("Inclination: ",kep[2][0],"deg") print("Argument of periapsis: ",kep[3][0],"deg") print("Longitude of Ascending Node:",kep[4][0],"deg") print("True Anomaly ",kep[5][0],"deg") # print data about residuals print() max_res = np.max(res,axis=0) min_res = np.min(res,axis=0) sum_res = np.sum(res,axis=0) avg_res = np.average(res,axis=0) std_res = np.std(res,axis=0) print("Printing data about residuals in each axis:") print("Max: ",max_res) print("Min: ",min_res) print("Sum: ",sum_res) print("Average: ",avg_res) print("Standard Deviation:",std_res)
[docs]def plot_kep(kep,data): """Plots the original data and the orbit defined by the keplerian elements. Args: kep(1x6 numpy array): keplerian elements data(nx3 numpy array): original data Returns: nothing """ a = kep[0] e = kep[1] inc = math.radians(kep[2]) t0 = math.radians(kep[3]) lan = math.radians(kep[4]) p_x = np.array([math.cos(lan), math.sin(lan), 0]) p_y = np.array([-math.sin(lan)*math.cos(inc), math.cos(lan)*math.cos(inc), math.sin(inc)]) # generate 1000 points on the ellipse theta = np.linspace(0,2*math.pi,1000) radii = a*(1-e**2)/(1+e*np.cos(theta-t0)) # convert to cartesian x_s = np.multiply(radii,np.cos(theta)) y_s = np.multiply(radii,np.sin(theta)) # convert to 3D mat = np.column_stack((p_x,p_y)) coords_3D = np.matmul(mat,[x_s,y_s]) fig = plt.figure() ax = Axes3D(fig) ax.axis('equal') # plot ax.plot3D(coords_3D[0],coords_3D[1],coords_3D[2],c = 'red',label='Fitted Ellipse') ax.scatter3D(data[:,0],data[:,1],data[:,2],c='black',label='Initial Data') # The Pale Blue Dot ax.scatter3D(0,0,0,c='blue',depthshade=False,label='Earth') ax.can_zoom() ax.legend() plt.show()
if __name__ == "__main__": args = __read_args() data = __read_file(args.file) kep, res = determine_kep(data) __print_kep(kep,res,args.units) plot_kep(kep,data)