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_ = name
28  # Generate an array of variables to use.
29  self.variables_variables_ = [sp.Symbol("x[{:d}]".format(i), real=True) for i in range(n)]
30  self.constraints_constraints_ = []
31 
32  def __getitem__(self, index):
33  """Return the index^th variable."""
34  return self.variables_variables_[index]
35 
36  def getVars(self):
37  """Create a variable vector."""
38  return sp.Matrix(self.variables_variables_)
39 
40  def getConstraints(self):
41  """Create a constraint function vector."""
42  return sp.Matrix(self.constraints_constraints_)
43 
44  def addConstraint(self, f):
45  """Add some symbolic function of variables to the list of constraints."""
46  self.constraints_constraints_.append(f)
47 
48  def jacobian(self):
49  """Compute the Jacobian of the current list of constraints."""
50  return self.getConstraintsgetConstraints().jacobian(self.variables_variables_)
51 
52  def funcCode(self):
53  ss = ""
54  for i in range(len(self.constraints_constraints_)):
55  ss += ' ' * 8
56  ss += sp.printing.cxxcode(sp.simplify(self.constraints_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.jacobianjacobian()
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_name_,
74  ambientDim=len(self.variables_variables_),
75  constraintDim=len(self.constraints_constraints_),
76  funcCode=self.funcCodefuncCode(),
77  jacCode=self.jacCodejacCode())
78 
79 if __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())