Add assert_plot_equal to compare a plot with a reference plot

README.md

@@ -78,6 +78,7 @@ Available helpers:
 - `assert_fn_equal(sandbox, func_name, args=..., expected=...)` — call + compare
 - `assert_fn_approx_equal(sandbox, func_name, args=..., expected=...)` — call + approx compare
 - `assert_true(condition)` / `assert_false(condition)` — boolean checks
+- `assert_plot_equal(plot, expected)` — compare a matplotlib plot with a reference plot
 
 ### Ungradable Submission Detection
 

quick_start.md

@@ -1017,6 +1017,10 @@ Same as `assert_fn_equal` but uses approximate comparison:
 assert_fn_approx_equal(sandbox, "compute_area", args=(3.0,), expected=28.274, rtol=1e-2)
 ```
 
+**`assert_plot_equal(plot, expected, *, name="your plot", rtol=1e-5, atol=1e-8, ..., msg=None)`**
+
+Compares a student's matplotlib plot with a reference plot. See [Comparing with a Reference Plot](#comparing-with-a-reference-plot).
+
 #### Combining with Partial Credit
 
 Friendly mode works seamlessly with `FeedbackFixture`:
@@ -1109,6 +1113,66 @@ def test_student_plot(sandbox: StudentFixture) -> None:
     pt.assert_title_contains(["Data Visualization"])
 ```
 
+#### Comparing with a Reference Plot
+
+`assert_plot_equal` checks a student's plot against a reference plot that the test draws itself:
+
+```python
+import numpy as np
+import pytest
+from matplotlib.figure import Figure
+from pytest_prairielearn_grader.assertions import assert_plot_equal
+from pytest_prairielearn_grader.fixture import StudentFixture
+from pytest_prairielearn_grader.utils import GradingOutputLevel
+
+
+@pytest.mark.output(level=GradingOutputLevel.Friendly)
+@pytest.mark.grading_data(name="plot", points=1)
+def test_plot(sandbox: StudentFixture) -> None:
+    reference = Figure()
+    axes = reference.subplots()
+    x = np.arange(0, 11)
+    axes.semilogy(x, np.exp(x))
+
+    assert_plot_equal(sandbox.query("plot"), reference)
+```
+
+Draw the reference on a `Figure` made directly, as here, rather than with `plt`, so that it doesn't
+share pyplot's current figure with anything else in the test. Both plots can be a `Figure` or an
+`Axes`, and subplots are compared in order.
+
+By default, it checks:
+
+- **Subplots:** the number of axes.
+- **Scales:** the x and y axis scales, such as linear or log.
+- **Data:** the lines, scatter plots and bar plots, which must have the same number of each kind. Data is compared with `numpy.allclose`, using `rtol` and `atol`.
+
+The series can be in any order, and a line also matches when its points are in reverse order, since
+it looks the same. The points in a scatter plot and the bars in a bar plot can be in any order.
+
+Keyword arguments change what's checked:
+
+| Argument       | Default | Checks                                                                                                     |
+| -------------- | ------- | ---------------------------------------------------------------------------------------------------------- |
+| `check_scales` | `True`  | The x and y axis scales.                                                                                   |
+| `check_data`   | `True`  | The data in lines, scatter plots and bar plots.                                                            |
+| `ordered`      | `False` | Requires lines, scatter plots, bar plots and legend labels to be in the reference's order.                 |
+| `allow_extra`  | `False` | Allows more lines, scatter plots or bar plots than the reference, such as a student's extra reference line. |
+| `check_labels` | `False` | The x and y axis labels, exactly.                                                                          |
+| `check_titles` | `False` | The title of each subplot and the figure, exactly.                                                         |
+| `check_legend` | `False` | The legend's labels and title, or that neither plot has a legend.                                          |
+| `check_limits` | `False` | The x and y axis limits, within `rtol` and `atol`.                                                         |
+| `check_styles` | `False` | Line colors, styles and markers, scatter point colors, and bar colors.                                     |
+
+A failure describes the first difference found, such as `Expected a log scale on the y-axis of your
+plot, but found a linear scale.` or `No line in your plot matches the line labeled 'exp'.` Lines
+are named by their labels when the reference has them. Use `name` to call the plot something else
+in messages, such as `name="plot"` for a variable, or `msg` to show your own message instead.
+
+Anything `assert_plot_equal` doesn't cover, such as checking that a label contains a word, can be
+checked directly on the `Figure`, or with `matplotcheck` as in the earlier example. Both can be
+used in the same test.
+
 #### Showing Plots to Students
 
 To show the student's plot with a test's results, pass it to `feedback.add_image`:

src/pytest_prairielearn_grader/assertions.py

@@ -188,6 +188,103 @@ def assert_false(
     raise AssertionError(msg or "Expected the condition to be false, but it was true.")
 
 
+def assert_plot_equal(
+    plot: Any,
+    expected: Any,
+    *,
+    name: str = "your plot",
+    rtol: float = 1e-5,
+    atol: float = 1e-8,
+    check_scales: bool = True,
+    check_data: bool = True,
+    ordered: bool = False,
+    allow_extra: bool = False,
+    check_labels: bool = False,
+    check_titles: bool = False,
+    check_legend: bool = False,
+    check_limits: bool = False,
+    check_styles: bool = False,
+    msg: str | None = None,
+) -> None:
+    """Assert that a student's matplotlib plot matches a reference plot.
+
+    Both plots can be a ``Figure`` or an ``Axes``.  ``sandbox.query`` returns
+    a student's plot as a ``Figure``; the reference is usually drawn by the
+    test itself.  Subplots are compared in order.
+
+    By default this checks the number of subplots, the axis scales, and the
+    data in each line, scatter plot and bar plot, which must have the same
+    number of each kind.  Series can be in any order, and a line also matches
+    when its points are in reverse order.  Scatter points and bars can be in
+    any order.  Numbers are compared with ``numpy.allclose``.
+
+    Parameters
+    ----------
+    plot:
+        The student's plot.  Anything other than a ``Figure`` or ``Axes``
+        fails the assertion.
+    expected:
+        The reference plot.
+    name:
+        What to call the student's plot in failure messages.
+    rtol, atol:
+        Tolerances for comparing data and axis limits.
+    check_scales:
+        Compare the x and y axis scales, such as ``"linear"`` and ``"log"``.
+    check_data:
+        Compare the data in lines, scatter plots and bar plots.
+    ordered:
+        Require lines, scatter plots, bar plots and legend labels to be in
+        the same order as in the reference.
+    allow_extra:
+        Allow the student's plot to have more lines, scatter plots or bar
+        plots than the reference.
+    check_labels:
+        Compare the x and y axis labels exactly.
+    check_titles:
+        Compare the titles of each subplot and the figure exactly.
+    check_legend:
+        Compare the legend's labels and title, or that neither plot has a
+        legend.
+    check_limits:
+        Compare the x and y axis limits.
+    check_styles:
+        Compare line colors, line styles and markers, scatter point colors,
+        and bar colors and orientation.
+    msg:
+        Optional message to show instead of the description of the
+        difference.
+
+    Raises
+    ------
+    AssertionError
+        Describing the first difference found.
+    TypeError
+        When *expected* is not a ``Figure`` or ``Axes``.
+    """
+    # matplotlib is an optional dependency, so only import it when a plot is checked.
+    from .plot_checks import plot_mismatch  # noqa: PLC0415
+
+    mismatch = plot_mismatch(
+        plot,
+        expected,
+        name=name,
+        rtol=rtol,
+        atol=atol,
+        check_scales=check_scales,
+        check_data=check_data,
+        ordered=ordered,
+        allow_extra=allow_extra,
+        check_labels=check_labels,
+        check_titles=check_titles,
+        check_legend=check_legend,
+        check_limits=check_limits,
+        check_styles=check_styles,
+    )
+    if mismatch is not None:
+        raise AssertionError(msg or mismatch)
+
+
 # ---------------------------------------------------------------------------
 # Sandbox-integrated assertions (call student function + compare result)
 # ---------------------------------------------------------------------------