diff --git a/mlair/data_handler/iterator.py b/mlair/data_handler/iterator.py index 39e20020f4f80a872428681d53e2ec9f1a3dd3f7..18466fd9ede1e666f52b49fa461585a7e38410dd 100644 --- a/mlair/data_handler/iterator.py +++ b/mlair/data_handler/iterator.py @@ -33,13 +33,18 @@ class StandardIterator(Iterator): class DataCollection(Iterable): - def __init__(self, collection: list = None): + def __init__(self, collection: list = None, name: str = None): if collection is None: collection = [] assert isinstance(collection, list) self._collection = collection self._mapping = {} self._set_mapping() + self._name = name + + @property + def name(self): + return self._name def __len__(self): return len(self._collection) diff --git a/mlair/plotting/postprocessing_plotting.py b/mlair/plotting/postprocessing_plotting.py index dc33219055db83012449be56c4071df2177af6e2..2ffeb536cfd1a7fc5094ffa72c23ee8819a324d6 100644 --- a/mlair/plotting/postprocessing_plotting.py +++ b/mlair/plotting/postprocessing_plotting.py @@ -10,6 +10,7 @@ from typing import Dict, List, Tuple import matplotlib import matplotlib.patches as mpatches +import matplotlib.lines as mlines import matplotlib.pyplot as plt import matplotlib.dates as mdates import numpy as np @@ -279,16 +280,22 @@ class PlotStationMap(AbstractPlotClass): import cartopy.crs as ccrs if generators is not None: + legend_elements = [] for element in generators: data_collection, plot_opts = self._get_collection_and_opts(element) marker = plot_opts.get("marker", "s") - ms = plot_opts.get("markersize", 6) + ms = plot_opts.get("ms", 6) mec = plot_opts.get("mec", "k") mfc = plot_opts.get("mfc", "b") + name = data_collection.name or "unknown" + legend_elements.append( + mlines.Line2D([], [], mfc=mfc, mec=mec, marker=marker, ms=ms, linestyle='None', label=name)) for station in data_collection: coords = station.get_coordinates() IDx, IDy = coords["lon"], coords["lat"] self._ax.plot(IDx, IDy, mfc=mfc, mec=mec, marker=marker, ms=ms, transform=ccrs.PlateCarree()) + if len(legend_elements) > 0: + self._ax.legend(handles=legend_elements, loc='best') @staticmethod def _get_collection_and_opts(element): diff --git a/mlair/run_modules/post_processing.py b/mlair/run_modules/post_processing.py index 0b9393e0ef902d4256bded0e6466d00a6d79910d..300d2e1dbd2ded279ee289ee64a2acd5f5c36fbc 100644 --- a/mlair/run_modules/post_processing.py +++ b/mlair/run_modules/post_processing.py @@ -277,8 +277,8 @@ class PostProcessing(RunEnvironment): logging.warning( f"Skip 'PlotStationMap` because running on a hpc node: {self.data_store.get('hostname')}") else: - gens = [(self.train_val_data, {"mfc": "r", "marker": 8}), - (self.test_data, {"mfc": "b", "marker": 9})] + gens = [(self.train_val_data, {"mfc": "r", "marker": 8, "ms": 10}), + (self.test_data, {"mfc": "b", "marker": 9, "ms": 10})] PlotStationMap(generators=gens, plot_folder=self.plot_path) if "PlotMonthlySummary" in plot_list: PlotMonthlySummary(self.test_data.keys(), path, r"forecasts_%s_test.nc", self.target_var, diff --git a/mlair/run_modules/pre_processing.py b/mlair/run_modules/pre_processing.py index c9e92e7deb22dcff12e9d4ab982f14289f764a97..bd3b9ec6fd471dcb6e794a3ba9b498e18ad76a37 100644 --- a/mlair/run_modules/pre_processing.py +++ b/mlair/run_modules/pre_processing.py @@ -257,7 +257,7 @@ class PreProcessing(RunEnvironment): logging.info("setup transformation using train data exclusively") self.transformation(data_handler, set_stations) # start station check - collection = DataCollection() + collection = DataCollection(name=set_name) valid_stations = [] kwargs = self.data_store.create_args_dict(data_handler.requirements(), scope=set_name)