diff --git a/src/moveroplot/daytime_scores.py b/src/moveroplot/daytime_scores.py index a85b291..7d5e8f3 100644 --- a/src/moveroplot/daytime_scores.py +++ b/src/moveroplot/daytime_scores.py @@ -80,7 +80,8 @@ def _daytime_scores_pipeline( def _daytime_score_transformation(df, header): - df["hh"] = df["hh"].astype(int) + # Use assign to avoid mutating the (potentially cached) input DataFrame + df = df.assign(hh=df["hh"].astype(int)) df = df.replace(float(header["Missing value code"][0]), np.NaN) return df diff --git a/src/moveroplot/load_files.py b/src/moveroplot/load_files.py index 2f1d7d2..2e3cdba 100644 --- a/src/moveroplot/load_files.py +++ b/src/moveroplot/load_files.py @@ -10,6 +10,28 @@ from .utils.atab import Atab +# Module-level cache for raw ATAB parse results, keyed by resolved file path. +# This avoids re-reading the same file from disk when multiple pipelines +# (e.g. total_scores and ensemble_scores) need the same underlying data. +_atab_cache: dict = {} + + +def _load_atab_cached(file_path: Path, sep: str = " "): + """Return (header, data) for an ATAB file, using cache when possible. + + The raw header dict and the cached DataFrame are returned directly. + Callers that apply a transform_func are expected to create a new + DataFrame (e.g. via ``df.replace``) rather than mutating in place. + When no transform is applied, ``load_relevant_files`` copies the + DataFrame to protect the cache. + """ + key = str(file_path.resolve()) + if key not in _atab_cache: + loaded_atab = Atab(file=file_path, sep=sep) + _atab_cache[key] = (loaded_atab.header, loaded_atab.data) + return _atab_cache[key] + + # pylint: disable=too-many-arguments,too-many-locals def is_valid_data(header): try: @@ -51,12 +73,13 @@ def load_relevant_files( in_lt_ranges = lt_range in lt_ranges if in_lt_ranges: - # extract header & dataframe - loaded_atab = Atab(file=file_path, sep=" ") - header = loaded_atab.header - df = loaded_atab.data + # extract header & dataframe (cached to avoid re-parsing) + header, df = _load_atab_cached(file_path, sep=" ") if transform_func: df = transform_func(df, header) + else: + # Copy to protect the cache when no transform is applied + df = df.copy() if is_valid_data(header): # add information to dict first_key, second_key = ( diff --git a/src/moveroplot/station_scores.py b/src/moveroplot/station_scores.py index ef03672..8bd3e93 100644 --- a/src/moveroplot/station_scores.py +++ b/src/moveroplot/station_scores.py @@ -1,13 +1,14 @@ # pylint: skip-file # relevant imports for parsing pipeline # Standard library +from functools import lru_cache from datetime import datetime from pathlib import Path # Third-party -import cartopy import cartopy.crs as ccrs import cartopy.feature as cfeature +import matplotlib.axes as _mpl_axes import matplotlib.colors as mcolors # relevant imports for plotting pipeline @@ -16,9 +17,6 @@ from cartopy.mpl.gridliner import LATITUDE_FORMATTER from cartopy.mpl.gridliner import LONGITUDE_FORMATTER -from matplotlib import cm -from matplotlib.colors import LinearSegmentedColormap -from matplotlib.colors import ListedColormap from netCDF4 import Dataset # First-party @@ -34,6 +32,128 @@ from .utils.FBI_scores_settings import param_score_range_fbi, _forward, _inverse, _forward_spec, _inverse_spec, fbi_custom_ticks +def _get_topo_rgba(topo_path: str): + """Load topography in rotated pole coordinates from a NetCDF file + and pre-compute the RGBA array.""" + try: + ds = Dataset(topo_path) + topo_x = ds["x_1"][:].data.copy() + topo_y = ds["y_1"][:].data.copy() + topo_hsurf = ds["HSURF"][0, ...].data.copy() + ds.close() + except Exception as exc: # pylint: disable=broad-except + print( + f"Warning: Failed to read topography file '{topo_path}': {exc}. " + "Skipping topography plotting." + ) + return None + + levels = np.arange(0, 6000, 300) + cmap = plt.get_cmap("gray_r") + norm = mcolors.BoundaryNorm(levels, cmap.N) + + # Map height values to RGBA + rgba = cmap(norm(topo_hsurf)) + + extent = ( + float(topo_x.min()), float(topo_x.max()), + float(topo_y.min()), float(topo_y.max()), + ) + return rgba, extent + + +def _get_cached_geometries( + category: str, name: str, scale: str +) -> list: + """Return (and cache) the shapely geometries for a NaturalEarth feature.""" + feature = cfeature.NaturalEarthFeature(category=category, name=name, scale=scale) + return list(feature.geometries()) + + +# Map features to draw on every station-score plot, with their styling. +# Mirrors the original add_feature() calls but with cached 10m geometries. +_MAP_FEATURES: list[tuple[cfeature.NaturalEarthFeature, dict]] = [ + (cfeature.NaturalEarthFeature("physical", "coastline", "10m"), { + "edgecolor": "black", + "facecolor": "none", + "alpha": 0.5, + }), + (cfeature.NaturalEarthFeature("cultural", "admin_0_boundary_lines_land", "10m"), { + "edgecolor": "black", + "facecolor": "none", + "linestyle": "--", + "alpha": 1, + }), + (cfeature.NaturalEarthFeature("physical", "ocean", "10m"), { + "facecolor": cfeature.COLORS["water"], + "edgecolor": "face", + }), + (cfeature.NaturalEarthFeature("physical", "lakes", "10m"), { + "facecolor": cfeature.COLORS["water"], + "edgecolor": "face", + "alpha": 0.5, + }), + (cfeature.NaturalEarthFeature("physical", "rivers_lake_centerlines", "10m"), { + "edgecolor": cfeature.COLORS["water"], + "facecolor": "none", + "alpha": 0.5, + }), + (cfeature.NaturalEarthFeature("physical", "lakes_europe", "10m"), { + "color": "#97b6e1", + "alpha": 0.5, + }), +] + + +@lru_cache(maxsize=2) +def _get_map_background_rgba(extent_key: str): + """Pre-render all map features into a single cached RGBA image.""" + _EXTENTS_PC = { + "ch": [5.8, 10.6, 45.75, 47.8], + "alps": [0.7, 16.5, 42.3, 50], + } + if extent_key not in _EXTENTS_PC: + return None + + proj = ccrs.RotatedPole(pole_longitude=-170, pole_latitude=43) + render_dpi = 150 + + fig_bg = plt.figure(figsize=(7.3, 5), dpi=render_dpi) + fig_bg.patch.set_alpha(0) + ax_bg = fig_bg.add_axes([0, 0, 1, 1], projection=proj) + ax_bg.set_extent(_EXTENTS_PC[extent_key], crs=ccrs.PlateCarree()) + ax_bg.patch.set_alpha(0) + try: + ax_bg.spines["geo"].set_visible(False) + except (KeyError, AttributeError): + try: + ax_bg.outline_patch.set_visible(False) + except AttributeError: + pass + + for feature, style in _MAP_FEATURES: + ax_bg.add_geometries( + _get_cached_geometries(feature.category, feature.name, feature.scale), + feature.crs, + **style, + ) + + native_extent = ax_bg.get_extent() # (x0, x1, y0, y1) in rotated-pole CRS + + fig_bg.canvas.draw() + w, h = fig_bg.canvas.get_width_height() + rgba = ( + np.frombuffer(fig_bg.canvas.buffer_rgba(), dtype=np.uint8) + .reshape(h, w, 4) + .copy() + .astype(np.float32) + / 255.0 + ) + plt.close(fig_bg) + + return rgba, native_extent + + def _calculate_figsize(num_rows, num_cols, single_plot_size=(8, 6), padding=(2, 2)): """Calculate the figure size given the number of rows and columns of subplots. @@ -62,17 +182,20 @@ def _initialize_plots(labels: list, scores: list, plot_setup: dict, topography=N ), nrows=num_rows, ncols=num_cols, - tight_layout=True, figsize=figsize, dpi=100, squeeze=False, ) + model_id = plot_setup["model_versions"][0][0] + extent_key = None for ax in axes.ravel(): - if "ch" in plot_setup["model_versions"][0][0]: + if "ch" in model_id: ax.set_extent([5.8, 10.6, 45.75, 47.8], crs=ccrs.PlateCarree()) - if "alps" in plot_setup["model_versions"][0][0]: + extent_key = "ch" + if "alps" in model_id: ax.set_extent([0.7, 16.5, 42.3, 50], crs=ccrs.PlateCarree()) - _add_features(ax, topography) + extent_key = "alps" + _add_features(ax, topography, extent_key=extent_key) fig.tight_layout(w_pad=8, h_pad=2, rect=(0.05, 0.05, 0.90, 0.90)) plt.subplots_adjust(bottom=0.15) return fig, axes @@ -320,7 +443,7 @@ def _station_score_transformation(df, header): # PLOTTING PIPELINE FOR STATION SCORES PLOTS -def _add_features(ax, topography=None): +def _add_features(ax, topography=None, extent_key=None): """Add features to map. # # point cartopy to the folder containing the shapefiles for the features on the map @@ -346,57 +469,57 @@ def _add_features(ax, topography=None): gl.xlabel_style = {"rotation": 0} gl.ylabel_style = {"rotation": 0} - ax.add_feature(cfeature.COASTLINE, alpha=0.5, rasterized=True, zorder=10) - ax.add_feature( - cfeature.BORDERS, linestyle="--", alpha=1, rasterized=True, zorder=10 - ) - ax.add_feature(cfeature.OCEAN, rasterized=True, zorder=10) - ax.add_feature(cfeature.LAKES, alpha=0.5, rasterized=True, zorder=10) - ax.add_feature(cfeature.RIVERS, alpha=0.5, rasterized=True, zorder=10) - ax.add_feature( - cartopy.feature.NaturalEarthFeature( - category="physical", - name="lakes_europe", - scale="10m", - rasterized=True, - ), - alpha=0.5, - rasterized=True, - color="#97b6e1", - zorder=10, - ) + # Use pre-rendered map background if available, avoiding costly + # FeatureArtist redraws and shapely.intersects on every draw(). + _bg_used = False + if extent_key is not None: + bg = _get_map_background_rgba(extent_key) + if bg is not None: + bg_rgba, (x0, x1, y0, y1) = bg + _mpl_axes.Axes.imshow( + ax, + bg_rgba, + extent=(x0, x1, y0, y1), + origin="upper", + interpolation="bilinear", + aspect="auto", + zorder=10, + ) + _bg_used = True + if not _bg_used: + # Fallback: add features individually (slower, triggers FeatureArtist) + for feature, style in _MAP_FEATURES: + ax.add_geometries( + _get_cached_geometries(feature.category, feature.name, feature.scale), + feature.crs, + rasterized=True, + zorder=10, + **style, + ) # ax.add_image(ShadedReliefESRI(), 8) - # add ICON-CH1-EPS topography on COSMO-1E grid + # add ICON-CH1-EPS topography on COSMO-1E grid (cached RGBA image) if topography: topo_file = Path(topography) - if not topo_file.is_file(): print( - f"Warning: --topography file '{topography}' not found, skipping topography plotting." + f"Warning: --topography file '{topography}' not found, " + "skipping topography plotting." ) - topo_file = None - - if topo_file is not None and topo_file.exists(): - try: - icon_ch1_eps_topo = Dataset(str(topo_file)) - ax.contourf( - icon_ch1_eps_topo["x_1"][:].data, - icon_ch1_eps_topo["y_1"][:].data, - icon_ch1_eps_topo["HSURF"][0, ...].data, - cmap="gray_r", - levels=np.arange(0, 8400, 400), - extend="both", - alpha=0.4, - ) - except Exception as exc: # pylint: disable=broad-except - print( - f"Warning: Failed to read topography file '{topo_file}': {exc}. Skipping topography plotting." - ) else: - if topo_file is not None: - print( - f"Warning: Topo file '{topo_file}' not found, skipping topography plotting." + result = _get_topo_rgba(str(topo_file)) + if result is not None: + rgba, extent = result + # Use base-class imshow to bypass cartopy reprojection; + # the topo data is already in the native RotatedPole CRS. + _mpl_axes.Axes.imshow( + ax, + rgba, + extent=extent, + origin="lower", + interpolation="bilinear", + aspect="auto", + zorder=5, ) diff --git a/src/moveroplot/utils/atab.py b/src/moveroplot/utils/atab.py index 09bfc55..05f7f03 100644 --- a/src/moveroplot/utils/atab.py +++ b/src/moveroplot/utils/atab.py @@ -1,5 +1,6 @@ """Atab file support.""" # Standard library +import io from pathlib import Path from typing import Any from typing import Dict @@ -37,6 +38,7 @@ def __init__(self, file: Path, sep: str = ";") -> None: self.n_header_lines = 0 self.header: Dict[str, list[str]] = {} self.data: pd.DataFrame = pd.DataFrame() + self._data_lines: str = "" self._parse() @@ -50,13 +52,13 @@ def _parse(self) -> None: self._parse_header() # Parse the data section - args: Dict[str, Any] = {"skiprows": self.n_header_lines, "parse_dates": False} + args: Dict[str, Any] = {"parse_dates": False} if self.sep == " ": args["delim_whitespace"] = True else: args["sep"] = self.sep - self.data = pd.read_csv(self.file, **args) + self.data = pd.read_csv(io.StringIO(self._data_lines), **args) if self.data.empty: raise OSError("ERROR: Atab file is empty") @@ -112,6 +114,8 @@ def _parse_header(self): # Stop extraction of header information if line contains no ":" if len(elements) == 1: self.n_header_lines = idx + # Keep remaining lines (including current) for data parsing + self._data_lines = line + "".join(lines) break # Store header information