TST: Calculate RMS and diff image in C++

The current implementation is not slow, but uses a lot of memory per
image.

In `compare_images`, we have:

- one actual and one expected image as uint8 (2×image)
- both converted to int16 (though original is thrown away) (4×)

which adds up to 4× the image allocated in this function.

Then it calls `calculate_rms`, which has:

- a difference between them as int16 (2×)
- the difference cast to 64-bit float (8×)
- the square of the difference as 64-bit float (though possibly the
  original difference was thrown away) (8×)

which at its peak has 16× the image allocated in parallel.

If the RMS is over the desired tolerance, then `save_diff_image` is
called, which:

- loads the actual and expected images _again_ as uint8 (2× image)
- converts both to 64-bit float (throwing away the original) (16×)
- calculates the difference (8×)
- calculates the absolute value (8×)
- multiples that by 10 (in-place, so no allocation)
- clips to 0-255 (8×)
- casts to uint8 (1×)

which at peak uses 32× the image.

So at their peak, `compare_images`→`calculate_rms` will have 20× the
image allocated, and then `compare_images`→`save_diff_image` will have
36× the image allocated. This is generally not a problem, but on
resource-constrained places like WASM, it can sometimes run out of
memory just in `calculate_rms`.

This implementation in C++ always allocates the diff image, even when
not needed, but doesn't have all the temporaries, so it's a maximum of
3× the image size (plus a few scalar temporaries).

Author: QuLogic

lib/matplotlib/testing/compare.py

@@ -19,7 +19,7 @@
 from PIL import Image
 
 import matplotlib as mpl
-from matplotlib import cbook
+from matplotlib import cbook, _image
 from matplotlib.testing.exceptions import ImageComparisonFailure
 
 _log = logging.getLogger(__name__)
@@ -412,7 +412,7 @@ def compare_images(expected, actual, tol, in_decorator=False):
 
     The two given filenames may point to files which are convertible to
     PNG via the `!converter` dictionary. The underlying RMS is calculated
-    with the `.calculate_rms` function.
+    in a similar way to the `.calculate_rms` function.
 
     Parameters
     ----------
@@ -483,17 +483,12 @@ def compare_images(expected, actual, tol, in_decorator=False):
         if np.array_equal(expected_image, actual_image):
             return None
 
-    # convert to signed integers, so that the images can be subtracted without
-    # overflow
-    expected_image = expected_image.astype(np.int16)
-    actual_image = actual_image.astype(np.int16)
-
-    rms = calculate_rms(expected_image, actual_image)
+    rms, abs_diff = _image.calculate_rms_and_diff(expected_image, actual_image)
 
     if rms <= tol:
         return None
 
-    save_diff_image(expected, actual, diff_image)
+    Image.fromarray(abs_diff).save(diff_image, format="png")
 
     results = dict(rms=rms, expected=str(expected),
                    actual=str(actual), diff=str(diff_image), tol=tol)

src/_image_wrapper.cpp

@@ -1,6 +1,8 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/numpy.h>
 
+#include <algorithm>
+
 #include "_image_resample.h"
 #include "py_converters.h"
 
@@ -202,6 +204,70 @@ image_resample(py::array input_array,
 }
 
 
+// This is used by matplotlib.testing.compare to calculate RMS and a difference image.
+static py::tuple
+calculate_rms_and_diff(py::array_t<unsigned char> expected_image,
+                       py::array_t<unsigned char> actual_image)
+{
+    for (const auto & [image, name] : {std::pair{expected_image, "Expected"},
+                                       std::pair{actual_image, "Actual"}})
+    {
+        if (image.ndim() != 3) {
+            auto exceptions = py::module_::import("matplotlib.testing.exceptions");
+            auto ImageComparisonFailure = exceptions.attr("ImageComparisonFailure");
+            py::set_error(
+                ImageComparisonFailure,
+                "{name} image must be 3-dimensional, but is {ndim}-dimensional"_s.format(
+                    "name"_a=name, "ndim"_a=expected_image.ndim()));
+            throw py::error_already_set();
+        }
+    }
+
+    auto height = expected_image.shape(0);
+    auto width = expected_image.shape(1);
+    auto depth = expected_image.shape(2);
+
+    if (height != actual_image.shape(0) || width != actual_image.shape(1) ||
+            depth != actual_image.shape(2)) {
+        auto exceptions = py::module_::import("matplotlib.testing.exceptions");
+        auto ImageComparisonFailure = exceptions.attr("ImageComparisonFailure");
+        py::set_error(
+            ImageComparisonFailure,
+            "Image sizes do not match expected size: {expected_image.shape} "_s
+            "actual size {actual_image.shape}"_s.format(
+                "expected_image"_a=expected_image, "actual_image"_a=actual_image));
+        throw py::error_already_set();
+    }
+    auto expected = expected_image.unchecked<3>();
+    auto actual = actual_image.unchecked<3>();
+
+    py::ssize_t diff_dims[3] = {height, width, 3};
+    py::array_t<unsigned char> diff_image(diff_dims);
+    auto diff = diff_image.mutable_unchecked<3>();
+
+    double total = 0.0;
+    for (auto i = 0; i < height; i++) {
+        for (auto j = 0; j < width; j++) {
+            for (auto k = 0; k < depth; k++) {
+                auto pixel_diff = static_cast<double>(expected(i, j, k)) -
+                                  static_cast<double>(actual(i, j, k));
+
+                total += pixel_diff*pixel_diff;
+
+                if (k != 3) { // Hard-code a fully solid alpha channel by omitting it.
+                    diff(i, j, k) = static_cast<unsigned char>(std::clamp(
+                        abs(pixel_diff) * 10, // Expand differences in luminance domain.
+                        0.0, 255.0));
+                }
+            }
+        }
+    }
+    total = total / (width * height * depth);
+
+    return py::make_tuple(sqrt(total), diff_image);
+}
+
+
 PYBIND11_MODULE(_image, m, py::mod_gil_not_used())
 {
     py::enum_<interpolation_e>(m, "_InterpolationType")
@@ -234,4 +300,7 @@ PYBIND11_MODULE(_image, m, py::mod_gil_not_used())
         "norm"_a = false,
         "radius"_a = 1,
         image_resample__doc__);
+
+    m.def("calculate_rms_and_diff", &calculate_rms_and_diff,
+          "expected_image"_a, "actual_image"_a);
 }