diff --git a/opty-dev-env.yml b/opty-dev-env.yml index e2e073a0..4671cd21 100644 --- a/opty-dev-env.yml +++ b/opty-dev-env.yml @@ -2,6 +2,7 @@ name: opty-dev channels: - conda-forge dependencies: + - python <= 3.12 - coverage - cyipopt >=1.1.0 - cython >=0.29.19 diff --git a/opty/direct_collocation.py b/opty/direct_collocation.py index c598695e..327aeab4 100644 --- a/opty/direct_collocation.py +++ b/opty/direct_collocation.py @@ -153,7 +153,7 @@ def __init__(self, obj, obj_grad, equations_of_motion, state_symbols, states, unknown trajectories, unknown parameters, or unknown time interval to a 2-tuple of floats, the first being the lower bound and the second the upper bound for that free variable, e.g. - ``{x(t): (-1.0, 5.0)}``. + ``{x(t): (-1.0, 5.0)}``. Mandatory for timeshift parameters. """ @@ -165,7 +165,7 @@ def __init__(self, obj, obj_grad, equations_of_motion, state_symbols, self.collocator = ConstraintCollocator( equations_of_motion, state_symbols, num_collocation_nodes, - node_time_interval, known_parameter_map, known_trajectory_map, + node_time_interval, known_parameter_map, known_trajectory_map, bounds, instance_constraints, time_symbol, tmp_dir, integration_method, parallel, show_compile_output=show_compile_output) @@ -205,6 +205,9 @@ def _generate_bound_arrays(self): num_state_nodes = N*self.collocator.num_states num_non_par_nodes = N*(self.collocator.num_states + self.collocator.num_unknown_input_trajectories) + num_var_dur_node = num_non_par_nodes + \ + self.collocator.num_unknown_parameters + num_tshift_nodes = num_var_dur_node + int(self.collocator._variable_duration) state_syms = self.collocator.state_symbols unk_traj = self.collocator.unknown_input_trajectories unk_par = self.collocator.unknown_parameters @@ -230,8 +233,12 @@ def _generate_bound_arrays(self): ub[idx] = bounds[1] elif (self.collocator._variable_duration and var == self.collocator.time_interval_symbol): - lb[-1] = bounds[0] - ub[-1] = bounds[1] + lb[num_var_dur_node] = bounds[0] + ub[num_var_dur_node] = bounds[1] + elif var in self.collocator.unknown_tshift_parameters: + i = self.collocator.unknown_tshift_parameters.index(var) + lb[num_tshift_nodes+i] = bounds[0] + ub[num_tshift_nodes+i] = bounds[1] else: msg = 'Bound variable {} not present in free variables.' raise ValueError(msg.format(var)) @@ -408,12 +415,19 @@ def plot_trajectories(self, vector, axes=None): self.collocator.unknown_input_trajectories) trajectories = state_traj + + tshift_trajs = \ + [self.collocator._to_general_time(v[0]) for v in self.collocator.timeshift_traj_substitutes.values()] if self.collocator.num_known_input_trajectories > 0: for knw_sym in self.collocator.known_input_trajectories: - trajectories = np.vstack( - (trajectories, - self.collocator.known_trajectory_map[knw_sym])) + + traj = self.collocator.known_trajectory_map[knw_sym] + + if knw_sym in tshift_trajs: + offset = self.collocator.timeshift_traj_offsets[knw_sym] + traj = traj[offset:offset+self.collocator.num_collocation_nodes] + trajectories = np.vstack((trajectories,traj)) if self.collocator.num_unknown_input_trajectories > 0: # NOTE : input_traj should be in the same order as @@ -466,11 +480,13 @@ def plot_constraint_violations(self, vector, axes=None): bars_per_plot = None rotation = -45 + + num_tshift_constraints = (self.collocator.num_tshift_parameters * self.collocator.num_collocation_nodes) # find the number of bars per plot, so the bars per plot are # aproximately the same on each plot hilfs = [] - len_constr = self.collocator.num_instance_constraints + len_constr = self.collocator.num_instance_constraints - num_tshift_constraints for i in range(6, 11): hilfs.append((i, i - len_constr % i)) if len_constr % i == 0: @@ -490,11 +506,14 @@ def plot_constraint_violations(self, vector, axes=None): num_plots = 1 else: num_plots = len_constr // bars_per_plot + 1 + + # add timeshift plots + num_plots += self.collocator.num_tshift_parameters # ensure that len(axes) is correct, raise ValuError otherwise if axes is not None: len_axes = len(axes.ravel()) - len_constr = self.collocator.num_instance_constraints + len_constr = self.collocator.num_instance_constraints - num_tshift_constraints if (len_constr <= bars_per_plot) and (len_axes < 2): raise ValueError('len(axes) must be equal to 2') @@ -533,17 +552,29 @@ def plot_constraint_violations(self, vector, axes=None): axes[0].set_title('Constraint violations') axes[0].set_xlabel('Node Number') axes[0].set_ylabel('EoM violation') + + + #plot tshift constraint violations + num_viol = con_violations.size + tshift_con_nodes = range(0, N) + for itshift in range(self.collocator.num_tshift_parameters): + + tshift_violations = con_violations[num_viol-(N*(itshift+1)):num_viol-(N*itshift)] + + axes[-1-itshift].plot(tshift_con_nodes, tshift_violations) + axes[-1-itshift].set_ylabel(f'Timeshift violation {self.collocator.unknown_tshift_parameters[-1-itshift]}') + axes[-1-itshift].set_xlabel('Node Number') if self.collocator.instance_constraints is not None: # reduce the instance constrtaints to 2 digits after the decimal # point. give the time in tha variables with 2 digits after the # decimal point. if variable h is used, use the result for h in # the time. - num_inst_viols = self.collocator.num_instance_constraints + num_inst_viols = self.collocator.num_instance_constraints - num_tshift_constraints instance_constr_plot = [] a_before = '' a_before_before = '' - for exp1 in self.collocator.instance_constraints: + for exp1 in self.collocator.instance_constraints[:num_inst_viols]: for a in sm.preorder_traversal(exp1): if ((isinstance(a_before, sm.Integer) or isinstance(a_before, sm.Float)) and @@ -559,7 +590,7 @@ def plot_constraint_violations(self, vector, axes=None): a_before = a instance_constr_plot.append(exp1) - for i in range(num_plots): + for i in range(num_plots-self.collocator.num_tshift_parameters): num_ticks = bars_per_plot if i == num_plots - 1: beginn = i * bars_per_plot @@ -619,10 +650,12 @@ class ConstraintCollocator(object): - n(N - 1) + o : number of constraints """ + + INF = 10e19 def __init__(self, equations_of_motion, state_symbols, num_collocation_nodes, node_time_interval, - known_parameter_map={}, known_trajectory_map={}, + known_parameter_map={}, known_trajectory_map={}, bounds={}, instance_constraints=None, time_symbol=None, tmp_dir=None, integration_method='backward euler', parallel=False, show_compile_output=False): @@ -659,6 +692,12 @@ def __init__(self, equations_of_motion, state_symbols, free trajectories optimization variables. If solving a variable duration problem, note that the values here are fixed at each node and will not scale with a varying time interval. + bounds : dictionary, optional + This dictionary should contain a mapping from any of the symbolic + states, unknown trajectories, unknown parameters, or unknown time + interval to a 2-tuple of floats, the first being the lower bound + and the second the upper bound for that free variable, e.g. + ``{x(t): (-1.0, 5.0)}``. Mandatory for timeshift parameters. instance_constraints : iterable of SymPy expressions, optional These expressions are for constraints on the states at specific times. They can be expressions with any state instance and any of @@ -718,26 +757,35 @@ def __init__(self, equations_of_motion, state_symbols, self.known_trajectory_map = known_trajectory_map self.instance_constraints = instance_constraints + + self.bounds = bounds self.num_constraints = self.num_states * (num_collocation_nodes - 1) self.tmp_dir = tmp_dir self.parallel = parallel self.show_compile_output = show_compile_output - + + + self._substitute_timeshift_trajectories() self._sort_parameters() self._check_known_trajectories() + self._sort_trajectories() self.num_free = ((self.num_states + self.num_unknown_input_trajectories) * self.num_collocation_nodes + self.num_unknown_parameters + - int(self._variable_duration)) - + int(self._variable_duration) + + self.num_tshift_parameters) + self.integration_method = integration_method - if instance_constraints is not None: - self.num_instance_constraints = len(instance_constraints) + if instance_constraints is not None or self._timeshift_params: + if self._timeshift_params: + self._generate_timeshift_constraints() + self._precalc_timshift_input_derivatives() + self.num_instance_constraints = len(self.instance_constraints) self.num_constraints += self.num_instance_constraints self._identify_functions_in_instance_constraints() self._find_closest_free_index() @@ -832,6 +880,13 @@ def _sort_parameters(self): self.parameters = res[0] + res[2] self.num_parameters = len(self.parameters) + + if self._timeshift_params: + for tsparam in self.unknown_tshift_parameters: + if tsparam in self.parameters: + raise ValueError(f"Timeshift parameters cannot be direct parameters of the" + f"eoms: {tsparam}") + def _check_known_trajectories(self): """Raises and error if the known trajectories are not the correct @@ -839,10 +894,152 @@ def _check_known_trajectories(self): N = self.num_collocation_nodes + #tshift trajectories are already checked in _substiture_timeshift_trajectories() + tshift_trajs = \ + [self._to_general_time(v[0]) for v in self.timeshift_traj_substitutes.values()] + for k, v in self.known_trajectory_map.items(): - if len(v) != N: - msg = 'The known parameter {} is not length {}.' - raise ValueError(msg.format(k, N)) + if k not in tshift_trajs: + if len(v) != N: + msg = 'The known parameter {} is not length {}.' + raise ValueError(msg.format(k, N)) + + def _to_general_time(self, func): + """Replaces the argument of a function with the time symbol""" + + assert len(func.args) <= 1, (f"Trying to replace the argument of a function with more then " + f"one argument: {func}") + + return func.subs(func.args[0], self.time_symbol) + + + def _check_timeshift_parameter(self, param): + """Check: + - timeshift parameter mustn't be known. + - timeshift parameter must have bounds specified. + """ + + if param in self.known_parameter_map: + raise ValueError(f"Known timeshift parameters are not supportet: {param}") + + if param not in self.bounds: + raise ValueError(f"Must provide bounds for timeshift parameter {param}.") + + def _check_timeshift_trajectory(self, tr, arg, param): + """Check: + - the unshifted original of the timeshifted trajectory is provided as known + trajectory + - values for the unshifted original are available for + t in [-upper_bound, N*time_interval + lower_bound] + + Returns + ------- + idx_offset : int + Offset so that x[i + idx_offset] = x(t) for an array x[i] for i in [n_min, n_max]. + """ + + tr_general = self._to_general_time(tr) + if tr_general not in self.known_trajectory_map: + raise ValueError((f'The unshifted original of an input trajectory with unknown' + f' timeshift has to be provided in the known trajectory map.' + f' The following trajectory is missing in known_trajectory_' + f'map: {tr_general}')) + + t_min = +self.INF + t_max = -self.INF + + for t_val in [0, self.num_collocation_nodes*self.node_time_interval]: + + val_dict = {param: self.bounds[param][0], self.time_symbol: t_val} + ext1 = arg.evalf(subs=val_dict) + + val_dict = {param: self.bounds[param][1], self.time_symbol: t_val} + ext2 = arg.evalf(subs=val_dict) + + t_min = min(min(ext1, ext2), t_min) + t_max = max(max(ext1, ext2), t_max) + + n_min = int((t_min / self.node_time_interval).floor()) + n_max = int((t_max / self.node_time_interval).ceiling()) + N_timeshift = n_max - n_min + + if N_timeshift != self.known_trajectory_map[tr_general].size: + raise ValueError((f'Values of the timeshift original {tr_general} must be ' + f'available for all t in the maximum interval defined by the range ' + f'of t and the bounds of {param}: [{t_min}, {t_max}[. Provide a ' + f'trajectory with {N_timeshift} samples or change the bounds of ' + f'{param}!')) + idx_offset = - n_min + return idx_offset + + + def _substitute_timeshift_trajectories(self): + """Identify and substitute time-shifted trajectories in the eom. verifies that timeshift + trajectories are of type u(t+/-tau). Creates a dictionary linking the substitutes with the + original trajectories and the timeshift parameters. + """ + + t_set = {self.time_symbol} + + trajs = self.eom.atoms(sm.Function) + self.timeshift_traj_substitutes = {} + self.timeshift_traj_offsets = {} + self.unknown_tshift_parameters = [] + + def is_timeshift_function(func): + """Check if func is a timeshift trajectory. Timeshift trajectories must be of type + u(t +/- tau)""" + + # must be a function of t + if not len(func.free_symbols.intersection(t_set)) == 1: + return False + # must not have more then one argument + if len(func.args) > 1: + return False + # must be an addition + if not isinstance(func.args[0], sm.core.add.Add): + return False + # must have two free parameters + if not len(traj.free_symbols) == 2: + return False + + return True + + for traj in trajs: + + function_candidates = [traj] + + for tr in function_candidates: + for arg in tr.args: + funcs_in_arg = arg.atoms(sm.Function) + if len(funcs_in_arg) > 0: + function_candidates += list(funcs_in_arg) + + if is_timeshift_function(tr): + + #find timeshift parameter + timeshift_param = list(tr.free_symbols - t_set)[0] + + + #check that timeshift parameter is not known + self._check_timeshift_parameter(timeshift_param) + + #check timeshift trajectory + idx_offset = self._check_timeshift_trajectory(tr, arg, timeshift_param) + + #substitute timeshift trajectory + tr_subs = sm.symbols(tr.name+"_shift", cls=sm.Function)(self.time_symbol) + self.eom = self.eom.subs(tr, tr_subs) + + #pack info into dict + self.timeshift_traj_offsets[self._to_general_time(tr)] = idx_offset + self.timeshift_traj_substitutes[tr_subs] = (tr, timeshift_param) + self.unknown_tshift_parameters.append(timeshift_param) + + self.unknown_tshift_parameters = tuple(self.unknown_tshift_parameters) + self.num_tshift_parameters = len(self.timeshift_traj_substitutes) + self._timeshift_params = self.num_tshift_parameters > 0 + def _sort_trajectories(self): """Finds and counts all of the non-state, time varying parameters in @@ -867,6 +1064,9 @@ def _sort_trajectories(self): self.input_trajectories = res[0] + res[2] self.num_input_trajectories = len(self.input_trajectories) + + + def _discrete_symbols(self): """Instantiates discrete symbols for each time varying variable in the @@ -899,6 +1099,8 @@ def _discrete_symbols(self): inputs. """ + tshift_trajs = \ + [self._to_general_time(v[0]) for v in self.timeshift_traj_substitutes.values()] # The previus, current, and next states. self.previous_discrete_state_symbols = \ @@ -914,10 +1116,10 @@ def _discrete_symbols(self): # The current and next known input trajectories. self.current_known_discrete_specified_symbols = \ tuple([sm.Symbol(f.__class__.__name__ + 'i', real=True) - for f in self.known_input_trajectories]) + for f in self.known_input_trajectories if f not in tshift_trajs]) self.next_known_discrete_specified_symbols = \ tuple([sm.Symbol(f.__class__.__name__ + 'n', real=True) - for f in self.known_input_trajectories]) + for f in self.known_input_trajectories if f not in tshift_trajs]) # The current and next unknown input trajectories. self.current_unknown_discrete_specified_symbols = \ @@ -944,10 +1146,14 @@ def _discretize_eom(self): The column vector of the discretized equations of motion. """ + + tshift_trajs = \ + [self._to_general_time(v[0]) for v in self.timeshift_traj_substitutes.values()] + logging.info('Discretizing the equations of motion.') x = self.state_symbols xd = self.state_derivative_symbols - u = self.input_trajectories + u = tuple([t for t in self.input_trajectories if t not in tshift_trajs]) xp = self.previous_discrete_state_symbols xi = self.current_discrete_state_symbols @@ -971,7 +1177,47 @@ def _discretize_eom(self): x_sub = {d: (i + n) / 2 for d, i, n in zip(x, xi, xn)} u_sub = {d: (i + n) / 2 for d, i, n in zip(u, ui, un)} self.discrete_eom = me.msubs(self.eom, xdot_sub, x_sub, u_sub) - + + def _generate_timeshift_constraints(self): + """ Generates a set of instance constraints to link the timeshifted trajectories with their + substitutes""" + + timeshift_constraints = [] + + for traj_subs, traj in self.timeshift_traj_substitutes.items(): + for i in range(self.num_collocation_nodes): + timeshift_constraints.append(traj_subs.subs(self.time_symbol, i * self.node_time_interval) - + traj[0].subs(self.time_symbol, i * self.node_time_interval)) + + if self.instance_constraints is None: + self.instance_constraints = tuple(timeshift_constraints) + else: + self.instance_constraints = tuple(list(self.instance_constraints) + timeshift_constraints) + + + def _precalc_timshift_input_derivatives(self): + """Precalculates the derivatives of timeshifted input for the partials of the constraint + jacobian.""" + + def _diff(trajectory): + """Differentiates a trajectory using midpoint or backward euler and edge value padding. + """ + if self.integration_method == 'midpoint': + padded = np.pad(trajectory, (1,1), mode='edge') + return (padded[2:] - padded[:-2])/(2*self.node_time_interval) + elif self.integration_method == 'backward euler': + padded = np.pad(trajectory, (1,0), mode='edge') + return (padded[1:] - padded[:-1])/(self.node_time_interval) + + self.timeshift_traj_derivative_map = {} + + for traj_subs in self.timeshift_traj_substitutes: + traj = self.timeshift_traj_substitutes[traj_subs][0] + traj_general = traj.subs(traj.args[0], self.time_symbol) + diff_vals = _diff(self.known_trajectory_map[traj_general]) + self.timeshift_traj_derivative_map[traj_general] = diff_vals + + def _identify_functions_in_instance_constraints(self): """Instantiates a set containing all of the instance functions, i.e. x(1.0) in the instance constraints.""" @@ -987,16 +1233,35 @@ def _find_closest_free_index(self): """Instantiates a dictionary mapping the instance functions to the nearest index in the free variables vector.""" - def determine_free_index(time_index, state): - state_index = self.state_symbols.index(state) - return time_index + state_index * self.num_collocation_nodes + def determine_free_index(time_index, trajectory): + + traj_index = None + + if trajectory in self.state_symbols: + traj_index = self.state_symbols.index(trajectory) + elif trajectory in self.unknown_input_trajectories: + traj_index = self.num_states + self.unknown_input_trajectories.index(trajectory) + else: + raise ValueError(f"{trajectory} is neither a state symbol not an unknown input" + "trajectory") + + if traj_index is None: + raise ValueError((f"'{trajectory}' is neither a state, a timeshift trajectory nor" + " an unknown input trajectory")) + + return time_index + traj_index * self.num_collocation_nodes N = self.num_collocation_nodes h = self.node_time_interval duration = h * (N - 1) + known_trajectory_names = [traj.name for traj in self.known_trajectory_map] node_map = {} for func in self.instance_constraint_function_atoms: + + if func.name in known_trajectory_names: + continue + if self._variable_duration: if func.args[0] == 0: time_idx = 0 @@ -1014,27 +1279,90 @@ def determine_free_index(time_index, state): func, time_idx, self.num_collocation_nodes - 1)) else: time_value = func.args[0] - time_vector = np.linspace(0.0, duration, num=N) - time_idx = np.argmin(np.abs(time_vector - time_value)) - free_index = determine_free_index(time_idx, - func.__class__(self.time_symbol)) + time_idx = (time_value / h).round() + #for a in sm.preorder_traversal(time_idx): + # if isinstance(a, sm.Float): + # time_idx = time_idx.subs(a, a.round()) + + free_index = determine_free_index(time_idx, func.__class__(self.time_symbol)) node_map[func] = free_index self.instance_constraints_free_index_map = node_map + + def _make_free_matrix_symbols(self, valtype = "input"): + """Make matrix symbols representing the vector of free variables and (in case of timeshifts) + the timeshfited trajectories + """ + + #create matrix symbols for the free vector and all timshift input trajectories + free = sm.MatrixSymbol('FREE', self.num_free, 1) + + unshifted_input_sym = {} + unshifted_input_vals = [] + for val in self.timeshift_traj_substitutes.values(): + if valtype == 'dudt': + trajvals = np.array(self.timeshift_traj_derivative_map[val[0].subs(val[1],0)]) + else: + trajvals = np.array(self.known_trajectory_map[val[0].subs(val[1],0)]) + unshifted_input_sym[val[0].name] = sm.MatrixSymbol(val[0].name.upper(), trajvals.size, 1) + unshifted_input_vals.append(trajvals[:,np.newaxis]) + + return free, unshifted_input_sym, unshifted_input_vals + def _instance_constraints_func(self): """Returns a function that evaluates the instance constraints given the free optimization variables.""" - free = sm.DeferredVector('FREE') - def_map = {k: free[v] for k, v in - self.instance_constraints_free_index_map.items()} - subbed_constraints = [con.subs(def_map) for con in - self.instance_constraints] - f = sm.lambdify(([free] + list(self.known_parameter_map.keys())), - subbed_constraints, modules=[{'ImmutableMatrix': - np.array}, "numpy"]) - - return lambda free: f(free, *self.known_parameter_map.values()) + + free, unshifted_input_sym, unshifted_input_vals = self._make_free_matrix_symbols("input") + to_int = sm.Function('int') + + # make map from unknown parameters to their value in FREE + unknown_param_map = {} + n_traj_vals = (self.num_states + self.num_unknown_input_trajectories) \ + * self.num_collocation_nodes + for v in self.unknown_parameters: + unknown_param_map[v] = free[n_traj_vals + self.unknown_parameters.index(v),0] + for tau in self.unknown_tshift_parameters: + unknown_param_map[tau] = free[n_traj_vals + + self.num_unknown_parameters + + int(self._variable_duration) + + self.unknown_tshift_parameters.index(tau),0] + + # make instance constraints free (MatrixSymbol) map + self.instance_constraints_free_map = {} + for k,v in self.instance_constraints_free_index_map.items(): + self.instance_constraints_free_map[k] = free[v,0] + + # substitue matrix symbols into constraints + subbed_constraints = [None]*self.num_instance_constraints + for i in range(self.num_instance_constraints): + + #substitute unknown trajectory values + subbed_constraints[i] = me.msubs(self.instance_constraints[i], + self.instance_constraints_free_map, + unknown_param_map) + + #substitute known trajectory values + for func in subbed_constraints[i].atoms(sm.Function): + if func.name in unshifted_input_sym: + arg = (func.args[0]/self.node_time_interval) + + #correct the argument for arrays of values not beginning at t=0 + tshift_idx_offset = self.timeshift_traj_offsets[self._to_general_time(func)] + arg += tshift_idx_offset + + func_subs = unshifted_input_sym[func.name][to_int(arg),0] + subbed_constraints[i] = subbed_constraints[i].subs(func, func_subs) + + #lambdify + args = [free] \ + + list(unshifted_input_sym.values()) \ + + list(self.known_parameter_map.keys()) + f = sm.lambdify(args, subbed_constraints, modules=[{'ImmutableMatrix': np.array}, "numpy"]) + + return lambda free: f(free[:,np.newaxis], + *unshifted_input_vals, *self.known_parameter_map.values()) def _instance_constraints_jacobian_indices(self): """Returns the row and column indices of the non-zero values in the @@ -1042,13 +1370,30 @@ def _instance_constraints_jacobian_indices(self): idx_map = self.instance_constraints_free_index_map num_eom_constraints = self.num_states*(self.num_collocation_nodes - 1) - + unshifted_input_trajs = {v[0].subs(v[0].args[0], self.time_symbol) : v[1] + for v in self.timeshift_traj_substitutes.values()} + + offset_free_params = self.num_unknown_parameters + int(self._variable_duration) \ + + (self.num_states + self.num_unknown_input_trajectories) \ + * self.num_collocation_nodes + rows = [] cols = [] for i, con in enumerate(self.instance_constraints): funcs = con.atoms(sm.Function) - indices = [idx_map[f] for f in funcs] + indices = [] + for f in funcs: + if f in idx_map: + # partials w.r.t states and unknown input trajectories + indices.append(idx_map[f]) + else: + # partials w.r.t timeshift parameters + f_general = f.subs(f.args[0], self.time_symbol) + if f_general in unshifted_input_trajs: + indices.append(offset_free_params + self.unknown_tshift_parameters.index( + unshifted_input_trajs[f_general])) + row_idxs = num_eom_constraints + i * np.ones(len(indices), dtype=int) rows += list(row_idxs) @@ -1059,20 +1404,61 @@ def _instance_constraints_jacobian_indices(self): def _instance_constraints_jacobian_values_func(self): """Returns the non-zero values of the constraint Jacobian associated with the instance constraints.""" - free = sm.DeferredVector('FREE') - - def_map = {k: free[v] for k, v in - self.instance_constraints_free_index_map.items()} + + #create matrix symbols for the free vector and all timshift input trajectories + free, unshifted_input_sym, unshifted_input_vals = self._make_free_matrix_symbols("dudt") + to_int = sm.Function('int') + + # make map from unknown parameters to their value in FREE + unknown_param_map = {} + n_traj_vals = (self.num_states + self.num_unknown_input_trajectories) \ + * self.num_collocation_nodes + for v in self.unknown_parameters: + unknown_param_map[v] = free[n_traj_vals + self.unknown_parameters.index(v),0] + for tau in self.unknown_tshift_parameters: + unknown_param_map[tau] = free[n_traj_vals + + self.num_unknown_parameters + + int(self._variable_duration) + + self.unknown_tshift_parameters.index(tau),0] + + known_traj_set = [traj.name for traj in self.known_trajectory_map.keys()] + timeshift_param_set = set([v[1] for v in self.timeshift_traj_substitutes.values()]) funcs = [] num_vals_per_func = [] + for con in self.instance_constraints: - partials = list(con.atoms(sm.Function)) - num_vals_per_func.append(len(partials)) - jac = sm.Matrix([con]).jacobian(partials) - jac = jac.subs(def_map) - funcs.append(sm.lambdify(([free] + - list(self.known_parameter_map.keys())), + + #identify partials w.r.t unknown trajectories and timeshift parameters + jac = sm.Matrix([[]]) + + for traj in con.atoms(sm.Function): + if not traj.name in known_traj_set: + jac = jac.row_join(sm.Matrix([con]).jacobian([traj])) + else: + time_idx = to_int(me.msubs(traj.args[0], unknown_param_map) / self.node_time_interval) + + #correct time_idx for arrays of values not beginning at t=0 + idx_offset = self.timeshift_traj_offsets[self._to_general_time(traj)] + time_idx += idx_offset + + jac = jac.row_join(sm.Matrix([unshifted_input_sym[traj.name][time_idx,0]])) + + #calculate jacobian entries w.r.t unknown input trajectories. + #traj_jac = sm.Matrix([con]).jacobian(traj_partials) + + #concat jacobian + ##if len(timeshift_jac) > 0: + # jac = sm.Matrix(timeshift_jac).row_join(traj_jac) + #else: + # jac = traj_jac + + + num_vals_per_func.append(len(jac)) + jac = me.msubs(jac, self.instance_constraints_free_map) + funcs.append(sm.lambdify(([free] \ + + list(unshifted_input_sym.values()) \ + + list(self.known_parameter_map.keys())), jac, modules=[{'ImmutableMatrix': np.array}, "numpy"])) length = np.sum(num_vals_per_func) @@ -1081,7 +1467,7 @@ def wrapped(free): arr = np.zeros(length) j = 0 for i, (f, num) in enumerate(zip(funcs, num_vals_per_func)): - arr[j:j + num] = f(free, *self.known_parameter_map.values()) + arr[j:j + num] = f(free[:,np.newaxis], *unshifted_input_vals, *self.known_parameter_map.values()) j += num return arr @@ -1669,9 +2055,14 @@ def constraints(free): variable_duration=self._variable_duration) time_interval = self.node_time_interval - all_specified = self._merge_fixed_free(self.input_trajectories, + tshift_trajs = \ + [self._to_general_time(v[0]) for v in self.timeshift_traj_substitutes.values()] + input_trajectories = [traj for traj in self.input_trajectories + if not traj in tshift_trajs] + + all_specified = self._merge_fixed_free(input_trajectories, self.known_trajectory_map, - free_specified, 'traj') + free_specified, 'traj').squeeze() all_constants = self._merge_fixed_free(self.parameters, self.known_parameter_map, @@ -1679,6 +2070,7 @@ def constraints(free): eom_con_vals = func(free_states, all_specified, all_constants, time_interval) + if self.instance_constraints is not None: if typ == 'con': diff --git a/opty/tests/test_direct_collocation.py b/opty/tests/test_direct_collocation.py index adc2f3f8..0e80c6e3 100644 --- a/opty/tests/test_direct_collocation.py +++ b/opty/tests/test_direct_collocation.py @@ -1566,3 +1566,302 @@ def test_for_algebraic_eoms(): ) assert excinfo.type is ValueError + + +class TestConstraintCollocatorTimeshiftConstraints(): + + def setup_method(self): + + m, c, k, t, tau = sym.symbols('m, c, k, t, tau') + x, v, f, f_shift = sym.symbols('x, v, f, f_shift', cls=sym.Function) + + self.discrete_symbols = sym.symbols('xi, vi, xp, vp, xn, vn, f_shifti, f_shiftn', + real=True) + + self.state_symbols = (x(t), v(t)) + self.time_symbol = t + self.constant_symbols = (m, c, k) + self.specified_symbols = (f(t),) + self.unknown_input_trajectories = (f_shift(t),) + + self.state_values = np.array([[1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0]]) + self.specified_values = np.array([2.0, 3.0, 4.0, 3.0, 2.0, 1.0, 1.0, 2.0]) + self.constant_values = np.array([1.0, 2.0, 3.0]) + self.interval_value = 0.01 + self.free = np.array([1.0, 2.0, 3.0, 4.0, # x + 5.0, 6.0, 7.0, 8.0, # v + 2.0, 2.0, 2.0, 2.0, # f_shift + 3.0, # k + 0.01]) #tau + + self.eom = sym.Matrix([x(t).diff() - v(t), + m * v(t).diff() + c * v(t) + k * x(t) - f(t-tau)]) + + + par_map = OrderedDict(zip(self.constant_symbols[:2], + self.constant_values[:2])) + self.known_trajectory_map = {f(t): self.specified_values} + self.known_parameter_map = {m: 1.0, c: 2.0} + self.bounds = {tau: [-0.02 , 0.02]} + + instance_constraints = (x(0)-1, v(0)-5) + + + self.u_vals_diff = [0.5, 0.5, 0, 0] + + self.collocator = \ + ConstraintCollocator(equations_of_motion=self.eom, + state_symbols=self.state_symbols, + num_collocation_nodes=4, + node_time_interval=self.interval_value, + known_parameter_map=par_map, + known_trajectory_map=self.known_trajectory_map, + instance_constraints=instance_constraints, + bounds = self.bounds, + time_symbol=t) + + self.eom_subs = sym.Matrix([x(t).diff() - v(t), + m * v(t).diff() + c * v(t) + k * x(t) - f_shift(t)]) + self.timeshift_inputs = {f_shift(t): (f(t - tau), tau)} + self.timeshift_traj_offsets = {f(t): 2} + + fv = self.specified_values + dt = 2*self.interval_value + self.expected_timeshift_traj_derivative_map_midpoint = {f(t): \ + [(fv[1]-fv[0])/dt, (fv[2]-fv[0])/dt, (fv[3]-fv[1])/dt, (fv[4]-fv[2])/dt, + (fv[5]-fv[3])/dt, (fv[6]-fv[4])/dt, (fv[7]-fv[5])/dt, (fv[7]-fv[6])/dt,]} + dt = self.interval_value + self.expected_timeshift_traj_derivative_map_bweuler = {f(t): \ + [(fv[0]-fv[0])/dt, (fv[1]-fv[0])/dt, (fv[2]-fv[1])/dt, (fv[3]-fv[2])/dt, + (fv[4]-fv[3])/dt, (fv[5]-fv[4])/dt, (fv[6]-fv[5])/dt, (fv[7]-fv[6])/dt,]} + + def test_substitute_timeshift_trajectories(self): + assert self.collocator.eom == self.eom_subs + assert self.collocator.timeshift_traj_substitutes == self.timeshift_inputs + assert self.collocator.timeshift_traj_offsets == self.timeshift_traj_offsets + + def test_precalc_timshift_input_derivatives(self): + + self.collocator.integration_method = 'midpoint' + self.collocator._precalc_timshift_input_derivatives() + for k,v in self.collocator.timeshift_traj_derivative_map.items(): + np.testing.assert_allclose(v, self.expected_timeshift_traj_derivative_map_midpoint[k]) + + self.collocator.integration_method = 'backward euler' + self.collocator._precalc_timshift_input_derivatives() + for k,v in self.collocator.timeshift_traj_derivative_map.items(): + np.testing.assert_allclose(v, self.expected_timeshift_traj_derivative_map_bweuler[k]) + + def test_init(self): + + assert self.collocator.state_symbols == self.state_symbols + assert self.collocator.state_derivative_symbols == \ + tuple([s.diff() for s in self.state_symbols]) + assert self.collocator.num_states == 2 + assert self.collocator.num_collocation_nodes == 4 + + def test_integration_method(self): + with raises(ValueError): + self.collocator.integration_method = 'booger' + + def test_sort_parameters(self): + + self.collocator._sort_parameters() + + m, c, k = self.constant_symbols + + assert self.collocator.known_parameters == (m, c) + assert self.collocator.num_known_parameters == 2 + + assert self.collocator.unknown_parameters == (k,) + assert self.collocator.num_unknown_parameters == 1 + + def test_sort_trajectories(self): + + self.collocator._sort_trajectories() + + assert self.collocator.known_input_trajectories == \ + self.specified_symbols + assert self.collocator.num_known_input_trajectories == 1 + + assert self.collocator.unknown_input_trajectories == self.unknown_input_trajectories + assert self.collocator.num_unknown_input_trajectories == 1 + + def test_discrete_symbols(self): + + self.collocator._discrete_symbols() + + xi, vi, xp, vp, xn, vn, f_shifti, f_shiftn = self.discrete_symbols + + assert self.collocator.previous_discrete_state_symbols == (xp, vp) + assert self.collocator.current_discrete_state_symbols == (xi, vi) + assert self.collocator.next_discrete_state_symbols == (xn, vn) + + assert self.collocator.current_discrete_specified_symbols == (f_shifti, ) + assert self.collocator.next_discrete_specified_symbols == (f_shiftn, ) + + def test_discretize_eom_backward_euler(self): + + m, c, k = self.constant_symbols + xi, vi, xp, vp, xn, vn, fi, fn = self.discrete_symbols + h = self.collocator.time_interval_symbol + + expected = sym.Matrix([(xi - xp) / h - vi, + m * (vi - vp) / h + c * vi + k * xi - fi]) + + self.collocator._discretize_eom() + + zero = sym.simplify(self.collocator.discrete_eom - expected) + + assert zero == sym.Matrix([0, 0]) + + def test_discretize_eom_midpoint(self): + + m, c, k = self.constant_symbols + xi, vi, xp, vp, xn, vn, fi, fn = self.discrete_symbols + + h = self.collocator.time_interval_symbol + + expected = sym.Matrix([(xn - xi) / h - (vi + vn) / 2, + m * (vn - vi) / h + c * (vi + vn) / 2 + + k * (xi + xn) / 2 - (fi + fn) / 2]) + + self.collocator.integration_method = 'midpoint' + self.collocator._discretize_eom() + + zero = sym.simplify(self.collocator.discrete_eom - expected) + + assert zero == sym.Matrix([0, 0]) + + def test_gen_multi_arg_con_func_backward_euler(self): + + self.collocator.integration_method ='backward euler' + self.collocator._gen_multi_arg_con_func() + + # Make sure the parameters are in the correct order. + constant_values = \ + np.array([self.constant_values[self.constant_symbols.index(c)] + for c in self.collocator.parameters]) + + # TODO : Once there are more than one specified, they will need to + # be put in the correct order too. + + result = self.collocator._multi_arg_con_func(self.state_values, + self.specified_values[1:-3], + constant_values, + self.interval_value) + + m, c, k = self.constant_values + h = self.interval_value + + expected_dynamic = np.zeros(3) + expected_kinematic = np.zeros(3) + + for i in [1, 2, 3]: + + xi, vi = self.state_values[:, i] + xp, vp = self.state_values[:, i - 1] + fi = self.specified_values[i+1] + + expected_dynamic[i - 1] = m * (vi - vp) / h + c * vi + k * xi - fi + expected_kinematic[i - 1] = (xi - xp) / h - vi + + expected = np.hstack((expected_kinematic, expected_dynamic)) + + np.testing.assert_allclose(result, expected) + + def test_gen_multi_arg_con_func_midpoint(self): + + self.collocator.integration_method = 'midpoint' + self.collocator._gen_multi_arg_con_func() + + # Make sure the parameters are in the correct order. + constant_values = \ + np.array([self.constant_values[self.constant_symbols.index(c)] + for c in self.collocator.parameters]) + + # TODO : Once there are more than one specified, they will need to + # be put in the correct order too. + + result = self.collocator._multi_arg_con_func(self.state_values, + self.specified_values[1:-3], + constant_values, + self.interval_value) + + m, c, k = self.constant_values + h = self.interval_value + + expected_dynamic = np.zeros(3) + expected_kinematic = np.zeros(3) + + for i in [0, 1, 2]: + + xi, vi = self.state_values[:, i] + xn, vn = self.state_values[:, i + 1] + fi = self.specified_values[i + 1:i + 1 + 1][0] + fn = self.specified_values[i + 1 + 1:i + 2 + 1][0] + + expected_kinematic[i] = (xn - xi) / h - (vi + vn) / 2 + expected_dynamic[i] = (m * (vn - vi) / h + c * (vn + vi) / 2 + k + * (xn + xi) / 2 - (fi + fn) / 2) + + expected = np.hstack((expected_kinematic, expected_dynamic)) + + np.testing.assert_allclose(result, expected) + + + def test_generate_jacobian_function(self): + + self.collocator.integration_method = 'backward euler' + jacobian = self.collocator.generate_jacobian_function() + + for tau_test in [-0.01, 0, 0.01]: + free = self.free + free[-1] = tau_test + jac_vals = jacobian(free) + + row_idxs, col_idxs = self.collocator.jacobian_indices() + + jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs))) + + # jacobian of eom_vector wrt vi, xi, xp, vp, k + # [ vi, xi, vp, xp, k] + # x: [ -1, 1/h, 0, -1/h, 0] + # v: [c + m/h, k, -m/h, 0, xi] + + x = self.state_values[0] + m, c, k = self.constant_values + h = self.interval_value + c_tau = np.array(list(self.expected_timeshift_traj_derivative_map_bweuler.values())[0]) + c_tau = c_tau[2-int(tau_test/self.interval_value):-2-int(tau_test/self.interval_value)] + + expected_jacobian = np.array( + # x1, x2, x3, x4, v1, v2, v3, v4, f1, f2, f3, f4, k, tau + [[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0], #eom-node 0 + [ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0], #eom-node 1 + [ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0], #eom-node 2 + [ 0, k, 0, 0, -m / h, c + m / h, 0, 0, 0, -1, 0, 0, x[1], 0], #eom-node 3 + [ 0, 0, k, 0, 0, -m / h, c + m / h, 0, 0, 0, -1, 0, x[2], 0], #eom-node 4 + [ 0, 0, 0, k, 0, 0, -m /h, c + m / h, 0, 0, 0, -1, x[3], 0], #eom-node 5 + [ 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], #c0: x[0] = 0 + [ 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], #c1: v[0] = 5 + [ 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, c_tau[0]], #c_tshift_0 + [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, c_tau[1]], #c_tshift_1 + [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, c_tau[2]], #c_tshift_2 + [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, c_tau[3]]], #c_tshift_3 + dtype=float) + + np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian) + + + def test_lambdify_instance_constraints(self): + + f = self.collocator._instance_constraints_func() + + extra_constraints = f(self.free) + + expected = np.array([0.0, 0.0, -1.0, -2.0, -1.0, 0.0]) + + np.testing.assert_allclose(extra_constraints, expected) + \ No newline at end of file