Add title arg (#6)

Signed-off-by: jparisu <[email protected]>
jparisu · Feb 1, 2024 · cfbae56 · cfbae56
1 parent 8e70272
commit cfbae56
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/sIArena/terrain/plot/plot_2D.py b/src/sIArena/terrain/plot/plot_2D.py
@@ -8,7 +8,9 @@ def plot_terrain_2D(
             terrain: Terrain,
             paths: List[Path] = [],
             colors: List[str] = ['r', 'y', 'm', 'k', 'c', 'g', 'b'],
-            cmap: str = 'terrain'):
+            cmap: str = 'terrain',
+            title: str = 'Terrain',
+        ):
     """Plots the terrain and the given paths"""
 
     plt.clf()
@@ -27,6 +29,7 @@ def plot_terrain_2D(
 
     plt.xlabel('row')
     plt.ylabel('col')
+    plt.title(title)
 
     plt.colorbar()
     plt.show(block=False)

diff --git a/src/sIArena/terrain/plot/plot_3D.py b/src/sIArena/terrain/plot/plot_3D.py
@@ -11,7 +11,9 @@ def plot_terrain_3D(
             angles: List[tuple] = [(45, 45), (45, 225)],
             paths: List[Path] = [],
             colors: List[str] = ['r', 'y', 'm', 'k', 'c', 'g', 'b'],
-            cmap: str = 'terrain'):
+            cmap: str = 'terrain',
+            title: str = 'Terrain',
+        ):
     """Plots the terrain and the given paths"""
 
     fig = plt.figure(figsize=(18, 6))
@@ -56,4 +58,5 @@ def plot_terrain_3D(
         ax.set_xlabel('row')
         ax.set_ylabel('col')
 
+    fig.title(title)
     plt.show()