Loading...
Searching...
No Matches
ConstraintGeneration.py
1#!/usr/bin/python3
2
3import sympy as sp
4sp.init_printing()
5
6# Constraint class template that will be filled in.
7template = """class {name:s}Constraint : public ompl::base::Constraint
8{{
9public:
10 {name:s}Constraint() : ompl::base::Constraint({ambientDim:d}, {constraintDim:d})
11 {{
12 }}
13
14 void function(const Eigen::Ref<const Eigen::VectorXd> &x, Eigen::Ref<Eigen::VectorXd> out) const override
15 {{
16{funcCode:s} }}
17
18 void jacobian(const Eigen::Ref<const Eigen::VectorXd> &x, Eigen::Ref<Eigen::MatrixXd> out) const override
19 {{
20{jacCode:s} }}
21}};
22"""
23
24class Constraint:
25
26 def __init__(self, name, n):
27 self.name_ = name
28 # Generate an array of variables to use.
29 self.variables_ = [sp.Symbol("x[{:d}]".format(i), real=True) for i in range(n)]
30 self.constraints_ = []
31
32 def __getitem__(self, index):
33 """Return the index^th variable."""
34 return self.variables_[index]
35
36 def getVars(self):
37 """Create a variable vector."""
38 return sp.Matrix(self.variables_)
39
40 def getConstraints(self):
41 """Create a constraint function vector."""
42 return sp.Matrix(self.constraints_)
43
44 def addConstraint(self, f):
45 """Add some symbolic function of variables to the list of constraints."""
46 self.constraints_.append(f)
47
48 def jacobian(self):
49 """Compute the Jacobian of the current list of constraints."""
50 return self.getConstraints().jacobian(self.variables_)
51
52 def funcCode(self):
53 ss = ""
54 for i in range(len(self.constraints_)):
55 ss += ' ' * 8
56 ss += sp.printing.cxxcode(sp.simplify(self.constraints_[i]),
57 assign_to="out[{:d}]".format(i))
58 ss += "\n"
59 return ss
60
61 def jacCode(self):
62 ss = ""
63 jac = self.jacobian()
64 for i in range(jac.shape[0]):
65 for j in range(jac.shape[1]):
66 ss += ' ' * 8
67 ss += sp.printing.cxxcode(sp.simplify(jac[i, j]),
68 assign_to="out({:d}, {:d})".format(i, j))
69 ss += "\n"
70 return ss
71
72 def toCode(self):
73 return template.format(name=self.name_,
74 ambientDim=len(self.variables_),
75 constraintDim=len(self.constraints_),
76 funcCode=self.funcCode(),
77 jacCode=self.jacCode())
78
79if __name__ == "__main__":
80 # Sphere constraint
81 s = Constraint("Sphere", 3)
82 s.addConstraint(s.getVars().norm() - 1)
83 print(s.toCode())
84
85 # Torus constraint
86 t = Constraint("Torus", 3)
87
88 outer_radius = 3
89 inner_radius = 1
90
91 c = t.getVars()
92 c[2] = 0
93 torus = (t.getVars() - outer_radius * c / c.norm()).norm() - inner_radius
94 t.addConstraint(torus)
95
96 print(t.toCode())