| 573 | |
| 574 | |
| 575 | class ConstraintSolver(object): |
| 576 | |
| 577 | opti: ca.Opti |
| 578 | variables: List[Tuple[ca.MX, ca.MX]] |
| 579 | starting_points: List[Tuple[ca.MX, ca.MX]] |
| 580 | constraints: List[Tuple[Tuple[int, ...], Constraint]] |
| 581 | locked: List[int] |
| 582 | ne: int |
| 583 | nc: int |
| 584 | scale: float |
| 585 | |
| 586 | def __init__( |
| 587 | self, |
| 588 | entities: List[Location], |
| 589 | constraints: List[Tuple[Tuple[int, ...], Constraint]], |
| 590 | locked: List[int] = [], |
| 591 | scale: float = 1, |
| 592 | ): |
| 593 | |
| 594 | self.scale = scale |
| 595 | self.opti = opti = ca.Opti() |
| 596 | self.variables = [ |
| 597 | (scale * opti.variable(NDOF_V), opti.variable(NDOF_Q)) |
| 598 | if i not in locked |
| 599 | else (opti.parameter(NDOF_V), opti.parameter(NDOF_Q)) |
| 600 | for i, _ in enumerate(entities) |
| 601 | ] |
| 602 | self.start_points = [ |
| 603 | (opti.parameter(NDOF_V), opti.parameter(NDOF_Q)) for _ in entities |
| 604 | ] |
| 605 | |
| 606 | # initialize, add the unit quaternion constraints and handle locked |
| 607 | for i, ((T, R), (T0, R0), loc) in enumerate( |
| 608 | zip(self.variables, self.start_points, entities) |
| 609 | ): |
| 610 | T0val, R0val = self._locToDOF6(loc) |
| 611 | |
| 612 | opti.set_value(T0, T0val) |
| 613 | opti.set_value(R0, R0val) |
| 614 | |
| 615 | if i in locked: |
| 616 | opti.set_value(T, (0, 0, 0)) |
| 617 | opti.set_value(R, (0, 0, 0)) |
| 618 | else: |
| 619 | opti.set_initial(T, (0.0, 0.0, 0.0)) |
| 620 | opti.set_initial(R, (1e-2, 1e-2, 1e-2)) |
| 621 | |
| 622 | self.constraints = constraints |
| 623 | |
| 624 | # additional book-keeping |
| 625 | self.ne = len(entities) |
| 626 | self.locked = locked |
| 627 | self.nc = len(self.constraints) |
| 628 | |
| 629 | @staticmethod |
| 630 | def _locToDOF6(loc: Location) -> DOF6: |
| 631 | |
| 632 | Tr = loc.wrapped.Transformation() |