Merge pull request #26 from InfoMusCP/Equilibrium

feat: add Equilibrium class with elliptical balance evaluation and in…

Author: simoneghisio

core/equilibrium.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+class Equilibrium:
+    """
+    Elliptical equilibrium evaluation between two feet and a barycenter.
+
+    This class defines an elliptical region of interest (ROI) aligned with the
+    line connecting the left and right foot. The ellipse is scaled by a margin
+    in millimeters and can be weighted along the Y-axis to emphasize forward–
+    backward sway more than lateral sway. A barycenter is evaluated against
+    this ellipse to compute a normalized equilibrium value.
+
+    Parameters
+    ----------
+    margin_mm : float, optional
+        Extra margin in millimeters added around the rectangle spanned by the
+        two feet (default: 100).
+    y_weight : float, optional
+        Weighting factor applied to the ellipse height along the Y-axis.
+        A value < 1 shrinks the ellipse in the forward/backward direction,
+        emphasizing sway in that axis (default: 0.5).
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> eq = Equilibrium(margin_mm=120, y_weight=0.6)
+    >>> left = np.array([0, 0, 0])
+    >>> right = np.array([400, 0, 0])
+    >>> barycenter = np.array([200, 50, 0])
+    >>> value, angle = eq(left, right, barycenter)
+    >>> round(value, 2)
+    0.91
+    >>> round(angle, 1)
+    0.0
+    """
+
+    def __init__(self, margin_mm=100, y_weight=0.5):
+        self.margin = margin_mm
+        self.y_weight = y_weight
+
+    def __call__(self, left_foot: np.ndarray, right_foot: np.ndarray, barycenter: np.ndarray) -> tuple[float, float]:
+        """
+        Evaluate the equilibrium value and ellipse angle.
+
+        Parameters
+        ----------
+        left_foot : numpy.ndarray, shape (3,)
+            3D coordinates (x, y, z) of the left foot in millimeters.
+            Only the x and y components are used.
+        right_foot : numpy.ndarray, shape (3,)
+            3D coordinates (x, y, z) of the right foot in millimeters.
+            Only the x and y components are used.
+        barycenter : numpy.ndarray, shape (3,)
+            3D coordinates (x, y, z) of the barycenter in millimeters.
+            Only the x and y components are used.
+
+        Returns
+        -------
+        value : float
+            Equilibrium value in [0, 1].
+            - 1 means the barycenter is perfectly at the ellipse center.
+            - 0 means the barycenter is outside the ellipse.
+        angle : float
+            Orientation of the ellipse in degrees, measured counter-clockwise
+            from the X-axis (line connecting left and right foot).
+
+        Notes
+        -----
+        - The ellipse is aligned with the line connecting the two feet.
+        - The ellipse width corresponds to the horizontal foot span + margin.
+        - The ellipse height corresponds to the vertical span + margin,
+          scaled by `y_weight`.
+        """
+        ps = np.array(left_foot)[:2]
+        pd = np.array(right_foot)[:2]
+        bc = np.array(barycenter)[:2]
+
+        min_xy = np.minimum(ps, pd) - self.margin
+        max_xy = np.maximum(ps, pd) + self.margin
+
+        center = (min_xy + max_xy) / 2
+        half_sizes = (max_xy - min_xy) / 2
+
+        a = half_sizes[0]
+        b = half_sizes[1] * self.y_weight
+
+        dx, dy = pd - ps
+        angle = np.arctan2(dy, dx)
+
+        rel = bc - center
+
+        rot_matrix = np.array([
+            [np.cos(-angle), -np.sin(-angle)],
+            [np.sin(-angle),  np.cos(-angle)]
+        ])
+        rel_rot = rot_matrix @ rel
+
+        norm = (rel_rot[0] / a) ** 2 + (rel_rot[1] / b) ** 2
+
+        if norm <= 1.0:
+            value = 1.0 - np.sqrt(norm)
+        else:
+            value = 0.0
+
+        return max(0.0, value), np.degrees(angle)

examples/test_equilibrium.py

@@ -0,0 +1,123 @@
+"""
+Interactive equilibrium simulation using matplotlib.
+
+This script provides a visual demonstration of the `Equilibrium` class.
+It creates an interactive 2D plot where the user can move the mouse to
+simulate the barycenter position, and the system evaluates whether the
+barycenter is within the equilibrium ellipse defined by two feet.
+
+The ellipse is dynamically updated in position, orientation, and color:
+- Green if the barycenter is within the ellipse.
+- Red if the barycenter is outside.
+
+The equilibrium value (in [0, 1]) is displayed in real-time.
+
+Examples
+--------
+Run the script:
+
+    $ python demo_equilibrium.py
+
+Then move the mouse cursor inside the figure window to simulate barycenter
+movements.
+
+Notes
+-----
+- Requires `matplotlib` for visualization.
+- Uses `numpy` for vector operations.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Ellipse
+import sys, os
+if os.getcwd() not in sys.path:
+    sys.path.append(os.getcwd())
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from core.equilibrium import Equilibrium
+
+# --- Setup test parameters ---
+left_foot = np.array([120, 200, 0])
+"""numpy.ndarray: 3D coordinates (x, y, z) of the left foot in millimeters."""
+
+right_foot = np.array([800, 600, 0])
+"""numpy.ndarray: 3D coordinates (x, y, z) of the right foot in millimeters."""
+
+eq = Equilibrium(margin_mm=100, y_weight=0.5)
+"""Equilibrium: instance of the equilibrium evaluator."""
+
+# --- Setup matplotlib figure ---
+fig, ax = plt.subplots()
+ax.set_xlim(-200, 1000)
+ax.set_ylim(-200, 800)
+ax.set_aspect('equal')
+ax.set_title("Equilibrium. Move the barycenter with mouse.")
+
+# Plot the two feet as blue points
+ax.plot(left_foot[0], left_foot[1], 'bo', markersize=8)
+ax.plot(right_foot[0], right_foot[1], 'bo', markersize=8)
+
+# Ellipse ROI (initial placeholder)
+roi_ellipse = Ellipse((0, 0), 0, 0,
+                      fill=True, facecolor='green', alpha=0.2,
+                      edgecolor='green', linewidth=2)
+ax.add_patch(roi_ellipse)
+
+# Red point for barycenter
+baricentro_plot, = ax.plot([], [], 'ro', markersize=8)
+
+# Text showing equilibrium value
+eq_text = ax.text(0.02, 1.02, "", transform=ax.transAxes, fontsize=12)
+
+def on_move(event):
+    """
+    Handle mouse movement and update equilibrium visualization.
+
+    Parameters
+    ----------
+    event : matplotlib.backend_bases.MouseEvent
+        The mouse event containing the current cursor position.
+        Only `event.xdata` and `event.ydata` are used.
+
+    Behavior
+    --------
+    - Computes equilibrium value and ellipse angle based on the
+      simulated barycenter.
+    - Updates ellipse position, size, orientation, and color.
+    - Updates barycenter position on the plot.
+    - Updates equilibrium value text.
+    """
+    if not event.inaxes:
+        return
+
+    baricentro = np.array([event.xdata, event.ydata, 0])
+
+    value, angle = eq(left_foot, right_foot, baricentro)
+
+    ps = np.array(left_foot)[:2]
+    pd = np.array(right_foot)[:2]
+    min_xy = np.minimum(ps, pd) - eq.margin
+    max_xy = np.maximum(ps, pd) + eq.margin
+    center = (min_xy + max_xy) / 2
+    half_sizes = (max_xy - min_xy) / 2
+    a = half_sizes[0]
+    b = half_sizes[1] * eq.y_weight
+
+    roi_ellipse.set_center(center)
+    roi_ellipse.width = 2 * a
+    roi_ellipse.height = 2 * b
+    roi_ellipse.angle = angle
+    roi_ellipse.set_facecolor('green' if value > 0 else 'red')
+    roi_ellipse.set_edgecolor('green' if value > 0 else 'red')
+
+    baricentro_plot.set_data(baricentro[0], baricentro[1])
+
+    eq_text.set_text(f"Equilibrium value = {value:.2f}")
+
+    fig.canvas.draw_idle()
+
+# Connect mouse motion event
+fig.canvas.mpl_connect('motion_notify_event', on_move)
+
+# Run interactive visualization
+plt.show()