ConstraintGeneration.py
1 #!/usr/bin/python3
2 
3 import sympy as sp
4 sp.init_printing()
5 
6 # Constraint class template that will be filled in.
7 template = """class {name:s}Constraint : public ompl::base::Constraint
8 {{
9 public:
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 
24 class 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]), assign_to="out[{:d}]".format(i))
57  ss += "\n"
58  return ss
59 
60  def jacCode(self):
61  ss = ""
62  jac = self.jacobian()
63  for i in range(jac.shape[0]):
64  for j in range(jac.shape[1]):
65  ss += ' ' * 8
66  ss += sp.printing.cxxcode(sp.simplify(jac[i, j]), assign_to="out({:d}, {:d})".format(i, j))
67  ss += "\n"
68  return ss
69 
70  def toCode(self):
71  return template.format(name = self.name_,
72  ambientDim = len(self.variables_),
73  constraintDim = len(self.constraints_),
74  funcCode = self.funcCode(),
75  jacCode = self.jacCode())
76 
77 if __name__ == "__main__":
78  # Sphere constraint
79  s = Constraint("Sphere", 3)
80  s.addConstraint(s.getVars().norm() - 1)
81  print(s.toCode())
82 
83  # Torus constraint
84  t = Constraint("Torus", 3)
85 
86  outer_radius = 3
87  inner_radius = 1
88 
89  c = t.getVars()
90  c[2] = 0
91  torus = (t.getVars() - outer_radius * c / c.norm()).norm() - inner_radius
92  t.addConstraint(torus)
93 
94  print(t.toCode())