mirror of
https://github.com/pim-n/pg-rad
synced 2026-03-22 21:48:11 +01:00
134 lines
3.8 KiB
Python
134 lines
3.8 KiB
Python
import logging
|
|
|
|
from matplotlib import pyplot as plt
|
|
from matplotlib.patches import Circle
|
|
import numpy as np
|
|
|
|
from pg_rad.path import Path
|
|
from pg_rad.objects import PointSource
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Landscape:
|
|
"""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,
|
|
size: int | tuple[int, int, int] = 500,
|
|
scale: str = 'meters'
|
|
):
|
|
if isinstance(size, int):
|
|
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.
|
|
|
|
Returns:
|
|
fig, ax: Matplotlib figure objects.
|
|
"""
|
|
x_lim, y_lim, _ = self.world.shape
|
|
|
|
fig, ax = plt.subplots()
|
|
ax.set_xlim(right=x_lim)
|
|
ax.set_ylim(top=y_lim)
|
|
ax.set_xlabel(f"X [{self.scale}]")
|
|
ax.set_ylabel(f"Y [{self.scale}]")
|
|
|
|
if self.path is not None:
|
|
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.
|
|
|
|
Args:
|
|
*sources (pg_rad.sources.PointSource): One or more sources,
|
|
passed as Source1, Source2, ...
|
|
Raises:
|
|
ValueError: If the source is outside the boundaries of the
|
|
landscape.
|
|
"""
|
|
max_x, max_y, max_z = self.world.shape[:3]
|
|
|
|
if any(
|
|
not (0 <= source.x < max_x and
|
|
0 <= source.y < max_y and
|
|
0 <= source.z < max_z)
|
|
for source in sources
|
|
):
|
|
raise ValueError("One or more sources are outside the landscape!")
|
|
|
|
self.sources.extend(sources)
|
|
|
|
def set_path(self, path: Path):
|
|
"""
|
|
Set the path in the landscape.
|
|
"""
|
|
self.path = path
|
|
|
|
|
|
def create_landscape_from_path(path: Path, max_z: float | int = 50):
|
|
"""Generate a landscape from a path, using its dimensions to determine
|
|
the size of the landscape.
|
|
|
|
Args:
|
|
path (Path): A Path object describing the trajectory.
|
|
max_z (int, optional): Height of the world. Defaults to 50 meters.
|
|
|
|
Returns:
|
|
landscape (pg_rad.landscape.Landscape): A landscape with dimensions
|
|
based on the provided Path.
|
|
"""
|
|
max_x = np.ceil(max(path.x_list))
|
|
max_y = np.ceil(max(path.y_list))
|
|
|
|
landscape = Landscape(size=(max_x, max_y, max_z))
|
|
landscape.path = path
|
|
return landscape
|