From 3dee03eb490309cc1da202de8417113695f6b233 Mon Sep 17 00:00:00 2001
From: Hongkai Dai <daihongkai@gmail.com>
Date: Sun, 31 Dec 2023 20:30:31 -0800
Subject: [PATCH] Write ellipsoid as an affine transformation of a ball. (#29)

---
 compatible_clf_cbf/ellipsoid_utils.py | 52 +++++++++++++++++++++++++++
 tests/test_ellipsoid_utils.py         | 23 ++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/compatible_clf_cbf/ellipsoid_utils.py b/compatible_clf_cbf/ellipsoid_utils.py
index e41dcca..e7a493d 100644
--- a/compatible_clf_cbf/ellipsoid_utils.py
+++ b/compatible_clf_cbf/ellipsoid_utils.py
@@ -4,6 +4,8 @@
 
 from typing import List, Optional, Tuple, Union
 
+import matplotlib.axes
+import matplotlib.lines
 import numpy as np
 
 import pydrake.solvers as solvers
@@ -319,3 +321,53 @@ def in_ellipsoid(
     else:
         assert pts.shape[1] == dim
         return (np.sum(pts * (pts @ S), axis=1) + pts @ b + c <= 0).tolist()
+
+
+def to_affine_ball(
+    S: np.ndarray, b: np.ndarray, c: float
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Convert the ellipsoid in the form {x | xᵀSx+bᵀx+c ≤ 0} to an alternative
+    form as an affine transformation of a ball {A(y+d) | |y|₂ ≤ 1 }
+
+    Args:
+      S: A symmetric matrix that parameterizes the ellipsoid.
+      b: A vector that parameterizes the ellipsoid.
+      c: A float that parameterizes the ellipsoid.
+    Returns:
+      A: The affine transformation of the ball.
+      d: The center of the ellipsoid.
+    """
+
+    # The ellipsoid can be written as
+    # {x | xᵀA⁻ᵀA⁻¹x − 2dᵀA⁻¹x + dᵀd−1 ≤ 0}. To match it with xᵀSx+bᵀx+c ≤ 0,
+    # we introduce a scalar k, with the constraint
+    # A⁻ᵀA⁻¹ = kS
+    # −2A⁻ᵀd = kb
+    # dᵀd−1 = kc
+    # To solve these equations, I define A̅ = A * √k, d̅ = d/√k
+    # Hence I know A̅⁻ᵀA̅⁻¹ = S
+    bar_A_inv = np.linalg.cholesky(S).T
+    bar_A = np.linalg.inv(bar_A_inv)
+    # −2A̅⁻ᵀd̅=b
+    bar_d = bar_A.T @ b / -2
+    # kd̅ᵀd̅−1 = kc
+    k = 1 / (bar_d.dot(bar_d) - c)
+    sqrt_k = np.sqrt(k)
+    A = bar_A / sqrt_k
+    d = bar_d * sqrt_k
+    return (A, d)
+
+
+def draw_ellipsoid2d(
+    ax: matplotlib.axes.Axes, A: np.ndarray, d: np.ndarray, **kwargs
+) -> List[matplotlib.lines.Line2D]:
+    """
+    Draw a 2D ellipsoid {A(y+d) | |y|₂ ≤ 1 }
+    """
+    theta = np.linspace(0, 2 * np.pi, 100)
+    xy_pts = A @ (
+        np.concatenate([np.cos(theta).reshape((1, -1)), np.sin(theta).reshape((1, -1))])
+        + d.reshape((-1, 1))
+    )
+    return ax.plot(xy_pts[0], xy_pts[1], **kwargs)
diff --git a/tests/test_ellipsoid_utils.py b/tests/test_ellipsoid_utils.py
index 546d610..58fb454 100644
--- a/tests/test_ellipsoid_utils.py
+++ b/tests/test_ellipsoid_utils.py
@@ -2,6 +2,7 @@
 
 import jax.numpy as jnp
 import jax
+import matplotlib.pyplot as plt
 import numpy as np
 import pytest  # noqa
 
@@ -311,3 +312,25 @@ def test_in_ellipsoid():
         c,
         (center + np.array([[0.5, 0, 0.2], [1.1, 0, 0], [1.2, 0, -1]])) @ A.T,
     ) == [True, False, False]
+
+
+def test_to_affine_ball():
+    def check(S, b, c):
+        A, d = mut.to_affine_ball(S, b, c)
+
+        # The ellipsoid is {x | xᵀA⁻ᵀA⁻¹x − 2dᵀA⁻¹x + dᵀd−1 ≤ 0}
+        ratio = (d.dot(d) - 1) / c
+        np.testing.assert_allclose(np.linalg.inv(A @ A.T), S * ratio)
+        np.testing.assert_allclose(-2 * np.linalg.solve(A.T, d), b * ratio)
+
+    check(np.eye(2), np.zeros(2), 1)
+    check(np.diag(np.array([1.0, 2.0, 3.0])), np.array([2.0, 3.0, 4.0]), -10)
+
+
+def test_draw_ellipsoid2d():
+    A = np.array([[2, 1], [0, 1.0]])
+    d = np.array([1, 3])
+    fig = plt.figure()
+    ax = fig.add_subplot()
+    lines = mut.draw_ellipsoid2d(ax, A, d, color="k")
+    assert isinstance(lines, list)