Improve test suite: ridgeplot._hist

docs/reference/changelog.md

@@ -9,12 +9,20 @@ Unreleased changes
 
 - Implement a new `template` parameter to allow users to specify a Plotly figure template ({gh-pr}`389`)
 
+### Bug fixes
+
+- Raise a clear `ValueError` when an empty samples array is passed to either density estimation method (KDE or histogram binning), instead of failing downstream with obscure errors ({gh-pr}`365`)
+- Raise clear `ValueError`s for invalid `sample_weights` values (negative or all-zero weights), instead of failing downstream with obscure errors ({gh-pr}`365`)
+- Improve the error messages raised when the length of a `sample_weights` array does not match the length of the corresponding samples array, or when the number of rows in a shallow per-row attributes array (e.g., `sample_weights` or `labels`) does not match the number of rows in the samples array ({gh-pr}`365`)
+
 ### Documentation
 
 - Add a section on theming with Plotly templates to the getting started guide ({gh-pr}`389`)
+- Fix the description of the `sample_weights` parameter in the `ridgeplot()` docstring, which incorrectly stated that the weights only applied to KDE, and document the accepted forms and their broadcasting semantics ({gh-pr}`365`)
 
 ### CI/CD
 
+- Improve the unit test suite for the `ridgeplot._hist` module ({gh-pr}`365`)
 - Review the coverage configuration in light of `covdefaults`, adopting its `assert_never` exclusion and `skip_covered` report setting, and raising all package coverage gates to 100% ({gh-pr}`390`)
 - Adopt pytest's strict mode, following the recommendations from pytest's "Good Integration Practices" guide ({gh-pr}`387`)
 - Make the e2e path sanity check independent of the checkout directory's name, so tests can run from git worktrees ({gh-pr}`391`)

src/ridgeplot/_hist.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from ridgeplot._kde import normalize_sample_weights
+from ridgeplot._kde import normalize_sample_weights, validate_trace_samples_and_weights
 
 if TYPE_CHECKING:
     from ridgeplot._types import (
@@ -25,16 +25,8 @@ def bin_trace_samples(
     nbins: int,
     weights: SampleWeights = None,
 ) -> DensityTrace:
-    trace_samples = np.asarray(trace_samples, dtype=float)
-    if not np.isfinite(trace_samples).all():
-        raise ValueError("The samples array should not contain any infs or NaNs.")
-    if weights is not None:
-        weights = np.asarray(weights, dtype=float)
-        if len(weights) != len(trace_samples):
-            raise ValueError("The weights array should have the same length as the samples array.")
-        if not np.isfinite(weights).all():
-            raise ValueError("The weights array should not contain any infs or NaNs.")
-    hist_counts, hist_edges = np.histogram(trace_samples, bins=nbins, weights=weights)
+    samples_array, weights_array = validate_trace_samples_and_weights(trace_samples, weights)
+    hist_counts, hist_edges = np.histogram(samples_array, bins=nbins, weights=weights_array)
     bin_centers = 0.5 * (hist_edges[:-1] + hist_edges[1:])
     return [(float(x), float(y)) for x, y in zip(bin_centers, hist_counts, strict=True)]
 

src/ridgeplot/_kde.py

@@ -98,6 +98,40 @@ def normalize_sample_weights(
     return sample_weights
 
 
+def validate_trace_samples_and_weights(
+    trace_samples: SamplesTrace,
+    weights: SampleWeights,
+) -> tuple[npt.NDArray[np.floating[Any]], npt.NDArray[np.floating[Any]] | None]:
+    """Validate a trace's samples and weights and coerce them to NumPy arrays.
+
+    This validation logic is shared between the KDE
+    (:func:`estimate_density_trace`) and histogram binning
+    (:func:`ridgeplot._hist.bin_trace_samples`) paths.
+    """
+    samples_array = np.asarray(trace_samples, dtype=float)
+    if len(samples_array) == 0:
+        raise ValueError("The samples array should not be empty.")
+    if not np.isfinite(samples_array).all():
+        raise ValueError("The samples array should not contain any infs or NaNs.")
+    if weights is None:
+        return samples_array, None
+    weights_array = np.asarray(weights, dtype=float)
+    if len(weights_array) != len(samples_array):
+        raise ValueError(
+            f"The weights array should have the same length as the samples array "
+            f"(got {len(weights_array)} weights for {len(samples_array)} samples). "
+            f"Note that a single flat array of weights is broadcast to all traces "
+            f"in the samples array."
+        )
+    if not np.isfinite(weights_array).all():
+        raise ValueError("The weights array should not contain any infs or NaNs.")
+    if np.any(weights_array < 0):
+        raise ValueError("The weights array should not contain negative values.")
+    if not np.any(weights_array):
+        raise ValueError("The weights array should not be all zeros.")
+    return samples_array, weights_array
+
+
 def estimate_density_trace(
     trace_samples: SamplesTrace,
     points: KDEPoints,
@@ -110,16 +144,14 @@ def estimate_density_trace(
     For a given set of sample values, computes the kernel densities (KDE) at
     the given points.
     """
-    trace_samples = np.asarray(trace_samples, dtype=float)
-    if not np.isfinite(trace_samples).all():
-        raise ValueError("The samples array should not contain any infs or NaNs.")
+    samples_array, weights_array = validate_trace_samples_and_weights(trace_samples, weights)
     if isinstance(points, int):
         # By default, we'll use a 'hard' KDE span. That is, we'll
         # evaluate the densities and N equally spaced points
         # over the range [min(samples), max(samples)]
         density_x = np.linspace(
-            start=min(trace_samples),
-            stop=max(trace_samples),
+            start=min(samples_array),
+            stop=max(samples_array),
             num=points,
         )
     else:
@@ -130,15 +162,9 @@ def estimate_density_trace(
                 f"The 'points' at which KDE is computed should be represented by a "
                 f"one-dimensional array, got an array of shape {density_x.shape} instead."
             )
-    if weights is not None:
-        weights = np.asarray(weights, dtype=float)
-        if len(weights) != len(trace_samples):
-            raise ValueError("The weights array should have the same length as the samples array.")
-        if not np.isfinite(weights).all():
-            raise ValueError("The weights array should not contain any infs or NaNs.")
 
     # ref: https://github.com/tpvasconcelos/ridgeplot/issues/116
-    dens = sm.nonparametric.KDEUnivariate(trace_samples)
+    dens = sm.nonparametric.KDEUnivariate(samples_array)
 
     # I'm hard-coding `fft=kernel == "gau" and weights is not None`
     # to avoid exposing yet another KDE parameter in ridgeplot()
@@ -150,9 +176,9 @@ def estimate_density_trace(
     # directly to the ridgeplot() figure factory.
     dens.fit(
         kernel=kernel,
-        fft=kernel == "gau" and weights is None,
+        fft=kernel == "gau" and weights_array is None,
         bw=bandwidth,  # pyright: ignore[reportArgumentType]
-        weights=weights,
+        weights=weights_array,
     )
     density_y = dens.evaluate(density_x)
     density_y = _validate_densities(x=density_x, y=density_y, kernel=kernel)

src/ridgeplot/_ridgeplot.py

@@ -266,9 +266,20 @@ def ridgeplot(
         .. versionadded:: 0.3.0
 
     sample_weights : SampleWeightsArray or ShallowSampleWeightsArray or SampleWeights or None
-        An (optional) array of KDE weights corresponding to each sample. The
-        weights should have the same shape as the samples array. If not
-        specified (default), all samples will be weighted equally.
+        An (optional) array of weights corresponding to each sample. The
+        weights are used by both density estimation methods (i.e., KDE and
+        histogram binning). If not specified (default), all samples will be
+        weighted equally. The following forms are accepted:
+
+        - A single flat array of weights, which will be broadcast to all
+          traces in the :paramref:`.samples` array. This requires all traces
+          to have the same number of samples as there are weights.
+        - A shallow array of weights with one entry per row, where each entry
+          is either a flat array of weights or ``None``. Each entry will be
+          applied to all traces in the corresponding row.
+        - A full (nested) array of weights with the same shape as the
+          :paramref:`.samples` array, specifying the weights for each trace
+          individually.
 
     norm : NormalisationOption or None
         The normalisation option to use when normalising the densities. If not

src/ridgeplot/_utils.py

@@ -232,14 +232,19 @@ def normalise_row_attrs(
     Raises
     ------
     ValueError
-        If the number of traces does not match the number of attributes for a
-        row.
+        If the number of rows in ``attrs`` does not match the number of rows
+        in ``l2_target``, or if the number of traces does not match the number
+        of attributes for a row.
 
     Examples
     --------
     >>> densities = [[[(0, 0), (1, 1), (2, 0)]]]  # Single row, single trace
     >>> normalise_row_attrs(["A"], densities)
     [['A']]
+    >>> normalise_row_attrs([["A"], ["B"]], densities)
+    Traceback (most recent call last):
+    ...
+    ValueError: Mismatch between number of rows in attrs (2) and samples/densities (1).
     >>> normalise_row_attrs([["A", "B"]], densities)
     Traceback (most recent call last):
     ...
@@ -311,6 +316,12 @@ def normalise_row_attrs(
     if not is_collection_l2(attrs):
         attrs = [attrs] if len(l2_target) == 1 else [[attr] for attr in attrs]
 
+    if len(attrs) != len(l2_target):
+        raise ValueError(
+            f"Mismatch between number of rows in attrs ({len(attrs)}) "
+            f"and samples/densities ({len(l2_target)})."
+        )
+
     result: list[list[_V]] = []
     for i, (row_attrs, row_traces) in enumerate(zip(attrs, l2_target, strict=True)):
         n_traces = len(row_traces)