rendezvous_sim.py 7.88 KB
Newer Older
navidmokh's avatar
navidmokh committed
1
2
3
4
5
6
7
###########################
######### Imports #########
###########################
import matlab.engine
import pandas as pd
import numpy as np
import datetime as dt
8
# import matplotlib.pyplot as plt
navidmokh's avatar
navidmokh committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

import os

###########################
##### Cython Imports ######
###########################
import pyximport; pyximport.install()
from acceleration import calculate_double_dots

###########################
####### Variables #########
###########################
COLUMNS         = ['time', 'x', 'y', 'x_dot', 'y_dot']
FULL_COLUMNS    = ['time', 'x', 'y', 'x_dot', 'y_dot', 'x_double_dot', 'y_double_dot', 'force_x', 'force_y']

mu          = 3.986e14 # m^3 / s^2
r_0         = 7.1e6 # m
n           = np.sqrt(mu/r_0**3)
m_chaser    = 500 # kg
m_target    = 2000 # kg

###########################
# compact_trace:    @input: arr (numpy.ndarray), compact factor (int)
#                   @output: compacted trace (numpy.ndarray)
# - This function takes in a trace and a compact factor. It uses
# - numpy splice notation to select each entry from the array with a
# - step size of compact_factor and returns the new array.
###########################
def compact_trace(arr, compact_factor):
    return arr[::compact_factor]

###########################
# generate_trace:   @input: curLabel (str), initCondition (list), transiteTime (int)
#                   @output: trace (numpy.ndarray)
# - This function starts a matlab session using the python matlab engine.
# - It then calls the matlab.engine.Simulate function. This function is
# - defined where the matlab.engine.addpath call points. It is defined by
# - the name of the file. The inputs to this function need to be altered to
# - meet the matlab format. The matlab code will generate a traces consiting
# - of a series of vectors containing [x, y, x_dot, y_dot] data. Since only
# - 4 values are calculated, only the first 4 initial conditions are passed.
# - The trace is returned as a numpy.ndarray
###########################
def generate_trace(curLabel, initCondition, transiteTime):
    eng = matlab.engine.start_matlab()
    path = os.path.abspath(os.path.dirname(__file__))+'/matlabcode'

    eng.addpath(path)
    sim_data = eng.Simulate(    curLabel,
                                matlab.double(initCondition[0:4]),
                                float(transiteTime)
    )
    return np.array(sim_data)

###########################
# add_variables:    @input: arr (numpy.ndarray)
#                   @output: arr (numpy.ndarray)
# - This function creates a new numpy.ndarray of zeros with same number of rows
# - as the input array and 4 columns to represent computed variables. The final
# - data will include: [x, y, x_dot. y_dot, x_double_dot, y_double_dot, force_x, force_y].
# - The new array is appended to the input array and returned.
###########################
def add_variables(arr):
    temp = np.zeros((arr.shape[0], 4))
    return np.append(arr, temp, axis=1)

###########################
# calculate_acceleration:   @input: arr (numpy.ndarray)
#                           @output: arr (numpy.ndarray)
# This function calculates the acceleration (double_dot) values for each
# timestep of an input array. It calls python code that has been compiled
# down to C for speed. It outputs the same array with double_dot values
# filled in.
###########################
def calculate_acceleration(arr):
    return calculate_double_dots(arr)

###########################
# find_min_index:   @input: arr (numpy.ndarray)
#                   @output: min_index (int)
# - This function finds the minimum index for the position of chaser spacecraft
# - in relation to the target spacecraft. This calculation is referred to as "rho"
# - which is defined as rho(x,y) = sqrt(x^2 + y^2). For each timestep in the trace
# - the function computes rho and compares it to the current minimum value seen.
# - Once it goes through the entire trace, it returns the index where the
# - minimum value occured.
###########################
def find_min_index(arr):
    min_val = float('Inf')
    min_index = None
    for i in range(arr.shape[0]):
        rho = np.sqrt(arr[i][1]**2 + arr[i][2]**2)
        if rho < min_val:
            min_val = rho
            min_index = i
    return min_index

###########################
# calculate_forces:     @input: arr (numpy.ndarray)
#                       @output: arr (numpy.ndarray)
# - This function take in a trace in the form of a numpy.ndarray. It works in two
# - steps: before rendezvous and after rendezvous. First, the min_index is found
# - using the find_min_index function. Next, for each timestep in the trace,
# - there is a check to see if it is the minimum index. If not, the trace is in
# - before rendezvous state. It calculates the force in the x and y direction as
# - it was stated in the research paper. If the index is the minimum found, the
# - chaser and target are as close as they will ever be and it can be assumed that
# - the rendezvous has occured. Since we no longer want to check this, the check
# - variable is set to false. We then add the mass of the target to the mass used
# - in calculating the force. Finally we use F = ma to calculate the force.
# - The array with filled in values for force is returned.
###########################
def calculate_forces(arr):
    m = m_chaser
    min_index = find_min_index(arr)
    check = True
    for i in range(arr.shape[0]):
        if check:
            if i == min_index:
                check = False
                m += m_target
                arr[i][7] = m * arr[i][5]
                arr[i][8] = m * arr[i][6]
            else:
                arr[i][7] = m * (arr[i][5] - 2*n*arr[i][4] - 3*n*n*arr[i][1])
                arr[i][8] = m * (arr[i][6] + 2*n*arr[i][3])
        else:
            arr[i][7] = m * arr[i][5]
            arr[i][8] = m * arr[i][6]
    return arr

###########################
# plot_total_force:     @input: arr (numpy.ndarray)
#                       @output: None
# - This function plots the total force vs time as a scatter plot. It was
# - used as debugging tool.
###########################
def plot_total_force(arr):
    df = pd.DataFrame(arr, columns=FULL_COLUMNS)
    force = np.sqrt(df['force_x']*df['force_x'] + df['force_y']*df['force_y'])
    plt.scatter(df['time'], force, s=1)
    plt.show()

###########################
# TC_Simulate:      @input: curLabel (str), initCondition (list), transiteTime (int)
#                   @output: arr (numpy.ndarray)
# - This is the function called by DryVR. It contains print statements and timing
# - that is provide to the user of DryVR. The main steps taken in this function:
# -     1) generate_trace
# -     2) add_variables (that are values of zero i.e. placeholders)
# -     3) calculate_acceleration
# -     4) compact_trace (done before adding force for runtime efficiencies)
# -     5) calculate_forces
# -     6) return trace
# - Each step is timed so that it is obvious what the limiting functions are.
###########################
def TC_Simulate(curLabel, initCondition, transiteTime):
li213's avatar
li213 committed
166
    print("\n~~~ Starting Simulation ~~~")
navidmokh's avatar
navidmokh committed
167
168
169
170
    start = dt.datetime.now()

    time = dt.datetime.now()
    arr = generate_trace(curLabel, initCondition, transiteTime)
li213's avatar
li213 committed
171
    print('\tGenerated Trace:', (dt.datetime.now() - time).total_seconds() * 1000, 'ms')
navidmokh's avatar
navidmokh committed
172
173
174

    time = dt.datetime.now()
    arr = add_variables(arr)
li213's avatar
li213 committed
175
    print('\tAdded Variables:', (dt.datetime.now() - time).microseconds / 1000, 'ms')
navidmokh's avatar
navidmokh committed
176
177
178

    time = dt.datetime.now()
    arr = calculate_acceleration(arr)
li213's avatar
li213 committed
179
    print('\tCalculated Acceleration:', (dt.datetime.now() - time).microseconds / 1000, 'ms')
navidmokh's avatar
navidmokh committed
180
181
182

    time = dt.datetime.now()
    arr = compact_trace(arr, 100)
li213's avatar
li213 committed
183
    print('\tCompacted Trace:', (dt.datetime.now() - time).microseconds / 1000, 'ms')
navidmokh's avatar
navidmokh committed
184
185
186

    time = dt.datetime.now()
    arr = calculate_forces(arr)
li213's avatar
li213 committed
187
    print('\tCalculated Forces:', (dt.datetime.now() - time).microseconds / 1000, 'ms')
navidmokh's avatar
navidmokh committed
188

li213's avatar
li213 committed
189
190
    print("Completed:", (dt.datetime.now() - start).total_seconds() * 1000, 'ms')
    print("Trace Length:", len(arr))
navidmokh's avatar
navidmokh committed
191
    return arr