def M2g(t, y):
vec = y[:-1, -1]
out = eom_func(vec, *args)
g = rn(dim + 1)
g.set_vector(out)
return g
def __call__(self, eom_func, quad_func, tspan, y0, q0, *args, **kwargs):
r"""
Propagates the differential equations over a defined time interval.
:param eom_func: Function representing the equations of motion.
:param quad_func: Function representing the quadratures.
:param tspan: Independent time interval.
:param y0: Initial state position.
:param q0: Initial quad position.
:param args: Additional arguments required by EOM files.
:param kwargs: Unused.
:return: A full reconstructed trajectory, :math:`\gamma`.
"""
y0 = np.array(y0, dtype=np.float64)
if self.program == 'scipy':
if self.variable_step is True:
int_sol = solve_ivp(lambda t, y: eom_func(y, *args),
[tspan[0], tspan[-1]],
y0,
rtol=self.reltol,
atol=self.abstol,
max_step=self.maxstep,
method=self.stepper)
else:
T = np.arange(tspan[0], tspan[-1], self.maxstep)
if T[-1] != tspan[-1]:
T = np.hstack((T, tspan[-1]))
int_sol = solve_ivp(lambda t, y: eom_func(y, *args),
[tspan[0], tspan[-1]],
y0,
rtol=self.reltol,
atol=self.abstol,
method=self.stepper,
t_eval=T)
gamma = Trajectory(int_sol.t, int_sol.y.T)
elif self.program == 'lie':
dim = y0.shape[0]
g = rn(dim + 1)
g.set_vector(y0)
y = HManifold(RN(dim + 1, exp(g)))
vf = VectorField(y)
vf.set_equationtype('general')
def M2g(t, y):
vec = y[:-1, -1]
out = eom_func(vec, *args)
g = rn(dim + 1)
g.set_vector(out)
return g
vf.set_M2g(M2g)
if self.method == 'RKMK':
ts = RKMK()
else:
raise NotImplementedError
ts.setmethod(self.stepper)
f = Flow(ts, vf, variablestep=self.variable_step)
ti, yi = f(y, tspan[0], tspan[-1], self.maxstep)
gamma = Trajectory(ti, np.vstack([_[:-1, -1] for _ in yi
])) # Hardcoded assuming RN
else:
raise NotImplementedError
if quad_func is not None and len(q0) is not 0:
if self.quick_reconstruct:
qf = integrate_quads(quad_func, tspan, gamma, *args)
gamma.q = np.vstack((q0, np.zeros(
(len(gamma.t) - 2, len(q0))), qf + q0))
else:
gamma = reconstruct(quad_func, gamma, q0, *args)
return gamma
def m2g(_, _y):
vec = _y[:-1, -1]
out = eom_func(vec, *args)
_g = rn(dim+1)
_g.set_vector(out)
return _g