Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Huawei wang homotopy direct collocation #40

Open
wants to merge 15 commits into
base: master
Choose a base branch
from
Prev Previous commit
Next Next commit
Update homotopy codes and unit test codes
1. reverse homotopy in direct collocation according Jason's comments: make homotopy_control and tracing_dynamic_control the keyword arguments, but didn't set the default value yet.

2. reverse the definition of k_matrix to make it a diagram matrix with k in the diagram line.

3. added the homotopy test cases in unit test. Only changed a few test case at the beginning. Update to see the correctness of the codes
HuaweiWang committed May 23, 2016
commit 7298aa68e3d9dc9845f7dc53c475994c09a196ac
1,195 changes: 1,195 additions & 0 deletions direct_collocation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,1195 @@
#!/usr/bin/env python

import numpy as np
import sympy as sym
from sympy.physics import mechanics as me
import ipopt

from .utils import ufuncify_matrix, parse_free


class Problem(ipopt.problem):

def __init__(self, obj, obj_grad, *args, **kwargs):
"""
Parameters
----------
obj : function
Returns the value of the objective function.
obj_grad : function
Returns the gradient of the objective function.
bounds : dictionary
This dictionary should contain a mapping from any of the
symbolic states, unknown trajectories, or unknown parameters 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)}.
"""

self.bounds = kwargs.pop('bounds', None)

self.collocator = ConstraintCollocator(*args, **kwargs)

self.obj = obj
self.obj_grad = obj_grad
self.con = self.collocator.generate_constraint_function()
self.con_jac = self.collocator.generate_jacobian_function()

self.con_jac_rows, self.con_jac_cols = \
self.collocator.jacobian_indices()

self.num_free = self.collocator.num_free
self.num_constraints = self.collocator.num_constraints

self._generate_bound_arrays()

# All constraints are expected to be equal to zero.
con_bounds = np.zeros(self.num_constraints)

super(Problem, self).__init__(n=self.num_free,
m=self.num_constraints,
lb=self.lower_bound,
ub=self.upper_bound,
cl=con_bounds,
cu=con_bounds)

self.obj_value = []

def _generate_bound_arrays(self):
INF = 10e19
lb = -INF * np.ones(self.num_free)
ub = INF * np.ones(self.num_free)

N = self.collocator.num_collocation_nodes
num_state_nodes = N * self.collocator.num_states
num_non_par_nodes = N * (self.collocator.num_states +
self.collocator.num_unknown_input_trajectories)
state_syms = self.collocator.state_symbols
unk_traj = self.collocator.unknown_input_trajectories
unk_par = self.collocator.unknown_parameters

if self.bounds is not None:
for var, bounds in self.bounds.items():
if var in state_syms:
i = state_syms.index(var)
start = i * N
stop = start + N
lb[start:stop] = bounds[0] * np.ones(N)
ub[start:stop] = bounds[1] * np.ones(N)
elif var in unk_traj:
i = unk_traj.index(var)
start = num_state_nodes + i * N
stop = start + N
lb[start:stop] = bounds[0] * np.ones(N)
ub[start:stop] = bounds[1] * np.ones(N)
elif var in unk_par:
i = unk_par.index(var)
idx = num_non_par_nodes + i
lb[idx] = bounds[0]
ub[idx] = bounds[1]

self.lower_bound = lb
self.upper_bound = ub

def objective(self, free):
return self.obj(free)

def gradient(self, free):
# This should return a column vector.
return self.obj_grad(free)

def constraints(self, free):
# This should return a column vector.
return self.con(free)

def jacobianstructure(self):
return (self.con_jac_rows, self.con_jac_cols)

def jacobian(self, free):
return self.con_jac(free)

def intermediate(self, *args):
self.obj_value.append(args[2])


class ConstraintCollocator(object):
"""This class is responsible for generating the constraint function and
the sparse Jacobian of the constraint function using direct collocation
methods for a non-linear programming problem where the essential
constraints are defined from the equations of motion of the system.
Notes
-----
N : number of collocation nodes
N - 1 : number of constraints
n : number of states
m : number of input trajectories
p : number of parameters
q : number of unknown input trajectories
r : number of unknown parameters
"""

time_interval_symbol = sym.Symbol('h', real=True)

def __init__(self, equations_of_motion, state_symbols,
num_collocation_nodes, node_time_interval,
known_parameter_map={}, known_trajectory_map={},
homotopy_control={}, tracing_dynamic_control={},
instance_constraints=None, time_symbol=None, tmp_dir=None,
integration_method='backward euler'):
"""
Parameters
----------
equations_of_motion : sympy.Matrix, shape(n, 1)
A column matrix of SymPy expressions defining the right hand
side of the equations of motion when the left hand side is zero,
e.g. 0 = x'(t) - f(x(t), u(t), p) or 0 = f(x'(t), x(t), u(t),
p). These should be in first order form but not necessairly
explicit.
state_symbols : iterable
An iterable containing all of the SymPy functions of time which
represent the states in the equations of motion.
num_collocation_nodes : integer
The number of collocation nodes, N. All known trajectory arrays
should be of this length.
node_time_interval : float
The time interval between collocation nodes.
known_parameter_map : dictionary, optional
A dictionary that maps the SymPy symbols representing the known
constant parameters to floats. Any parameters in the equations
of motion not provided in this dictionary will become free
optimization variables.
known_trajectory_map : dictionary, optional
A dictionary that maps the non-state SymPy functions of time to
ndarrays of floats of shape(N,). Any time varying parameters in
the equations of motion not provided in this dictionary will
become free trajectories optimization variables.
homotopy_control: float, optional
A parameter in homotopy method that controls the change of
motion equations. The default value of it is 0, which means
the homotopy method does not apply.
tracing_dynamic_control: float, optional
A parameter in homotopy method that adjust the dynamics of extra
'data tracing' term. The default value of it is 1.
integration_method : string, optional
The integration method to use, either `backward euler` or
`midpoint`.
instance_constraints : iterable of SymPy expressions
These expressions are for constraints on the states at specific
time points. They can be expressions with any state instance and
any of the known parameters found in the equations of motion.
All states should be evaluated at a specific instant of time.
For example, the constraint x(0) = 5.0 would be specified as
x(0) - 5.0 and the constraint x(0) = x(5.0) would be specified
as x(0) - x(5.0). Unknown parameters and time varying
parameters other than the states are currently not supported.
time_symbol : SymPy Symbol, optional
The symbol representating time in the equations of motion. If
not given, it is assumed to be the default stored in
dynamicsymbols._t.
tmp_dir : string, optional
If you want to see the generated Cython and C code for the
constraint and constraint Jacobian evaluations, pass in a path
to a directory here.
"""

if time_symbol is not None:
self.time_symbol = time_symbol
else:
self.time_symbol = me.dynamicsymbols._t

self.state_symbols = tuple(state_symbols)
self.state_derivative_symbols = tuple([s.diff(self.time_symbol) for
s in state_symbols])
self.num_states = len(self.state_symbols)

self.num_collocation_nodes = num_collocation_nodes
self.node_time_interval = node_time_interval

self.known_parameter_map = known_parameter_map
self.known_trajectory_map = known_trajectory_map

self.lamda = homotopy_control
self.dynamic = tracing_dynamic_control

self.num_know_trajectory_map = len(self.known_trajectory_map)

self.instance_constraints = instance_constraints

self.num_constraints = self.num_states * (num_collocation_nodes - 1)

self.tmp_dir = tmp_dir

Lamda_matrix = sym.zeros(self.num_states)
Lamda_matrix[self.num_states/2:, self.num_states/2:] = self.lamda * sym.eye(self.num_states/2)

k_matrix = self.dynamic * sym.eye(self.num_states)

if homotopy_control == 0:
self.eom = equations_of_motion

else:
self.eom = (1-Lamda_matrix)*equations_of_motion + Lamda_matrix*(self.state_derivative_symbols - k_matrix
* (known_trajectory_map[self.num_know_trajectory_map
- self.num_states:]-self.state_symbols))

"""
add extra term in equations of motion, the goal is to have the following formula:
homotopy motion equations = [1 0 0 0 ] [f1(x, xd, p, r) ]
0 1 0 0 * f2(x, xd, p, r)
0 0 1-la 0 f3(x, xd, p, r)
[0 0 0 1-la] [f4(x, xd, p, r) ]
[0 0 0 0 ] [x1_dot] [e1] [x1]
+ 0 0 0 0 * ( x2_dot - K*( e2 - x2 ) )
0 0 la 0 x3_dot e3 x3
[0 0 0 la] [x4_dot] [e4] [x4]
"""

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)

self.integration_method = integration_method

if instance_constraints is not None:
self.num_instance_constraints = len(instance_constraints)
self.num_constraints += self.num_instance_constraints
self._identify_functions_in_instance_constraints()
self._find_closest_free_index()
self.eval_instance_constraints = \
self._instance_constraints_func()
self.eval_instance_constraints_jacobian_values = \
self._instance_constraints_jacobian_values_func()

@property
def integration_method(self):
return self._integration_method

@integration_method.setter
def integration_method(self, method):
if method not in ['backward euler', 'midpoint']:
msg = ("{} is not a valid integration method.")
raise ValueError(msg.format(method))
else:
self._integration_method = method
self._discrete_symbols()
self._discretize_eom()

@staticmethod
def _parse_inputs(all_syms, known_syms):
"""Returns sets of symbols and their counts, based on if the known
symbols exist in the set of all symbols.
Parameters
----------
all_syms : sequence
A set of SymPy symbols or functions.
known_syms : sequence
A set of SymPy symbols or functions.
Returns
-------
known : tuple
The set of known symbols.
num_known : integer
The number of known symbols.
unknown : tuple
The set of unknown symbols in all_syms.
num_unknown :integer
The number of unknown symbols.
"""
all_syms = set(all_syms)
known_syms = known_syms

def sort_sympy(seq):
seq = list(seq)
try: # symbols
seq.sort(key=lambda x: x.name)
except AttributeError: # functions
seq.sort(key=lambda x: x.__class__.__name__)
return seq

if not all_syms: # if empty sequence
if known_syms:
msg = '{} are not in the provided equations of motion.'
raise ValueError(msg.format(known_syms))
else:
known = tuple()
num_known = 0
unknown = tuple()
num_unknown = 0
else:
if known_syms:
known = tuple(known_syms) # don't sort known syms
num_known = len(known)
unknown = tuple(sort_sympy(all_syms.difference(known)))
num_unknown = len(unknown)
else:
known = tuple()
num_known = 0
unknown = tuple(sort_sympy(all_syms))
num_unknown = len(unknown)

return known, num_known, unknown, num_unknown

def _sort_parameters(self):
"""Finds and counts all of the parameters in the equations of motion
and categorizes them based on which parameters the user supplies.
The unknown parameters are sorted by name."""

parameters = self.eom.free_symbols.copy()
parameters.remove(self.time_symbol)

res = self._parse_inputs(parameters,
self.known_parameter_map.keys())

self.known_parameters = res[0]
self.num_known_parameters = res[1]
self.unknown_parameters = res[2]
self.num_unknown_parameters = res[3]

self.parameters = res[0] + res[2]
self.num_parameters = len(self.parameters)

def _check_known_trajectories(self):
"""Raises and error if the known tracjectories are not the correct
length."""

N = self.num_collocation_nodes

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))

def _sort_trajectories(self):
"""Finds and counts all of the non-state, time varying parameters in
the equations of motion and categorizes them based on which
parameters the user supplies. The unknown parameters are sorted by
name."""

states = set(self.state_symbols)
states_derivatives = set(self.state_derivative_symbols)

time_varying_symbols = me.find_dynamicsymbols(self.eom)
state_related = states.union(states_derivatives)
non_states = time_varying_symbols.difference(state_related)

res = self._parse_inputs(non_states,
self.known_trajectory_map.keys())

self.known_input_trajectories = res[0]
self.num_known_input_trajectories = res[1]
self.unknown_input_trajectories = res[2]
self.num_unknown_input_trajectories = res[3]

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 euqations of motion.
Instantiates
------------
previous_discrete_state_symbols : tuple of sympy.Symbols
The n symbols representing the system's (ith - 1) states.
current_discrete_state_symbols : tuple of sympy.Symbols
The n symbols representing the system's ith states.
next_discrete_state_symbols : tuple of sympy.Symbols
The n symbols representing the system's (ith + 1) states.
current_known_discrete_specified_symbols : tuple of sympy.Symbols
The symbols representing the system's ith known input
trajectories.
next_known_discrete_specified_symbols : tuple of sympy.Symbols
The symbols representing the system's (ith + 1) known input
trajectories.
current_unknown_discrete_specified_symbols : tuple of sympy.Symbols
The symbols representing the system's ith unknown input
trajectories.
next_unknown_discrete_specified_symbols : tuple of sympy.Symbols
The symbols representing the system's (ith + 1) unknown input
trajectories.
current_discrete_specified_symbols : tuple of sympy.Symbols
The m symbols representing the system's ith specified inputs.
next_discrete_specified_symbols : tuple of sympy.Symbols
The m symbols representing the system's (ith + 1) specified
inputs.
"""

# The previus, current, and next states.
self.previous_discrete_state_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'p', real=True)
for f in self.state_symbols])
self.current_discrete_state_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'i', real=True)
for f in self.state_symbols])
self.next_discrete_state_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'n', real=True)
for f in self.state_symbols])

# The current and next known input trajectories.
self.current_known_discrete_specified_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'i', real=True)
for f in self.known_input_trajectories])
self.next_known_discrete_specified_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'n', real=True)
for f in self.known_input_trajectories])

# The current and next unknown input trajectories.
self.current_unknown_discrete_specified_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'i', real=True)
for f in self.unknown_input_trajectories])
self.next_unknown_discrete_specified_symbols = \
tuple([sym.Symbol(f.__class__.__name__ + 'n', real=True)
for f in self.unknown_input_trajectories])

self.current_discrete_specified_symbols = \
(self.current_known_discrete_specified_symbols +
self.current_unknown_discrete_specified_symbols)
self.next_discrete_specified_symbols = \
(self.next_known_discrete_specified_symbols +
self.next_unknown_discrete_specified_symbols)

def _discretize_eom(self):
"""Instantiates the constraint equations in a discretized form using
backward Euler discretization.
Instantiates
------------
discrete_eoms : sympy.Matrix, shape(n, 1)
The column vector of the discretized equations of motion.
"""
x = self.state_symbols
xd = self.state_derivative_symbols
u = self.input_trajectories

xp = self.previous_discrete_state_symbols
xi = self.current_discrete_state_symbols
xn = self.next_discrete_state_symbols
ui = self.current_discrete_specified_symbols
un = self.next_discrete_specified_symbols

h = self.time_interval_symbol

if self.integration_method == 'backward euler':

deriv_sub = {d: (i - p) / h for d, i, p in zip(xd, xi, xp)}

func_sub = dict(zip(x + u, xi + ui))

self.discrete_eom = me.msubs(self.eom, deriv_sub, func_sub)

elif self.integration_method == 'midpoint':

xdot_sub = {d: (n - i) / h for d, i, n in zip(xd, xi, xn)}
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 _identify_functions_in_instance_constraints(self):
"""Instantiates a set containing all of the instance functions, i.e.
x(1.0) in the instance constraints."""

all_funcs = set()

for con in self.instance_constraints:
all_funcs = all_funcs.union(con.atoms(sym.Function))

self.instance_constraint_function_atoms = all_funcs

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

N = self.num_collocation_nodes
h = self.node_time_interval
duration = h * (N - 1)

time_vector = np.linspace(0.0, duration, num=N)

node_map = {}
for func in self.instance_constraint_function_atoms:
time_value = func.args[0]
time_index = np.argmin(np.abs(time_vector - time_value))
free_index = determine_free_index(time_index,
func.__class__(self.time_symbol))
node_map[func] = free_index

self.instance_constraints_free_index_map = node_map

def _instance_constraints_func(self):
"""Returns a function that evaluates the instance constraints given
the free optimization variables."""
free = sym.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 = sym.lambdify(([free] + self.known_parameter_map.keys()),
subbed_constraints, modules=[{'ImmutableMatrix':
np.array}, "numpy"])
def wrapped(free):
return f(free, *self.known_parameter_map.values())

return wrapped

def _instance_constraints_jacobian_indices(self):
"""Returns the row and column indices of the non-zero values in the
Jacobian of the constraints."""
idx_map = self.instance_constraints_free_index_map

num_eom_constraints = self.num_states * (self.num_collocation_nodes - 1)

rows = []
cols = []

for i, con in enumerate(self.instance_constraints):
funcs = con.atoms(sym.Function)
indices = [idx_map[f] for f in funcs]
row_idxs = num_eom_constraints + i * np.ones(len(indices),
dtype=int)
rows += list(row_idxs)
cols += indices

return np.array(rows), np.array(cols)

def _instance_constraints_jacobian_values_func(self):
"""Retruns the non-zero values of the constraint Jacobian associated
with the instance constraints."""
free = sym.DeferredVector('FREE')

def_map = {k: free[v] for k, v in
self.instance_constraints_free_index_map.items()}

funcs = []
num_vals_per_func = []
for con in self.instance_constraints:
partials = list(con.atoms(sym.Function))
num_vals_per_func.append(len(partials))
jac = sym.Matrix([con]).jacobian(partials)
jac = jac.subs(def_map)
funcs.append(sym.lambdify(([free] +
self.known_parameter_map.keys()),
jac,
modules=[{'ImmutableMatrix': np.array},
"numpy"]))
l = np.sum(num_vals_per_func)

def wrapped(free):
arr = np.zeros(l)
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())
j += num
return arr

return wrapped

def _gen_multi_arg_con_func(self):
"""Instantiates a function that evaluates the constraints given all
of the arguments of the functions, i.e. not just the free
optimization variables.
Instantiates
------------
_multi_arg_con_func : function
A function which returns the numerical values of the constraints
at collocation nodes 2,...,N.
Notes
-----
args:
all current states (x1i, ..., xni)
all previous states (x1p, ... xnp)
all current specifieds (s1i, ..., smi)
parameters (c1, ..., cb)
time interval (h)
args: (x1i, ..., xni, x1p, ... xnp, s1i, ..., smi, c1, ..., cb, h)
n: num states
m: num specified
b: num parameters
The function should evaluate and return an array:
[con_1_2, ..., con_1_N, con_2_2, ...,
con_2_N, ..., con_n_2, ..., con_n_N]
for n states and N-1 constraints at the time points.
"""
xi_syms = self.current_discrete_state_symbols
xp_syms = self.previous_discrete_state_symbols
xn_syms = self.next_discrete_state_symbols
si_syms = self.current_discrete_specified_symbols
sn_syms = self.next_discrete_specified_symbols
h_sym = self.time_interval_symbol
constant_syms = self.known_parameters + self.unknown_parameters

if self.integration_method == 'backward euler':

args = [x for x in xi_syms] + [x for x in xp_syms]
args += [s for s in si_syms] + list(constant_syms) + [h_sym]

current_start = 1
current_stop = None
adjacent_start = None
adjacent_stop = -1

elif self.integration_method == 'midpoint':

args = [x for x in xi_syms] + [x for x in xn_syms]
args += [s for s in si_syms] + [s for s in sn_syms]
args += list(constant_syms) + [h_sym]

current_start = None
current_stop = -1
adjacent_start = 1
adjacent_stop = None

f = ufuncify_matrix(args, self.discrete_eom,
const=constant_syms + (h_sym,),
tmp_dir=self.tmp_dir)

def constraints(state_values, specified_values, constant_values,
interval_value):
"""Returns a vector of constraint values given all of the
unknowns in the equations of motion over the 2, ..., N time
steps.
Parameters
----------
states : ndarray, shape(n, N)
The array of n states through N time steps.
specified_values : ndarray, shape(m, N) or shape(N,)
The array of m specifieds through N time steps.
constant_values : ndarray, shape(b,)
The array of b parameters.
interval_value : float
The value of the discretization time interval.
Returns
-------
constraints : ndarray, shape(N-1,)
The array of constraints from t = 2, ..., N.
[con_1_2, ..., con_1_N, con_2_2, ...,
con_2_N, ..., con_n_2, ..., con_n_N]
"""

if state_values.shape[0] < 2:
raise ValueError('There should always be at least two states.')

assert state_values.shape == (self.num_states,
self.num_collocation_nodes)

x_current = state_values[:, current_start:current_stop] # n x N - 1
x_adjacent = state_values[:, adjacent_start:adjacent_stop] # n x N - 1

# 2n x N - 1
args = [x for x in x_current] + [x for x in x_adjacent]

# 2n + m x N - 1
if len(specified_values.shape) == 2:
assert specified_values.shape == \
(self.num_input_trajectories,
self.num_collocation_nodes)
si = specified_values[:, current_start:current_stop]
args += [s for s in si]
if self.integration_method == 'midpoint':
sn = specified_values[:, adjacent_start:adjacent_stop]
args += [s for s in sn]
elif len(specified_values.shape) == 1 and specified_values.size != 0:
assert specified_values.shape == \
(self.num_collocation_nodes,)
si = specified_values[current_start:current_stop]
args += [si]
if self.integration_method == 'midpoint':
sn = specified_values[adjacent_start:adjacent_stop]
args += [sn]

args += [c for c in constant_values]
args += [interval_value]

num_constraints = state_values.shape[1] - 1

# TODO : Move this to an attribute of the class so that it is
# only initialized once and just reuse it on each evaluation of
# this function.
result = np.empty((num_constraints, state_values.shape[0]))

return f(result, *args).T.flatten()

self._multi_arg_con_func = constraints

def jacobian_indices(self):
"""Returns the row and column indices for the non-zero values in the
constraint Jacobian.
Returns
-------
jac_row_idxs : ndarray, shape(2 * n + q + r,)
The row indices for the non-zero values in the Jacobian.
jac_col_idxs : ndarray, shape(n,)
The column indices for the non-zero values in the Jacobian.
"""

N = self.num_collocation_nodes
n = self.num_states

num_constraint_nodes = N - 1

if self.integration_method == 'backward euler':

num_partials = n * (2 * n + self.num_unknown_input_trajectories
+ self.num_unknown_parameters)

elif self.integration_method == 'midpoint':

num_partials = n * (2 * n + 2 *
self.num_unknown_input_trajectories +
self.num_unknown_parameters)

num_non_zero_values = num_constraint_nodes * num_partials

if self.instance_constraints is not None:
ins_row_idxs, ins_col_idxs = \
self._instance_constraints_jacobian_indices()
num_non_zero_values += len(ins_row_idxs)

jac_row_idxs = np.empty(num_non_zero_values, dtype=int)
jac_col_idxs = np.empty(num_non_zero_values, dtype=int)

"""
The symbolic derivative matrix for a single constraint node follows
these patterns:
Backward Euler
--------------
i: ith, p: ith-1
For example:
x1i = the first state at the ith constraint node
uqi = the qth state at the ith constraint node
uqn = the qth state at the ith+1 constraint node
[x1] [x1i, ..., xni, x1p, ..., xnp, u1i, .., uqi, p1, ..., pr]
[. ]
[. ]
[. ]
[xn]
Midpoint
--------
i: ith, n: ith+1
[x1] [x1i, ..., xni, x1n, ..., xnn, u1i, .., uqi, u1n, ..., uqn, p1, ..., pp]
[. ]
[. ]
[. ]
[xn]
Each of these matrices are evaulated at N-1 constraint nodes and
then the 3D matrix is flattened into a 1d array. The backward euler
uses nodes 1 <= i <= N-1 and the midpoint uses 0 <= i <= N - 2. So
the flattened arrays looks like:
M = N-1
P = N-2
Backward Euler
--------------
i=1 x1 | [x11, ..., xn1, x10, ..., xn0, u11, .., uq1, p1, ..., pr,
x2 | x11, ..., xn1, x10, ..., xn0, u11, .., uq1, p1, ..., pr,
... | ...,
xn | x11, ..., xn1, x10, ..., xn0, u11, .., uq1, p1, ..., pr,
i=2 x1 | x12, ..., xn2, x11, ..., xn1, u12, .., uq2, p1, ..., pr,
x2 | x12, ..., xn2, x11, ..., xn1, u12, .., uq2, p1, ..., pr,
... | ...,
xn | x12, ..., xn2, x11, ..., xn1, u12, .., uq2, p1, ..., pr,
| ...,
i=M x1 | x1M, ..., xnM, x1P, ..., xnP, u1M, .., uqM, p1, ..., pr,
x2 | x1M, ..., xnM, x1P, ..., xnP, u1M, .., uqM, p1, ..., pr,
... | ...,
xn | x1M, ..., xnM, x1P, ..., xnP, u1M, .., uqM, p1, ..., pr]
Midpoint
--------
i=0 x1 | [x10, ..., xn0, x11, ..., xn1, u10, .., uq0, u11, .., uq1, p1, ..., pr,
x2 | x10, ..., xn0, x11, ..., xn1, u10, .., uq0, u11, .., uq1, p1, ..., pr,
... | ...,
xn | x10, ..., xn0, x11, ..., xn1, u10, .., uq0, u11, .., uq1, p1, ..., pr,
i=1 x1 | x11, ..., xn1, x12, ..., xn2, u11, .., uq1, u12, .., uq2, p1, ..., pr,
x2 | x11, ..., xn1, x12, ..., xn2, u11, .., uq1, u12, .., uq2, p1, ..., pr,
... | ...,
xn | x11, ..., xn1, x12, ..., xn2, u11, .., uq1, u12, .., uq2, p1, ..., pr,
... | ...,
i=P x1 | x1P, ..., xnP, x1M, ..., xnM, u1P, .., uqP, u1M, .., uqM, p1, ..., pr,
x2 | x1P, ..., xnP, x1M, ..., xnM, u1P, .., uqP, u1M, .., uqM, p1, ..., pr,
... | ...,
xn | x1P, ..., xnP, x1M, ..., xnM, u1P, .., uqP, u1M, .., uqM, p1, ..., pr]
These two arrays contain of the non-zero values of the sparse
Jacobian[#]_.
.. [#] Some of the partials can be equal to zero and could be
excluded from the array. These could be a significant number.
Now we need to generate the triplet format indices of the full
sparse Jacobian for each one of the entries in these arrays. The
format of the Jacobian matrix is:
Backward Euler
--------------
[x10, ..., x1N-1, ..., xn0, ..., xnN-1, u10, ..., u1N-1, ..., uq0, ..., uqN-1, p1, ..., pr]
[x11]
[x12]
[...]
[x1M]
[...]
[xn1]
[xn2]
[...]
[xnM]
Midpoint
--------
[x10, ..., x1N-1, ..., xn0, ..., xnN-1, u10, ..., u1N-1, ..., uq0, ..., uqN-1, p1, ..., pr]
[x10]
[x11]
[...]
[x1P]
[...]
[xn0]
[xn1]
[...]
[xnP]
"""
for i in range(num_constraint_nodes):

# n : number of states
# m : number of input trajectories
# p : number of parameters
# q : number of unknown input trajectories
# r : number of unknown parameters

# the states repeat every N - 1 constraints
# row_idxs = [0 * (N - 1), 1 * (N - 1), 2 * (N - 1), ..., n * (N - 1)]

# This gives the Jacobian row indices matching the ith
# constraint node for each state. ith corresponds to the loop
# indice.
row_idxs = [j * (num_constraint_nodes) + i for j in range(n)]

# first row, the columns indices mapping is:
# [1, N + 1, ..., N - 1] : [x1p, x1i, 0, ..., 0]
# [0, N, ..., 2 * (N - 1)] : [x2p, x2i, 0, ..., 0]
# [-p:] : p1,..., pp the free constants

# i=0: [1, ..., n * N + 1, 0, ..., n * N + 0, n * N:n * N + p]
# i=1: [2, ..., n * N + 2, 1, ..., n * N + 1, n * N:n * N + p]
# i=2: [3, ..., n * N + 3, 2, ..., n * N + 2, n * N:n * N + p]

if self.integration_method == 'backward euler':

col_idxs = [j * N + i + 1 for j in range(n)]
col_idxs += [j * N + i for j in range(n)]
col_idxs += [n * N + j * N + i + 1 for j in
range(self.num_unknown_input_trajectories)]
col_idxs += [(n + self.num_unknown_input_trajectories) * N + j
for j in range(self.num_unknown_parameters)]

elif self.integration_method == 'midpoint':

col_idxs = [j * N + i for j in range(n)]
col_idxs += [j * N + i + 1 for j in range(n)]
col_idxs += [n * N + j * N + i for j in
range(self.num_unknown_input_trajectories)]
col_idxs += [n * N + j * N + i + 1 for j in
range(self.num_unknown_input_trajectories)]
col_idxs += [(n + self.num_unknown_input_trajectories) * N + j
for j in range(self.num_unknown_parameters)]

row_idx_permutations = np.repeat(row_idxs, len(col_idxs))
col_idx_permutations = np.array(list(col_idxs) * len(row_idxs),
dtype=int)

start = i * num_partials
stop = (i + 1) * num_partials
jac_row_idxs[start:stop] = row_idx_permutations
jac_col_idxs[start:stop] = col_idx_permutations

if self.instance_constraints is not None:
jac_row_idxs[-len(ins_row_idxs):] = ins_row_idxs
jac_col_idxs[-len(ins_col_idxs):] = ins_col_idxs

return jac_row_idxs, jac_col_idxs

def _gen_multi_arg_con_jac_func(self):
"""Instantiates a function that evaluates the Jacobian of the
constraints.
Instantiates
------------
_multi_arg_con_jac_func : function
A function which returns the numerical values of the constraints
at time points 2,...,N.
"""
xi_syms = self.current_discrete_state_symbols
xp_syms = self.previous_discrete_state_symbols
xn_syms = self.next_discrete_state_symbols
si_syms = self.current_discrete_specified_symbols
sn_syms = self.next_discrete_specified_symbols
ui_syms = self.current_unknown_discrete_specified_symbols
un_syms = self.next_unknown_discrete_specified_symbols
h_sym = self.time_interval_symbol
constant_syms = self.known_parameters + self.unknown_parameters

if self.integration_method == 'backward euler':

# The free parameters are always the n * (N - 1) state values,
# the unknown input trajectories, and the unknown model
# constants, so the base Jacobian needs to be taken with respect
# to the ith, and ith - 1 states, and the free model constants.
wrt = (xi_syms + xp_syms + ui_syms + self.unknown_parameters)

# The arguments to the Jacobian function include all of the free
# Symbols/Functions in the matrix expression.
args = xi_syms + xp_syms + si_syms + constant_syms + (h_sym,)

current_start = 1
current_stop = None
adjacent_start = None
adjacent_stop = -1

elif self.integration_method == 'midpoint':

wrt = (xi_syms + xn_syms + ui_syms + un_syms +
self.unknown_parameters)

# The arguments to the Jacobian function include all of the free
# Symbols/Functions in the matrix expression.
args = (xi_syms + xn_syms + si_syms + sn_syms + constant_syms +
(h_sym,))

current_start = None
current_stop = -1
adjacent_start = 1
adjacent_stop = None

# This creates a matrix with all of the symbolic partial derivatives
# necessary to compute the full Jacobian.
symbolic_partials = self.discrete_eom.jacobian(wrt)

# This generates a numerical function that evaluates the matrix of
# partial derivatives. This function returns the non-zero elements
# needed to build the sparse constraint Jacobian.
eval_partials = ufuncify_matrix(args, symbolic_partials,
const=constant_syms + (h_sym,),
tmp_dir=self.tmp_dir)

result = np.empty((self.num_collocation_nodes - 1,
symbolic_partials.shape[0] *
symbolic_partials.shape[1]))

def constraints_jacobian(state_values, specified_values,
parameter_values, interval_value):
"""Returns the values of the sparse constraing Jacobian matrix
given all of the values for each variable in the equations of
motion over the N - 1 nodes.
Parameters
----------
states : ndarray, shape(n, N)
The array of n states through N time steps. There are always
at least two states.
specified_values : ndarray, shape(m, N) or shape(N,)
The array of m specified inputs through N time steps.
parameter_values : ndarray, shape(p,)
The array of p parameter.
interval_value : float
The value of the discretization time interval.
Returns
-------
constraint_jacobian_values : ndarray, shape(see below,)
backward euler: shape((N - 1) * n * (2*n + q + r),)
midpoint: shape((N - 1) * n * (2*n + 2*q + r),)
The values of the non-zero entries of the constraints
Jacobian. These correspond to the triplet formatted indices
returned from jacobian_indices.
Notes
-----
N : number of collocation nodes
n : number of states
m : number of input trajectories
p : number of parameters
q : number of unknown input trajectories
r : number of unknown parameters
n * (N - 1) : number of constraints
"""
if state_values.shape[0] < 2:
raise ValueError('There should always be at least two states.')

# Each of these arrays are shape(n, N - 1). The x_adjacent is
# either the previous value of the state or the next value of
# the state, depending on the integration method.
x_current = state_values[:, current_start:current_stop]
x_adjacent = state_values[:, adjacent_start:adjacent_stop]

# 2n x N - 1
args = [x for x in x_current] + [x for x in x_adjacent]

# 2n + m x N - 1
if len(specified_values.shape) == 2:
si = specified_values[:, current_start:current_stop]
args += [s for s in si]
if self.integration_method == 'midpoint':
sn = specified_values[:, adjacent_start:adjacent_stop]
args += [s for s in sn]
elif len(specified_values.shape) == 1 and specified_values.size != 0:
si = specified_values[current_start:current_stop]
args += [si]
if self.integration_method == 'midpoint':
sn = specified_values[adjacent_start:adjacent_stop]
args += [sn]

args += [c for c in parameter_values]
args += [interval_value]

# backward euler: shape(N - 1, n, 2*n + q + r)
# midpoint: shape(N - 1, n, 2*n + 2*q + r)
non_zero_derivatives = eval_partials(result, *args)

return non_zero_derivatives.ravel()

self._multi_arg_con_jac_func = constraints_jacobian

@staticmethod
def _merge_fixed_free(syms, fixed, free, typ):
"""Returns an array with the fixed and free values combined. This
just takes the known and unknown values and combines them for the
function evaluation.
This assumes that you have the free constants in the correct order.
Parameters
----------
syms : iterable of SymPy Symbols or Functions
fixed : dictionary
A mapping from Symbols to floats or Functions to 1d ndarrays.
free : ndarray, (N,) or shape(n,N)
An array
typ : string
traj or par
"""

merged = []
n = 0
# syms is order as known (fixed) then unknown (free)
for i, s in enumerate(syms):
if s in fixed.keys():
merged.append(fixed[s])
else:
if typ == 'traj' and len(free.shape) == 1:
merged.append(free)
else:
merged.append(free[n])
n += 1
return np.array(merged)

def _wrap_constraint_funcs(self, func, typ):
"""Returns a function that evaluates all of the constraints or
Jacobian of the constraints given the free optimization variables.
Parameters
----------
func : function
A function that takes the full parameter set and evaluates the
constraint functions or the Jacobian of the contraint functions.
i.e. the output of _gen_multi_arg_con_func or
_gen_multi_arg_con_jac_func.
Returns
-------
func : function
A function which returns constraint values given the system's
free optimization variables.
"""

def constraints(free):

free_states, free_specified, free_constants = \
parse_free(free, self.num_states,
self.num_unknown_input_trajectories,
self.num_collocation_nodes)

all_specified = self._merge_fixed_free(self.input_trajectories,
self.known_trajectory_map,
free_specified, 'traj')

all_constants = self._merge_fixed_free(self.parameters,
self.known_parameter_map,
free_constants, 'par')

eom_con_vals = func(free_states, all_specified, all_constants,
self.node_time_interval)

if self.instance_constraints is not None:
if typ == 'con':
ins_con_vals = self.eval_instance_constraints(free)
elif typ == 'jac':
ins_con_vals = self.eval_instance_constraints_jacobian_values(free)
return np.hstack((eom_con_vals, ins_con_vals))
else:
return eom_con_vals

intro, second = func.__doc__.split('Parameters')
params, returns = second.split('Returns')
new_doc = '{}Parameters\n----------\nfree : ndarray, shape()\n\nReturns\n{}'
constraints.__doc__ = new_doc.format(intro, returns)

return constraints

def generate_constraint_function(self):
"""Returns a function which evaluates the constraints given the
array of free optimization variables."""
self._gen_multi_arg_con_func()
return self._wrap_constraint_funcs(self._multi_arg_con_func, 'con')

def generate_jacobian_function(self):
"""Returns a function which evaluates the Jacobian of the
constraints given the array of free optimization variables."""
self._gen_multi_arg_con_jac_func()
return self._wrap_constraint_funcs(self._multi_arg_con_jac_func, 'jac')
1,117 changes: 1,117 additions & 0 deletions test_direct_collocation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,1117 @@
#!/usr/bin/env python

from collections import OrderedDict

import numpy as np
import sympy as sym
from scipy import sparse
from nose.tools import raises

from ..direct_collocation import Problem, ConstraintCollocator


def test_Problem():

m, c, k, t = sym.symbols('m, c, k, t')
x, v, f = [s(t) for s in sym.symbols('x, v, f', cls=sym.Function)]

state_symbols = (x, v)

interval_value = 0.01

eom = sym.Matrix([x.diff() - v,
m * v.diff() + c * v + k * x - f])

prob = Problem(lambda x: 1.0,
lambda x: x,
eom,
state_symbols,
2,
interval_value,
bounds={x: (-10.0, 10.0),
f: (-8.0, 8.0),
m: (-1.0, 1.0),
c: (-0.5, 0.5)})

INF = 10e19
expected_lower = np.array([-10.0, -10.0,
-INF, -INF,
-8.0, -8.0,
-0.5, -INF, -1.0])
np.testing.assert_allclose(prob.lower_bound, expected_lower)
expected_upper = np.array([10.0, 10.0,
INF, INF,
8.0, 8.0,
0.5, INF, 1.0])
np.testing.assert_allclose(prob.upper_bound, expected_upper)


class TestConstraintCollocator():

def setup(self):

m, c, k, la, dy, t = sym.symbols('m, c, k, t')
la, dy = sym.symbols('la, dy')
x, v, f, xm, vm = [s(t) for s in sym.symbols('x, v, f, xm, vm', cls=sym.Function)]

self.state_symbols = (x, v)
self.constant_symbols = (m, c, k)
self.lamda = la
self.dynamic = dy
self.specified_symbols = (f, xm, vm)
self.discrete_symbols = sym.symbols('xi, vi, xp, vp, xn, vn, fi, fn',
real=True)

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, 2.0, 2.0, 2.0],
[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0]])
self.constant_values = np.array([1.0, 2.0, 3.0])
self.lamda = 0.0
self.dynamic = 5.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
3.0]) # k

lamda_matrix = sym.zeros(2)
lamda_matrix[1:, 1:] = self.lamda * sym.eye(1)

#k_matrix = self.dynamic * sym.eye(2)

self.eom = (1-lamda_matrix) * sym.Matrix([x.diff() - v,
m * v.diff() + c * v + k * x - f]) + lamda_matrix*sym.Matrix([x.diff() - self.dynamic*(xm - x),
v.diff() - self.dynamic*(vm - v)])

par_map = OrderedDict(zip(self.constant_symbols[:2],
self.constant_values[:2]))
traj_map = OrderedDict(zip([f], [self.specified_values]))

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=traj_map,
homotopy_control=self.lamda,
tracing_dynamic_control=self.dynamic,
tmp_dir='test_ufuncs')

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

@raises(ValueError)
def test_integration_method(self):
self.collocator.integration_method = 'booger'

def test_sort_parameters(self):

# TODO : Added checks for the other cases.

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):

# TODO : Added checks for the other cases.

self.collocator._sort_trajectories()

assert self.collocator.known_input_trajectories == \
self.specified_symbols
assert self.collocator.num_known_input_trajectories == 3

assert self.collocator.unknown_input_trajectories == tuple()
assert self.collocator.num_unknown_input_trajectories == 0

def test_homotopy_parameters(self):

# TODO : Added checks for the other cases.

self.collocator.lamda()

la = self.lamda
dy = self.dynamic

assert self.collocator.lamda == la
assert self.collocator.dynamic == dy

def test_discrete_symbols(self):

self.collocator._discrete_symbols()

xi, vi, xp, vp, xn, vn, fi, fn, xmi, xmn, vmi, vmn = 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 == (fi, xmi, vmi)
assert self.collocator.next_discrete_specified_symbols == (fn, xmn, vmn)

def test_discretize_eom_backward_euler(self):

m, c, k = self.constant_symbols
xi, vi, xp, vp, xn, vn, fi, fn, xmi, xmn, vmi, vmn = self.discrete_symbols
h = self.collocator.time_interval_symbol

lamda_matrix = sym.zeros(2)
lamda_matrix[1:, 1:] = self.lamda * sym.eye(1)
#expected = sym.Matrix([(xi - xp) / h - vi,
# m * (vi - vp) / h + c * vi + k * xi - fi])

expected = (1-lamda_matrix)*sym.Matrix([(xi - xp) / h - vi,
m * (vi - vp) / h + c * vi + k * xi - fi]) + lamda_matrix*sym.Matrix([(xi - xp) / h - self.dynamic*(xmi - xi),
(vi - vp) / h - self.dynamic*(vmi - vi)])

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, xmi, xmn, vmi, vmn = self.discrete_symbols

h = self.collocator.time_interval_symbol

lamda_matrix = sym.zeros(2)
lamda_matrix[1:, 1:] = self.lamda * sym.eye(1)

#expected = sym.Matrix([(xn - xi) / h - (vi + vn) / 2,
# m * (vn - vi) / h + c * (vi + vn) / 2 +
# k * (xi + xn) / 2 - (fi + fn) / 2])

expected = (1-lamda_matrix) * sym.Matrix([(xn - xi) / h - (vi + vn) / 2,
m * (vn - vi) / h + c * (vi + vn) / 2 + k * (xi + xn) / 2 - (fi + fn) / 2]) + lamda_matrix*sym.Matrix([(xn - xi) / h - self.dynamic*((xmn - xmi) / h - (xn - xi) / h),
(vn - vi) / h - self.dynamic*((vmn - vmi) / h - (vn - vi) / h)])

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._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,
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]

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,
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:i + 1]
fn = self.specified_values[i + 1:i + 2]

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_gen_multi_arg_con_jac_func_backward_euler(self):

self.collocator._gen_multi_arg_con_jac_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.

jac_vals = self.collocator._multi_arg_con_jac_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

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

expected_jacobian = np.array(
# x1, x2, x3, x4, v1, v2, v3, v4, k
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0],
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0],
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0],
[ 0, k, 0, 0, -m / h, c + m / h, 0, 0, x[1]],
[ 0, 0, k, 0, 0, -m / h, c + m / h, 0, x[2]],
[ 0, 0, 0, k, 0, 0, -m /h, c + m / h, x[3]]],
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)

def test_gen_multi_arg_con_jac_func_midpoint(self):

self.collocator.integration_method = 'midpoint'
self.collocator._gen_multi_arg_con_jac_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.

jac_vals = self.collocator._multi_arg_con_jac_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

row_idxs, col_idxs = self.collocator.jacobian_indices()

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

x = self.state_values[0]
m, c, k = self.constant_values
h = self.interval_value

part1 = np.array(
# x0, x1, x2, x3
[[ -1.0 / h, 1.0 / h, 0, 0],
[ 0, -1.0 / h, 1.0 / h, 0],
[ 0, 0, -1.0 / h, 1.0 / h],
[ k / 2.0, k / 2.0, 0, 0],
[ 0, k / 2.0, k / 2.0, 0],
[ 0, 0, k / 2.0, k / 2.0]],
dtype=float)

part2 = np.array(
# v0, v1, v2, v3, k i
[[ -1.0 / 2.0, -1.0 / 2.0, 0, 0, 0], # 0
[ 0, -1.0 / 2.0, -1.0 / 2.0, 0, 0], # 1
[ 0, 0, -1.0 / 2.0, -1.0 / 2.0, 0], # 2
[ -m / h + c / 2.0, m / h + c / 2.0, 0, 0, (x[1] + x[0]) / 2.0], # 0
[ 0, -m / h + c / 2.0, m / h + c / 2.0, 0, (x[2] + x[1]) / 2.0], # 1
[ 0, 0, -m / h + c / 2.0, m / h + c / 2.0, (x[3] + x[2]) / 2.0]], # 2
dtype=float)

expected_jacobian = np.hstack((part1, part2))

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)

def test_generate_constraint_function(self):

constrain = self.collocator.generate_constraint_function()

result = constrain(self.free)

expected_dynamic = np.zeros(3)
expected_kinematic = np.zeros(3)

m, c, k = self.constant_values
h = self.interval_value

for i in [1, 2, 3]:

xi, vi = self.state_values[:, i]
xp, vp = self.state_values[:, i - 1]
fi = self.specified_values[i]

expected_kinematic[i - 1] = (xi - xp) / h - vi
expected_dynamic[i - 1] = m * (vi - vp) / h + c * vi + k * xi - fi

expected = np.hstack((expected_kinematic, expected_dynamic))

np.testing.assert_allclose(result, expected)

def test_generate_jacobian_function(self):

jacobian = self.collocator.generate_jacobian_function()

jac_vals = jacobian(self.free)

row_idxs, col_idxs = self.collocator.jacobian_indices()

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

m, c, k = self.constant_values
h = self.interval_value
x = self.state_values[0]

expected_jacobian = np.array(
# x1, x2, x3, x4, v1, v2, v3, v4, k
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0],
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0],
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0],
[ 0, k, 0, 0, -m / h, c + m / h, 0, 0, x[1]],
[ 0, 0, k, 0, 0, -m / h, c + m / h, 0, x[2]],
[ 0, 0, 0, k, 0, 0, -m /h, c + m / h, x[3]]],
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)


class TestConstraintCollocatorUnknownTrajectories():

def setup(self):

# constant parameters
m, c, t = sym.symbols('m, c, t')
# time varying parameters
x, v, f, k = [s(t) for s in sym.symbols('x, v, f, k',
cls=sym.Function)]

self.state_symbols = (x, v)
self.constant_symbols = (m, c)
self.specified_symbols = (f, k)
self.discrete_symbols = sym.symbols('xi, vi, xp, vp, fi, ki',
real=True)

# The state order should match the eom order and state symbol order.
self.state_values = np.array([[1.0, 2.0, 3.0, 4.0], # x
[5.0, 6.0, 7.0, 8.0]]) # v
# This must be ordered as the known first then the unknown.
self.specified_values = np.array([[9.0, 10.0, 11.0, 12.0], # f
[13.0, 14.0, 15.0, 16.0]]) # k
# This must be ordered as the known first then the unknown.
self.constant_values = np.array([1.0, # m
2.0]) # c
self.interval_value = 0.01
# The free optimization variable array.
self.free = np.array([1.0, 2.0, 3.0, 4.0, # x
5.0, 6.0, 7.0, 8.0, # v
13.0, 14.0, 15.0, 16.0, # k
2.0]) # c

self.eom = sym.Matrix([x.diff() - v,
m * v.diff() + c * v + k * x - f])

par_map = OrderedDict(zip([m], [1.0]))
traj_map = OrderedDict(zip([f], [self.specified_values[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=traj_map)

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

def test_sort_parameters(self):

# TODO : Added checks for the other cases.

self.collocator._sort_parameters()

m, c = self.constant_symbols

assert self.collocator.known_parameters == (m,)
assert self.collocator.num_known_parameters == 1

assert self.collocator.unknown_parameters == (c,)
assert self.collocator.num_unknown_parameters == 1

def test_sort_trajectories(self):

# TODO : Added checks for the other cases.

self.collocator._sort_trajectories()

f, k = self.specified_symbols

assert self.collocator.known_input_trajectories == (f,)
assert self.collocator.num_known_input_trajectories == 1

assert self.collocator.unknown_input_trajectories == (k,)
assert self.collocator.num_unknown_input_trajectories == 1

def test_discrete_symbols(self):

self.collocator._discrete_symbols()

xi, vi, xp, vp, fi, ki = self.discrete_symbols

assert self.collocator.current_discrete_state_symbols == (xi, vi)
assert self.collocator.previous_discrete_state_symbols == (xp, vp)
# The following should be in order of known and unknown.
assert self.collocator.current_discrete_specified_symbols == (fi, ki)

def test_discretize_eom(self):

m, c = self.constant_symbols
xi, vi, xp, vp, fi, ki = self.discrete_symbols
h = self.collocator.time_interval_symbol

expected = sym.Matrix([(xi - xp) / h - vi,
m * (vi - vp) / h + c * vi + ki * xi - fi])

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(self):

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,
constant_values,
self.interval_value)

m, c = 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, ki = self.specified_values[:, i]

expected_dynamic[i - 1] = m * (vi - vp) / h + c * vi + ki * 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_jacobian_indices(self):

row_idxs, col_idxs = self.collocator.jacobian_indices()

expected_row_idxs = np.array([0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 1,
1, 1, 1, 1, 1, 4, 4, 4, 4, 4, 4, 2, 2,
2, 2, 2, 2, 5, 5, 5, 5, 5, 5])

expected_col_idxs = np.array([1, 5, 0, 4, 9, 12, 1, 5, 0, 4, 9, 12,
2, 6, 1, 5, 10, 12, 2, 6, 1, 5, 10,
12, 3, 7, 2, 6, 11, 12, 3, 7, 2, 6,
11, 12])

np.testing.assert_allclose(row_idxs, expected_row_idxs)
np.testing.assert_allclose(col_idxs, expected_col_idxs)

def test_gen_multi_arg_con_jac_func(self):

self.collocator._gen_multi_arg_con_jac_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.

jac_vals = self.collocator._multi_arg_con_jac_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

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, v = self.state_values
m, c = self.constant_values
f, k = self.specified_values
h = self.interval_value

expected_jacobian = np.array(
# x0, x1, x2, x3, v0, v1, v2, v3, k0, k1, k2, k3, c con
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0], # 1
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0], # 2
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0], # 3
[ 0, k[1], 0, 0, -m / h, c + m / h, 0, 0, 0, x[1], 0, 0, v[1]], # 1
[ 0, 0, k[2], 0, 0, -m / h, c + m / h, 0, 0, 0, x[2], 0, v[2]], # 2
[ 0, 0, 0, k[3], 0, 0, -m /h, c + m / h, 0, 0, 0, x[3], v[3]]], # 3
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)

def test_gen_multi_arg_con_jac_func_midpoint(self):

self.collocator.integration_method = 'midpoint'
self.collocator._gen_multi_arg_con_jac_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.

jac_vals = self.collocator._multi_arg_con_jac_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

row_idxs, col_idxs = self.collocator.jacobian_indices()

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

x, v = self.state_values
m, c = self.constant_values
f, k = self.specified_values
h = self.interval_value

part1 = np.array(
# x0, x1, x2, x3
[[ -1.0 / h, 1.0 / h, 0, 0],
[ 0, -1.0 / h, 1.0 / h, 0],
[ 0, 0, -1.0 / h, 1.0 / h],
[ (k[0] + k[1]) / 4.0, (k[0] + k[1]) / 4.0, 0, 0],
[ 0, (k[1] + k[2]) / 4.0, (k[1] + k[2]) / 4.0, 0],
[ 0, 0, (k[2] + k[3]) / 4.0, (k[2] + k[3]) / 4.0]],
dtype=float)

part2 = np.array(
# v0, v1, v2, v3 i
[[ -1.0 / 2.0, -1.0 / 2.0, 0, 0], # 0
[ 0, -1.0 / 2.0, -1.0 / 2.0, 0], # 1
[ 0, 0, -1.0 / 2.0, -1.0 / 2.0], # 2
[ -m / h + c / 2, m / h + c / 2, 0, 0], # 0
[ 0, -m / h + c / 2, m / h + c / 2, 0], # 1
[ 0, 0, -m / h + c / 2, m / h + c / 2]], # 2
dtype=float)

part3 = np.array(
# k0, k1, k2, k3, c i
[[ 0, 0, 0, 0, 0], # 0
[ 0, 0, 0, 0, 0], # 1
[ 0, 0, 0, 0, 0], # 2
[ (x[1] + x[0]) / 4.0, (x[1] + x[0]) / 4.0, 0, 0, (v[1] + v[0]) / 2], # 0
[ 0, (x[2] + x[1]) / 4.0, (x[2] + x[1]) / 4.0, 0, (v[2] + v[1]) / 2], # 1
[ 0, 0, (x[3] + x[2]) / 4.0, (x[3] + x[2]) / 4.0, (v[3] + v[2]) / 2]], # 2
dtype=float)

expected_jacobian = np.hstack((part1, part2, part3))

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)

def test_generate_constraint_function(self):

constrain = self.collocator.generate_constraint_function()

result = constrain(self.free)

m, c = 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, ki = self.specified_values[:, i]

expected_dynamic[i - 1] = m * (vi - vp) / h + c * vi + ki * 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_generate_jacobian_function(self):

jacobian = self.collocator.generate_jacobian_function()

jac_vals = jacobian(self.free)

row_idxs, col_idxs = self.collocator.jacobian_indices()

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

x, v = self.state_values
m, c = self.constant_values
f, k = self.specified_values
h = self.interval_value

expected_jacobian = np.array(
# x0, x1, x2, x3, v0, v1, v2, v3, k0, k1, k2, k3, c con
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0], # 1
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0], # 2
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0], # 3
[ 0, k[1], 0, 0, -m / h, c + m / h, 0, 0, 0, x[1], 0, 0, v[1]], # 1
[ 0, 0, k[2], 0, 0, -m / h, c + m / h, 0, 0, 0, x[2], 0, v[2]], # 2
[ 0, 0, 0, k[3], 0, 0, -m /h, c + m / h, 0, 0, 0, x[3], v[3]]], # 3
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)


def test_merge_fixed_free_parameters():

m, c, k = sym.symbols('m, c, k')
all_syms = m, c, k
known = {m: 1.0, c: 2.0}
unknown = np.array([3.0])

merged = ConstraintCollocator._merge_fixed_free(all_syms, known,
unknown, 'par')

expected = np.array([1.0, 2.0, 3.0])

np.testing.assert_allclose(merged, expected)

a, b, c, d = sym.symbols('a, b, c, d')
all_syms = a, b, c, d
known = {a: 1.0, b: 2.0}
unknown = np.array([3.0, 4.0])

merged = ConstraintCollocator._merge_fixed_free(all_syms, known,
unknown, 'par')

expected = np.array([1.0, 2.0, 3.0, 4.0])

np.testing.assert_allclose(merged, expected)


def test_merge_fixed_free_trajectories():

t = sym.symbols('t')
f, k = sym.symbols('f, k', cls=sym.Function)
f = f(t)
k = k(t)
all_syms = f, k
known = {f: np.array([1.0, 2.0])}
unknown = np.array([3.0, 4.0])

merged = ConstraintCollocator._merge_fixed_free(all_syms, known,
unknown, 'traj')

expected = np.array([[1.0, 2.0],
[3.0, 4.0]])

np.testing.assert_allclose(merged, expected)

t = sym.symbols('t')
all_syms = [g(t) for g in sym.symbols('a, b, c, d', cls=sym.Function)]
a, b, c, d = all_syms
known = {a: np.array([1.0, 2.0]),
b: np.array([3.0, 4.0])}
unknown = np.array([[5.0, 6.0],
[7.0, 8.0]])

merged = ConstraintCollocator._merge_fixed_free(all_syms, known,
unknown, 'traj')

expected = np.array([[1.0, 2.0],
[3.0, 4.0],
[5.0, 6.0],
[7.0, 8.0]])

np.testing.assert_allclose(merged, expected)


class TestConstraintCollocatorInstanceConstraints():

def setup(self):

I, m, g, d, t = sym.symbols('I, m, g, d, t')
theta, omega, T = [f(t) for f in sym.symbols('theta, omega, T',
cls=sym.Function)]

self.state_symbols = (theta, omega)
self.constant_symbols = (I, m, g, d)
self.specified_symbols = (T,)
self.discrete_symbols = \
sym.symbols('thetai, omegai, thetap, omegap, Ti', real=True)

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([9.0, 10.0, 11.0, 12.0])
self.constant_values = np.array([1.0, 1.0, 9.81, 1.0])
self.interval_value = 0.01
self.time = np.array([0.0, 0.01, 0.02, 0.03])
self.free = np.array([1.0, 2.0, 3.0, 4.0, # theta
5.0, 6.0, 7.0, 8.0, # omega
9.0, 10.0, 11.0, 12.0, # T
1.0, # I
1.0, # m
9.81, # g
1.0]) # d

self.eom = sym.Matrix([theta.diff() - omega,
I * omega.diff() + m * g * d * sym.sin(theta) - T])

par_map = OrderedDict(zip(self.constant_symbols,
self.constant_values))

# Additional node equality contraints.
theta, omega = sym.symbols('theta, omega', cls=sym.Function)
instance_constraints = (2.0 * theta(0.0),
3.0 * theta(0.03) - sym.pi,
4.0 * omega(0.0),
5.0 * omega(0.03))

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,
instance_constraints=instance_constraints)

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

def test_sort_parameters(self):

self.collocator._sort_parameters()

assert self.collocator.known_parameters == self.constant_symbols
assert self.collocator.num_known_parameters == 4

assert self.collocator.unknown_parameters == tuple()
assert self.collocator.num_unknown_parameters == 0

def test_sort_trajectories(self):

self.collocator._sort_trajectories()

assert self.collocator.known_input_trajectories == tuple()
assert self.collocator.num_known_input_trajectories == 0

assert self.collocator.unknown_input_trajectories == self.specified_symbols
assert self.collocator.num_unknown_input_trajectories == 1

def test_discrete_symbols(self):

self.collocator._discrete_symbols()

thetai, omegai, thetap, omegap, Ti = self.discrete_symbols

assert self.collocator.current_discrete_state_symbols == (thetai,
omegai)
assert self.collocator.previous_discrete_state_symbols == (thetap,
omegap)
assert self.collocator.current_discrete_specified_symbols == (Ti,)

def test_discretize_eom(self):

I, m, g, d = self.constant_symbols
thetai, omegai, thetap, omegap, Ti = self.discrete_symbols
h = self.collocator.time_interval_symbol

expected = sym.Matrix([(thetai - thetap) / h - omegai,
I * (omegai - omegap) / h + m * g * d * sym.sin(thetai) - Ti])

self.collocator._discretize_eom()

zero = sym.simplify(self.collocator.discrete_eom - expected)

assert zero == sym.Matrix([0, 0])

def test_identify_function_in_instance_constraints(self):

self.collocator._identify_functions_in_instance_constraints()

theta, omega = sym.symbols('theta, omega', cls=sym.Function)

expected = set((theta(0.0), theta(0.03), omega(0.0), omega(0.03)))

assert self.collocator.instance_constraint_function_atoms == expected

def test_find_free_index(self):

theta, omega = sym.symbols('theta, omega', cls=sym.Function)

self.collocator._identify_functions_in_instance_constraints()
self.collocator._find_closest_free_index()

expected = {theta(0.0): 0,
theta(0.03): 3,
omega(0.0): 4,
omega(0.03): 7}

assert self.collocator.instance_constraints_free_index_map == expected

def test_lambdify_instance_constraints(self):

f = self.collocator._instance_constraints_func()

extra_constraints = f(self.free)

expected = np.array([2.0, 12.0 - np.pi, 20.0, 40.0])

np.testing.assert_allclose(extra_constraints, expected)

def test_instance_constraints_jacobian_indices(self):

rows, cols = self.collocator._instance_constraints_jacobian_indices()

np.testing.assert_allclose(rows, np.array([6, 7, 8, 9]))
np.testing.assert_allclose(cols, np.array([0, 3, 4, 7]))

def test_instance_constraints_jacobian_values(self):

f = self.collocator._instance_constraints_jacobian_values_func()

vals = f(self.free)

np.testing.assert_allclose(vals, np.array([2.0, 3.0, 4.0, 5.0]))

def test_gen_multi_arg_con_func(self):

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])

result = self.collocator._multi_arg_con_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

I, m, g, d = self.constant_values
h = self.interval_value

expected_dynamic = np.zeros(3)
expected_kinematic = np.zeros(3)

for i in [1, 2, 3]:

thetai, omegai = self.state_values[:, i]
thetap, omegap = self.state_values[:, i - 1]
Ti = self.specified_values[i]

expected_kinematic[i - 1] = (thetai - thetap) / h - omegai
expected_dynamic[i - 1] = (I * (omegai - omegap) / h + m * g * d
* sym.sin(thetai) - Ti)

expected = np.hstack((expected_kinematic, expected_dynamic))

np.testing.assert_allclose(result, expected)

def test_jacobian_indices(self):

row_idxs, col_idxs = self.collocator.jacobian_indices()

expected_row_idxs = np.array([0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 1, 1, 1,
1, 1, 4, 4, 4, 4, 4, 2, 2, 2, 2, 2, 5,
5, 5, 5, 5, 6, 7, 8, 9])

expected_col_idxs = np.array([1, 5, 0, 4, 9, 1, 5, 0, 4, 9, 2, 6, 1,
5, 10, 2, 6, 1, 5, 10, 3, 7, 2, 6, 11,
3, 7, 2, 6, 11, 0, 3, 4, 7])

np.testing.assert_allclose(row_idxs, expected_row_idxs)
np.testing.assert_allclose(col_idxs, expected_col_idxs)

def test_gen_multi_arg_con_jac_func(self):

self.collocator._gen_multi_arg_con_jac_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])

jac_vals = self.collocator._multi_arg_con_jac_func(self.state_values,
self.specified_values,
constant_values,
self.interval_value)

row_idxs, col_idxs = self.collocator.jacobian_indices()
row_idxs = row_idxs[:-4]
col_idxs = col_idxs[:-4]

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

# thetai, omegai, thetap, omegap, Ti
# theta : [ 1/h, -1, -1/h, 0, 0],
# omega : [g*m*cos(thetai), I/h, 0, -I/h, -1]])

theta = self.state_values[0]
I, m, g, d = self.constant_values
h = self.interval_value

expected_jacobian = np.array(
# theta0, theta1, theta2, theta3, omega0, omega1, omega2, omega3, T0, T1, T2, T3
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0], # 1
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0], # 2
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0], # 3
[ 0, d * g * m * np.cos(theta[1]), 0, 0, -I / h, I / h, 0, 0, 0, -1, 0, 0], # 1
[ 0, 0, d * g * m * np.cos(theta[2]), 0, 0, -I / h, I / h, 0, 0, 0, -1, 0], # 2
[ 0, 0, 0, d * g * m * np.cos(theta[3]), 0, 0, -I / h, I / h, 0, 0, 0, -1]], # 3
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)

def test_generate_constraint_function(self):

constrain = self.collocator.generate_constraint_function()

result = constrain(self.free)

expected_dynamic = np.zeros(3)
expected_kinematic = np.zeros(3)

I, m, g, d = self.constant_values
h = self.interval_value

expected_dynamic = np.zeros(3)
expected_kinematic = np.zeros(3)

for i in [1, 2, 3]:

thetai, omegai = self.state_values[:, i]
thetap, omegap = self.state_values[:, i - 1]
Ti = self.specified_values[i]

expected_kinematic[i - 1] = (thetai - thetap) / h - omegai
expected_dynamic[i - 1] = (I * (omegai - omegap) / h + m * g * d
* sym.sin(thetai) - Ti)
theta_values = self.state_values[0]
omega_values = self.state_values[1]

expected_node_constraints = np.array([2.0 * theta_values[0],
3.0 * theta_values[3] - np.pi,
4.0 * omega_values[0],
5.0 * omega_values[3]])

expected = np.hstack((expected_kinematic, expected_dynamic,
expected_node_constraints))

np.testing.assert_allclose(result, expected)

def test_generate_jacobian_function(self):

jacobian = self.collocator.generate_jacobian_function()

jac_vals = jacobian(self.free)

row_idxs, col_idxs = self.collocator.jacobian_indices()

jacobian_matrix = sparse.coo_matrix((jac_vals, (row_idxs, col_idxs)))

theta = self.state_values[0]
I, m, g, d = self.constant_values
h = self.interval_value

expected_jacobian = np.array(
# theta0, theta1, theta2, theta3, omega0, omega1, omega2, omega3, T0, T1, T2, T3
[[-1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0], # 1
[ 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0, 0], # 2
[ 0, 0, -1 / h, 1 / h, 0, 0, 0, -1, 0, 0, 0, 0], # 3
[ 0, d * g * m * np.cos(theta[1]), 0, 0, -I / h, I / h, 0, 0, 0, -1, 0, 0], # 1
[ 0, 0, d * g * m * np.cos(theta[2]), 0, 0, -I / h, I / h, 0, 0, 0, -1, 0], # 2
[ 0, 0, 0, d * g * m * np.cos(theta[3]), 0, 0, -I / h, I / h, 0, 0, 0, -1], # 3
[ 2.0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 3.0, 0, 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 4.0, 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0, 5.0, 0, 0, 0, 0]],
dtype=float)

np.testing.assert_allclose(jacobian_matrix.todense(), expected_jacobian)