Skip to content

Commit 704fdec

Browse files
murrayrmmurrayrm
committed
add trajectory_method='collocation'
1 parent 0862b5a commit 704fdec

2 files changed

Lines changed: 200 additions & 69 deletions

File tree

‎control/optimal.py‎

Lines changed: 100 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
from .timeresp import TimeResponseData
2222

2323
# Define module default parameter values
24+
_optimal_trajectory_methods = {'shooting', 'collocation'}
2425
_optimal_defaults = {
26+
'optimal.trajectory_method': 'shooting',
2527
'optimal.minimize_method': None,
2628
'optimal.minimize_options': {},
2729
'optimal.minimize_kwargs': {},
@@ -65,9 +67,11 @@ class OptimalControlProblem():
6567
inputs should either be a 2D vector of shape (ninputs, horizon)
6668
or a 1D input of shape (ninputs,) that will be broadcast by
6769
extension of the time axis.
68-
method : string, optional
69-
Method to use for carrying out the optimization. Currently only
70-
'shooting' is supported.
70+
trajectory_method : string, optional
71+
Method to use for carrying out the optimization. Currently supported
72+
methods are 'shooting' and 'collocation' (continuous time only). The
73+
default value is 'shooting' for discrete time systems and
74+
'collocation' for continuous time systems
7175
log : bool, optional
7276
If `True`, turn on logging messages (using Python logging module).
7377
Use ``logging.basicConfig`` to enable logging output (e.g., to a file).
@@ -133,7 +137,7 @@ class OptimalControlProblem():
133137
def __init__(
134138
self, sys, timepts, integral_cost, trajectory_constraints=[],
135139
terminal_cost=None, terminal_constraints=[], initial_guess=None,
136-
method='shooting', basis=None, log=False, **kwargs):
140+
trajectory_method=None, basis=None, log=False, **kwargs):
137141
"""Set up an optimal control problem."""
138142
# Save the basic information for use later
139143
self.system = sys
@@ -144,11 +148,15 @@ def __init__(
144148
self.basis = basis
145149

146150
# Keep track of what type of method we are using
147-
if method not in {'shooting', 'collocation'}:
151+
if trajectory_method is None:
152+
# TODO: change default
153+
# trajectory_method = 'collocation' if sys.isctime() else 'shooting'
154+
trajectory_method = 'shooting' if sys.isctime() else 'shooting'
155+
elif trajectory_method not in _optimal_trajectory_methods:
148156
raise NotImplementedError(f"Unkown method {method}")
149-
else:
150-
self.shooting = method in {'shooting'}
151-
self.collocation = method in {'collocation'}
157+
158+
self.shooting = trajectory_method in {'shooting'}
159+
self.collocation = trajectory_method in {'collocation'}
152160

153161
# Process keyword arguments
154162
self.solve_ivp_kwargs = {}
@@ -249,20 +257,21 @@ def __init__(
249257
self.eqconst_value))
250258
if self.collocation:
251259
# Add the collocation constraints
252-
colloc_zeros = np.zeros((self.timepts.size - 1) * sys.nstates)
260+
colloc_zeros = np.zeros(sys.nstates * self.timepts.size)
261+
self.colloc_vals = np.zeros((sys.nstates, self.timepts.size))
253262
self.constraints.append(sp.optimize.NonlinearConstraint(
254263
self._collocation_constraint, colloc_zeros, colloc_zeros))
255264

265+
# Initialize run-time statistics
266+
self._reset_statistics(log)
267+
256268
# Process the initial guess
257269
self.initial_guess = self._process_initial_guess(initial_guess)
258270

259271
# Store states, input, used later to minimize re-computation
260272
self.last_x = np.full(self.system.nstates, np.nan)
261273
self.last_coeffs = np.full(self.initial_guess.shape, np.nan)
262274

263-
# Reset run-time statistics
264-
self._reset_statistics(log)
265-
266275
# Log information
267276
if log:
268277
logging.info("New optimal control problem initailized")
@@ -468,10 +477,6 @@ def _eqconst_function(self, coeffs):
468477
# Checked above => we should never get here
469478
raise TypeError("unknown constraint type {ctype}")
470479

471-
# Add the collocation constraints
472-
if self.collocation:
473-
raise NotImplementedError("collocation not yet implemented")
474-
475480
# Update statistics
476481
self.eqconst_evaluations += 1
477482
if self.log:
@@ -489,11 +494,32 @@ def _eqconst_function(self, coeffs):
489494
# Return the value of the constraint function
490495
return np.hstack(value)
491496

492-
def _collocation_constraints(self, coeffs):
497+
def _collocation_constraint(self, coeffs):
493498
# Compute the states and inputs
494499
states, inputs = self._compute_states_inputs(coeffs)
495500

496-
raise NotImplementedError("collocation not yet implemented")
501+
if self.system.isctime():
502+
# Compute the collocation constraints
503+
for i, t in enumerate(self.timepts):
504+
if i == 0:
505+
# Initial condition constraint (self.x = initial point)
506+
self.colloc_vals[:, 0] = states[:, 0] - self.x
507+
fk = self.system._rhs(
508+
self.timepts[0], states[:, 0], inputs[:, 0])
509+
continue
510+
511+
# M. Kelly, SIAM Review (2017), equation (3.2), i = k+1
512+
# x[k+1] - x[k] = 0.5 hk (f(x[k+1], u[k+1] + f(x[k], u[k]))
513+
fkp1 = self.system._rhs(t, states[:, i], inputs[:, i])
514+
self.colloc_vals[:, i] = states[:, i] - states[:, i-1] - \
515+
0.5 * (self.timepts[i] - self.timepts[i-1]) * (fkp1 + fk)
516+
fk = fkp1
517+
else:
518+
raise NotImplementedError(
519+
"collocation not yet implemented for discrete time systems")
520+
521+
# Return the value of the constraint function
522+
return self.colloc_vals.reshape(-1)
497523

498524
#
499525
# Initial guess
@@ -509,40 +535,61 @@ def _collocation_constraints(self, coeffs):
509535
# TODO: figure out how to modify this for collocation
510536
#
511537
def _process_initial_guess(self, initial_guess):
512-
if self.shooting and initial_guess is not None:
513-
# Convert to a 1D array (or higher)
514-
initial_guess = np.atleast_1d(initial_guess)
515-
516-
# See whether we got entire guess or just first time point
517-
if initial_guess.ndim == 1:
518-
# Broadcast inputs to entire time vector
519-
try:
520-
initial_guess = np.broadcast_to(
521-
initial_guess.reshape(-1, 1),
522-
(self.system.ninputs, self.timepts.size))
523-
except ValueError:
524-
raise ValueError("initial guess is the wrong shape")
525-
526-
elif initial_guess.shape != \
527-
(self.system.ninputs, self.timepts.size):
528-
raise ValueError("initial guess is the wrong shape")
538+
# Sort out the input guess and the state guess
539+
if self.collocation and initial_guess is not None and \
540+
isinstance(initial_guess, tuple):
541+
state_guess, input_guess = initial_guess
542+
else:
543+
state_guess, input_guess = None, initial_guess
544+
545+
# Process the input guess
546+
if input_guess is not None:
547+
input_guess = self._broadcast_initial_guess(
548+
input_guess, (self.system.ninputs, self.timepts.size))
529549

530550
# If we were given a basis, project onto the basis elements
531551
if self.basis is not None:
532-
initial_guess = self._inputs_to_coeffs(initial_guess)
552+
input_guess = self._inputs_to_coeffs(input_guess)
553+
else:
554+
input_guess = np.zeros(
555+
self.system.ninputs *
556+
(self.timepts.size if self.basis is None else self.basis.N))
557+
558+
# Process the state guess
559+
if self.collocation:
560+
if state_guess is None:
561+
# Run a simulation to get the initial guess
562+
inputs = input_guess.reshape(self.system.ninputs, -1)
563+
state_guess = self._simulate_states(
564+
np.zeros(self.system.nstates), inputs)
565+
else:
566+
state_guess = self._broadcast_initial_guess(
567+
state_guess, (self.system.nstates, self.timepts.size))
533568

534569
# Reshape for use by scipy.optimize.minimize()
535-
return initial_guess.reshape(-1)
570+
return np.hstack([
571+
input_guess.reshape(-1), state_guess.reshape(-1)])
572+
else:
573+
# Reshape for use by scipy.optimize.minimize()
574+
return input_guess.reshape(-1)
575+
576+
def _broadcast_initial_guess(self, initial_guess, shape):
577+
# Convert to a 1D array (or higher)
578+
initial_guess = np.atleast_1d(initial_guess)
579+
580+
# See whether we got entire guess or just first time point
581+
if initial_guess.ndim == 1:
582+
# Broadcast inputs to entire time vector
583+
try:
584+
initial_guess = np.broadcast_to(
585+
initial_guess.reshape(-1, 1), shape)
586+
except ValueError:
587+
raise ValueError("initial guess is the wrong shape")
536588

537-
elif self.collocation and initial_guess is not None:
538-
raise NotImplementedError("collocation not yet implemented")
589+
elif initial_guess.shape != shape:
590+
raise ValueError("initial guess is the wrong shape")
539591

540-
# Default is no initial guess
541-
input_size = self.system.ninputs * \
542-
(self.timepts.size if self.basis is None else self.basis.N)
543-
state_size = 0 if not self.collocation else \
544-
(self.timepts.size - 1) * self.sys.nstates
545-
return np.zeros(input_size + state_size)
592+
return initial_guess
546593

547594
#
548595
# Utility function to convert input vector to coefficient vector
@@ -642,9 +689,10 @@ def _compute_states_inputs(self, coeffs):
642689
#
643690
if self.collocation:
644691
# States are appended to end of (input) coefficients
645-
states = coeffs[-self.sys.nstates * self.timepts.size:0]
646-
coeffs = coeffs[0:-self.sys.nstates * self.timepts.size]
647-
692+
states = coeffs[-self.system.nstates * self.timepts.size:].reshape(
693+
self.system.nstates, -1)
694+
coeffs = coeffs[:-self.system.nstates * self.timepts.size]
695+
648696
# Compute input at time points
649697
if self.basis:
650698
inputs = self._coeffs_to_inputs(coeffs)
@@ -855,11 +903,8 @@ def __init__(
855903
# Remember the optimal control problem that we solved
856904
self.problem = ocp
857905

858-
# Compute input at time points
859-
if ocp.basis:
860-
inputs = ocp._coeffs_to_inputs(res.x)
861-
else:
862-
inputs = res.x.reshape((ocp.system.ninputs, -1))
906+
# Parse the optimization variables into states and inputs
907+
states, inputs = ocp._compute_states_inputs(res.x)
863908

864909
# See if we got an answer
865910
if not res.success:
@@ -877,6 +922,7 @@ def __init__(
877922

878923
if return_states and inputs.shape[1] == ocp.timepts.shape[0]:
879924
# Simulate the system if we need the state back
925+
# TODO: this is probably not needed due to compute_states_inputs()
880926
_, _, states = ct.input_output_response(
881927
ocp.system, ocp.timepts, inputs, ocp.x, return_x=True,
882928
solve_ivp_kwargs=ocp.solve_ivp_kwargs, t_eval=ocp.timepts)
@@ -1005,7 +1051,7 @@ def solve_ocp(
10051051
'optimal', 'return_x', kwargs, return_states, pop=True)
10061052

10071053
# Process (legacy) method keyword
1008-
if kwargs.get('method'):
1054+
if kwargs.get('method') and method not in optimal_methods:
10091055
if kwargs.get('minimize_method'):
10101056
raise ValueError("'minimize_method' specified more than once")
10111057
kwargs['minimize_method'] = kwargs.pop('method')

0 commit comments

Comments
 (0)