Merge pull request #398 from hiddenSymmetries/ml/linkingnum_speedup

Speed up LinkingNumber

src/simsopt/geo/curveobjectives.py

@@ -479,48 +479,45 @@ class MinimumDistance(CurveCurveDistance):
 
 class LinkingNumber(Optimizable):
 
-    def __init__(self, curves):
+    def __init__(self, curves, downsample=1):
         Optimizable.__init__(self, depends_on=curves)
         self.curves = curves
+        for curve in curves:
+            assert np.mod(len(curve.quadpoints), downsample) == 0, f"Downsample {downsample} does not divide the number of quadpoints {len(curve.quadpoints)}."
+
+        self.downsample = downsample
+        self.dphis = np.array([(c.quadpoints[1] - c.quadpoints[0]) * downsample for c in self.curves])
+
         r"""
-        Compute the Linking number of a set of curves (whether the curves 
-        are interlocked or not).
+        Compute the Gauss linking number of a set of curves, i.e. whether the curves
+        are interlocked or not.
 
-        The value is 1 if the are interlocked, 0 if not.
+        The value is an integer, >= 1 if the curves are interlocked, 0 if not. For each pair
+        of curves, the contribution to the linking number is
         
         .. math::
-            Link(c1,c2) = \frac{1}{4\pi} \oint_{c1}\oint_{c2}\frac{\textbf{R1} - \textbf{R2}}{|\textbf{R1}-\textbf{R2}|^3} (d\textbf{R1} \times d\textbf{R2})
+            Link(c_1, c_2) = \frac{1}{4\pi} \left| \oint_{c_1}\oint_{c_2}\frac{\textbf{r}_1 - \textbf{r}_2}{|\textbf{r}_1 - \textbf{r}_2|^3} (d\textbf{r}_1 \times d\textbf{r}_2) \right|
             
-        where :math:`c1` is the first curve and :math:`c2` is the second curve, 
-        :math:`\textbf{R1}` is the radius vector of the first curve, and 
-        :math:`\textbf{R2}` is the radius vector of the second curve
+        where :math:`c_1` is the first curve, :math:`c_2` is the second curve,
+        :math:`\textbf{r}_1` is the position vector along the first curve, and
+        :math:`\textbf{r}_2` is the position vector along the second curve.
 
         Args:
-            curves: the set of curves on which the linking number should be computed.
-        
+            curves: the set of curves for which the linking number should be computed.
+            downsample: integer factor by which to downsample the quadrature
+                points when computing the linking number. Setting this parameter to
+                a value larger than 1 will speed up the calculation, which may
+                be useful if the set of coils is large, though it may introduce
+                inaccuracy if ``downsample`` is set too large.
         """
 
     def J(self):
-        ncoils = len(self.curves)
-        linkNum = np.zeros([ncoils + 1, ncoils + 1])
-        i = 0
-        for c1 in self.curves[:(ncoils + 1)]:
-            j = 0
-            i = i + 1
-            for c2 in self.curves[:(ncoils + 1)]:
-                j = j + 1
-                if i < j:
-                    R1 = c1.gamma()
-                    R2 = c2.gamma()
-                    dS = c1.quadpoints[1] - c1.quadpoints[0]
-                    dT = c2.quadpoints[1] - c1.quadpoints[0]
-                    dR1 = c1.gammadash()
-                    dR2 = c2.gammadash()
-
-                    integrals = sopp.linkNumber(R1, R2, dR1, dR2) * dS * dT
-                    linkNum[i-1][j-1] = 1/(4*np.pi) * (integrals)
-        linkNumSum = sum(sum(abs(linkNum)))
-        return round(linkNumSum)
+        return sopp.compute_linking_number(
+            [c.gamma() for c in self.curves],
+            [c.gammadash() for c in self.curves],
+            self.dphis,
+            self.downsample,
+        )
 
     @derivative_dec
     def dJ(self):

src/simsoptpp/python_distance.cpp

@@ -173,29 +173,45 @@ void init_distance(py::module_ &m){
 
     m.def("get_pointclouds_closer_than_threshold_within_collection", &get_close_candidates_pdist, "In a list of point clouds, get all pairings that are closer than threshold to each other.", py::arg("pointClouds"), py::arg("threshold"), py::arg("num_base_curves"));
     m.def("get_pointclouds_closer_than_threshold_between_two_collections", &get_close_candidates_cdist, "Between two lists of pointclouds, get all pairings that are closer than threshold to each other.", py::arg("pointCloudsA"), py::arg("pointCloudsB"), py::arg("threshold"));
-    m.def("linkNumber", [](const PyArray& curve1, const PyArray& curve2, const PyArray& curve1dash, const PyArray& curve2dash) {
-        int linknphi1 = curve1.shape(0);
-        int linknphi2 = curve2.shape(0);
-        int linkntheta1 = curve1.shape(1);
-        int linkntheta2 = curve2.shape(1);
-        const double *curve1_ptr = curve1.data();
-        const double *curve2_ptr = curve2.data();
-        const double *curve1dash_ptr = curve1dash.data();
-        const double *curve2dash_ptr = curve2dash.data();
-        double difference[3] = { 0 };
-        double total = 0;
-        for(int i=0; i < linknphi1; i++){
-            for(int j=0; j < linknphi2; j++){
-                difference[0] = (curve1_ptr[3*i+0] - curve2_ptr[3*j+0]);
-                difference[1] = (curve1_ptr[3*i+1] - curve2_ptr[3*j+1]);
-                difference[2] = (curve1_ptr[3*i+2] - curve2_ptr[3*j+2]);
-                double denom = pow(std::sqrt(difference[0]*difference[0] + difference[1]*difference[1] + difference[2]*difference[2]), 3);
-                double det = curve1dash_ptr[3*i+0]*(curve2dash_ptr[3*j+1]*difference[2] - curve2dash_ptr[3*j+2]*difference[1]) - curve1dash_ptr[3*i+1]*(curve2dash_ptr[3*j+0]*difference[2] - curve2dash_ptr[3*j+2]*difference[0]) + curve1dash_ptr[3*i+2]*(curve2dash_ptr[3*j+0]*difference[1] - curve2dash_ptr[3*j+1]*difference[0]);
-                double r = det/denom;
-                total += r;
+    m.def("compute_linking_number", [](const vector<PyArray>& gammas, const vector<PyArray>& gammadashs, const PyArray& dphis, const double downsample) {
+        int ncurves = gammas.size();
+        // assert(dphis.size() == ncurves);
+        // assert(gammadashs.size() == ncurves);
+
+        int linking_number = 0;
+        #pragma omp parallel for reduction(+:linking_number)
+        for (int p = 1; p < ncurves; p++) {
+            int linknphi1 = gammas[p].shape(0);
+            const double *curve1_ptr = gammas[p].data();
+            const double *curve1dash_ptr = gammadashs[p].data();
+            // assert(gammas[p].size() == gammadashs[p].size());
+            for (int q = 0; q < p; q++) {
+                // assert(gammas[q].size() == gammadashs[q].size());
+                int linknphi2 = gammas[q].shape(0);
+                const double *curve2_ptr = gammas[q].data();
+                const double *curve2dash_ptr = gammadashs[q].data();
+                double difference[3] = { 0 };
+                double total = 0;
+                double dr, det;
+                for (int i=0; i < linknphi1; i += downsample){
+                    for (int j=0; j < linknphi2; j += downsample){
+                        difference[0] = (curve1_ptr[3*i+0] - curve2_ptr[3*j+0]);
+                        difference[1] = (curve1_ptr[3*i+1] - curve2_ptr[3*j+1]);
+                        difference[2] = (curve1_ptr[3*i+2] - curve2_ptr[3*j+2]);
+                        dr = std::sqrt(difference[0]*difference[0] + difference[1]*difference[1] + difference[2]*difference[2]);
+                        det = curve1dash_ptr[3*i+0]*(curve2dash_ptr[3*j+1]*difference[2] 
+                            - curve2dash_ptr[3*j+2]*difference[1]) 
+                            - curve1dash_ptr[3*i+1]*(curve2dash_ptr[3*j+0]*difference[2] 
+                            - curve2dash_ptr[3*j+2]*difference[0]) 
+                            + curve1dash_ptr[3*i+2]*(curve2dash_ptr[3*j+0]*difference[1] 
+                            - curve2dash_ptr[3*j+1]*difference[0]);
+                        total += det / (dr * dr * dr);
+                    }
+                }
+                linking_number += std::round(std::abs(total * dphis[p] * dphis[q]) / (4 * M_PI));
             }
         }
-        return total;
+        return linking_number;
     });
 
 }

tests/geo/test_curve_objectives.py

@@ -358,34 +358,33 @@ def test_curve_surface_distance(self):
             err = err_new
 
     def test_linking_number(self):
-
-        curves1 = create_equally_spaced_curves(2, 1, stellsym=True, R0=1, R1=0.5, order=5, numquadpoints=128)
-        curve1 = CurveXYZFourier(200, 3)
-        coeffs = curve1.dofs_matrix
-        coeffs[1][0] = 1.
-        coeffs[1][1] = 0.5
-        coeffs[2][2] = 0.5
-        curve1.set_dofs(np.concatenate(coeffs))
-
-        curve2 = CurveXYZFourier(150, 3)
-        coeffs = curve2.dofs_matrix
-        coeffs[1][0] = 0.5
-        coeffs[1][1] = 0.5
-        coeffs[0][0] = 0.1
-        coeffs[0][1] = 0.5
-        coeffs[0][2] = 0.5
-        curve2.set_dofs(np.concatenate(coeffs))
-        curves2 = [curve1, curve2]
-        Object1 = LinkingNumber(curves1)
-        Object2 = LinkingNumber(curves2)
-
-        fullArray = Object1.J()
-        fullArray2 = Object2.J()
-        print("Link Number Testing (should be 0, 1)")
-        print(fullArray)
-        print(fullArray2)
-        self.assertAlmostEqual(fullArray, 0)
-        self.assertAlmostEqual(fullArray2, 1)
+        for downsample in [1, 2, 5]:
+            curves1 = create_equally_spaced_curves(2, 1, stellsym=True, R0=1, R1=0.5, order=5, numquadpoints=120)
+            curve1 = CurveXYZFourier(200, 3)
+            coeffs = curve1.dofs_matrix
+            coeffs[1][0] = 1.
+            coeffs[1][1] = 0.5
+            coeffs[2][2] = 0.5
+            curve1.set_dofs(np.concatenate(coeffs))
+
+            curve2 = CurveXYZFourier(150, 3)
+            coeffs = curve2.dofs_matrix
+            coeffs[1][0] = 0.5
+            coeffs[1][1] = 0.5
+            coeffs[0][0] = 0.1
+            coeffs[0][1] = 0.5
+            coeffs[0][2] = 0.5
+            curve2.set_dofs(np.concatenate(coeffs))
+            curves2 = [curve1, curve2]
+            curves3 = [curve2, curve1]
+            objective1 = LinkingNumber(curves1, downsample)
+            objective2 = LinkingNumber(curves2, downsample)
+            objective3 = LinkingNumber(curves3, downsample)
+
+            print("Linking number testing (should be 0, 1, 1):", objective1.J(), objective2.J(), objective3.J())
+            np.testing.assert_allclose(objective1.J(), 0, atol=1e-14, rtol=1e-14)
+            np.testing.assert_allclose(objective2.J(), 1, atol=1e-14, rtol=1e-14)
+            np.testing.assert_allclose(objective3.J(), 1, atol=1e-14, rtol=1e-14)
 
 
 if __name__ == "__main__":