add save and show options to Plotter

This commit is contained in:
Pim Nelissen
2026-02-17 09:41:54 +01:00
parent f49b2a5f8a
commit 3fd2eafb2a

View File

@ -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,6 +34,14 @@ class LandscapeSlicePlotter:
self._draw_point_sources(ax, landscape)
ax.set_aspect("equal")
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):
@ -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",