docs(example): Adds Confidence Interval Ellipses

pyproject.toml

@@ -98,6 +98,7 @@ doc = [
     "sphinx_copybutton",
     "sphinx-design",
     "scipy",
+    "scipy-stubs; python_version>=\"3.10\"",
 ]
 
 [tool.altair.vega]
@@ -458,8 +459,10 @@ module = [
     "ipykernel.*",
     "ibis.*",
     "vegafusion.*",
+    "scipy.*"
 ]
 ignore_missing_imports = true
+disable_error_code = ["import-untyped"]
 
 [tool.pyright]
 enableExperimentalFeatures=true

sphinxext/utils.py

@@ -58,7 +58,7 @@ def create_generic_image(
 """
 
 
-def _parse_source_file(filename: str) -> tuple[ast.Module | None, str]:
+def _parse_source_file(filename: str | Path) -> tuple[ast.Module | None, str]:
     """
     Parse source file into AST node.
 
@@ -88,7 +88,7 @@ def _parse_source_file(filename: str) -> tuple[ast.Module | None, str]:
     return node, content
 
 
-def get_docstring_and_rest(filename: str) -> tuple[str, str | None, str, int]:
+def get_docstring_and_rest(filename: str | Path) -> tuple[str, str | None, str, int]:
     """
     Separate ``filename`` content between docstring and the rest.
 
@@ -160,7 +160,7 @@ def get_docstring_and_rest(filename: str) -> tuple[str, str | None, str, int]:
         if (
             node.body
             and isinstance(node.body[0], ast.Expr)
-            and isinstance(node.body[0].value, (ast.Str, ast.Constant))
+            and isinstance(node.body[0].value, ast.Constant)
         ):
             docstring_node = node.body[0]
             docstring = docstring_node.value.s  # pyright: ignore[reportAttributeAccessIssue]

tests/__init__.py

@@ -95,6 +95,10 @@ def windows_has_tzdata() -> bool:
     https://github.com/vega/vegafusion
 """
 
+skip_requires_scipy: pytest.MarkDecorator = pytest.mark.skipif(
+    find_spec("scipy") is None, reason="`scipy` not installed."
+)
+
 
 def skip_requires_pyarrow(
     fn: Callable[..., Any] | None = None, /, *, requires_tzdata: bool = False
@@ -229,6 +233,7 @@ def _distributed_examples(
                 "wind_vector_map": slow,
                 r"\.point_map\.py": slow,
                 "line_chart_with_color_datum": slow,
+                "deviation_ellipses": skip_requires_scipy,
             },
         )
     ),

tests/examples_arguments_syntax/deviation_ellipses.py

@@ -0,0 +1,91 @@
+"""
+Confidence Interval Ellipses
+----------------------------
+This example shows bivariate deviation ellipses of petal length and width of three iris species.
+
+Inspired by `ggplot2.stat_ellipse`_ and directly based on `Deviation ellipses example`_ by `@essicolo`_
+
+.. _ggplot2.stat_ellipse:
+    https://ggplot2.tidyverse.org/reference/stat_ellipse.html#ref-examples
+.. _Deviation ellipses example:
+    https://github.com/vega/altair/pull/514
+.. _@essicolo:
+    https://github.com/essicolo
+"""
+
+# category: distributions
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy.stats import f as F
+
+import altair as alt
+from vega_datasets import data
+
+
+def np_ellipse(
+    arr: np.ndarray[tuple[int, int], np.dtype[np.float64]],
+    conf_level: float = 0.95,
+    segments: int = 50,
+):
+    """
+    Calculate confidence interval ellipse.
+
+    Parameters
+    ----------
+    arr
+        numpy array with 2 columns
+    conf_level
+        lower tail probability
+    segments
+        number of points describing the ellipse.
+    """
+    n_elements = len(arr)
+    # Degrees of freedom of the chi-squared distribution in the **numerator**
+    dfn = 2
+    # Degrees of freedom of the chi-squared distribution in the **denominator**
+    dfd = n_elements - 1
+    # Percent point function at `conf_level` of an F continuous random variable
+    quantile = F.ppf(conf_level, dfn=dfn, dfd=dfd)
+    radius = np.sqrt(2 * quantile)
+    angles = np.arange(0, segments) * 2 * np.pi / segments
+    circle = np.column_stack((np.cos(angles), np.sin(angles)))
+    center = np.mean(arr, axis=0)
+    cov_mat = np.cov(arr, rowvar=False)
+    return center + radius * (circle @ np.linalg.cholesky(cov_mat).T)
+
+
+def pd_ellipse(
+    df: pd.DataFrame, col_x: str, col_y: str, col_group: str
+) -> pd.DataFrame:
+    cols = col_x, col_y
+    groups = []
+    # TODO: Rewrite in a more readable way
+    categories = df[col_group].unique()
+    for category in categories:
+        sliced = df.loc[df[col_group] == category, cols]
+        ell_df = pd.DataFrame(np_ellipse(sliced.to_numpy()), columns=cols) # type: ignore
+        ell_df[col_group] = category
+        groups.append(ell_df)
+    return pd.concat(groups).reset_index()
+
+
+col_x = "petalLength"
+col_y = "petalWidth"
+col_group = "species"
+x = alt.X(col_x, scale=alt.Scale(zero=False))
+y = alt.Y(col_y, scale=alt.Scale(zero=False))
+color = alt.Color(col_group)
+
+source = data.iris()
+ellipse = pd_ellipse(source, col_x=col_x, col_y=col_y, col_group=col_group)
+points = alt.Chart(source).mark_circle(size=50, tooltip=True).encode(x, y, color)
+lines = (
+    alt.Chart(ellipse)
+    .mark_line(filled=True, fillOpacity=0.2)
+    .encode(x, y, color, order="index")
+)
+
+chart = (lines + points).properties(height=500, width=500)
+chart