From 3fd2eafb2a146d0bdf3ece01c81e4aab8d38a2b5 Mon Sep 17 00:00:00 2001
From: Pim Nelissen <pim.nelissen@protonmail.com>
Date: Tue, 17 Feb 2026 09:41:54 +0100
Subject: [PATCH] add save and show options to Plotter

---
 src/pg_rad/plotting/landscape_plotter.py | 34 +++++++++++++++++-------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/pg_rad/plotting/landscape_plotter.py b/src/pg_rad/plotting/landscape_plotter.py
index 9c5552b..49693c9 100644
--- a/src/pg_rad/plotting/landscape_plotter.py
+++ b/src/pg_rad/plotting/landscape_plotter.py
@@ -5,17 +5,24 @@ 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):
+    def plot(
+        self,
+        landscape: Landscape,
+        z: int = 0,
+        show: bool = True,
+        save: bool = False
+    ):
         """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.
+            show (bool, optional): Show the plot. Defaults to True.
+            save (bool, optional): Save the plot. Defaults to False.
         """        """
 
         """
@@ -27,7 +34,15 @@ class LandscapeSlicePlotter:
         self._draw_point_sources(ax, landscape)
 
         ax.set_aspect("equal")
-        plt.show()
+
+        if save:
+            name = landscape.name.lower().replace(' ', '_')
+            plt.savefig(
+                f"{name}_z{self.z}.png"
+            )
+
+        if show:
+            plt.show()
 
     def _draw_base(self, ax, landscape):
         width, height = landscape.size[:2]
@@ -38,7 +53,7 @@ class LandscapeSlicePlotter:
         ax.set_title(f"Landscape (top-down, z = {self.z})")
 
     def _draw_path(self, ax, landscape):
-        if landscape.path.z < self.z:
+        if landscape.path.z <= self.z:
             ax.plot(landscape.path.x_list, landscape.path.y_list, 'bo-')
         else:
             logger.warning(
@@ -48,18 +63,19 @@ class LandscapeSlicePlotter:
 
     def _draw_point_sources(self, ax, landscape):
         for s in landscape.point_sources:
-            if s.z <= self.z:
+            x, y, z = s.pos
+            if z <= self.z:
                 dot = Circle(
-                        (s.x, s.y),
+                        (x, y),
                         radius=5,
                         color=s.color,
                         zorder=5
                 )
 
                 ax.text(
-                    s.x + 0.06,
-                    s.y + 0.06,
-                    s.name,
+                    x + 0.06,
+                    y + 0.06,
+                    s.name+", z="+str(z),
                     color=s.color,
                     fontsize=10,
                     ha="left",