mirror of
https://github.com/pim-n/pg-rad
synced 2026-03-11 19:58:11 +01:00
add save and show options to Plotter
This commit is contained in:
@ -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",
|
||||
|
||||
Reference in New Issue
Block a user