fixed plot for loop

flatironinstitute · Oct 27, 2024 · dffec5a · dffec5a
1 parent c65383c
commit dffec5a
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/nemos/basis.py b/src/nemos/basis.py
@@ -953,11 +953,11 @@ def evaluate_on_grid(self, *n_samples: int) -> Tuple[Tuple[NDArray], NDArray]:
         >>> from nemos.basis import MSplineBasis
         >>> mspline_basis = MSplineBasis(n_basis_funcs=4, order=3)
         >>> sample_points, basis_values = mspline_basis.evaluate_on_grid(100)
-        >>> p = plt.plot(sample_points, basis_values, label=f'Function {i+1}')
-        >>> plt.title('M-Spline Basis Functions');
-        >>> plt.xlabel('Domain');
-        >>> plt.ylabel('Basis Function Value');
-        >>> l = plt.legend()
+        >>> p = plt.plot(sample_points, basis_values)
+        >>> _ = plt.title('M-Spline Basis Functions')
+        >>> _ = plt.xlabel('Domain')
+        >>> _ = plt.ylabel('Basis Function Value')
+        >>> _ = plt.legend([f'Function {i+1}' for i in range(4)]);
         """
         self._check_input_dimensionality(n_samples)