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
