def test_shooting_4():
# This problem contains a quad and tests if the prob solver correctly
# integrates the quadfun. Also tests multiple shooting.
def odefun(x, _, __):
return -x[1], x[0]
def quadfun(x, _, __):
return x[0]
def bcfun(y0, _, __, qf, ___, ____, _____):
return y0[0], y0[1] - 1, qf[0] - 1.0
algo = Shooting(odefun, quadfun, bcfun, num_arcs=4)
solinit = Trajectory()
solinit.t = np.linspace(0, np.pi / 2, 2)
solinit.y = np.array([[1, 0], [1, 0]])
solinit.q = np.array([[0], [0]])
solinit.k = np.array([])
out = algo.solve(solinit)['traj']
assert (out.y[0, 0] - 0) < tol
assert (out.y[0, 1] - 1) < tol
assert (out.q[0, 0] - 2) < tol
assert (out.q[-1, 0] - 1) < tol
def inv_map(self, sol: Solution) -> Solution:
qinv = np.ones_like(sol.t)
pinv = np.ones_like(sol.t)
for ii, t in enumerate(sol.t):
qinv[ii] = self.fn_q_inv(sol.y[ii], sol.lam[ii], sol.q[ii], sol.p,
sol.k)
pinv[ii] = self.fn_p_inv(sol.y[ii], sol.lam[ii], sol.p, sol.k)
state = pinv
qval = qinv
if self.remove_parameter_dict['location'] == 'states':
sol.y = np.column_stack(
(sol.y[:, :self.remove_parameter_dict['index']], state,
sol.y[:, self.remove_parameter_dict['index']:]))
elif self.remove_parameter_dict['location'] == 'costates':
sol.lam = np.column_stack(
(sol.lam[:, :self.remove_parameter_dict['index']], state,
sol.lam[:, self.remove_parameter_dict['index']:]))
if self.remove_symmetry_dict['location'] == 'states':
sol.y = np.column_stack(
(sol.y[:, :self.remove_symmetry_dict['index']], qval,
sol.y[:, self.remove_symmetry_dict['index']:]))
elif self.remove_symmetry_dict['location'] == 'costates':
sol.lam = np.column_stack(
(sol.lam[:, :self.remove_symmetry_dict['index']], qval,
sol.lam[:, self.remove_symmetry_dict['index']:]))
sol.q = np.delete(sol.q, np.s_[-1], axis=1)
sol.p = sol.p[:-1]
return sol
def map(self, sol: Solution) -> Solution:
cval = self.fn_p(sol.y[0], sol.lam[0], sol.p, sol.k)
qval = np.ones_like(sol.t)
sol.p = np.hstack((sol.p, cval))
for ii, t in enumerate(sol.t):
qval[ii] = self.fn_q(sol.y[ii], sol.lam[ii], sol.p, sol.k)
if self.remove_parameter_dict['location'] == 'states':
sol.y = np.delete(sol.y,
np.s_[self.remove_parameter_dict['index']],
axis=1)
elif self.remove_parameter_dict['location'] == 'costates':
sol.lam = np.delete(sol.lam,
np.s_[self.remove_parameter_dict['index']],
axis=1)
if self.remove_symmetry_dict['location'] == 'states':
sol.y = np.delete(sol.y,
np.s_[self.remove_symmetry_dict['index']],
axis=1)
elif self.remove_symmetry_dict['location'] == 'costates':
sol.lam = np.delete(sol.lam,
np.s_[self.remove_symmetry_dict['index']],
axis=1)
sol.q = np.column_stack((sol.q, qval))
return sol
def scale_sol(sol: Trajectory, scale_factors, inv=False):
sol = copy.deepcopy(sol)
if inv:
op = np.multiply
else:
op = np.divide
sol.t = op(sol.t, scale_factors[0])
sol.y = op(sol.y, scale_factors[1])
sol.q = op(sol.q, scale_factors[2])
if sol.u.size > 0:
sol.u = op(sol.u, scale_factors[3])
sol.p = op(sol.p, scale_factors[4])
sol.nu = op(sol.nu, scale_factors[5])
sol.k = op(sol.k, scale_factors[6])
return sol
def test_r18(const):
def odefun(y, _, k):
return -y[0] / k[0]
def quadfun(y, _, __):
return y[0]
def bcfun(_, q0, __, qf, ___, ____, k):
return q0[0] - 1, qf[0] - np.exp(-1 / k[0])
algo = SPBVP(odefun, quadfun, bcfun)
solinit = Trajectory()
solinit.t = np.linspace(0, 1, 2)
solinit.y = np.array([[1], [1]])
solinit.q = np.array([[0], [0]])
solinit.k = np.array([const])
sol = algo.solve(solinit)['traj']
e1 = np.exp(-sol.t / sol.k[0])
e2 = -1 / (sol.k[0] * np.exp(sol.t / sol.k[0]))
assert all(e1 - sol.q[:, 0] < tol)
assert all(e2 - sol.y[:, 0] < tol)
def test_r8(const):
def odefun(y, _, k):
return -y[0] / k[0]
def quadfun(y, _, __):
return y[0]
def bcfun(_, q0, __, qf, ___, ____, _____):
return q0[0] - 1, qf[0] - 2
algo = Shooting(odefun, quadfun, bcfun)
solinit = Trajectory()
solinit.t = np.linspace(0, 1, 2)
solinit.y = np.array([[1], [1]])
solinit.q = np.array([[0], [0]])
solinit.k = np.array([const])
sol = algo.solve(solinit)['traj']
e1 = (1.e0 - np.exp((sol.t - 1.e0) / sol.k)) / (1.e0 - np.exp(-1.e0 / sol.k))
e2 = np.exp((sol.t - 1) / sol.k) / (sol.k * (1 / np.exp(1 / sol.k) - 1))
assert all(e1 - sol.q[:, 0] < tol)
assert all(e2 - sol.y[:, 0] < tol)
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: FunctionComponent representing the equations of motion.
:param quad_func: FunctionComponent 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=beluga.DTYPE)
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:
ti = np.arange(tspan[0], tspan[-1], self.maxstep)
if ti[-1] != tspan[-1]:
ti = np.hstack((ti, 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=ti)
gamma = Trajectory(t=int_sol.t, y=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(_, _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) != 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