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