diff --git a/src/pg_rad/main.py b/src/pg_rad/main.py index 19f4161..7d24b6f 100644 --- a/src/pg_rad/main.py +++ b/src/pg_rad/main.py @@ -2,6 +2,7 @@ import argparse from pg_rad.logger import setup_logger from pg_rad.landscape import LandscapeDirector +from pg_rad.plotting import LandscapeSlicePlotter def main(): @@ -20,13 +21,19 @@ def main(): default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], ) + parser.add_argument( + "--saveplot", + action="store_true", + help="Save the plot or not." + ) args = parser.parse_args() setup_logger(args.loglevel) if args.test: landscape = LandscapeDirector().build_test_landscape() - print(landscape.name) + plotter = LandscapeSlicePlotter() + plotter.plot(landscape, save=args.saveplot) if __name__ == "__main__": diff --git a/src/pg_rad/plotting/landscape_plotter.py b/src/pg_rad/plotting/landscape_plotter.py index 49693c9..d54ca9a 100644 --- a/src/pg_rad/plotting/landscape_plotter.py +++ b/src/pg_rad/plotting/landscape_plotter.py @@ -36,10 +36,10 @@ class LandscapeSlicePlotter: ax.set_aspect("equal") if save: - name = landscape.name.lower().replace(' ', '_') - plt.savefig( - f"{name}_z{self.z}.png" - ) + landscape_name = landscape.name.lower().replace(' ', '_') + filename = f"{landscape_name}_z{self.z}.png" + plt.savefig(filename) + logger.info("Plot saved to file: "+filename) if show: plt.show()