From f4150b62d34e6c5bd9d1faca6157e112465b1eb5 Mon Sep 17 00:00:00 2001
From: adam-urbanczyk <adam.jan.urbanczyk@gmail.com>
Date: Thu, 24 Sep 2020 22:27:47 +0200
Subject: [PATCH] Major rework of the solver

* numerical jacobian by hand taking into account the sparsity pattern
* better code structure
* scaling of the dir constraint
---
 cadquery/occ_impl/solver.py | 93 +++++++++++++++++++++++++++++++++----
 1 file changed, 85 insertions(+), 8 deletions(-)

diff --git a/cadquery/occ_impl/solver.py b/cadquery/occ_impl/solver.py
index 573f8ee6..d50004db 100644
--- a/cadquery/occ_impl/solver.py
+++ b/cadquery/occ_impl/solver.py
@@ -1,7 +1,7 @@
 from typing import Tuple, Union, Any, Callable, List, Optional
 from nptyping import NDArray as Array
 
-from numpy import array
+from numpy import array, eye, zeros
 from scipy.optimize import minimize
 
 from OCP.gp import gp_Vec, gp_Dir, gp_Pnt, gp_Trsf, gp_Quaternion
@@ -13,6 +13,8 @@ ConstraintMarker = Union[gp_Dir, gp_Pnt]
 Constraint = Tuple[Tuple[ConstraintMarker, ...], Tuple[Optional[ConstraintMarker], ...]]
 
 NDOF = 6
+DIR_SCALING = 1e3
+DIFF_EPS = 1e-8
 
 
 class ConstraintSolver(object):
@@ -78,7 +80,22 @@ class ConstraintSolver(object):
 
         return rv
 
-    def _cost(self) -> Callable[[Array[(Any,), float]], float]:
+    def _cost(
+        self,
+    ) -> Tuple[
+        Callable[[Array[(Any,), float]], float],
+        Callable[[Array[(Any,), float]], Array[(Any,), float]],
+    ]:
+        def pt_cost(m1: gp_Pnt, m2: gp_Pnt, t1: gp_Trsf, t2: gp_Trsf) -> float:
+
+            return (m1.Transformed(t1).XYZ() - m2.Transformed(t2).XYZ()).SquareModulus()
+
+        def dir_cost(
+            m1: gp_Dir, m2: gp_Dir, t1: gp_Trsf, t2: gp_Trsf, val: float = -1
+        ) -> float:
+
+            return DIR_SCALING * (val - m1.Transformed(t1).Dot(m2.Transformed(t2))) ** 2
+
         def f(x):
 
             constraints = self.constraints
@@ -96,25 +113,85 @@ class ConstraintSolver(object):
 
                 for m1, m2 in zip(ms1, ms2):
                     if isinstance(m1, gp_Pnt):
-                        rv += (
-                            m1.Transformed(t1).XYZ() - m2.Transformed(t2).XYZ()
-                        ).Modulus() ** 2
+                        rv += pt_cost(m1, m2, t1, t2)
                     elif isinstance(m1, gp_Dir):
-                        rv += (-1 - m1.Transformed(t1).Dot(m2.Transformed(t2))) ** 2
+                        rv += dir_cost(m1, m2, t1, t2)
                     else:
                         raise NotImplementedError(f"{m1,m2}")
 
             return rv
 
-        return f
+        def jac(x):
+
+            constraints = self.constraints
+            ne = self.ne
+
+            delta = DIFF_EPS * eye(NDOF)
+
+            rv = zeros(NDOF * ne)
+
+            transforms = [
+                self._build_transform(*x[NDOF * i : NDOF * (i + 1)]) for i in range(ne)
+            ]
+
+            transforms_delta = [
+                self._build_transform(*(x[NDOF * i : NDOF * (i + 1)] + delta[j, :]))
+                for i in range(ne)
+                for j in range(NDOF)
+            ]
+
+            for i, ((k1, k2), (ms1, ms2)) in enumerate(constraints):
+                t1 = transforms[k1] if k1 not in self.locked else gp_Trsf()
+                t2 = transforms[k2] if k2 not in self.locked else gp_Trsf()
+
+                for m1, m2 in zip(ms1, ms2):
+                    if isinstance(m1, gp_Pnt):
+                        tmp = pt_cost(m1, m2, t1, t2)
+
+                        for j in range(NDOF):
+
+                            t1j = transforms_delta[k1 * NDOF + j]
+                            t2j = transforms_delta[k2 * NDOF + j]
+
+                            if k1 not in self.locked:
+                                tmp1 = pt_cost(m1, m2, t1j, t2)
+                                rv[k1 * NDOF + j] += (tmp1 - tmp) / DIFF_EPS
+
+                            if k2 not in self.locked:
+                                tmp2 = pt_cost(m1, m2, t1, t2j)
+                                rv[k2 * NDOF + j] += (tmp2 - tmp) / DIFF_EPS
+
+                    elif isinstance(m1, gp_Dir):
+                        tmp = dir_cost(m1, m2, t1, t2)
+
+                        for j in range(NDOF):
+
+                            t1j = transforms_delta[k1 * NDOF + j]
+                            t2j = transforms_delta[k2 * NDOF + j]
+
+                            if k1 not in self.locked:
+                                tmp1 = dir_cost(m1, m2, t1j, t2)
+                                rv[k1 * NDOF + j] += (tmp1 - tmp) / DIFF_EPS
+
+                            if k2 not in self.locked:
+                                tmp2 = dir_cost(m1, m2, t1, t2j)
+                                rv[k2 * NDOF + j] += (tmp2 - tmp) / DIFF_EPS
+                    else:
+                        raise NotImplementedError(f"{m1,m2}")
+
+            return rv
+
+        return f, jac
 
     def solve(self) -> List[Location]:
 
         x0 = array([el for el in self.entities]).ravel()
+        f, jac = self._cost()
 
         res = minimize(
-            self._cost(),
+            f,
             x0,
+            jac=jac,
             method="L-BFGS-B",
             options=dict(disp=True, ftol=1e-9, maxiter=1000),
         )