feat: make a FieldPlotter for CalculusPatch

Author: alexfikl

sumpy/point_calculus.py

@@ -29,6 +29,8 @@
 
 from pytools import memoize_method, obj_array
 
+from sumpy.visualization import FieldPlotter
+
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
@@ -345,6 +347,9 @@ def norm(self,
         else:
             raise ValueError("unsupported norm")
 
+    def make_field_plotter(self) -> FieldPlotter:
+        return FieldPlotter(self.center, self.h, points=self._points_shaped)
+
     def plot_nodes(self) -> None:
         if self.dim != 2:
             raise ValueError(f"cannot plot {self.dim}d fields")
@@ -357,6 +362,11 @@ def plot_nodes(self) -> None:
             "o")
 
     def plot(self, f: Array1D[np.floating[Any]]) -> None:
+        from warnings import warn
+        warn(f"Calling '{type(self).__name__}.plot' is deprecated. Use "
+             f"'{type(self).__name__}.make_field_plotter' instead, which also "
+             "supports 3d fields.", DeprecationWarning, stacklevel=2)
+
         if self.dim != 2:
             raise ValueError(f"cannot plot {self.dim}d fields")
 

sumpy/visualization.py

@@ -96,52 +96,59 @@ def make_field_plotter_from_bbox(
 
 class FieldPlotter:
     """
+    .. autoattribute:: dimensions
+    .. autoattribute:: npoints
+    .. autoattribute:: points
+
     .. automethod:: set_matplotlib_limits
     .. automethod:: show_scalar_in_matplotlib
     .. automethod:: show_scalar_in_mayavi
     .. automethod:: write_vtk_file
     """
 
     dimensions: int
+    npoints: int
+    points: onp.Array2D[np.floating[Any]]
+
     a: onp.Array1D[np.floating[Any]]
     b: onp.Array1D[np.floating[Any]]
-
-    nd_points: onp.Array2D[np.floating[Any]]
-    points: onp.Array2D[np.floating[Any]]
-    npoints: int
+    nd_points: onp.ArrayND[np.floating[Any]]
 
     def __init__(self,
                  center: onp.ToArray1D[np.floating[Any]],
                  extent: float | onp.Array1D[np.floating[Any]] = 1,
-                 npoints: int | tuple[int, ...] = 1000) -> None:
+                 npoints: int | tuple[int, ...] = 1000,
+                 points: onp.ArrayND[np.floating[Any]] | None = None) -> None:
         center = np.asarray(center)
         dim, = cast("tuple[int]", center.shape)
 
         self.dimensions = dim
         self.a = a = center - 0.5 * extent
         self.b = b = center + 0.5 * extent
 
-        from numbers import Number
-        if isinstance(npoints, (int, Number)):
-            npoints = dim*(npoints,)
+        if points is None:
+            from numbers import Number
+            if isinstance(npoints, (int, Number)):
+                npoints = dim*(npoints,)
+            else:
+                if len(npoints) != dim:
+                    raise ValueError("length of npoints must match dimension")
+
+            for i in range(dim):
+                if npoints[i] == 1:
+                    a[i] = center[i]
+
+            mgrid_index = tuple(
+                    slice(a[i], b[i], 1j*npoints[i])
+                    for i in range(dim))
+            mgrid = np.mgrid[mgrid_index]
         else:
-            if len(npoints) != dim:
-                raise ValueError("length of npoints must match dimension")
-
-        for i in range(dim):
-            if npoints[i] == 1:
-                a[i] = center[i]
-
-        mgrid_index = tuple(
-                slice(a[i], b[i], 1j*npoints[i])
-                for i in range(dim))
-        mgrid = np.mgrid[mgrid_index]
+            mgrid = points
 
         # (axis, point x idx, point y idx, ...)
         self.nd_points = mgrid
-
-        self.points = self.nd_points.reshape(dim, -1).copy()
-        self.npoints = np.prod(npoints)
+        self.points = mgrid.reshape(dim, -1).copy()
+        self.npoints = mgrid.size
 
     def _get_nontrivial_dims(self) -> onp.Array1D[np.bool_]:
         return np.array(self.nd_points.shape[1:]) != 1