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.dual[ii], sol.q[ii], sol.dynamical_parameters, sol.const) pinv[ii] = self.fn_p_inv(sol.y[ii], sol.dual[ii], sol.dynamical_parameters, sol.const) 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.dual = np.column_stack( (sol.dual[:, :self.remove_parameter_dict['index']], state, sol.dual[:, 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.dual = np.column_stack( (sol.dual[:, :self.remove_symmetry_dict['index']], qval, sol.dual[:, self.remove_symmetry_dict['index']:])) sol.q = np.delete(sol.q, np.s_[-1], axis=1) sol.dynamical_parameters = sol.dynamical_parameters[:-1] return sol
def map(self, sol: Solution) -> Solution: cval = self.fn_p(sol.y[0], sol.dual[0], sol.dynamical_parameters, sol.const) qval = np.ones_like(sol.t) sol.dynamical_parameters = np.hstack((sol.dynamical_parameters, cval)) for ii, t in enumerate(sol.t): qval[ii] = self.fn_q(sol.y[ii], sol.dual[ii], sol.dynamical_parameters, sol.const) 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.dual = np.delete(sol.dual, 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.dual = np.delete(sol.dual, 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.dynamical_parameters = op(sol.dynamical_parameters, scale_factors[4]) sol.nondynamical_parameters = op(sol.nondynamical_parameters, scale_factors[5]) sol.const = op(sol.const, 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.const = np.array([const]) sol = algo.solve(solinit)['sol'] e1 = np.exp(-sol.t / sol.const[0]) e2 = -1 / (sol.const[0] * np.exp(sol.t / sol.const[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 = 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.const = np.array([const]) sol = algo.solve(solinit)['sol'] e1 = (1.e0 - np.exp( (sol.t - 1.e0) / sol.const)) / (1.e0 - np.exp(-1.e0 / sol.const)) e2 = np.exp((sol.t - 1) / sol.const) / (sol.const * (1 / np.exp(1 / sol.const) - 1)) assert all(e1 - sol.q[:, 0] < tol) assert all(e2 - sol.y[:, 0] < tol)
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.const = np.array([]) out = algo.solve(solinit)['sol'] 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 solve(self, solinit, **kwargs): solinit = copy.deepcopy(solinit) nstates = solinit.y.shape[1] nquads = 0 def return_nil(*_, **__): return np.array([]) if solinit.q.size > 0: nquads = solinit.q.shape[1] else: nquads = 0 self.quadrature_function = return_nil ndyn = solinit.dynamical_parameters.size nnondyn = solinit.nondynamical_parameters.size empty_array = np.array([]) if nquads == 0: # TODO: Try to vectorize def _fun(t, y, params=empty_array, const=solinit.const): return np.vstack([self.derivative_function(yi[:nstates], params[:ndyn], const) for yi in y.T]).T def _bc(ya, yb, params=empty_array, const=solinit.const): return self.boundarycondition_function(ya, yb, params[:ndyn], params[ndyn:ndyn + nnondyn], const) else: def _fun(t, y, params=empty_array, const=solinit.const): y = y.T o1 = np.vstack([self.derivative_function(yi[:nstates], params[:ndyn], const) for yi in y]) o2 = np.vstack([self.quadrature_function(yi[:nstates], params[:ndyn], const) for yi in y]) return np.hstack((o1, o2)).T def _bc(ya, yb, params=np.array([]), const=solinit.const): return self.boundarycondition_function(ya[:nstates], ya[nstates:nstates+nquads], yb[:nstates], yb[nstates:nstates+nquads], params[:ndyn], params[ndyn:ndyn+nnondyn], const) if self.derivative_function_jac is not None: def _fun_jac(t, y, params=np.array([]), const=solinit.const): y = y.T df_dy = np.zeros((y[0].size, y[0].size, t.size)) df_dp = np.zeros((y[0].size, ndyn+nnondyn, t.size)) for ii, yi in enumerate(y): df_dy[:, :, ii], _df_dp = self.derivative_function_jac(yi, params[:ndyn], const) if nstates > 1 and len(_df_dp.shape) == 1: _df_dp = np.array([_df_dp]).T df_dp[:, :, ii] = np.hstack((_df_dp, np.zeros((nstates, nnondyn)))) if ndyn + nnondyn == 0: return df_dy else: return df_dy, df_dp else: _fun_jac = None if self.boundarycondition_function_jac is not None: if nquads > 0: def _bc_jac(ya, yb, params=np.array([]), const=solinit.const): dbc_dya, dbc_dyb, dbc_dp = \ self.boundarycondition_function_jac(ya[:nstates], ya[nstates:nstates+nquads], yb[:nstates], yb[nstates:nstates+nquads], params[:ndyn], params[ndyn:ndyn+nnondyn], const) return dbc_dya, dbc_dyb, dbc_dp else: def _bc_jac(ya, yb, params=np.array([]), const=solinit.const): dbc_dya, dbc_dyb, dbc_dp = \ self.boundarycondition_function_jac(ya, yb, params[:ndyn], params[ndyn:ndyn+nnondyn], const) return dbc_dya, dbc_dyb, dbc_dp else: _bc_jac = None if nquads > 0: opt = solve_bvp(_fun, _bc, solinit.t, np.hstack((solinit.y, solinit.q)).T, np.hstack((solinit.dynamical_parameters, solinit.nondynamical_parameters)), max_nodes=self.max_nodes, fun_jac=_fun_jac, bc_jac=_bc_jac) else: opt = solve_bvp(_fun, _bc, solinit.t, solinit.y.T, np.hstack((solinit.dynamical_parameters, solinit.nondynamical_parameters)), max_nodes=self.max_nodes, fun_jac=_fun_jac, bc_jac=_bc_jac) sol = Trajectory(solinit) sol.t = opt['x'] sol.y = opt['y'].T[:, :nstates] sol.q = opt['y'].T[:, nstates:nstates+nquads] sol.dual = np.zeros_like(sol.y) if opt['p'] is not None: sol.dynamical_parameters = opt['p'][:ndyn] sol.nondynamical_parameters = opt['p'][ndyn:ndyn+nnondyn] else: sol.dynamical_parameters = np.array([]) sol.nondynamical_parameters = np.array([]) sol.converged = opt['success'] out = BVPResult(sol=sol, success=opt['success'], message=opt['message'], rms_residuals=opt['rms_residuals'], niter=opt['niter']) return out
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: 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) != 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