""" North and south polar-stereographic views of a latitude/longitude field. """ import os import numpy as np import matplotlib.pyplot as plt import matplotlib.path as mpath import cartopy.crs as ccrs from ..figure import (attach_format_coord, display_figure, is_video, name_window, save_figure, still_path) from ..topography import overlay_topography from . import layers as layer_draw from .movie import add_frame_slider, frame_label, write_movie def plot_polar_views(lon2d, lat2d, data2d, colormap, varname, units=None, topo_overlay=True, output_path=None, norm=None, frames=None, frame_dim=None, frame_axis=None, fps=12, interactive=True, layers=(), overlay_style='blend', interpolate=True): """ Plot two polar‐stereographic views (north & south) of the same data. If output_path is given, write '_north' and '_south' instead of opening windows. With `frames`, each hemisphere steps through the animated dimension: two movies when the output names one, a slider under each figure otherwise. `layers` are the `--overlay` variables the map has already prepared - the same arrays, so the two views cannot end up disagreeing about the seam or about which longitude a column belongs to. """ figs = [] # collect (pole, figure, mesh) so an animation can re-fill them for pole in ("north", "south"): # Choose projection and extent for each pole if pole == "north": proj = ccrs.NorthPolarStereo(central_longitude=180) extent = [-180, 180, 60, 90] else: proj = ccrs.SouthPolarStereo(central_longitude=180) extent = [-180, 180, -90, -60] # Create figure and GeoAxes fig = plt.figure(figsize=(8, 6)) ax = fig.add_subplot(1, 1, 1, projection=proj, aspect=True) ax.set_global() ax.set_extent(extent, ccrs.PlateCarree()) # Draw circular boundary theta = np.linspace(0, 2 * np.pi, 100) center, radius = [0.5, 0.5], 0.5 verts = np.vstack([np.sin(theta), np.cos(theta)]).T circle = mpath.Path(verts * radius + center) ax.set_boundary(circle, transform=ax.transAxes) # Add meridians/parallels gl = ax.gridlines( draw_labels=True, color='k', xlocs=range(-180, 181, 30), ylocs=range(-90, 91, 10), linestyle='--', linewidth=0.5 ) # Plot data in PlateCarree projection underlay = layer_draw.paint_blank_cells(ax, lon2d, lat2d, data2d, colormap, transform=ccrs.PlateCarree()) # The cells cartopy would warn about here are in the *other* hemisphere, # which this view never shows; see `layers.quiet_wrap_warning`. with layer_draw.quiet_wrap_warning(): cf = ax.pcolormesh( lon2d, lat2d, data2d, shading=layer_draw.mesh_shading(interpolate, lon2d, lat2d, data2d), cmap=colormap, norm=norm, transform=ccrs.PlateCarree() ) cf._dispnc_blank_underlay = underlay # Pass the coordinates in data order so the hover readout is not flipped # for files whose latitudes run north-to-south attach_format_coord(ax, data2d, lon2d[0, :], lat2d[:, 0], 'lon', 'lat', varname) # Optionally overlay topography if topo_overlay: overlay_topography(ax, transform=ccrs.PlateCarree(), levels=20) # The layers go on before the bars, so the base bar is placed knowing # how many are coming if layers: # Same far-hemisphere cells as the base mesh above with layer_draw.quiet_wrap_warning(): layer_draw.draw(fig, ax, layers, lon2d[0, :], lat2d[:, 0], style=overlay_style, transform=ccrs.PlateCarree(), interpolate=interpolate) # Colorbar and title cbar = fig.colorbar(cf, ax=ax, pad=0.1) label = varname + (f" ({units})" if units else "") cbar.set_label(label) title = f"{varname} — {pole.capitalize()} polar region" ax.set_title(title, pad=20, y=1.05, fontsize=12, fontweight='bold') name_window(fig, title) figs.append((pole, fig, ax, cf, title)) if frames is not None: _animate(figs, frames, frame_dim, frame_axis, output_path, fps, interactive, lon2d[0, :], lat2d[:, 0], varname) return if output_path: stem, ext = os.path.splitext(still_path(output_path)) for pole, fig, _ax, _cf, _title in figs: save_figure(fig, f"{stem}_{pole}{ext or '.png'}") else: # Show both figures for _pole, fig, _ax, _cf, _title in figs: display_figure(fig) def _animate(figs, frames, frame_dim, frame_axis, output_path, fps, interactive, lons, lats, varname): """ Step both hemispheres through the animated dimension. Each pole gets its own writer, so the pair comes out as '_north.mp4' and '_south.mp4' - the same naming the stills already use. """ stem, ext = os.path.splitext(output_path or '') for pole, fig, ax, mesh, title in figs: # The frame label goes below the disc rather than under the title: a # still is saved with bbox_inches='tight' and can absorb an overrunning # title, but a movie frame has a fixed canvas, and up there the second # line lands on the 180 degree gridline label. stamp = fig.text(0.5, 0.02, '', ha='center', fontsize=11) def on_frame(index, ax=ax, mesh=mesh, stamp=stamp): mesh.set_array(frames[index]) layer_draw.refill_blank_cells(mesh, frames[index]) stamp.set_text(frame_label(frame_axis, frame_dim, index)) # Keep the hover readout on the frame actually being shown try: attach_format_coord(ax, frames[index], lons, lats, 'lon', 'lat', varname) except ValueError: pass return (mesh, stamp) on_frame(0) if is_video(output_path): write_movie(fig, len(frames), on_frame, f"{stem}_{pole}{ext}", fps=fps) elif output_path: suffix = os.path.splitext(still_path(output_path))[1] or '.png' save_figure(fig, f"{stem}_{pole}{suffix}") elif interactive: add_frame_slider(fig, len(frames), on_frame, label=frame_dim or 'frame') if not output_path: for _pole, fig, _ax, _mesh, _title in figs: display_figure(fig)