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 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 inv_map(self, sol: Solution) -> Solution:
        sol = copy.deepcopy(sol)

        if self.delta_ind_idx is None:
            self.delta_ind_idx = np.shape(sol.dynamical_parameters)[0] - 1

        sol.t = sol.t * sol.dynamical_parameters[self.delta_ind_idx]
        sol.dynamical_parameters = np.delete(sol.dynamical_parameters,
                                             self.delta_ind_idx)

        return sol
    def map(self, sol: Solution) -> Solution:
        sol = copy.deepcopy(sol)

        if self.delta_ind_idx is None:
            self.delta_ind_idx = np.shape(sol.dynamical_parameters)[0]

        delta_t = sol.t[-1] - sol.t[0]
        sol.dynamical_parameters = np.insert(sol.dynamical_parameters,
                                             self.delta_ind_idx, delta_t)

        sol.t = (sol.t - sol.t[0]) / delta_t

        return sol
def test_spbvp_3():
    # This problem contains a parameter, but it is not explicit in the BCs.
    # Since time is buried in the ODEs, this tests if the BVP solver calculates
    # sensitivities with respect to parameters.
    def odefun(_, p, __):
        return 1 * p[0]

    def bcfun(y0, yf, _, __, ___):
        return y0[0] - 0, yf[0] - 2

    algo = SPBVP(odefun, None, bcfun)
    solinit = Trajectory()
    solinit.t = np.linspace(0, 1, 4)
    solinit.y = np.array([[0], [0], [0], [0]])
    solinit.dynamical_parameters = np.array([1])
    solinit.const = np.array([])
    out = algo.solve(solinit)['sol']
    assert abs(out.dynamical_parameters - 2) < tol
        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 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