def compile_deriv_jac_func(self, use_control_arg=False):
if self.prob.func_jac['df_dy'] is None:
return None
elif use_control_arg:
df_dy = self.lambdify(self._dynamic_args_w_controls,
np.array(self.prob.func_jac['df_dy']))
df_dp = self.lambdify(self._dynamic_args_w_controls,
np.array(self.prob.func_jac['df_dp']))
def deriv_func_jac(_y, _u, _p, _k):
return np.array(df_dy(_y, _u, _p,
_k)), np.array(df_dp(_y, _u, _p, _k))
self.deriv_func_jac = jit_compile_func(
deriv_func_jac,
self._dynamic_args_w_controls,
func_name='deriv_func_jac')
elif self.compute_u is not None:
df_dy = self.lambdify(self._dynamic_args_w_controls,
np.array(self.prob.func_jac['df_dy']))
df_dp = self.lambdify(self._dynamic_args_w_controls,
np.array(self.prob.func_jac['df_dp']))
compute_u = self.compute_u
def deriv_func_jac(_y, _p, _k):
_u = compute_u(_y, _p, _k)
return np.array(df_dy(_y, _u, _p,
_k)), np.array(df_dp(_y, _u, _p, _k))
self.deriv_func_jac = jit_compile_func(deriv_func_jac,
self._dynamic_args,
func_name='deriv_func_jac')
else:
df_dy = self.lambdify(self._dynamic_args,
np.array(self.prob.func_jac['df_dy']))
df_dp = self.lambdify(self._dynamic_args,
np.array(self.prob.func_jac['df_dp']))
def deriv_func_jac(_y, _p, _k):
return np.array(df_dy(_y, _p, _k)), np.array(df_dp(_y, _p, _k))
self.deriv_func_jac = jit_compile_func(deriv_func_jac,
self._dynamic_args,
func_name='deriv_func_jac')
return self.deriv_func_jac
def compile_control(control_options,
args,
ham_func,
lambdify_func=jit_lambdify):
num_options = len(control_options)
if num_options == 0:
return None
elif num_options == 1:
compiled_option = lambdify_func(args, control_options[0])
def calc_u(_y, _p, _k):
return np.array(compiled_option(_y, _p, _k))
else:
compiled_options = lambdify_func(args, control_options)
def calc_u(_y, _p, _k):
u_set = np.array(compiled_options(_y, _p, _k))
u = u_set[0, :]
ham = ham_func(_y, u, _p, _k)
for n in range(1, num_options):
ham_i = ham_func(_y, u_set[n, :], _p, _k)
if ham_i < ham:
u = u_set[n, :]
ham = ham_i
return u
return jit_compile_func(calc_u, args, func_name='control_function')
def __init__(self,
name: str,
func: Callable,
units: str,
arg_units: Collection,
dim_consistent=False,
local_compiler=None):
self.arg_types = ['scalar' for _ in arg_units]
if not dim_consistent:
raise NotImplementedError(
'Dimensionally inconsistant functions not yet implemented')
else:
self.func = jit_compile_func(func, self.arg_types)
self.sym = custom_functions.CustomFunctionGenerator(
self.func,
name=name,
arg_types=self.arg_types,
local_compiler=local_compiler)
local_compiler.add_symbolic_local(name, self.sym)
# local_compiler.add_function_local(name, self.func)
super(FunctionStruct, self).__init__(name,
units,
arg_units,
dim_consistent=dim_consistent,
local_compiler=local_compiler)
def compile_cost(self, use_quad_arg=True, use_control_arg=False):
if use_quad_arg:
_bc_args = self._bc_args_w_quads
else:
_bc_args = self._bc_args
self.initial_cost_func = self.lambdify(_bc_args,
self.prob.cost.initial)
self.terminal_cost_func = self.lambdify(_bc_args,
self.prob.cost.terminal)
if self.compute_u is None:
if use_control_arg:
_dyn_args = self._dynamic_args_w_controls
else:
_dyn_args = self._dynamic_args
self.path_cost_func = self.lambdify(_dyn_args, self.prob.cost.path)
else:
_compute_u = self.compute_u
_compute_path_cost = self.lambdify(self._dynamic_args_w_controls,
self.prob.cost.path)
def path_cost_func(_y, _p, _k):
_u = _compute_u(_y, _p, _k)
return np.array(_compute_path_cost(_y, _u, _p, _k))
self.path_cost_func = jit_compile_func(path_cost_func,
self._dynamic_args)
return self.initial_cost_func, self.path_cost_func, self.terminal_cost_func
def compile_deriv_func(self, use_control_arg=False):
sym_eom = getattr_from_list(self.prob.states, 'eom')
sym_eom_q = getattr_from_list(self.prob.quads, 'eom')
if self.compute_u is None:
if use_control_arg:
_args = self._dynamic_args_w_controls
else:
_args = self._dynamic_args
self.compute_y_dot = self.lambdify(_args, sym_eom)
if len(sym_eom_q) > 0:
self.compute_q_dot = self.lambdify(_args, sym_eom_q)
self.deriv_func, self.quad_func = self.compute_y_dot, self.compute_q_dot
else:
self.compute_y_dot = self.lambdify(self._dynamic_args_w_controls,
sym_eom)
compute_u = self.compute_u
compute_y_dot = self.compute_y_dot
def deriv_func(_y, _p, _k):
_u = compute_u(_y, _p, _k)
return np.array(compute_y_dot(_y, _u, _p, _k))
self.deriv_func = jit_compile_func(deriv_func, self._dynamic_args)
if len(sym_eom_q) > 0:
self.compute_q_dot = self.lambdify(
self._dynamic_args_w_controls, sym_eom_q)
compute_q_dot = self.compute_q_dot
def quad_func(_y, _p, _k):
_u = compute_u(_y, _p, _k)
return np.array(compute_q_dot(_y, _u, _p, _k))
self.quad_func = jit_compile_func(quad_func,
self._dynamic_args)
return self.deriv_func, self.quad_func
def __new__(cls, base_name, base_func, arg_list):
obj = super(CustomFunctionMeta, cls).__new__(cls, *arg_list)
obj.nargs = (len(arg_list), )
obj.arg_list = tuple(arg_list)
obj.base_name = base_name
if type(base_func) is CPUDispatcher:
obj.base_func = base_func
else:
obj.base_func = jit_compile_func(base_func, arg_list)
return obj
def __init__(self, prob: Problem):
num_options = len(prob.control_law)
_args = \
[prob.independent_variable.sym, extract_syms(prob.states), extract_syms(prob.costates),
extract_syms(prob.parameters), extract_syms(prob.constants)]
_args_w_control = copy.copy(_args)
_args_w_control.insert(3, extract_syms(prob.controls))
# self.compute_u = compile_control(prob.control_law, _args, prob.hamiltonian.expr, lambdify_func=prob.lambdify)
if num_options == 0:
raise RuntimeError
elif num_options == 1:
compiled_option = prob.lambdify(_args, prob.control_law[0])
def calc_u(_t, _y, _lam, _p, _k):
return np.array(compiled_option(_t, _y, _lam, _p, _k))
else:
compiled_options = prob.lambdify(_args, prob.control_law)
ham_func = prob.lambdify(_args_w_control, prob.hamiltonian.expr)
def calc_u(_t, _y, _lam, _p, _k):
u_set = np.array(compiled_options(_t, _y, _lam, _p, _k))
u = u_set[0, :]
ham = ham_func(_t, _y, _lam, u, _p, _k)
for n in range(1, num_options):
ham_i = ham_func(_t, _y, _lam, u_set[n, :], _p, _k)
if ham_i < ham:
u = u_set[n, :]
ham = ham_i
return u
self.compute_u = jit_compile_func(calc_u,
_args,
func_name='control_function')
def __init__(self,
base_func,
name=None,
arg_types=None,
local_compiler=None,
deriv_list=None,
order=None,
num_deriv_type='c_diff'):
if arg_types is None:
self.arg_types = [
'scalar' for _ in inspect.signature(base_func).parameters
]
else:
self.arg_types = arg_types
if order is None:
self.order = tuple([0] * len(self.arg_types))
else:
self.order = order
self.deriv_list = deriv_list
self.local_compiler = local_compiler
self.num_deriv_type = num_deriv_type
if type(base_func) is CPUDispatcher:
self.base_func = base_func
else:
self.base_func = jit_compile_func(base_func,
self.arg_types,
complex_numbers=False)
if name is None:
self.name = 'Sym(' + self.base_func.__name__ + ')'
else:
self.name = name
def compile_bc_jac_func(self, use_quad_arg=False):
if self.prob.bc_jac['initial']['dbc_dy'] is None:
return None
elif use_quad_arg:
_args = self._bc_args_w_quads
else:
_args = self._bc_args
num_bc_0, num_states = self.prob.bc_jac['initial']['dbc_dy'].shape
num_bc_f, num_parameters = self.prob.bc_jac['terminal']['dbc_dp'].shape
calc_dbc_0_dy = self.lambdify(
_args, np.array(self.prob.bc_jac['initial']['dbc_dy']))
calc_dbc_f_dy = self.lambdify(
_args, np.array(self.prob.bc_jac['terminal']['dbc_dy']))
calc_dbc_0_dp = self.lambdify(
_args, np.array(self.prob.bc_jac['initial']['dbc_dp']))
calc_dbc_f_dp = self.lambdify(
_args, np.array(self.prob.bc_jac['terminal']['dbc_dp']))
empty_dbc_0_dy = np.zeros((num_bc_0, num_states))
empty_dbc_f_dy = np.zeros((num_bc_f, num_states))
if use_quad_arg:
_, num_quads = self.prob.bc_jac['initial']['dbc_dq'].shape
empty_dbc_0_dq = np.zeros((num_bc_0, num_quads))
empty_dbc_f_dq = np.zeros((num_bc_f, num_quads))
calc_dbc_0_dq = self.lambdify(
_args, np.array(self.prob.bc_jac['initial']['dbc_dq']))
calc_dbc_f_dq = self.lambdify(
_args, np.array(self.prob.bc_jac['terminal']['dbc_dq']))
def bc_func_jac(_y0, _q0, _yf, _qf, _p, _p_con, _k):
dbc_dy0 = np.vstack(
(np.array(calc_dbc_0_dy(_y0, _q0, _p, _p_con,
_k)), empty_dbc_f_dy))
dbc_dyf = np.vstack(
(empty_dbc_0_dy,
np.array(calc_dbc_f_dy(_yf, _qf, _p, _p_con, _k))))
dbc_dp = np.vstack(
(np.array(calc_dbc_0_dp(_y0, _q0, _p, _p_con, _k)),
np.array(calc_dbc_f_dp(_yf, _qf, _p, _p_con, _k))))
dbc_dq0 = np.vstack(
(np.array(calc_dbc_0_dq(_y0, _q0, _p, _p_con,
_k)), empty_dbc_f_dq))
dbc_dqf = np.vstack(
(empty_dbc_0_dq,
np.array(calc_dbc_f_dq(_yf, _qf, _p, _p_con, _k))))
return dbc_dy0, dbc_dyf, dbc_dp, dbc_dq0, dbc_dqf
_combined_args = \
([self.prob.independent_variable.sym] + self._state_syms + self._quad_syms) * 2 \
+ self._parameter_syms + self._constraint_parameters_syms + self._constant_syms
else:
def bc_func_jac(_y0, _yf, _p, _p_con, _k):
dbc_dy0 = np.vstack(
(np.array(calc_dbc_0_dy(_y0, _p, _p_con,
_k)), empty_dbc_f_dy))
dbc_dyf = np.vstack(
(empty_dbc_0_dy,
np.array(calc_dbc_f_dy(_yf, _p, _p_con, _k))))
dbc_dp = np.vstack(
(np.array(calc_dbc_0_dp(_y0, _p, _p_con, _k)),
np.array(calc_dbc_f_dp(_yf, _p, _p_con, _k))))
return dbc_dy0, dbc_dyf, dbc_dp
_combined_args = [
self._state_syms, self._state_syms, self._parameter_syms,
self._constraint_parameters_syms, self._constant_syms
]
self.bc_func_jac = jit_compile_func(bc_func_jac,
_combined_args,
func_name='bc_func_jac')
return self.bc_func_jac
def compile_bc(self, use_quad_arg=False):
sym_initial_bc = getattr_from_list(self.prob.constraints['initial'],
'expr')
sym_terminal_bc = getattr_from_list(self.prob.constraints['terminal'],
'expr')
if use_quad_arg:
if self.compute_u is None:
_args = self._bc_args_w_quads
compute_initial_bc = self.lambdify(_args, sym_initial_bc)
compute_terminal_bc = self.lambdify(_args, sym_terminal_bc)
def bc_func(_y0, _q0, _yf, _qf, _p, _p_con, _k):
bc_0 = np.array(
compute_initial_bc(_y0, _q0, _p, _p_con, _k))
bc_f = np.array(
compute_terminal_bc(_yf, _qf, _p, _p_con, _k))
return np.concatenate((bc_0, bc_f))
else:
_args = self._bc_args_w_quads_controls
compute_initial_bc = self.lambdify(_args, sym_initial_bc)
compute_terminal_bc = self.lambdify(_args, sym_terminal_bc)
compute_u = self.compute_u
def bc_func(_y0, _q0, _yf, _qf, _p, _p_con, _k):
_u0 = compute_u(_y0, _p, _k)
_uf = compute_u(_yf, _p, _k)
bc_0 = np.array(
compute_initial_bc(_y0, _q0, _u0, _p, _p_con, _k))
bc_f = np.array(
compute_terminal_bc(_yf, _qf, _uf, _p, _p_con, _k))
return np.concatenate((bc_0, bc_f))
_combined_args = \
([self.prob.independent_variable.sym] + self._state_syms + self._quad_syms) * 2 \
+ self._parameter_syms + self._constraint_parameters_syms + self._constant_syms
self.bc_func = jit_compile_func(bc_func, _combined_args)
self.compute_initial_bc, self.compute_terminal_bc = compute_initial_bc, compute_terminal_bc
else:
if self.compute_u is None:
_args = self._bc_args
compute_initial_bc = self.lambdify(_args, sym_initial_bc)
compute_terminal_bc = self.lambdify(_args, sym_terminal_bc)
def bc_func(_y0, _yf, _p, _p_con, _k):
bc_0 = np.array(compute_initial_bc(_y0, _p, _p_con, _k))
bc_f = np.array(compute_terminal_bc(_yf, _p, _p_con, _k))
return np.concatenate((bc_0, bc_f))
else:
_args = self._bc_args_w_controls
compute_initial_bc = self.lambdify(_args, sym_initial_bc)
compute_terminal_bc = self.lambdify(_args, sym_terminal_bc)
compute_u = self.compute_u
def bc_func(_y0, _yf, _p, _p_con, _k):
_u0 = compute_u(_y0, _p, _k)
_uf = compute_u(_yf, _p, _k)
bc_0 = np.array(
compute_initial_bc(_y0, _u0, _p, _p_con, _k))
bc_f = np.array(
compute_terminal_bc(_yf, _uf, _p, _p_con, _k))
return np.concatenate((bc_0, bc_f))
_combined_args = [self._state_syms] * 2 + [self._parameter_syms] + [self._constraint_parameters_syms] \
+ [self._constant_syms]
self.bc_func = jit_compile_func(bc_func, _combined_args)
self.compute_initial_bc, self.compute_terminal_bc = compute_initial_bc, compute_terminal_bc
return self.bc_func