snapshot 2025-01-13 10:15:22.811433

gridlock/__init__.py

@@ -17,6 +17,8 @@ Dependencies:
 """
 from .error import GridError as GridError
 from .grid import Grid as Grid
+from .draw import XYZExtent
+
 
 __author__ = 'Jan Petykiewicz'
 __version__ = '1.2'

gridlock/draw.py

@@ -21,6 +21,73 @@ foreground_callable_t = Callable[[NDArray, NDArray, NDArray], NDArray]
 foreground_t = float | foreground_callable_t
 
 
+def _ctrspan(
+        xyz: str,
+        cmin: float | None,
+        cctr: float | None,
+        cmax: float | None,
+        lc: float | None,
+        ) -> tuple[float, float]:
+    if sum(cc is None for cc in (cmin, cctr, cmax, lc)) != 2:
+        raise GridError(f'Exactly two of {xyz}min, {xyz}ctr, {xyz}max, {xyz}x must be None!')
+
+    if lc is None:
+        if cctr is None:
+            # assert cmin is not None
+            # assert cmax is not None
+            cctr = 0.5 * (cmax + cmin)
+            lc = cmax - cmin
+        elif cmax is None:
+            # assert cmin is not None
+            # assert cctr is not None
+            lc = 2 * (cctr - cmin)
+        elif xmin is None:
+            # assert cctr is not None
+            # assert cmax is not None
+            lc = 2 * (cmax - cctr)
+    else:                        # noqa: PLR5501
+        if cctr is not None:
+            pass
+        elif cmax is None:
+            # assert xmin is not None
+            # assert lx is not None
+            cctr = cmin + 0.5 * lc
+        elif cmin is None:
+            # assert xmax is not None
+            # assert lx is not None
+            cctr = cmax - 0.5 * lc
+
+    return cctr, lc
+
+
+class XYZExtent:
+    center: NDArray[numpy.float64]
+    span: NDArray[numpy.float64]
+
+    def __init__(
+            self,
+            xmin: float | None = None,
+            xctr: float | None = None,
+            xmax: float | None = None,
+            lx: float | None = None,
+            ymin: float | None = None,
+            yctr: float | None = None,
+            ymax: float | None = None,
+            ly: float | None = None,
+            zmin: float | None = None,
+            zctr: float | None = None,
+            zmax: float | None = None,
+            lz: float | None = None,
+            ) -> None:
+
+        cx, lx = _ctrspan('x', cmin=xmin, cctr=xctr, cmax=xmax, lc=lx)
+        cy, ly = _ctrspan('y', cmin=ymin, cctr=yctr, cmax=ymax, lc=ly)
+        cz, lz = _ctrspan('z', cmin=zmin, cctr=zctr, cmax=zmax, lc=lz)
+        self.center = numpy.array((cx, cy, cz), dtype=float)
+        self.span = numpy.array((lx, ly, lz), dtype=float)
+
+
+
 class GridDrawMixin(GridPosMixin):
     def draw_polygons(
             self,
@@ -279,9 +346,9 @@ class GridDrawMixin(GridPosMixin):
     def draw_cuboid(
             self,
             cell_data: NDArray,
-            center: ArrayLike,
-            dimensions: ArrayLike,
             foreground: Sequence[foreground_t] | foreground_t,
+            *,
+            extent: XYZExtent | None = None
             ) -> None:
         """
         Draw an axis-aligned cuboid
@@ -293,7 +360,9 @@ class GridDrawMixin(GridPosMixin):
                  sizes of the cuboid
             foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
         """
-        dimensions = numpy.asarray(dimensions)
+        center = extent.center
+        dimensions = extent.span
+
         p = numpy.array([[-dimensions[0], +dimensions[1]],
                          [+dimensions[0], +dimensions[1]],
                          [+dimensions[0], -dimensions[1]],

gridlock/read.py

@@ -93,7 +93,7 @@ class GridReadMixin(GridPosMixin):
             sample_period: int = 1,
             finalize: bool = True,
             pcolormesh_args: dict[str, Any] | None = None,
-            ) -> tuple['matplotlib.axes.Axes', 'matplotlib.figure.Figure']:
+            ) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']:
         """
         Visualize a slice of a grid.
         Interpolates if given a position between two planes.
@@ -145,7 +145,7 @@ class GridReadMixin(GridPosMixin):
             sample_period: int = 1,
             show_edges: bool = True,
             finalize: bool = True,
-            ) -> tuple['matplotlib.axes.Axes', 'matplotlib.figure.Figure']:
+            ) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']:
         """
         Draw an isosurface plot of the device.