Skip to content

Commit ce5e67a

Browse files
murrayrmmurrayrm
committed
add trajectory_method='collocation'
1 parent d9ec3ec commit ce5e67a

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 = {}
@@ -257,20 +265,21 @@ def __init__(
257265
self.eqconst_value))
258266
if self.collocation:
259267
# Add the collocation constraints
260-
colloc_zeros = np.zeros((self.timepts.size - 1) * sys.nstates)
268+
colloc_zeros = np.zeros(sys.nstates * self.timepts.size)
269+
self.colloc_vals = np.zeros((sys.nstates, self.timepts.size))
261270
self.constraints.append(sp.optimize.NonlinearConstraint(
262271
self._collocation_constraint, colloc_zeros, colloc_zeros))
263272

273+
# Initialize run-time statistics
274+
self._reset_statistics(log)
275+
264276
# Process the initial guess
265277
self.initial_guess = self._process_initial_guess(initial_guess)
266278

267279
# Store states, input, used later to minimize re-computation
268280
self.last_x = np.full(self.system.nstates, np.nan)
269281
self.last_coeffs = np.full(self.initial_guess.shape, np.nan)
270282

271-
# Reset run-time statistics
272-
self._reset_statistics(log)
273-
274283
# Log information
275284
if log:
276285
logging.info("New optimal control problem initailized")
@@ -476,10 +485,6 @@ def _eqconst_function(self, coeffs):
476485
# Checked above => we should never get here
477486
raise TypeError("unknown constraint type {ctype}")
478487

479-
# Add the collocation constraints
480-
if self.collocation:
481-
raise NotImplementedError("collocation not yet implemented")
482-
483488
# Update statistics
484489
self.eqconst_evaluations += 1
485490
if self.log:
@@ -497,11 +502,32 @@ def _eqconst_function(self, coeffs):
497502
# Return the value of the constraint function
498503
return np.hstack(value)
499504

500-
def _collocation_constraints(self, coeffs):
505+
def _collocation_constraint(self, coeffs):
501506
# Compute the states and inputs
502507
states, inputs = self._compute_states_inputs(coeffs)
503508

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

506532
#
507533
# Initial guess
@@ -517,40 +543,61 @@ def _collocation_constraints(self, coeffs):
517543
# TODO: figure out how to modify this for collocation
518544
#
519545
def _process_initial_guess(self, initial_guess):
520-
if self.shooting and initial_guess is not None:
521-
# Convert to a 1D array (or higher)
522-
initial_guess = np.atleast_1d(initial_guess)
523-
524-
# See whether we got entire guess or just first time point
525-
if initial_guess.ndim == 1:
526-
# Broadcast inputs to entire time vector
527-
try:
528-
initial_guess = np.broadcast_to(
529-
initial_guess.reshape(-1, 1),
530-
(self.system.ninputs, self.timepts.size))
531-
except ValueError:
532-
raise ValueError("initial guess is the wrong shape")
533-
534-
elif initial_guess.shape != \
535-
(self.system.ninputs, self.timepts.size):
536-
raise ValueError("initial guess is the wrong shape")
546+
# Sort out the input guess and the state guess
547+
if self.collocation and initial_guess is not None and \
548+
isinstance(initial_guess, tuple):
549+
state_guess, input_guess = initial_guess
550+
else:
551+
state_guess, input_guess = None, initial_guess
552+
553+
# Process the input guess
554+
if input_guess is not None:
555+
input_guess = self._broadcast_initial_guess(
556+
input_guess, (self.system.ninputs, self.timepts.size))
537557

538558
# If we were given a basis, project onto the basis elements
539559
if self.basis is not None:
540-
initial_guess = self._inputs_to_coeffs(initial_guess)
560+
input_guess = self._inputs_to_coeffs(input_guess)
561+
else:
562+
input_guess = np.zeros(
563+
self.system.ninputs *
564+
(self.timepts.size if self.basis is None else self.basis.N))
565+
566+
# Process the state guess
567+
if self.collocation:
568+
if state_guess is None:
569+
# Run a simulation to get the initial guess
570+
inputs = input_guess.reshape(self.system.ninputs, -1)
571+
state_guess = self._simulate_states(
572+
np.zeros(self.system.nstates), inputs)
573+
else:
574+
state_guess = self._broadcast_initial_guess(
575+
state_guess, (self.system.nstates, self.timepts.size))
541576

542577
# Reshape for use by scipy.optimize.minimize()
543-
return initial_guess.reshape(-1)
578+
return np.hstack([
579+
input_guess.reshape(-1), state_guess.reshape(-1)])
580+
else:
581+
# Reshape for use by scipy.optimize.minimize()
582+
return input_guess.reshape(-1)
583+
584+
def _broadcast_initial_guess(self, initial_guess, shape):
585+
# Convert to a 1D array (or higher)
586+
initial_guess = np.atleast_1d(initial_guess)
587+
588+
# See whether we got entire guess or just first time point
589+
if initial_guess.ndim == 1:
590+
# Broadcast inputs to entire time vector
591+
try:
592+
initial_guess = np.broadcast_to(
593+
initial_guess.reshape(-1, 1), shape)
594+
except ValueError:
595+
raise ValueError("initial guess is the wrong shape")
544596

545-
elif self.collocation and initial_guess is not None:
546-
raise NotImplementedError("collocation not yet implemented")
597+
elif initial_guess.shape != shape:
598+
raise ValueError("initial guess is the wrong shape")
547599

548-
# Default is no initial guess
549-
input_size = self.system.ninputs * \
550-
(self.timepts.size if self.basis is None else self.basis.N)
551-
state_size = 0 if not self.collocation else \
552-
(self.timepts.size - 1) * self.sys.nstates
553-
return np.zeros(input_size + state_size)
600+
return initial_guess
554601

555602
#
556603
# Utility function to convert input vector to coefficient vector
@@ -650,9 +697,10 @@ def _compute_states_inputs(self, coeffs):
650697
#
651698
if self.collocation:
652699
# States are appended to end of (input) coefficients
653-
states = coeffs[-self.sys.nstates * self.timepts.size:0]
654-
coeffs = coeffs[0:-self.sys.nstates * self.timepts.size]
655-
700+
states = coeffs[-self.system.nstates * self.timepts.size:].reshape(
701+
self.system.nstates, -1)
702+
coeffs = coeffs[:-self.system.nstates * self.timepts.size]
703+
656704
# Compute input at time points
657705
if self.basis:
658706
inputs = self._coeffs_to_inputs(coeffs)
@@ -863,11 +911,8 @@ def __init__(
863911
# Remember the optimal control problem that we solved
864912
self.problem = ocp
865913

866-
# Compute input at time points
867-
if ocp.basis:
868-
inputs = ocp._coeffs_to_inputs(res.x)
869-
else:
870-
inputs = res.x.reshape((ocp.system.ninputs, -1))
914+
# Parse the optimization variables into states and inputs
915+
states, inputs = ocp._compute_states_inputs(res.x)
871916

872917
# See if we got an answer
873918
if not res.success:
@@ -885,6 +930,7 @@ def __init__(
885930

886931
if return_states and inputs.shape[1] == ocp.timepts.shape[0]:
887932
# Simulate the system if we need the state back
933+
# TODO: this is probably not needed due to compute_states_inputs()
888934
_, _, states = ct.input_output_response(
889935
ocp.system, ocp.timepts, inputs, ocp.x, return_x=True,
890936
solve_ivp_kwargs=ocp.solve_ivp_kwargs, t_eval=ocp.timepts)
@@ -1017,7 +1063,7 @@ def solve_ocp(
10171063
'optimal', 'return_x', kwargs, return_states, pop=True)
10181064

10191065
# Process (legacy) method keyword
1020-
if kwargs.get('method'):
1066+
if kwargs.get('method') and method not in optimal_methods:
10211067
if kwargs.get('minimize_method'):
10221068
raise ValueError("'minimize_method' specified more than once")
10231069
kwargs['minimize_method'] = kwargs.pop('method')

0 commit comments

Comments
 (0)