Move plotting functionality out of Landscape to LandscapeSlicePlotter

src/pg_rad/landscape/landscape.py

@@ -13,77 +13,34 @@ logger = logging.getLogger(__name__)
 
 
 class Landscape:
-    """A generic Landscape that can contain a Path and sources.
+    """
-
+    A generic Landscape that can contain a Path and sources.
-    Args:
-        air_density (float, optional): Air density, kg/m^3. Defaults to 1.243.
-        size (int | tuple[int, int, int], optional): Size of the world.
-        Defaults to 500.
-        scale (str, optional): The scale of the size argument passed.
-        Defaults to 'meters'.
     """
     def __init__(
             self,
-            air_density: float = 1.243,
+            path: Path,
-            size: int | tuple[int, int, int] = 500,
+            point_sources: list[PointSource] = [],
-            scale: str = 'meters'
+            size: tuple[int, int, int] = [500, 500, 50],
+            air_density: float = 1.243
             ):
-        if isinstance(size, int):
+        """Initialize a landscape.
-            self.world = np.zeros((size, size, size))
-        elif isinstance(size, tuple) and len(size) == 3:
-            self.world = np.zeros(size)
-        else:
-            raise TypeError("size must be integer or a tuple of 3 integers.")
-
-        self.air_density = air_density
-        self.scale = scale
-        self.path: Path = None
-        self.sources: list[PointSource] = []
-        logger.debug("Landscape initialized.")
-
-    def plot(self, z: float | int = 0):
-        """Plot a slice of the world at a height `z`.
 
         Args:
-            z (int, optional): Height of slice. Defaults to 0.
+            path (Path): A Path object.
+            point_sources (list[PointSource], optional): List of point sources.
+            air_density (float, optional): Air density in kg/m^3. Defaults to 1.243.
+            size (tuple[int, int, int], optional): (x,y,z) dimensions of world in meters. Defaults to [500, 500, 50].
 
-        Returns:
+        Raises:
-            fig, ax: Matplotlib figure objects.
+            TypeError: _description_
         """
-        x_lim, y_lim, _ = self.world.shape
 
-        fig, ax = plt.subplots()
+        self.path = path
-        ax.set_xlim(right=x_lim)
+        self.point_sources = point_sources
-        ax.set_ylim(top=y_lim)
+        self.size = size
-        ax.set_xlabel(f"X [{self.scale}]")
+        self.air_density = air_density
-        ax.set_ylabel(f"Y [{self.scale}]")
 
-        if self.path is not None:
+        logger.debug("Landscape initialized.")
-            ax.plot(self.path.x_list, self.path.y_list, 'bo-')
-
-        for s in self.sources:
-            if np.isclose(s.z, z):
-                dot = Circle(
-                        (s.x, s.y),
-                        radius=5,
-                        color=s.color,
-                        zorder=5
-                )
-
-                ax.text(
-                    s.x + 0.06,
-                    s.y + 0.06,
-                    s.name,
-                    color=s.color,
-                    fontsize=10,
-                    ha="left",
-                    va="bottom",
-                    zorder=6
-                )
-
-                ax.add_patch(dot)
-
-        return fig, ax
 
     def add_sources(self, *sources: PointSource):
         """Add one or more point sources to the world.
@@ -150,3 +107,8 @@ def create_landscape_from_path(path: Path, max_z: float | int = 50):
     landscape = Landscape(size=(max_x, max_y, max_z))
     landscape.path = path
     return landscape
+
+
+class LandscapeBuilder:
+    def __init__(self):
+        pass

src/pg_rad/plotting/__init__.py

@@ -0,0 +1,7 @@
+# do not expose internal logger when running mkinit
+__ignore__ = ["logger"]
+from pg_rad.plotting import landscape_plotter
+
+from pg_rad.plotting.landscape_plotter import (LandscapeSlicePlotter,)
+
+__all__ = ['LandscapeSlicePlotter', 'landscape_plotter']

src/pg_rad/plotting/landscape_plotter.py

@@ -0,0 +1,75 @@
+import logging
+
+from matplotlib import pyplot as plt
+from matplotlib.patches import Circle
+
+from pg_rad.landscape import Landscape
+
+
+logger = logging.getLogger(__name__)
+
+
+class LandscapeSlicePlotter:
+    def plot(self, landscape: Landscape, z: int = 0):
+        """Plot a top-down slice of the landscape at a height z.
+
+        Args:
+            landscape (Landscape): the landscape to plot
+            z (int, optional): Height at which to plot slice. Defaults to 0.
+        """        """
+
+        """
+        self.z = z
+        fig, ax = plt.subplots()
+
+        self._draw_base(ax, landscape)
+        self._draw_path(ax, landscape)
+        self._draw_point_sources(ax, landscape)
+
+        ax.set_aspect("equal")
+        plt.show()
+
+    def _draw_base(self, ax, landscape):
+        width, height = landscape.size[:2]
+        ax.set_xlim(right=width)
+        ax.set_ylim(top=height)
+        ax.set_xlabel("X [m]")
+        ax.set_ylabel("Y [m]")
+        ax.set_title(f"Landscape (top-down, z = {self.z})")
+
+    def _draw_path(self, ax, landscape):
+        if landscape.path.z < self.z:
+            ax.plot(landscape.path.x_list, landscape.path.y_list, 'bo-')
+        else:
+            logger.warning(
+                "Path is above the slice height z."
+                "It will not show on the plot."
+                )
+
+    def _draw_point_sources(self, ax, landscape):
+        for s in landscape.point_sources:
+            if s.z <= self.z:
+                dot = Circle(
+                        (s.x, s.y),
+                        radius=5,
+                        color=s.color,
+                        zorder=5
+                )
+
+                ax.text(
+                    s.x + 0.06,
+                    s.y + 0.06,
+                    s.name,
+                    color=s.color,
+                    fontsize=10,
+                    ha="left",
+                    va="bottom",
+                    zorder=6
+                )
+
+                ax.add_patch(dot)
+            else:
+                logger.warning(
+                    f"Source {s.name} is above slice height z."
+                    "It will not show on the plot."
+                    )