I am trying to integrate orbits in Python, and when considering only the 6x1 state vector x (3 spatial dimensions + 3 velocities), this is fairly easy and even easier when passing it to something like scipy.integrate.solve_ivp() alongside the system of ODEs in question.
But in trajectory design situations, we also want to use a matrix called the STM, which is a 6x6, and whose ODE is dependent on both itself and the state vector:
phi_d(t, t_0) = A(t)phi(t, t_0)
where A(t) = f(x(t)).
So my problem is this: I can write a function for the derivatives that takes in a pseudo-state kind of object that is a tuple or list of (state_vector, phi) and returns (state_vector_dot, phi_dot) (I can't use numpy.hstack() because of the differing dimensions). But then I can't just pass this to solve_ivp() as the same issue as the hstack() error comes up:
def deriv(state_phi):
state, phi_t = state_phi
x, y, z, x_d, y_d, z_d = state
A_i = self.STM_A(x, y ,z)
phi_d_t = np.matmul(A_i, phi_t)
state_deriv = self.EOM_eqn(0, state)
return [state_deriv, phi_d_t]
def integrate(self,
state: np.ndarray,
deriv: callable,
dt: float,
method: str = 'DOP853'):
solver = solve_ivp(deriv, t_span=(0,dt), y0=state, method=method)
assert abs(solver.t[-1]-dt) < tol
return np.array(solver.y).T[-1]
new_phi = integrate([current_state, current_STM], deriv, dt)[1]
Traceback:
fun = <function <locals>.deriv at 0x0000025E0B87B060>
y0 = [array([x, y, z , x_d, y_d, z_d]), array([[1., 0., 0., 0., 0., 0.],
[0., 1., 0., 0., 0.,..., 1., 0., 0., 0.],
[0., 0., 0., 1., 0., 0.],
[0., 0., 0., 0., 1., 0.],
[0., 0., 0., 0., 0., 1.]])]
support_complex = True
def check_arguments(fun, y0, support_complex):
"""Helper function for checking arguments common to all solvers."""
> y0 = np.asarray(y0)
E ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (2, 6) + inhomogeneous part.
(This was for using the DOP853 method, but I get the same if I use RK45).
So, how do I go about integrating a term (phi) whose ODE depends on other variables that are of differing dimension (x) that also require integrating their own ODE every step and virtual step of the phi ODE? I had considered trying to compute the whole trajectory's worth of x terms, then go back and use them as effectively lookup terms for my phi ODE, but this both introduces error as I'll need to interpolate between them and also is a horrible way to solve this that introduces a whole lot more complexity than I think should be required for something as straightforward as this. Also, this method would remove the nuance of the DOP853 and RK45 methods' error checking and variable timestep.
Any assistance would be very much appreciated.
The answer was straightforward in the end: simply flatten the matrix and concatenate that with the state vector, then inside the derivative function, un-concatenate and unflatten to get the vector and matrix back, do the matrix maths, and then reflatten and concatenate as the return of the derivative function (see below). Scipy doesn't care what it gets fed or how they relate to each other, and it doesnt care what the deriv function black box is doing internally either. It just wants something to feed in and something to feed out.
def deriv(t, pseudostate):
state, pf = pseudostate[:6], pseudostate[6:]
phi_t = np.array([pf[:6], pf[6:12], pf[12:18], pf[18:24], pf[24:30], pf[30:]])
x, y, z, x_d, y_d, z_d = state
A_i = self.STM_A(x, y ,z)
phi_d_t: np.ndarray = np.matmul(A_i, phi_t)
state_deriv = self.EOM_eqn(0, state)
phi_d_t_flat = phi_d_t.flatten()
return np.hstack((state_deriv,phi_d_t_flat))
def integrate(self,
state: np.ndarray,
deriv: callable,
dt: float,
method: str = 'DOP853'):
solver = solve_ivp(deriv, t_span=(0,dt), y0=state, method=method)
assert abs(solver.t[-1]-dt) < tol
return np.array(solver.y).T[-1]
new_phi = integrate(np.hstack((current_state, current_STM.flatten)), deriv, dt)[1]