-
Notifications
You must be signed in to change notification settings - Fork 3
/
augmented_mp.py
52 lines (38 loc) · 1.82 KB
/
augmented_mp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
from pydrake.all import MathematicalProgram, SolverType
from pydrake.symbolic import sin, cos
import numpy as np
from collections import Iterable
class AugmentedMathematicalProgram(MathematicalProgram):
#generate state and input matrices
##################################
def set_max_iters(max_solver_iterations):
self.SetSolverOption(SolverType.kSnopt, 'MajorIterationsLimit', max_solver_iterations)
def var_stack(self, N, num_states, name):
return self.NewContinuousVariables(N, num_states, name)
def states(self, N, states_initial, states_final=None):
num_boats = states_initial.shape[0]
#intialize states for a boat
def boat_state(b):
S = self.var_stack(N, len(states_initial[b]), 'b_%d-s' % b)
self.add_equal_constraints(S[0], states_initial[b])
if states_final is not None:
self.add_equal_constraints(S[-1], states_final[b])
return S
boats_S = np.stack((boat_state(b) for b in range(num_boats)))
return boats_S
def inputs(self, N, num_inputs, num_boats=1):
boats_U = np.stack(self.var_stack(N, num_inputs, 'b_%d-u' % b) for b in range(num_boats))
return boats_U
#constraint functions
#####################
def add_equal_constraints(self, state, val, linear=False):
slack = 0.0001
def add_equal_constraint(x, y):
self.AddConstraint(x == y)
def add_equal_linear_constraint(x, y):
self.AddLinearConstraint(y-slack == x)
if isinstance(state, Iterable):
for x,y in zip(state, val):
add_equal_linear_constraint(x, y) if linear else add_equal_constraint(x, y)
else:
add_equal_constraint(state, val)