add function to get center of mass of path

brunobeltran · brunobeltran · commit cb37c8f88922 · 2020-04-09T14:28:04.000-07:00
diff --git a/lib/matplotlib/bezier.py b/lib/matplotlib/bezier.py
@@ -310,6 +310,73 @@ def arc_length(self, rtol=None, atol=None):
                     return curr_est
         return curr_est
 
+    def arc_center_of_mass(self):
+        r"""
+        Center of mass of the (even-odd-rendered) area swept out by the ray
+        from the origin to the path.
+
+        Summing this vector for each segment along a closed path will produce
+        that area's center of mass.
+
+        Returns
+        -------
+        r_cm : (2,) np.array<float>
+            the "arc's center of mass"
+
+        Notes
+        -----
+        A simple analytical form can be derived for general Bezier curves.
+        Suppose the curve was closed, so :math:`B(0) = B(1)`. Call the area
+        enclosed by :math:`B(t)` :math:`B_\text{int}`. The center of mass of
+        :math:`B_\text{int}` is defined by the expected value of the position
+        vector :math:`\vec{r}`
+
+        .. math::
+
+            \vec{R}_\text{cm} = \int_{B_\text{int}} \vec{r} \left( \frac{1}{
+            \int_{B_\text{int}}} d\vec{r} \right) d\vec{r}
+
+        where :math:`(1/\text{Area}(B_\text{int})` can be interpreted as a
+        probability density.
+
+        In order to compute this integral, we choose two functions
+        :math:`F_0(x,y) = [x^2/2, 0]` and :math:`F_1(x,y) = [0, y^2/2]` such
+        that :math:`[\div \cdot F_0, \div \cdot F_1] = \vec{r}`. Then, applying
+        the divergence integral (componentwise), we get that
+
+        .. math::
+            \vec{R}_\text{cm} &= \oint_{B(t)} F \cdot \vec{n} dt \\
+            &= \int_0^1 \left[ \begin{array}{1}
+                B^{(0)}(t) \frac{dB^{(1)}(t)}{dt}  \\
+              - B^{(1)}(t) \frac{dB^{(0)}(t)}{dt}  \end{array} \right] dt
+
+        After expanding in Berstein polynomials and moving the integral inside
+        all the sums, we get that
+
+        .. math::
+            \vec{R}_\text{cm} = \frac{1}{6} \sum_{i,j=0}^n\sum_{k=0}^{n-1}
+                \frac{{n \choose i}{n \choose j}{{n-1} \choose k}}
+                     {{3n - 1} \choose {i + j + k}}
+                \left(\begin{array}{1}
+                    P^{(0)}_i P^{(0)}_j (P^{(1)}_{k+1} - P^{(1)}_k)
+                  - P^{(1)}_i P^{(1)}_j (P^{(0)}_{k+1} - P^{(0)}_k)
+                \right) \end{array}
+
+        where :math:`P_i = [P^{(0)}_i, P^{(1)}_i]` is the :math:`i`'th control
+        point of the curve and :math:`n` is the degree of the curve.
+        """
+        n = self.degree
+        r_cm = np.zeros(2)
+        P = self.control_points
+        dP = np.diff(P, axis=0)
+        Pn = np.array([[1, -1]])*dP[:, ::-1]  # n = [y, -x]
+        for i in range(n + 1):
+            for j in range(n + 1):
+                for k in range(n):
+                    r_cm += _comb(n, i) * _comb(n, j) * _comb(n - 1, k) \
+                            * P[i]*P[j]*Pn[k] / _comb(3*n - 1, i + j + k)
+        return r_cm/6
+
     def arc_area(self):
         r"""
         (Signed) area swept out by ray from origin to curve.
@@ -406,6 +473,15 @@ def arc_area(self):
                         * (P[j, 0]*dP[k, 1] - P[j, 1]*dP[k, 0])
         return (1/4)*area
 
+    def center_of_mass(self):
+        """Return the center of mass of the curve (not the filled curve!)
+
+        Notes
+        -----
+        Computed as the mean of the control points.
+        """
+        return np.mean(self._cpoints, axis=0)
+
     @classmethod
     def differentiate(cls, B):
         """Return the derivative of a BezierSegment, itself a BezierSegment"""
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -700,10 +700,147 @@ def signed_area(self, **kwargs):
         # add final implied CLOSEPOLY, if necessary
         if start_point is not None \
                 and not np.all(np.isclose(start_point, prev_point)):
-            B = BezierSegment(np.array([prev_point, start_point]))
-            area += B.arc_area()
+            Bclose = BezierSegment(np.array([prev_point, start_point]))
+            area += Bclose.arc_area()
         return area
 
+    def center_of_mass(self, dimension=None, **kwargs):
+        r"""
+        Center of mass of the path, assuming constant density.
+
+        The center of mass is defined to be the expected value of a vector
+        located uniformly within either the filled area of the path
+        (:code:`dimension=2`) or the along path's edge (:code:`dimension=1`) or
+        along isolated points of the path (:code:`dimension=0`).  Notice in
+        particular that for this definition, if the filled area is used, then
+        any 0- or 1-dimensional components of the path will not contribute to
+        the center of mass. Similarly, for if *dimension* is 1, then isolated
+        points in the path (i.e.  "0-dimensional" strokes made up of only
+        :code:`Path.MOVETO`'s) will not contribute to the center of mass.
+
+        For the 2d case, the center of mass is computed using the same
+        filling strategy as `signed_area`. So, if a path is self-intersecting,
+        the drawing rule "even-odd" is used and only the filled area is
+        counted, and all sub paths are treated as if they had been closed. That
+        is, if there is a MOVETO without a preceding CLOSEPOLY, one is added.
+
+        For the 1d measure, the curve is averaged as-is (the implied CLOSEPOLY
+        is not added).
+
+        For the 0d measure, any non-isolated points are ignored.
+
+        Parameters
+        ----------
+        dimension : 2, 1, or 0 (optional)
+            Whether to compute the center of mass by taking the expected value
+            of a position uniformly distributed within the filled path
+            (2D-measure), the path's edge (1D-measure), or between the
+            discrete, isolated points of the path (0D-measure), respectively.
+            By default, the intended dimension of the path is inferred by
+            checking first if `Path.signed_area` is non-zero (implying a
+            *dimension* of 2), then if the `Path.length` is non-zero (implying
+            a *dimension* of 1), and finally falling back to the counting
+            measure (*dimension* of 0).
+        kwargs : Dict[str, object]
+            Passed thru to `Path.cleaned` via `Path.iter_bezier`.
+
+        Returns
+        -------
+        r_cm : (2,) np.array<float>
+            The center of mass of the path.
+
+        Raises
+        ------
+        ValueError
+            An empty path has no well-defined center of mass.
+
+            In addition, if a specific *dimension* is requested and that
+            dimension is not well-defined, an error is raised. This can happen
+            if::
+
+                1) 2D expected value was requested but the path has zero area
+                2) 1D expected value was requested but the path has only
+                `Path.MOVETO` directives
+                3) 0D expected value was requested but the path has NO
+                subsequent `Path.MOVETO` directives.
+
+            This error cannot be raised if the function is allowed to infer
+            what *dimension* to use.
+        """
+        area = None
+        cleaned = self.cleaned(**kwargs)
+        move_codes = cleaned.codes == Path.MOVETO
+        if len(cleaned.codes) == 0:
+            raise ValueError("An empty path has no center of mass.")
+        if dimension is None:
+            dimension = 2
+            area = cleaned.signed_area()
+            if not np.isclose(area, 0):
+                dimension -= 1
+            if np.all(move_codes):
+                dimension = 0
+        if dimension == 2:
+            # area computation can be expensive, make sure we don't repeat it
+            if area is None:
+                area = cleaned.signed_area()
+            if np.isclose(area, 0):
+                raise ValueError("2d expected value over empty area is "
+                                 "ill-defined.")
+            return cleaned._2d_center_of_mass(area)
+        if dimension == 1:
+            if np.all(move_codes):
+                raise ValueError("1d expected value over empty arc-length is "
+                                 "ill-defined.")
+            return cleaned._1d_center_of_mass()
+        if dimension == 0:
+            adjacent_moves = (move_codes[1:] + move_codes[:-1]) == 2
+            if len(move_codes) > 1 and not np.any(adjacent_moves):
+                raise ValueError("0d expected value with no isolated points "
+                                 "is ill-defined.")
+            return cleaned._0d_center_of_mass()
+
+    def _2d_center_of_mass(self, normalization=None):
+        #TODO: refactor this and signed_area (and maybe others, with
+        # close= parameter)?
+        if normalization is None:
+            normalization = self.signed_area()
+        r_cm = np.zeros(2)
+        prev_point = None
+        prev_code = None
+        start_point = None
+        for B, code in self.iter_bezier():
+            if code == Path.MOVETO:
+                if prev_code is not None and prev_code is not Path.CLOSEPOLY:
+                    Bclose = BezierSegment(np.array([prev_point, start_point]))
+                    r_cm += Bclose.arc_center_of_mass()
+                start_point = B.control_points[0]
+            r_cm += B.arc_center_of_mass()
+            prev_point = B.control_points[-1]
+            prev_code = code
+        # add final implied CLOSEPOLY, if necessary
+        if start_point is not None \
+                and not np.all(np.isclose(start_point, prev_point)):
+            Bclose = BezierSegment(np.array([prev_point, start_point]))
+            r_cm += Bclose.arc_center_of_mass()
+        return r_cm / normalization
+
+    def _1d_center_of_mass(self):
+        r_cm = np.zeros(2)
+        Bs = list(self.iter_bezier())
+        arc_lengths = np.array([B.arc_length() for B in Bs])
+        r_cms = np.array([B.center_of_mass() for B in Bs])
+        total_length = np.sum(arc_lengths)
+        return np.sum(r_cms*arc_lengths)/total_length
+
+    def _0d_center_of_mass(self):
+        move_verts = self.codes
+        isolated_verts = move_verts.copy()
+        if len(move_verts) > 1:
+            isolated_verts[:-1] = (move_verts[:-1] + move_verts[1:]) == 2
+            isolated_verts[-1] = move_verts[-1]
+        num_verts = np.sum(isolated_verts)
+        return np.sum(self.vertices[isolated_verts], axis=0)/num_verts
+
     def interpolated(self, steps):
         """
         Return a new path resampled to length N x steps.
diff --git a/lib/matplotlib/tests/test_bezier.py b/lib/matplotlib/tests/test_bezier.py
@@ -10,6 +10,11 @@
 _test_curves = [list(tc.path.iter_bezier())[-1][0] for tc in _test_curves]
 
 
+def _integral_center_of_mass(B):
+    return np.array([scipy.integrate.quad(lambda t: B(t)[0], 0, 1)[0],
+                     scipy.integrate.quad(lambda t: B(t)[1], 0, 1)[0]])
+
+
 def _integral_arc_length(B):
     dB = B.differentiate(B)
     def integrand(t):
@@ -27,6 +32,27 @@ def integrand(t):
     return scipy.integrate.quad(integrand, 0, 1)[0]
 
 
+def _integral_arc_com(B):
+    dB = B.differentiate(B)
+    def integrand(t):
+        dr = dB(t).T
+        n = np.array([dr[1], -dr[0]])
+        return B(t).T**2 * n / 2
+    def integrand_x(t):
+        return integrand(t)[0]
+    def integrand_y(t):
+        return integrand(t)[1]
+    return np.array([
+        scipy.integrate.quad(integrand_x, 0, 1)[0],
+        scipy.integrate.quad(integrand_y, 0, 1)[0]
+    ])
+
+
+def _integral_com(B):
+    return np.array([scipy.integrate.quad(lambda t: B(t)[0], 0, 1)[0],
+                     scipy.integrate.quad(lambda t: B(t)[1], 0, 1)[0]])
+
+
 @pytest.mark.parametrize("B", _test_curves)
 def test_area_formula(B):
     assert np.isclose(_integral_arc_area(B), B.arc_area())
@@ -37,3 +63,13 @@ def test_length_iteration(B):
     assert np.isclose(_integral_arc_length(B),
                       B.arc_length(rtol=1e-5, atol=1e-8),
                       rtol=1e-5, atol=1e-8)
+
+
+@pytest.mark.parametrize("B", _test_curves)
+def test_center_of_mass_1d(B):
+    assert np.all(np.isclose(B.center_of_mass(), _integral_center_of_mass(B)))
+
+
+@pytest.mark.parametrize("B", _test_curves)
+def test_center_of_mass_2d(B):
+    assert np.all(np.isclose(B.arc_center_of_mass(), _integral_arc_com(B)))
